StableMask: Refining Causal Masking in Decoder-only Transformer

要約

因果マスキングと相対位置エンコーディング (RPE) を備えたデコーダー専用の Transformer アーキテクチャは、言語モデリングにおける事実上の選択肢となっています。
さまざまなタスクにわたって優れたパフォーマンスを発揮しますが、2 つの制限があることを確認しました。1 つ目は、現在の埋め込みに十分な自己完結型情報がある場合でも、すべての注意スコアがゼロではなく、合計が 1 になる必要があることです。
これにより、モデルは特定のトークンに不釣り合いな過剰な注意を払うことを余儀なくされます。
第 2 に、RPE ベースのトランスフォーマは、絶対位置情報をエンコードする能力が限られているため、汎用の近似器ではなく、位置が重要なタスクでのアプリケーションが制限されます。
この研究では、因果マスクを改良することで両方の制限に対処するパラメータフリーの方法である StableMask を提案します。
これは、擬似注意値を導入して注意分布のバランスをとり、徐々に減少するマスク率を介して絶対位置情報をエンコードします。
StableMask の有効性は理論的にも経験的にも検証されており、さまざまなデータセットとエンコード方法にわたって 71M から 1.4B の範囲のパラメーター サイズで言語モデルが大幅に強化されていることが示されています。
さらに、(1) StreamingLLM などの特別なトリックを必要としない効率的な外挿、および (2) 既存のアテンション最適化技術との簡単な統合を自然にサポートしていることを示します。

要約(オリジナル)

The decoder-only Transformer architecture with causal masking and relative position encoding (RPE) has become the de facto choice in language modeling. Despite its exceptional performance across various tasks, we have identified two limitations: First, it requires all attention scores to be non-zero and sum up to 1, even if the current embedding has sufficient self-contained information. This compels the model to assign disproportional excessive attention to specific tokens. Second, RPE-based Transformers are not universal approximators due to their limited capacity at encoding absolute positional information, which limits their application in position-critical tasks. In this work, we propose StableMask: a parameter-free method to address both limitations by refining the causal mask. It introduces pseudo-attention values to balance attention distributions and encodes absolute positional information via a progressively decreasing mask ratio. StableMask’s effectiveness is validated both theoretically and empirically, showing significant enhancements in language models with parameter sizes ranging from 71M to 1.4B across diverse datasets and encoding methods. We further show that it naturally supports (1) efficient extrapolation without special tricks such as StreamingLLM and (2) easy integration with existing attention optimization techniques.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google