LLMアーキテクチャの最新動向：KV共有、mHC、圧縮注意機構

短い家族休暇の後、忙しい数週間のオープンウェイトLLMリリースに追いつくことができて嬉しく思います。最も印象的だったのは、新しいアーキテクチャが長コンテキスト効率に重点を置いていることです。

推論モデルとエージェントワークフローがより多くのトークンを（より長く）保持するにつれて、KVキャッシュサイズ、メモリトラフィック、注意コストが主要なボトルネックとなり、LLM開発者はこれらのコストを削減するためにますます多くのアーキテクチャ上の工夫を追加しています。

この記事で取り上げる主な例は、Gemma 4のKV共有とレイヤー別埋め込み、Laguna XS.2のレイヤー別注意予算、ZAYA1-8Bの圧縮畳み込み注意、DeepSeek V4のmHCと圧縮注意です。

これらの変更のほとんどは私のアーキテクチャ図では小さな調整に見えますが、中には詳細な議論に値する非常に複雑な設計変更もあります。

1. 層をまたぐKVテンソルの再利用（Gemma 4）

アーキテクチャの進歩のツアーでは、4月初めにGoogleがリリースした新しいオープンウェイトGemma 4モデルシリーズから始めます。これらは3つの大きなカテゴリに分類されます：モバイルおよび小型ローカルデバイス向けのGemma 4 E2BおよびE4Bモデル、効率的なローカル推論に最適化された26B混合専門家（MoE）モデル、そして最高品質とより便利な後処理を目指す31B密モデルです。

E2BおよびE4Bバリアントの最初の小さなアーキテクチャ調整は、共有KVキャッシュ方式を採用していることです。これにより、後続の層が前の層のキーと値の状態を再利用し、長コンテキストのメモリと計算を削減します。

このKV共有はGemma 4が発明したものではありません（例えば、Brandon et al., "Reducing Transformer Key-Value Cache Size with Cross-Layer Attention", NeurIPS 2024を参照）。しかし、私がこの概念が適用されたのを見た最初のポピュラーなアーキテクチャです。（クロスレイヤー注意をクロスアテンションと混同しないでください。）

KV共有をさらに説明する前に、動機について簡単に話しましょう。最近数か月で私が書いたり話したりしたように、LLMアーキテクチャ設計の最近の主要テーマの1つはKVキャッシュサイズの削減です。そして、KVキャッシュサイズ削減の背後にある動機は、必要なメモリを減らし、より長いコンテキストを扱えるようにすることであり、これは推論モデルとエージェントの時代に特に関連します。

実際、私が以前の記事で説明した人気のある注意バリアントのほとんどすべては、KVキャッシュサイズを削減するように設計されています。例えば、Gemma 4がまだ使用しているGrouped Query Attention（GQA）は、異なるクエリヘッド間でKVヘッドを共有することでKVキャッシュサイズを削減します。

前述のように、Gemma 4はGQAを使用しています。しかし、GQAの一部としてのクエリ間のKV共有に加えて、Gemma 4は各層の注意モジュールの一部として計算する代わりに、異なる層間でKV投影も共有します。このKV共有方式はクロスレイヤー注意とも呼ばれます。

例えば、Gemma 4 E2Bには35のトランスフォーマー層がありますが、最初の15層だけが独自のKV投影を計算します。残りの20層は、同じ注意タイプの最新の非共有層からKVテンソルを再利用します。これにより、KVキャッシュサイズが約半分になります。最小のE2Bモデルでは、128Kの長コンテキストで2.7 GB（bfloat16精度）の節約になります。

KV共有の欠点は、もちろん、それが本物の「近似」であること、つまりモデル容量を減らすことです。しかし、クロスレイヤー注意の論文によれば、影響は最小限である可能性があります（テストされた小規模モデルでは）。

2. レイヤー別埋め込みと「有効」サイズ（Gemma 4 E2B/E4B）

Gemma 4 E2BおよびE4Bバリアントには、レイヤー別埋め込み（PLE）と呼ばれる2つ目の効率指向の設計選択が含まれています。これは上記のKV共有方式とは別のものです。

KV共有はKVキャッシュを削減します。一方、PLEはパラメータ効率に関するもので、小型Gemma 4モデルが、メイントランスフォーマースタックを同じ総パラメータ数の密モデルと同じくらい高価にすることなく、より多くのトークン固有情報を使用できるようにします。

例えば、Gemma 4 E2Bは2.3B有効パラメータ、または埋め込みを含めると5.1Bパラメータとしてリストされています。つまり、メイントランスフォーマースタックの計算は小さい方の数に近く、大きい方の数には追加の埋め込みテーブル層が含まれます。

PLEパスでは、各トランスフォーマーブロックは通常の注意とフィードフォワードパスに加えて、層固有のトークンベクトルを受け取ります。このベクトルは、トークンIDのレイヤー別埋め込みルックアップと通常のトークン埋め込みの線形投影を組み合わせて作成され、スケーリングと再形成を経て、各層が独自のスライスを受け取ります。ブロック内部では、注意とフィードフォワードブランチが通常通り実行され、その後の隠れ状態が層固有のPLEベクトルをゲートし、ゲートされたPLEベクトルがモデルの隠れサイズに投影され、正規化されて追加の残差更新として追加されます。

PLEの利点は、表現能力を高める一方で、トランスフォーマースタック全体をより大きなパラメータ数に拡張するコストを回避できることです。

3. その他のアーキテクチャの進展

記事ではさらに、Laguna XS.2のレイヤー別注意予算（層の重要度に応じて注意ヘッド予算を動的に割り当てる）、ZAYA1の圧縮畳み込み注意（畳み込みカーネルで注意パターンを近似して計算を削減）、DeepSeek V4のmHC（マルチヘッド圧縮）と圧縮注意についても議論しています。

これらの設計は、LLMアーキテクチャが単純なスケーリングから、よりインテリジェントなリソース割り当てへと移行し、より良い長コンテキスト性能を実現していることを示しています。