Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routing

要約

大規模な言語モデルの最近の進歩は、自己触たちの過度の二次コストを強調しました。
重要な研究努力にもかかわらず、地下の注意方法は依然として実際に劣っているパフォーマンスに苦しんでいます。
私たちは、コンテンツベースの動的で学んだコンテンツベースのスパースがより効率的な注意メカニズムにつながる可能性があると仮定します。
私たちは、専門家（MOE）と専門家の選択ルーティングの混合に触発された斬新なアプローチである、Sparse Anterestion（MOSA）の混合物を提示します。
MOSAは、各注意ヘッドのトークンを動的に選択し、任意のまばらな注意パターンを可能にします。
長さの$ t $のシーケンスから$ k $トークンを選択することにより、MOSAは各注意ヘッドの計算複雑さを$ o（t^2）$から$ o（k^2 + t）$に減らします。
これにより、同じ計算予算内でより多くのヘッドを使用することができ、より高い専門化が可能になります。
テストされたまばらな注意バリアントの中で、モザは密なベースラインよりも優れている唯一のものであり、同一の計算予算で最大27％の困惑を伴うことがあることを示します。
MOSAは、密集した自己告発と比較して、リソースの使用を減らすこともできます。
最適化されたカーネルなしでトーチの実装を使用しているにもかかわらず、困惑したMOSAモデルは、壁2時間で同時に高速で、トレーニングのメモリが少なくなり、密な変圧器ベースラインと比較してKVキャッシュのサイズを大幅に削減します。

要約(オリジナル)

Recent advances in large language models highlighted the excessive quadratic cost of self-attention. Despite the significant research efforts, subquadratic attention methods still suffer from inferior performance in practice. We hypothesize that dynamic, learned content-based sparsity can lead to more efficient attention mechanisms. We present Mixture of Sparse Attention (MoSA), a novel approach inspired by Mixture of Experts (MoE) with expert choice routing. MoSA dynamically selects tokens for each attention head, allowing arbitrary sparse attention patterns. By selecting $k$ tokens from a sequence of length $T$, MoSA reduces the computational complexity of each attention head from $O(T^2)$ to $O(k^2 + T)$. This enables using more heads within the same computational budget, allowing higher specialization. We show that among the tested sparse attention variants, MoSA is the only one that can outperform the dense baseline, sometimes with up to 27% better perplexity for an identical compute budget. MoSA can also reduce the resource usage compared to dense self-attention. Despite using torch implementation without an optimized kernel, perplexity-matched MoSA models are simultaneously faster in wall-clock time, require less memory for training, and drastically reduce the size of the KV-cache compared to the dense transformer baselines.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google