技術要約: MoVE (Mixture of Value Embeddings)

1. 問題提起

現代の生成AIの基盤である自己回帰型シーケンスモデリングは、根本的な構造的ボトルネックに直面している。それは、パラメトリックメモリ（モデルが事実的知識や視覚的パターンを蓄積する能力）と計算コストの間の硬直した結合である。従来、モデルのメモリ容量を増やすには、ネットワークを深くするか、あるいは幅を広げる必要があり、これは訓練および推論時のアクティブなFLOPsの比例的な線形増加を招く。これにより、「メモリ密度が高い」モデル（膨大な百科事典的知識や視覚的知識を持つモデル）は、計算コストが極めて高くなるというトレードオフが生じる。本論文は、この結合を打破し、ネットワークの深さとは直交する新しい軸として容量をスケールさせる手法を模索する。

2. 手法: MoVE

著者らは、AttentionメカニズムのValueストリームを拡張することで、メモリと計算をデカップリング（分離）するメカニズムであるMoVE (Mixture of Value Embeddings) を提案する。

コアアーキテクチャ

• グローバル・バリュー・エンベディング・バンク (Global Value Embedding Bank): MoVEは、すべてのAttentionレイヤーからアクセス可能な「概念ベクトル」（例：言語的定義や視覚的属性）の共有リポジトリを表す、学習可能なグローバルテンソル $E \in \mathbb{R}^{N_{vocab} \times M \times d}$ を導入する。ここで、$M$ はトークンあたりのエンベディング・スロット数である。
• 役割の分離 (Decoupling Roles): 入力エンベディング ($W_I$) がアテンション・パターンと意味的内容の両方の役割を担う標準的なTransformerとは異なり、MoVEはこれらを分離する。$W_I$ はルーティングのためにコンパクトなまま維持され、バリュー・エンベディング・バンク $E$ がValueストリーム専用の高容量リポジトリを提供する。
• ソフトゲーティングによる動的な混合 (Dynamic Mixing via Soft Gating): 各シーケンスステップおよびアテンションヘッドに対して、モデルはバンク内の各スロットに対する微分可能なソフトゲーティング係数を計算する。最終的なバリューテンソル $V_S^{(h)}$ は、標準的な高密度投影と、取得されたグローバル・エンベディングの加重和となる：
  $$V_S^{(h)} = g_{t,0}^{(h)} \odot V^{(h)} + \sum_{i=1}^{M} g_{t,i}^{(h)} \odot M_{t,i}^{(h)}$$
  ゲート $g_{t,0}$ は標準的なパスを制御し、$g_{t,i}$ は特定のメモリ・スロットの寄与を制御する。これにより、モデルはローカルなコンテキストとグローバルなパラメトリック知識を動的にバランスさせることができる。
• レイヤー間での共有 (Shared Across Layers): 極めて重要な点として、エンベディング・バンク $E$ はすべての $L$ 個のアテンションレイヤー間で共有される。これにより「勾配のハイウェイ (gradient highways)」が形成され、概念が迅速かつ堅牢に学習される。これは、レイヤーごとの構成においてスパースになりがちなパラメータに対して、信号を高密度化させる効果がある。

効率的なアーキテクチャとの統合

本論文では、MoVEがMulti-Head Latent Attention (MLA) にシームレスに統合できることを示している。メモリをフルランクのValueベクトルに注入する（これはMLAの効率性を損なうことになる）のではなく、MoVEは圧縮された低ランク潜在空間 ($c_{KV}$) に直接グローバルメモリを注入する。これにより、MLAの推論効率（KVキャッシュ圧縮によるもの）を維持しながら、パラメトリックメモリ容量を拡張することができる。

3. 主な貢献

1. 新しいスケーリング軸: MoVEは、ネットワークの深さに依存せずにパラメトリックメモリをスケールさせる手法を確立した。エンベディング・スロット数 ($M$) を増やすことで、計算の深さやFLOPsを増やすことなく、モデルを「メモリ密度が高い」状態にすることができる。
2. モダリティに依存しない設計: このメカニズムはテキストおよび画像生成の両方で検証されており、多様な自己回帰型モダリティへの適用可能性を証明している。
3. MLAとの相乗効果: 本研究は、メモリ拡張と推論圧縮（MLA）が補完関係にあることを示している。MoVEは、高ランクのテンソルを実体化させることなく、圧縮された潜在空間の容量を高めることができる。
4. アブレーションによる洞察: 著者らは、グローバル共有バンクとゲーティング機構の寄与を分離して調査し、グローバルなアーキテクチャが効率性の主要な要因であり、標準的なパスをゲーティングすることが、取得された知識とローカルなコンテキストの間の仲裁に不可欠であることを発見した。

4. 実験結果

著者らは、テキストには nanochat フレームワークを、画像には LlamaGen を用いて、厳密に制御された実験を行い、MoVEを標準的なTransformerおよびレイヤー単位のバリュー・エンベディング (LaVE) と比較した。

• テキスト生成 (FineWeb-Edu):

  • MoVEは、モデルの深さ（D12, D20, D32）に関わらず、標準およびLaVEのベースラインを一貫して上回った。
  • スケーラビリティ: LaVEの性能はメモリがレイヤー数に結合されているために急速に飽和したが、MoVEは持続的な改善を示した。例えば、D12モデルにおいて、MoVEのスロットを$\times8$ (4L) に増やすと、検証時のBits Per Byte (BPB) は、標準の0.838から0.797へと減少（0.041の改善）した。
  • 効率性: MoVE-$\times1$（LaVE-$\times2$ の半分のスロット数）は、しばしばLaVE-$\times2$と同等以上の性能を示し、グローバル共有による優れたパラメータ利用率を示唆した。

• 画像生成 (ImageNet):

  • GPT-B (111M) および GPT-L (343M) モデルにおいて、MoVEは画像忠実度を大幅に向上させた。
  • GPT-Bにおいて、MoVEはFréchet Inception Distance (FID) を標準の6.53から5.62へと減少させた。
  • 特筆すべきは、LaVEベースラインがGPT-Lスケールで性能を低下させた（FID 3.77）一方で、MoVEは明確なリードを維持し、FID 3.10を達成した。

• MLAとの統合:

  • Multi-Head Latent Attentionに適用されたMoVEは、$\times32$のメモリ密度まで単調なスケーリングを示した。
  • MLAの圧縮により、$\times32$のMoVE設定は標準的な$\times1$のバンクと同等のパラメータ数を追加するだけで、大幅なBPBの利得（D12で0.0136）を達成しており、メモリのスケーリングがネットワークの深化を必要としないことを裏付けた。

5. 意義と主張

本論文は、MoVEが「メモリは安価であり、知能（計算）は固定されている」という独自の「スケーリングの経済性」を提供すると主張している。パラメトリックメモリを計算の深さから切り離すことで、MoVEは巨大なデンスネットワークの禁止的なコストを負うことなく、膨大な知識を備えたメモリ密度が高いモデルの構築を可能にする。

著者らは、MoVEを自己回帰型モデリングの汎用的なブースターとして位置づけており、混合コンポーネントを増やすだけで容量をスケールさせることができるとしている。また、制限事項として、MoVEはFLOPsを削減する一方でメモリ帯域幅のコストを導入すること、および現在の追加パラメータあたりのパラメータ効率は標準的なデンス・スケーリングよりも低いことを認めている。しかし、本研究はネットワークの深さと直交するスケーリングの新しい次元を確立しており、将来の生成AIシステムに向けた有望なアーキテクチャ・プリミティブを提示している。