LLMアーキテクチャの最近の動向: KV共有、mHC、そして圧縮アテンション

 (magazine.sebastianraschka.com)
6 ポイント 投稿者 GN⁺ 5 시간 전 | まだコメントはありません。 | WhatsAppで共有
  • 最近公開されたオープンウェイトLLMは長文コンテキスト効率に注力しており、KVキャッシュサイズ・メモリトラフィック・アテンションコストを減らすためのアーキテクチャ上の工夫が急速に増えている
  • Gemma 4は層間KV共有(cross-layer attention) と per-layer embeddings(PLE) により、KVキャッシュとパラメータ効率を同時に改善
  • Laguna XS.2はレイヤーごとにクエリヘッド数を変えて割り当てる layer-wise attention budgeting を導入
  • ZAYA1-8BはCompressed Convolutional Attention(CCA) により、圧縮された潜在空間で直接アテンション演算を実行し、KVキャッシュだけでなくアテンションFLOPsも削減
  • DeepSeek V4はmHC(Manifold-Constrained Hyper-Connections) で残差経路を拡張し、CSA/HCA でシーケンス長を圧縮することで、1MトークンコンテキストにおいてV3.2比でFLOPs・KVキャッシュを大幅に削減

概要: 長文コンテキスト効率に注力する最新アーキテクチャ

  • 推論(reasoning)モデルとエージェントワークフローが、より多くのトークンをより長く保持するようになり、KVキャッシュサイズ・メモリトラフィック・アテンションコストが主要な制約として浮上
  • 4月〜5月に公開された主要なオープンウェイトモデルで新たに登場した設計ポイント
    • Gemma 4: KV sharing と per-layer embeddings
    • Laguna XS.2: layer-wise attention budgeting
    • ZAYA1-8B: compressed convolutional attention
    • DeepSeek V4: mHC + compressed attention
  • 本文では、データミックス、学習スケジュール、ポストトレーニング、RLレシピ、ベンチマークは扱わず、トランスフォーマーブロック・残差ストリーム・KVキャッシュ・アテンション演算の内部変化に焦点を当てる

1. Gemma 4: 層間KV共有でキャッシュを縮小

  • Googleが4月初旬に公開したGemma 4ファミリーは3つのカテゴリで構成
    • Gemma 4 E2B/E4B: モバイル・組み込みデバイス(IoT)向けの小型モデル
    • Gemma 4 26B MoE: 効率的なローカル推論に最適化されたMoEモデル
    • Gemma 4 31B dense: 最高品質とポストトレーニングのしやすさを狙ったdenseモデル

  • KV共有(cross-layer attention)の導入

    • 後半のレイヤーは独自のK/Vプロジェクションを計算せず、同じアテンションタイプの直近の先行する非共有レイヤーのKVテンソルを再利用する
    • スライディングウィンドウレイヤーは前のスライディングウィンドウレイヤーと、full-attentionレイヤーは前のfull-attentionレイヤーとKVを共有
    • クエリプロジェクションは各レイヤーが独自に計算するため、レイヤーごとのアテンションパターンは維持される
    • Gemma 4 E2Bは35個のトランスフォーマーレイヤーのうち最初の15個だけが独自にKVを計算し、最後の20レイヤーは再利用
    • Gemma 4 E4Bは42レイヤーのうち24個だけが独自にKVを計算し、最後の18個を再利用

  • 削減効果

    • およそ半分のKVを共有し、KVキャッシュサイズを約半分削減
    • 128K長文コンテキスト(bfloat16)基準でE2Bは2.7 GB、E4Bは約6 GBを削減

  • 限界

    • KV共有は一種の近似であり、モデルcapacityの低下をもたらす
    • cross-layer attention論文によれば、(テストされた小型モデルでは)影響は最小限
  • この概念自体はBrandon et al., "Reducing Transformer Key-Value Cache Size with Cross-Layer Attention" (NeurIPS 2024)に基づくもので、Gemma 4はこれを広く知られたアーキテクチャで初めて適用した事例

