技術サマリー：N-vium – 高速な正確な生成のための混合出口トランスフォーマー

1. 問題定義

標準的な自己回帰型大規模言語モデル（LLM）は、デコーダ層のスタックを介してトークンを逐次的に処理します。このフォワードパスは、(1) 将来の注意機構のために各層でキー・バリュー（KV）表現を構築すること、および (2) 最終的な隠れ状態を投影して次のトークンを予測すること、という 2 つの結合した機能を提供します。その結果、予測分布がすでに収束しているかどうかに関わらず、すべてのトークンはすべての層を介した完全なフォワードパスを必要とします。

推論効率を向上させる既存のアプローチ、例えばEarly Exit（早期終了）手法は、中間層に予測ヘッダーを接続することでこれらの機能を分離しようと試みます。しかし、これらの手法は通常、中間予測を最終層の分布の近似として扱います。これにより、以下のいくつかの欠点が生じます。

• 近似誤差: 中間予測は、完全な深さの予測よりも弱くなることが多く、学習とテストの不一致を引き起こします。
• ヒューリスティックなルーティング: 終了のタイミングに関する決定は、しばしば較正を必要とする手作業で設計された信頼度閾値に依存します。
• 不完全な KV キャッシュ: 計算を節約するため、早期終了手法はバイパスされた層の KV キャッシュ計算をスキップすることが多く、近似コピーやスキップを必要とし、品質を低下させます。
• 品質の低下: 速度は頻繁に、パープレキシティや生成品質の犠牲によって達成されます。

本論文は、「logit lens」などの経験的観察（例えば、予測分布が最終層に到達する前に安定化することが多い）が示唆するところを踏まえつつも、現在の手法は品質の低下なしに正確な生成のためにこれを活用できていないと主張しています。

2. 手法：N-vium

著者らは、標準的なトランスフォーマーアーキテクチャを厳密に一般化するN-vium、すなわち混合出口フレームワークを導入します。N-vium は、早期終了を近似として扱うのではなく、モデルの出力を複数の出口にわたる学習された混合分布として再定義します。

アーキテクチャ

• 混合出口: モデルは、バックボーンの $L$ 層を $N_V$ ブロックに分割します。各ブロックの末尾（ジャンクション $k$）において、ルーターは、その層の分布 $\pi_k$ から次のトークンをサンプリングする確率を決定します。
• トークン適応型ルーティング: ルーター MLP は、ジャンクション $k$ における隠れ状態 $h^{(\ell_k)}_t$ を処理して確率 $w_k$ を出力します。ジャンクション $k$ での終了確率 $p_k$ は以下のように計算されます：
  $$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$$
  ここで、慣例により $w_{N_V} = 1$ とします。
• 混合分布: 最終的な次のトークン分布は、すべての出口分布の重み付き和です：
  $$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$$
• 共有ヘッダー: すべての出口は、最終的な言語モデリング（LM）投影ヘッダー $W_{lm}$ を共有し、中間隠れ状態をヘッダーの期待される幾何学形状に一致させるために軽量なアダプター MLP が変換を行います。

学習

• 目的関数: モデルは、個々のヘッダーを最終層に一致させるように個別に最適化するのではなく、結合混合分布 $\pi_{mix}$ のクロスエントロピー損失を最小化するように学習されます。
• 計算ペナルティ: ルーターがより早い出口に確率質量を割り当て、モデルを効率化へと誘導するために、補助的な損失項 $\mathcal{L}_{compute}$ が追加されます：
  $$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$$
  ここで、$\mathcal{L}_{compute}$ は期待される正規化された層の深さです。
• ウォームアップ: 「ルーターの崩壊」（すべてのトークンが最終層へルーティングされる現象）を防ぐために、ルーターのウォームアップ段階が採用されます。これは、初期段階で一様分布からの逸脱を罰することで実現されます。

推論：ピギーバック

正確な KV キャッシュの回復を計算のスキップなしに保証するためのピギーバック技術が、重要な革新です。

1. 遅延計算: トークンがジャンクション $k$ で終了した場合、層 $k+1$ から $L$ までの計算は遅延されます。
2. 並列バッチ処理: 将来のトークンが処理されるにつれて、それらのフォワードパスは遅延された層に到達します。これらの深さにおいて、遅延されたトークンの隠れ状態は、現在のトークンの状態と連結され、並列に処理されます。
3. 結果: これにより、すべての遅延されたトークンの欠落している KV 状態が、単一の並列パスで埋められ、KV キャッシュが完全で正確かつサンプリングされたトークンと整合性が取れていることが保証されます。計算がスキップされるのではなく、単に逐次的な実行から並列バッチ実行へと再順序付けされます。

