Transformersには3つの投影が必要か？ QKVバリアントの系統的な研究

要点：投影の共有がメモリオーバーヘッドを削減する

Transformersは、Query (Q)、Key (K)、およびValue (V) 行列間で重み投影を共有することで、パフォーマンスを維持しながらメモリ要件を大幅に削減できます。具体的には、KeyとValueの投影を共有する（K=V）ことで、Multi-Query Attention (MQA) と組み合わせた場合に最大96.9%のKVキャッシュ削減が可能となり、オンデバイスおよびエッジAI展開において非常に実行可能な戦略となります。

QKV投影バリアントの評価

標準的なTransformerアーキテクチャは、アテンションのために3つの異なる投影を使用します。本研究では、これらの投影が冗長であるかどうかを判断するために、3つの代替案を系統的に評価しています。

  • Shared Key-Value (K=V): KeyとValueの投影が結合されています。
  • Shared Query-Key (Q=K): QueryとKeyの投影が結合され、対称的なアテンションマップが生成されます。
  • Single Projection (Q=K=V): 3つの投影すべてが同じ重みを共有し、これも対称的なアテンションマップを生成します。

対称的なアテンションマップ（Q=K または Q=K=V）における方向性の喪失を軽減するため、研究者たちは2D positional encodingsを使用して非対称性を導入することを検討しました。

タスクおよびモデルスケール全体でのパフォーマンス

合成タスク、ビジョン (MNIST, CIFAR, TinyImageNet)、および言語モデリング（10B tokensでトレーニングされた300Mおよび1.2Bパラメータのモデル）にわたる実験の結果、投影共有Transformerは、標準的なQKV Transformerと同等、あるいは時として標準よりも優れたパフォーマンスを示します。

言語モデリングとKVキャッシュの効率性

言語モデリングタスクにおいて、K=Vバリアントは最も大きな実用的な利点を提供します：

  • 直接的な影響: KとVの投影を共有することで、パープレキシティの低下をわずか3.1%に抑えつつ、KVキャッシュを50%削減できます。
  • ヘッド共有との相乗効果: 投影共有は、Grouped-Query Attention (GQA) および Multi-Query Attention (MQA) と補完的です。
    • K=V + GQA-4: 87.5%のキャッシュ削減を実現します。
    • K=V + MQA: 96.9%のキャッシュ削減を実現します。

理論的洞察：なぜK=Vが機能するのか

本研究は、K=Vがモデルの品質を維持できる理由として、KeyとValueがしばしば同様の表現空間を占めることを示唆しています。さらに、アテンションは通常、低ランク領域で動作するため、重要な情報を失うことなくKとVの間の冗長性を活用できる可能性があります。逆に、QueryとKey (Q=K) の投影を共有することは、アテンションに必要な方向性を損なう傾向があり、KとVの共有よりもパフォーマンスをより大幅に低下させます。

コミュニティ分析と技術的批判

エッジ展開において結果は有望ですが、技術的な議論では、これらの知見の汎用性に関していくつかの注意点（caveats）が指摘されています。

トレーニングスケールと過学習（Over-training）

批判的な意見では、本研究で使用された言語モデル（最大1.2Bパラメータ）は10B tokensでトレーニングされており、これはChinchilla-optimalな計算量閾値を大幅に下回っています。

"Modern overtrained 1B LLMs are trained on the order of 10T tokens (1000x more). This is important because... simplifications and alternatives to standard attention can look fine in the under-trained regime but lag after over-training."

誘導バイアスとスケーリング

アテンションにおける強力な誘導バイアス（inductive bias）の欠如は、標準的なQKVアーキテクチャの表現力が、より大規模なスケールや、より多くのトレーニングデータを用いる場合にのみ明らかになる可能性があるという懸念があります。スケーリング曲線がないため、これらの簡略化が大規模スケールにおいてパレート・エッジ（Pareto envelope）を押し広げるのか、それとも小規模スケールにおいて単に「それなりに」機能するだけなのかを判断するのは困難であると主張する観察者もいます。

実装上の注意

論文の読者は、ソーステキストにおける「Q-K=V」という表記が、数学的な減算操作ではなく、共有されたKey-Value投影（K=V）を指すために使用されていることに注意してください。

Sources