2. Gemma 4 E2B/E4BのPer-Layer Embeddings (PLE)と"Effective"サイズ

  • PLEはKV共有とは別個の効率化設計であり、パラメータ効率に焦点を当てている

  • "E"はeffectiveを意味

    • Gemma 4 E2B: 2.3B effective parameters、埋め込みを含めると5.1B
    • Gemma 4 E4B: 4.5B effective parameters、埋め込みを含めると8B
    • メインのトランスフォーマースタックの計算量は小さい数値に近く、大きい数値には追加の埋め込みテーブルレイヤーが含まれる

  • PLEの構造

    • PLEベクトルは繰り返されるトランスフォーマーブロックの外側で準備される
    • トークンIDはper-layer embedding lookupを通り、通常のトークン埋め込みは同じPLE空間へ線形射影される
    • 2つの結果を足し合わせ、スケーリング・reshapeして各レイヤー用のスライスを1つずつ持つテンソルを生成
    • 各レイヤー l は自身のスライス(ple_l)だけを受け取る

  • トランスフォーマーブロック内部の動作

    • アテンションとフィードフォワードの残差更新は通常どおり実行
    • 2回目の残差addの後、hidden state z がレイヤー別PLEベクトルをゲーティングする
    • ゲートされたPLEベクトルをモデルのhidden sizeへ再射影・正規化したうえで、追加の残差更新として加える

  • PLEの目的

    • 高コストなトランスフォーマーブロックは小さな"effective"サイズに近いまま維持
    • 追加のcapacityはper-layer embeddingテーブルに格納し、lookupベースのため、アテンション・FFN重みを追加するよりはるかに低コスト
    • 単純にdenseモデルを小さくする代替案と異なり、メイン計算部のcapacityを犠牲にしない
  • PLEは原理的には小型モデルに限られないが、大型モデルはすでに十分なcapacityを持ち、MoEでcapacity拡張も可能

3. Laguna XS.2: Layer-Wise Attention Budgeting

  • Lagunaは、コーディング用途LLMに注力する欧州拠点の企業Poolsideによる最初のオープンウェイトモデル

  • 基本構成

    • 合計40レイヤーで、そのうち30レイヤーがスライディングウィンドウアテンション、10レイヤーがglobal/full attention
    • スライディングウィンドウレイヤーのウィンドウサイズ: 512トークン
    • スライディングウィンドウ + globalの混合パターン自体はGemma 4など他のアーキテクチャでも使われている

  • 新しい点: レイヤーごとのクエリヘッド数の差別化

    • Hugging Faceのconfig.jsonにある num_attention_heads_per_layer 設定で、レイヤーごとに異なるクエリヘッド数を指定でき、KVキャッシュ形状との互換性は維持される
    • スライディングウィンドウレイヤー: KVヘッドあたり8個のクエリヘッド
    • Full attentionレイヤー: KVヘッドあたり6個のクエリヘッド
    • KVヘッドは8個で固定

  • 設計意図

    • 全レイヤーに同じアテンション予算を与えるのではなく、有用な箇所にアテンションcapacityを集中する
    • full-attentionレイヤーはコンテキスト全体を見るため高コストであり、クエリヘッドを少なめに割り当てる
  • レイヤーごとのcapacity差別化という発想自体は、少なくともAppleの2024年のOpenELMまでさかのぼれ、Laguna XS.2はプロダクション級オープンモデルにおける最も目立つ最近の事例
  • 付随してLagunaはper-head attention-output gating も適用している (Qwen3-Nextなどに類似)