3. 主要な貢献

1. 正確なサンプリング混合: N-vium は、標準的な早期終了手法に固有の近似誤差を回避し、学習済み分布からの正確なサンプリングを可能にするトークン適応型深さルーティングを実現する最初の手法です。
2. 完全な KV 回復: 上位層の計算を遅延およびバッチ化することにより、N-vium は完全かつ正確な KV キャッシュを回復します。これは、他の適応的計算手法ではしばしば失われる機能です。
3. 学習フレームワーク: 著者らは、学習可能な計算ペナルティを通じてモデルを効率化へと誘導する戦略を含む、混合目的を直接最適化する実用的な学習パイプラインを提供します。
4. スケーラビリティ: この手法は、最大15 億パラメータのスケールで検証されており、このアプローチが小規模モデルに限定されないことを実証しています。

4. 実験結果

著者らは、LLaMA 2 アーキテクチャを使用して C4 データセット上で N-vium モデル（具体的には 4 つの出口を持つ「Quadrivium」）を事前学習し、パラメータ数とデータが一致する密なベースラインと比較しました。

• 速度向上対パープレキシティ:
  • 最大モデル（15 億パラメータ、48 層、隠れ幅 1536）は、密なベースラインに対して57.9% のウォールクロック速度向上を達成し、パープレキシティの低下はありませんでした（0% の差）。
  • 小規模モデルは、23% から 41% の一貫した速度向上を示し、わずかなパープレキシティの改善が見られました（例：36 層モデルは 36.4% の速度向上と 3.7% のパープレキシティ改善を示しました）。
• パレートフロンティア: 品質 - 速度プロットにおいて、N-vium は「右上」の象限（より高速で、同等またはそれ以上の品質）を占めます。一方、CALM（CALM* として再実装）などの手法は「右下」（より高速だが品質が悪い）に位置し、LayerSkip はベースラインより下の一定のパープレキシティ線上に位置します。
• ドメイン適応後のパフォーマンス: Tulu 3 混合データセットでの教師あり微調整（SFT）後、N-vium モデルは MMLU、ARC、HellaSwag などの下流ベンチマークにおいて密なベースラインと同等かそれ以上のパフォーマンスを示し、モデルが誘導可能であり、タスクパフォーマンスを犠牲にしてクロスエントロピー目的に過剰適合していないことを確認しました。
• ハードウェアの堅牢性: A100、RTX A6000、RTX 3090 GPU 上での速度向上測定は、一貫した結果（約 57%）をもたらしましたが、ネイティブな bfloat16 サポートを欠く古い Turing アーキテクチャ GPU は、より低い利益を示しました。

5. 意義と主張

本論文は、N-vium が推論効率におけるパラダイムシフトを表すと主張しています。トークンあたりの FLOPs を削減する（多くの場合近似を介して）ことで秒あたりのトークン数を向上させるほとんどの手法とは異なり、N-vium は秒あたりの FLOPs を増加させることで効率を向上させます。これは、逐次的な計算を並列バッチパスに再編成することで達成され、計算を破棄することなくハードウェアの並列性をより効果的に活用します。

著者らは、N-vium が標準的なトランスフォーマーの厳密な一般化であることを強調しています。中間出口のルーティング重みがゼロの場合、モデルは標準的なトランスフォーマーと完全に同じように動作します。この手法は、著者が他の既存の手法が同時に満たしていないと主張する 4 つの重要な基準を満たします。

1. 正確な推論: 近似なしに学習済み分布からのサンプリング。
2. 完全な KV キャッシュ回復: スキップされた層や近似キャッシュの埋め込みがない。
3. 学習されたルーティング: トークンごとの計算割り当てはヒューリスティックではなく、学習されたルーターによって行われる。
4. デコード時の速度向上: 自己回帰生成中のウォールクロックレイテンシを削減する。

論文は結論として、N-vium の学習は（複数のジャンクションでヘッダーを評価することによる）トークンあたりのより高い FLOP コストを伴うが、推論コストが支配的な本番環境での展開にはこのオーバーヘッドが正当化され、モデルサイズがスケールするにつれて相対的なオーバーヘッドは減少すると述べています。