RATTENTION: Towards the Minimal Sliding Window Size in Local-Global Attention Models

要約

ローカルグローバルの注意モデルは最近、標準的な変圧器の説得力のある代替品として浮上しており、トレーニングと推論効率の両方の改善を約束しています。
ただし、ウィンドウサイズの重要な選択はパレートのトレードオフを示します。大きなウィンドウは、完全な注意と同様のパフォーマンスを維持しますが、短いコンテキストシナリオでは最小限の効率の向上を提供しますが、ウィンドウの小さなウィンドウはパフォーマンスの低下につながる可能性があります。
Gemma2やMistralなどの現在のモデルは、パフォーマンスを維持するために保守的なウィンドウサイズ（たとえば、8192の事前トレーニング長のうち4096）を採用しています。
この作業は、このパレートフロンティアをシフトする戦略を調査し、ローカルグローバルモデルが短いコンテキスト体制でさえ効率の向上を達成できるようにします。
私たちの核となる動機は、地元の注意の本質的な制限に対処することです。これは、定義されたウィンドウの外側のトークンに対する完全な無視です。
これらのウィンドウ外のトークンから情報をキャプチャするように設計された特殊な線形注意メカニズムと統合されたローカルな注意のバリアントであるRattentionを探ります。
3Bおよび12Bスケールでの事前除去実験は、格子がパフォーマンスと効率の間の優れたパレートトレードオフを達成することを示しています。
スイートスポットとして、わずか512のウィンドウサイズのラテンショニングは、多様な設定にわたるフルアテンションモデルのパフォーマンスと一貫して一致します。
さらに、ラッティングの線形注意成分に固有の再発性の性質は、ルーラーベンチマークで検証されているように、長いコンテキストのパフォーマンスの向上に貢献します。
重要なことに、これらの改善はトレーニング効率を損なうものではありません。
特殊なカーネルの実装とウィンドウサイズの削減により、ラテンションは既存の最先端のアプローチに匹敵するトレーニング速度を維持します。

要約(オリジナル)

Local-global attention models have recently emerged as compelling alternatives to standard Transformers, promising improvements in both training and inference efficiency. However, the crucial choice of window size presents a Pareto tradeoff: larger windows maintain performance akin to full attention but offer minimal efficiency gains in short-context scenarios, while smaller windows can lead to performance degradation. Current models, such as Gemma2 and Mistral, adopt conservative window sizes (e.g., 4096 out of an 8192 pretraining length) to preserve performance. This work investigates strategies to shift this Pareto frontier, enabling local-global models to achieve efficiency gains even in short-context regimes. Our core motivation is to address the intrinsic limitation of local attention — its complete disregard for tokens outside the defined window. We explore RATTENTION, a variant of local attention integrated with a specialized linear attention mechanism designed to capture information from these out-of-window tokens. Pretraining experiments at the 3B and 12B scales demonstrate that RATTENTION achieves a superior Pareto tradeoff between performance and efficiency. As a sweet spot, RATTENTION with a window size of just 512 consistently matches the performance of full-attention models across diverse settings. Furthermore, the recurrent nature inherent in the linear attention component of RATTENTION contributes to enhanced long-context performance, as validated on the RULER benchmark. Crucially, these improvements do not compromise training efficiency; thanks to a specialized kernel implementation and the reduced window size, RATTENTION maintains training speeds comparable to existing state-of-the-art approaches.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google