技術概要：自己回帰的推測デコーディングにおけるアテンションドリフト

問題定義

推測デコーディングは、軽量な「ドラフター」モデルを用いて将来のトークンを予測し、それをより大きな「ターゲット」モデルで検証することで、大規模言語モデル（LLM）の推論を加速する。効果的ではあるが、現在のドラフターモデル（特に EAGLE-3 およびマルチトークン予測ヘッド）は、テンプレート摂動や長文脈入力といった展開環境の変化に対して著しい脆弱性を示す。実務家は、パフォーマンスを維持するために特定の推論エンジンやプロンプトに合わせてドラフターを再学習させることが多く、これは根本的なアーキテクチャの不安定さを示唆している。

本論文は、アテンションドリフトと呼ばれる以前報告されていない現象を特定した。ドラフターが推測トークンの系列を生成するにつれて、そのアテンション機構は「アテンションシンク」（通常、プロンプトの開始付近にある安定したアンカートークン）から次第に離れ、自身が直近に生成したトークンに集中するようになる。このドリフトは、特に長文脈シナリオや入力テンプレートが訓練分布から逸脱する場合に、受入率の低下と相関する。

手法と分析

著者らは、複数のモデルファミリー（Llama 3.1、Qwen 3/3.5、GPT-oss）における経験的観察とアーキテクチャ分析を組み合わせ、アテンションドリフトの根本原因を調査した。

1. 現象学的観察

アテンションヒートマップと定量的指標を用いて、著者らは以下のことを実証した：

• シンクの移動: 強力なアテンションシンクを持つモデル（Llama、Qwen など）において、推測チェーンが深くなるにつれてドラフターのアテンションはシンクトークンから減少し、直近にドラフトされたトークンへのアテンションが増加する。
• 普遍性: このドリフトは、凍結されたターゲット上で事後学習された EAGLE-3 ドラフターと、ターゲットと共同学習された MTP ヘッドの両方で観察される。これは、自身の出力を消費する自己回帰的デザインに固有の性質であることを示唆している。
• シンクの不在: シンクを抑制するように設計されたモデル（ゲート付きアテンションや学習可能なバイアスを備えた GPT-oss など）であっても、プロンプトの開始ではなく、繰り返されるテンプレートトークンにアテンションが集中するドリフトの形態が発生する。

2. 根本原因分析：隠れ状態の大きさ

本論文は、標準的なプレノーマンアーキテクチャにおける推測ステップ間の正規化されていない残差パスにアテンションドリフトの根源を特定した：

• 単調な増大: プレノーマン・トランスフォーマーでは、残差結合が正規化されずに大きさを蓄積する。著者らは、ドラフターの隠れ状態の大きさ（$\|h^*_k\|$）が推測深度（$k$）とともに単調に増大することを示した。
• 層のスタッキング対自己回帰: この増大により、ドラフターは単一の自己回帰予測器として振る舞うのではなく、ターゲットの上に積み重ねられた追加のトランスフォーマー層の系列として振る舞うようになる。ステップ $k$ において、ドラフターは安定した分布に基づいて次のトークンを予測するのではなく、実質的に深度 $N+k$ におけるターゲットモデルを近似する。
• 不均衡な入力: ターゲットモデルから融合された隠れ状態（$h_{FC}$）は、しばしばより大きな大きさを持つ深い層によって支配されており、これがドラフターに入力される際のスケールミスマッチをさらに悪化させる。

3. 提案される介入

大きさの蓄積に対抗し、ドラフティングプロセスを安定化させるため、著者らは 2 つのアーキテクチャ修正を提案する：

1. ドラフター隠れ状態へのポストノーマ: 残差結合の後（推測ステップ間）にドラフターの隠れ状態に対して RMSNorm を適用することで、大きさの単調な増大を防ぐ。これにより、ドラフターは深度固有の微調整ではなく、深度不変の自己回帰則を学習することを強制する。
2. 隠れ状態ごとの RMSNorm: ターゲットの隠れ状態（$h_{low}, h_{mid}, h_{high}$）をドラフター入力に融合する前に RMSNorm レイヤーを追加し、入力信号の大きさをバランスさせる。

著者らはまた、固定されたシンクへの依存を抑制するゲート付きアテンションを別の介入として評価した。その結果、ゲート化はシンクを排除するが、大きさの増大を防ぐことはできず、ゲート化とポストノーマを組み合わせると過剰正則化（エントロピーの崩壊）を招くことが判明した。

主要な結果

提案されたポストノーマアーキテクチャは、さまざまなベンチマークおよび条件下で、最先端のプレノーマ EAGLE-3 と比較評価された：

• テンプレート摂動: テンプレート摂動下（チャットテンプレートや文頭トークンの削除など）において、ポストノーマドラフターはプレノーマと比較して受入長を最大2 倍改善した。プレノーマモデルは壊滅的な低下（最大 52% の損失）を被ったのに対し、ポストノーマモデルは耐性が高く（最大 5% の損失）、安定していた。
• 長文脈タスク: 長文脈タスク（LongBench、スライディングウィンドウアテンションを伴うマルチターンチャット）において、ポストノーマは受入長を1.18 倍改善した。フルアテンションの長文脈設定では、プレノーマモデルは受入率がほぼゼロに崩壊したが、ポストノーマは実用的なパフォーマンスを維持した。
• 標準ベンチマーク: 7 つの標準ベンチマーク（マルチターンチャット、数学、コーディング）全体で、ポストノーマは受入長において一貫して1.10 倍の改善を示した。
• 学習効率: ポストノーマドラフターは、訓練時の視野を超えた推論深度に汎化する。これにより、より短い「訓練時テスト（TTT）」深度（例：8 ステップから 4 ステップへの削減）での学習が可能となり、パフォーマンスを犠牲にすることなく訓練時間を約 3 分の 1 に削減できる。
• ノイズ耐性: ポストノーマアーキテクチャは、隠れ状態におけるノイズに対して著しく高い耐性を示す（$\alpha=0.5$ のノイズで 58% の受入率維持に対し、プレノーマは 5%）。これは、量子化や分布シフトに対するより優れた堅牢性を示唆している。

意義と主張

本論文は、アテンションドリフトが、正規化されていない残差パスと自己回帰的な自己フィードバックの相互作用に起因する、推測デコーディングにおける重大かつ以前報告されていない故障モードであると主張する。

• 診断: 著者らは、アテンションドリフトはより深い問題の症状であると論じる。すなわち、隠れ状態の大きさの蓄積により、ドラフターが安定したトークン予測ではなく、深度依存の変換を学習してしまうことである。
• 解決策: 彼らは、単にアテンションシンクを修正する（例えばゲート化を通じて）だけでは不十分であり、根本的な大きさのダイナミクスに対処する必要があると主張する。提案されたポストノーマアーキテクチャは、ドラフティングプロセスを安定化させるシンプルかつ効果的な解決策として提示される。
• 汎化: この知見は、特にさまざまなプロンプトや長文脈に対する堅牢性が求められる生産システムにおいて、ポストノーマを推測デコーディング用ドラフターのデフォルトのアーキテクチャ選択として考慮すべきであることを示唆している。

本論文は、「アテンションドリフト」という名称を付け、その特性を記述することがドラフターの堅牢性に対する有用な診断基準を提供し、大きさの蓄積を標的としたアーキテクチャ介入が信頼性の高い推測デコーディングに不可欠であると結論づけている。