技術概要：D-PACE

問題定義

Speculative decoding（SD）は、より小さなドラフトモデルがトークンを提案し、より大きなターゲットモデルがそれを並列に検証することで、大規模言語モデル（LLM）の推論を加速します。最近の並列ドラフターである DFlash は、単一のフォワードパスで $B$ 個のトークン全体を予測することで効率を向上させていますが、その学習目的には固定された位置依存の重み付けスケジュール（例：ヘッド依存の重みまたは指数関数的減衰）に依存しています。

著者らは、これらの固定スケジュールに重大な限界があることを特定しました。すなわち、それらはドラフティングプロセスの動的な性質に適応しないという点です。学習中、ドラフトブロックの受諾を制限する位置は変化します。

• 初期の位置が非常に正確であれば、ボトルネックは後続の位置へシフトし、それらにはより多くの学習シグナルを与えるべきです。
• 初期の位置が誤っていれば、その後の位置の品質に関わらずそれらは拒否され、それらに割り当てられた学習シグナルは無効になります。

固定スケジュールは、特定の例の現在の受諾ボトルネックではなく、トークンの位置のみに依存するため、学習シグナルの配分を誤ります。

手法：D-PACE

これに対処するため、著者らは**D-PACE（Dynamic Position-Aware Cross-Entropy）**を提案しました。これは、期待される受諾ドラフト長に基づいて位置ごとの重みを動的に導出する学習目的です。

1. 受諾長の代理変数

核心的な洞察は、期待される受諾長 $E[X]$ が、滑らかで微分可能な代理変数 $\tilde{S}$ によって近似できるという点です。

• $q_i$ を、位置 $i$ においてターゲットモデルが選択したトークンに対するドラフトモデルの信頼度とします。
• 位置 $k$ までのプレフィックスが受諾される確率は、累積積 $\prod_{i=1}^k q_i$ によって近似されます。
• 非負整数に対するテール和の恒等式を用いると、期待される受諾長は以下のように近似されます：
  $$\tilde{S}(\psi) = \sum_{k=1}^B \prod_{i=1}^k q_i$$
  この代理変数は、学習中のブロックの受諾長の代理として機能します。

2. 適応的重みの導出

D-PACE は、各位置が総受諾長にどのように寄与するかを決定するために、ドラフトの信頼度に対する代理変数 $\tilde{S}$ の勾配を計算します。

• $\tilde{S}$ を $q_j$ に関して微分することで、著者らは位置 $j$ に対する閉形式の重み $w_j$ を導出します：
  $$w_j = \left( \prod_{i \le j} q_i \right) \cdot f_j$$
  ここで、$f_j$ は「継続値」であり、プレフィックスが受諾されたという条件のもとで、位置 $j$ およびそれ以降のすべての位置が寄与する期待される受諾長を表します。
• この重みは、累積信頼度（位置 $j$ に到達する確率）と継続値（残りのブロックの価値）を組み合わせます。

3. 非対称な重み平滑化

生身の累積積 $\prod q_i$ は数値的に不安定です。初期の信頼度が低い場合、後続の位置の重みは消滅し、それらの学習シグナルが抑制されます。

• 解決策：著者らは非対称な平滑化を適用します。重みの計算時のみ信頼度を平滑化します：$\tilde{q}_i = (1-\alpha)q_i + \alpha$。
• クロスエントロピー損失項自体は、平滑化されていない確率 $-\log q_j$ のまま残ります。これにより、標準的なトークンレベルの勾配を保持しつつ、初期の低信頼度予測が後続の位置の重みをゼロにするのを防ぎます。

4. D-PACE 損失

最終的な目的関数は、重み $\bar{w}_j$ が勾配から切り離され（定数として扱われる）、重み付けされたクロスエントロピー損失です：
$$L_{\text{D-PACE}} = \sum_{j=1}^B \bar{w}_j (-\log q_j)$$
このアプローチは、ドラフターのアーキテクチャや推論手順に変更を加えることを必要としません。

主要な貢献

1. 適応的な位置ごとの重み：固定された減衰スケジュールを、特定の例に対して現在受諾を制限しているブロックの位置に適応する重みに置き換える、重み付けされたクロスエントロピー目的関数の導出。
2. 非対称平滑化による安定した学習：重みの計算のみを平滑化することで累積積の重みを安定化させ、クロスエントロピー項の標準的な勾配特性を保持する手法。
3. 効率性：この手法は過剰なオーバーヘッド（学習時間の 2.3% 増加）を導入せず、アーキテクチャの変更も必要としません。

実験結果

著者らは、D-PACE を 6 つのベンチマーク（数学、コード、チャット）、2 つのターゲットモデル（Qwen3-4B、Qwen3-8B、Llama-3.1-8B）、およびさまざまなドラフト深度（3L、5L）と温度条件下で評価しました。

• 性能向上：D-PACE は、DFlash ベースラインと比較して、ウォールクロックの速度向上と平均放出長さ（$\tau$）の両方を一貫して改善しました。
  • Qwen3-4B（3L）において、D-PACE は平均放出長さを**+8.5% から +10.7%、速度向上を+8.0% から +9.7%**改善しました。
  • Qwen3-8B および Llama-3.1-8B において、放出長さの改善はそれぞれ約**+12.8%および+12.5%**でした。
• アブレーション研究：
  • コンポーネント分析：「累積のみ」のバリアント（プレフィックス確率による重み付け）が向上の大部分を提供し、「継続のみ」のバリアントは限定的な改善しか提供しませんでした。両方を組み合わせた完全な D-PACE が最良の結果をもたらしました。
  • 平滑化の感度：性能は平滑化パラメータ $\alpha \in [0.3, 0.7]$ 全体で安定していました。平滑化なし（$\alpha=0$）の場合、性能は崩壊しました。
  • ブロックサイズ：D-PACE はブロックサイズ $B=8, 12, 16$ 全体で一貫した改善を示し、ブロックサイズが増加するにつれて利益がわずかに増加しました。
• 一般化：この手法は、コアロジックの再調整なしに、異なるターゲットモデル（Llama-3.1-8B、Qwen3-8B）およびデータセットへ効果的に転移しました。

意義と主張

本論文は、D-PACE が静的な位置ベースの学習スケジュールから動的な貢献ベースの重み付けへの転換を表すと主張しています。学習シグナルを期待される受諾長に直接整合させることで、D-PACE はドラフターが並列デコーディングプロセスを現在制約している位置に学習努力を集中させることを保証します。

著者らは、この改善が以下の点においてなしで達成されていることを強調しています：

• ドラフターのアーキテクチャの変更。
• 推論手順の修正。
• 顕著な計算オーバーヘッドの導入。

この研究は、並列ブロックドラフターにとって、受諾の微分可能な代理変数から導出された重みで固定スケジュールを置き換えることが、Speculative decoding の効率を向上させるための一般化可能な戦略であることを示唆しています。論文は明示的に、この導出は 1 回のパスですべての位置を生成する並列ブロックドラフターに適用可能であり、状態の蓄積が固定プレフィックスの因数分解を破る逐次ドラフターには現時点では拡張されていないと述べています。