Scalable-Softmax Is Superior for Attention

要約

ソフトマックス関数によるベクトル出力の最大要素は、入力ベクトルサイズが大きくなるとゼロに近づきます。
トランスベースの言語モデルは、SoftMaxに依存して注意スコアを計算し、コンテキストサイズが大きくなるにつれて注意分布が平らになります。
これにより、キー情報に効果的に優先順位を付けるモデルの能力が低下し、その長さの一般化が潜在的に制限されます。
この問題に対処するために、スケーラブルソフトマックス（SSMAX）を提案します。これは、入力ベクトルサイズが異なるシナリオでSoftMaxを置き換えます。
SSMAXは、既存の変圧器ベースのアーキテクチャにシームレスに統合できます。
言語モデリングの実験結果は、SSMAXを使用したモデルが前oraining前の速度の損失削減を達成するだけでなく、長いコンテキストと主要な情報検索でのパフォーマンスを大幅に改善することを示しています。
さらに、注意スコアの分析により、SSMAXにより、長いコンテキストでもモデルが重要な情報に注意を向けることができることが明らかになりました。
さらに、事前トレーニングの開始からSSMAXを使用するモデルは、より良い長さの一般化を実現しますが、すでに事前に脱出を開始したモデルは、前脱出中または前後のいずれかでSSMAXをSSMAXに置き換えることにより、この能力の一部を獲得できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google