4. ZAYA1-8B: Compressed Convolutional Attention (CCA)

  • Zyphraが開発したオープンウェイトモデルで、NVIDIA GPUやGoogle TPUではなくAMD GPUで学習された点が特徴

  • 構造

    • config.json上では80個の交互レイヤー項目があり、CCA/GQAアテンションとMoEフィードフォワードが交互に現れる (見た目には40個のアテンション+MoEペアとして表現される)
    • 4:1 GQAレイアウトとともにCCAを使用
    • MoEは非常に疎な設定で、トークンごとにルーティングされるexpertは1つだけ有効

  • CCAの核心

    • MLAと同様に、アテンションブロックに圧縮された潜在表現を導入
    • 違いは、MLAが潜在表現を主にKVキャッシュ縮小用に使い、実際のアテンションはアテンションヘッド空間へ再射影して行う点
    • CCAはQ、K、Vをすべて圧縮したうえで、圧縮された潜在空間で直接アテンション演算を行い、得られたアテンションベクトルを再びup-projectionする
    • その結果、KVキャッシュだけでなくprefill・学習時のアテンションFLOPsも削減する

  • Convolutional Mixing

    • "Convolutional"という名称は、圧縮されたK、Q表現に追加のconvolutional mixing が入るため
    • 圧縮によってQ、K、Vは細くなり、計算量・キャッシュは削減されるが、アテンション表現力が低下する可能性がある
    • convolutionは圧縮されたQ、Kに局所コンテキストを低コストで追加する手段
    • Vには適用しない — Q、Kはアテンションスコアを決定し、Vはそのスコアで平均化されるコンテンツだからである
    • sequence mixingに加えて、channel mixingコンポーネントも存在する

  • 性能

    • CCAは、ZAYA1-8B技術レポートより前の別論文 "Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Space" (2025年10月) で導入された
    • CCA論文の実験では、同じ圧縮設定でMLAより優れる結果が報告されている

