Transformerの支配的地位を揺るがしたMambaの著者の一人であるTri Daoが、新たに作品を発表しました。
推論に特化した2つのアテンションメカニズムを提案しています。
モデルの性能を維持しつつ、デコード速度とスループットを最大2倍に向上させ、モデルの長文脈推論能力を大幅に最適化しました。
この研究の3人の著者は全員プリンストン大学出身で、論文には2つの主要な貢献があります：
一つ目は、Grouped-Tied Attention（GTA）を提案すること。これはLLaMA 3にすでに統合されているアテンションメカニズムGQAと同等の品質でありながら、KVキャッシュの使用量を約50%削減します。
二つ目は、Grouped Latent Attention（GLA）を提案すること。これはDeepSeekが使用するアテンションメカニズムMLAと品質が一致しながらも、デコード速度が速く、場合によってはFlashMLAよりも2倍高速です。
著者の一人であるテッド・ザドゥーリのまとめによると：
GTAはGQAの有効な代替品であり、GLAはMLAの実用的な代替品です。
簡潔に言えば、アテンションメカニズムのメモリ使用量と計算ロジックを最適化することで、モデルの生成品質を犠牲にすることなく、大規模言語モデルの推論効率とハードウェアリソース利用率を大幅に向上させることができ、特に長文脈シナリオでの利点が際立っています。
関連論文の公開後、多くの研究者が祝福に駆けつけました〜
では、この研究は具体的に何を述べているのでしょうか？
推論認識アテンションメカニズムの導入
概括すると、論文の核は推論認識アテンションメカニズムの導入です。これは、モデルの推論段階におけるメモリ冗長性、計算非効率性、長文脈ボトルネックなどの問題に対処するためにアテンションメカニズムを再設計するものです。
トライ・ダオによると、この研究は次のような発想から始まりました：
推論駆動型AIが発展する時代において、「理想的な」アーキテクチャはどのようなものになるべきか？
特に長文脈推論の場合、現在のLLM（大規模言語モデル）はメモリアクセスボトルネックと並列処理の制限という二つの大きな課題に直面しています。
つまり、モデルがテキストを生成する際、毎回大量の「履歴」をメモリから読み込む必要があり、これにより各単語の生成が遅くなるだけでなく、順番にしか生成できず、複数のチップが同時に作業を行うことができません。
これに対し、チームは2つの方向からアテンションメカニズムを再設計することを計画しました：
より高いハードウェア効率："メモリロードあたりの計算量"（算術強度）を増やすことで、メモリ帯域幅への依存を減らす。
並列スケーラビリティの維持：モデルの並列学習/推論能力を犠牲にすることなく、デコード速度を最適化する。
そして最終的に提案されたGTAとGLAは、KVキャッシュの使用量を削減しつつ、既存のソリューションと同等のモデル品質を維持し、デコード速度を大幅に向上させます。
ここで言及されている「既存のソリューション」とは、主に学術界で知られている2つの方法を指します：
一つはGrouped Query Attention（GQA）メカニズムで、KVキャッシュをグループ化して共有することでメモリ使用量を削減します。ビジョントランスフォーマー（ViT）などのタスクで良好なパフォーマンスを示し、大規模データ処理に適しており、現在Llama 3などのオープンソースモデルに適用されています。
二つ目はMulti-headed Latent Attention（MLA）メカニズムで、その起源は「Attention Is All You Need」論文にまで遡り、後にDeepSeekによって再び注目を集めました。これは、異なる層間でアテンション情報をどのように統合するかを重視し、各層の冗長な計算を削減することができます。
ただし、GQAは各クエリヘッドグループに独立したKVを格納する必要があり、MLAは並列最適化が不十分であるため、さらなる改善が必要です。
以下に、チームが提案する新しい手法であるGTAとGLAについてそれぞれ説明します。
Grouped-Tied Attention（GTA）
GTAの核心的な設計思想は、異なるクエリヘッドのキー（Key）とバリュー（Value）の状態を組み合わせ、再利用することで、メモリ転送回数を削減することです。
具体的には（右図）、マルチヘッドアテンションのヘッドをいくつかのグループに分け、各グループ内のヘッドは同じKeyとValueパラメータを共有します。計算時には、同じグループ内のヘッドは同じKVキャッシュを使用し、クエリ（Query）パラメータのみが独立しています。
対照的に、中央の従来のマルチヘッドアテンションメカニズム（MHA）では、各クエリヘッドが独立したキーとバリューを持つため、共有がないことから、すべてのキーとバリューを格納するためにより多くのメモリが必要となります。
さらにGQA（左図）と比較すると、GQAはグループごとにKVを共有しますが、各グループは依然として独立して格納しています。それに対し、GTAはパラメータの結合により、KVのより徹底的な再利用を実現します。
Grouped Latent Attention（GLA）
一方、GLAの設計は2層構造を採用しています：
潜在層（Latent Layer）：固定数の潜在トークンを導入し、グローバルコンテキストの圧縮表現として、元のトークンの一部のKVキャッシュを置き換えます。
グループ化されたヘッドメカニズム：クエリヘッドをグループ化し、各グループは潜在トークンのKVを共有しつつ、元のトークンとの相互作用を維持します。
デコードプロセスでは、MLA（左図）と比較して、GLAは共有された結合潜在表現を介して、各デバイスがロードする必要のあるKVキャッシュの量を削減し、それによってメモリアクセス量を減少させます。
そして、各デバイス上のKVキャッシュ量が減少したことにより、より多くのリクエストを同時に処理できるようになります。
「GQAとMLA」の有効な代替品
それでは、GTAとGLAの効果は一体どうなのでしょうか？
チームは4つの規模のモデルで実験を行いました。これには、小型（183M）、中型（433M）、大型（876M）、XL（1471M）が含まれます。これらのモデルはFineWeb-Edu-100Bデータセットで学習され、GPT-3アーキテクチャとLlama 3トークナイザを使用しています。
テスト指標は主に2つのカテゴリに分けられます：
品質指標：パープレキシティ（Perplexity）、下流タスクの精度（Winogrande、SciQなど7つのベンチマーク）；
効率指標：トークンあたりのデコード遅延、スループット、KVキャッシュ占有量。
実験では、GQA、MLA、FlashMLA、従来のMHAなど、様々なアテンションメカニズムが比較されました。
パープレキシティの実験では、GTAが中・大型モデルでGQAを上回ることが示され、GTAがモデルのさらなる拡張により適している可能性を示唆しています。一方、GLAはほとんどのシナリオでMLAと同等であり、GLAの設計が合理的であることを示しています。これは並列計算とモデル品質の間で良好なバランスを見つけることができることを意味します。
いくつかの手法を下流タスク（一般的な常識推論、論理推論、知識問答などのシナリオをカバー）で比較すると、全体的なパフォーマンスに大きな差はありませんでした。
しかし、変化の傾向を見ると（下図は中型から大型への移行を示す）、GTAとGLAは、中型からXLサイズへの下流タスクのパフォーマンスを維持または向上させることができます。
KVキャッシュに関して、モデルの品質を犠牲にすることなく、GTAはGQAと比較してKVキャッシュを約50%削減し、「パラメータ結合＋グループ再利用」の有効性を検証しました。
同時に、クエリ長が1の場合、MLAは計算ボトルネック（610 TFLOPS/sに達する）に近づいているのに対し、GLAはまだ計算リソースを飽和させていません（360 TFLOPS/s）。
また、シーケンス長が1Kから64Kに増加するにつれて、GLAのデコード速度はFlashMLAの2倍高速です。
さらに、リアルタイムサーバーのパフォーマンステストでは、64の同時リクエストに対する出力スループット（高いほど良い）において、同じ並列方式の下でGLAはMLAよりも優れたパフォーマンスを示しました。
次に、チームはDeepSeek Coder V2 Base（236B）モデルで、FP8精度を使用した場合の、異なるプリフィル長とデコード長における出力スループットを比較しました。
その結果、プリフィル長が32Kと64Kの場合、GLA-8の出力スループットがMLAよりも明らかに高かった。これは、長文脈処理において、GLAがスループットでMLAを上回ることを示しています。
GLA-8は、不均衡な負荷を処理する際にも、より高い出力スループットを示しました。これは、GLAが異なる長さのリクエストを処理する際に、リソースをより効率的に利用し、全体的なパフォーマンスを向上させることができることを示しています。
上記の実験はすべて、論文著者の主張である「GTAとGLA」が「GQAとMLA」の有効な代替品であることを証明しています。
論文著者は全員プリンストン大学出身
論文の著者はトライ・ダオを含め3名で、全員がプリンストン大学出身です。
テッド・ザドゥーリは現在、プリンストン大学の博士課程学生で、機械学習を研究しています。
以前はインテルで2回のインターンシップ経験（深層学習の研究）があり、AIスタートアップCohereで研究員として短期間勤務していました。
ヒューバート・シュトラウスはプリンストン大学の研究エンジニアで、機械学習とモデルの深層学習を研究しています。
フランスの有名な工学学校Arts et Métiersを卒業後、ジョージア工科大学でオペレーションズリサーチの修士号を取得しました。
卒業後はいくつかのインターンシップや職務経験を積み、プリンストン大学のエンジニアになる前は、ある企業で機械学習エンジニアとしてモデルのトレーニングとTransformerの最適化を担当していました。
トライ・ダオは現在、プリンストン大学のコンピュータサイエンス助教授であり、生成AIスタートアップTogether AIの主任科学者です。
彼はTransformerモデルのアテンションメカニズムを最適化する一連の研究で学界に名を馳せています。
その中でも最も影響力があるのは、彼が共同著者の一人として提案したMambaアーキテクチャです。このアーキテクチャは、言語、音声、ゲノミクスなど様々なモダリティでSOTA（State-of-the-Art）性能を達成しました。
特に言語モデリングの分野では、Mamba-3Bモデルは事前学習と下流評価の両方で同規模のTransformerモデルを上回り、その2倍の規模のTransformerモデルに匹敵します。
また、彼はFlashAttention 1-3バージョンの発表にも携わっており、FlashAttentionはTransformerの高速化に広く利用され、アテンション速度を4〜8倍向上させました。
いずれにせよ、この研究に戻ると、論文著者のテッド・ザドゥーリは率直に述べています：
これはテスト時推論の「理想的な」アーキテクチャへの第一歩に過ぎません！
論文：https://arxiv.org/abs/2505.21487
コード：https://github.com/Dao-AILab/grouped-latent-attention
参考文献：
[1]https://x.com/tri_dao/status/1928170648863473892
[2]https://x.com/gm8xx8/status/1927572103806554262
[3]https://x.com/tedzadouri/status/1928167296821854363
— 終わり —