スケーラブルソフトマックスはアテンションにおいて優れています。

Scalable-Softmax Is Superior for Attention

January 31, 2025

著者: Ken M. Nakanishi

cs.AI

要旨

Softmax 関数によって出力されるベクトルの最大要素は、入力ベクトルのサイズが増加するにつれてゼロに近づく。Transformer ベースの言語モデルは、Softmax を使用して注意スコアを計算し、コンテキストのサイズが大きくなると注意分布が平坦化する。これにより、モデルが重要な情報を効果的に優先する能力が低下し、長さの一般化が制限される可能性がある。この問題に対処するために、私たちは Scalable-Softmax (SSMax) を提案する。SSMax は、入力ベクトルのサイズが異なるシナリオで Softmax を置き換えるものであり、既存の Transformer ベースのアーキテクチャにシームレスに統合できる。言語モデリングにおける実験結果は、SSMax を使用するモデルが、事前学習中により速い損失削減を達成するだけでなく、長いコンテキストと重要な情報の取得においても性能が著しく向上することを示している。さらに、注意スコアの分析により、SSMax を使用することで、モデルが長いコンテキストでも重要な情報に焦点を当てることが可能になることが明らかになった。また、事前学習の初めから SSMax を使用するモデルが長さの一般化をより良く達成する一方、既に事前学習を開始したモデルでも、事前学習中または事前学習後に注意層の Softmax を SSMax に置き換えることで、この能力の一部を獲得することができる。

English

The maximum element of the vector output by the Softmax function approaches zero as the input vector size increases. Transformer-based language models rely on Softmax to compute attention scores, causing the attention distribution to flatten as the context size grows. This reduces the model's ability to prioritize key information effectively and potentially limits its length generalization. To address this problem, we propose Scalable-Softmax (SSMax), which replaces Softmax in scenarios where the input vector size varies. SSMax can be seamlessly integrated into existing Transformer-based architectures. Experimental results in language modeling show that models using SSMax not only achieve faster loss reduction during pretraining but also significantly improve performance in long contexts and key information retrieval. Furthermore, an analysis of attention scores reveals that SSMax enables the model to focus attention on key information even in long contexts. Additionally, although models that use SSMax from the beginning of pretraining achieve better length generalization, those that have already started pretraining can still gain some of this ability by replacing Softmax in the attention layers with SSMax, either during or after pretraining.