5. DeepSeek V4: CSA/HCA、mHC、圧縮アテンションキャッシュ

  • DeepSeek V4は今年もっとも大きな話題とモデル規模を示したリリースであり、DeepSeek V4-Proはactive parameter比率の観点で最も疎なMoE

  • 本文では、従来アーキテクチャと比べた2つの新しい核心に注目する

    • mHC: より広い残差経路
    • CSA/HCA: 長文コンテキスト向けのアテンション圧縮・疎化

  • 5.1 Manifold-Constrained Hyper-Connections (mHC)

    • 2025年12月31日のDeepSeekチーム論文 "mHC: Manifold-Constrained Hyper-Connections" に基づくもので、当時は27B規模でしか実験されていなかったが、今回のフラッグシップで本格適用された
    • トランスフォーマーブロック内の残差接続設計を現代化することが目的 — アテンション/正規化/MoEに集中していた変更とは一線を画す

    • Hyper-Connections (HC) の背景

      • Zhu et al. (2024) "Hyper-connections" に基づく
      • 単一の残差ストリームを複数の並列残差ストリームと学習済みマッピングに置き換える
      • アテンション・MoEレイヤーが通常のhidden sizeで動作するために、Pre Mapping(並列ストリーム → 1つのhidden vector) と Post Mapping(レイヤー出力 → 並列ストリームへの分配) を追加
      • 残差経路の表現力を高めつつ、アテンション・MoE自体は広げない
      • 7B OLMo MoE実験では、トークンあたりFLOPsは13.36G → 13.38Gで実質的に変化せず、学習トークン数が約半分でbaseline性能に到達

    • HC → mHCの変更点

      • 通常のHCのRes Mappingは学習可能な行列であり、複数層を通るうちに信号の増幅・縮小が予測不能になる
      • mHCは残差マッピングをdoubly stochastic matricesのマニフォールドに射影する — すべての要素が非負で、各行・各列の和が1
      • 残差混合はストリーム間の安定した情報再配分のように機能する
      • Pre Mapping、Post Mappingも非負かつ有界に制約し、拡張された残差状態の読み書き時に相殺が起きるのを防ぐ
      • より深いモデルほど重要になるスケーリング安定性を確保

    • コスト

      • 27Bモデル実験では、DeepSeekチームの最適化実装(fusion, recomputation, pipeline scheduling)により、n=4残差ストリーム使用時の学習時間オーバーヘッドは6.7%

  • 5.2 CSAとHCAによる圧縮アテンション

    • 非常に長いコンテキストにおいて、アテンションスコア計算だけでなくKVキャッシュがシーケンス長に比例して大きくなる問題の解決が目的
    • DeepSeek V4は2種類の圧縮アテンションをハイブリッドに使用: Compressed Sparse Attention (CSA) と Heavily Compressed Attention (HCA)

    • MLAとの違い

      • DeepSeek V2/V3のMLAはトークンごとのKV表現を圧縮するが、トークンごとに1つの潜在KVエントリは保持する
      • CSA/HCAはシーケンス次元に沿って圧縮し、複数トークングループをより少ない圧縮KVエントリに要約するため、キャッシュ自体が短くなる
      • トークン単位の情報を一部手放す代わりに、長文コンテキストのコストを大幅に削減する

    • CSA vs HCA

      • CSA: 弱い圧縮率(m=4) + DeepSeek Sparse Attention (DSA)スタイルのtop-k選択
      • HCA: 強い圧縮(m'=128、128トークンを圧縮KVエントリ1個に) + 短くなったキャッシュ上でのdense attention
      • どちらの方式も、直近の非圧縮トークンのために128トークンのスライディングウィンドウ分岐を維持
      • CSAはより細部を残しつつ疎に選択し、HCAはエントリ数を大きく減らしてdense attentionを可能にするため、相補的であり、DeepSeek V4では両レイヤーを交互に配置

    • 効率結果 (1Mトークンコンテキスト、DeepSeek V3.2比)

      • DeepSeek V4-Pro: 単一トークン推論FLOPs 27%、KVキャッシュサイズ 10%
      • DeepSeek V4-Flash: FLOPs 10%、KVキャッシュサイズ 7%

    • 評価時の注意点

      • CSA/HCAがMLAより一般的に"優れている"と断定するのは難しく、より攻めた長文コンテキスト設計であり、より複雑
      • 論文にはablation studyがない
      • DeepSeek V4-Flash-Baseが多くのbaseベンチマークでV3.2-Baseを上回り、1Mトークンretrievalでも強い結果を示したが、これはより良いデータ・Muonベース最適化・mHC・精度/保存最適化・学習/推論システム変更を含む全体レシピの結果である

6. 結論

  • 今年の新しいオープンウェイトモデルに共通するパターンは、総パラメータ数を減らさずに長文コンテキスト推論コストを下げること
    • Gemma 4: cross-layer KV sharing でKVキャッシュを縮小し、per-layer embeddings でcapacityを追加
    • Laguna XS.2: レイヤーごとのアテンションcapacityを差別化
    • ZAYA1-8B: アテンションを圧縮潜在空間へ移す
    • DeepSeek V4: 制約付き残差ストリーム混合 + 圧縮長文アテンション
  • トランスフォーマーブロックは今なお変化を続けているが、狙いが明確な形での修正であり、基本骨格はGPT decoder-onlyアーキテクチャを維持している
  • 定性的なモデリング性能は主としてデータ品質・量と学習レシピによって牽引される
  • 現時点ではトランスフォーマーがSOTAアーキテクチャの現状(status quo) を維持しており、diffusionモデルなどの代替も存在する
  • 基本的なトランスフォーマーブロックはPyTorchで50〜100行程度で実装可能だったが、最近はアテンション変種などによりコードの複雑さが約10倍に増した
  • 複雑さの増大自体はランタイムコストを減らすため、必ずしも悪いことではないが、個々のコンポーネントとその相互作用を明確に理解することはますます難しくなっている
  • 学習に推奨されるアプローチ: 元祖decoder-style LLM(GPT/GPT-2)から出発し、新しいコンポーネントを1つずつ追加しながら学ぶ

関連記事

まだコメントはありません。

まだコメントはありません。