DBLAST: Dependent Block Drafting for Stochastic Speculative Decoding

技術要約: DBlast – 確率的投機的デコーディングのための依存ブロックドラフティング

問題提起

投機的デコーディング（Speculative Decoding）は、軽量なドラフター（下書きモデル）を用いて複数の将来のトークンを提案し、ターゲットモデルがそれらを並列に検証することで、大規模言語モデル（LLM）の推論を加速させる。近年のブロックベースのドラフター（例：DFlash）は、一度のパスで複数の位置を予測できるが、通常は独立したブロックサンプリングに依存している。このアプローチは、ドラフトブロック内のトークンが、接頭辞（prefix）を与えられた条件下で条件付き独立であるという仮定に基づいている。

この仮定は、貪欲（greedy）または低エントロピーのデコーディングにおいては十分に成立するが、対話、クリエイティブ・ライティング、強化学習による探索といったオープンエンドなタスクに不可欠な確率的（non-greedy）デコーディングのレジームにおいては脆弱になる。確率的な設定では、ターゲット分布は意図的に多様化されるため、複数の継続パターンが妥当となり得る。本論文は、根本的なミスマッチを特定している。すなわち、ドラフターは周辺的なターゲット分布に一致するようにトークンを独立にサンプリングする一方で、ターゲットモデルは、以前に受理されたトークンに条件付けられた状態で、それらを**逐次的（sequential）**に検証するという点である。その結果、ブロック内の各位置の周辺分布が正確であっても、一貫した条件付き軌跡を形成できず、ターゲットのサンプリング分布のエントロピーが増大するにつれて、受理されるドラフト長が著しく低下するという問題が生じる。

手法

著者らは、並列予測の効率性を損なうことなく、ブロック拡散ドラフターに依存関係を注入するフレームワークであるDBlast（Dependent Block Drafting）を提案する。この手法は、以下の2つのコアコンポーネントで構成される。

1. 潜在混合による依存ブロックサンプリング:
   提案分布を各位置で分解する（$q(x_{1:b}|y) = \prod q_i(x_i|y)$）代わりに、DBlastはブロック内の相関を誘導するために、カテゴリカルな潜在変数 $z$（「分岐」または「モード」）を導入する。結合分布は、低ランクの正準CP（Canonical Polyadic）混合としてモデル化される：
   $$q_{dep}(x_{1:b} | y) = \sum_{z=1}^{K} q(z | y) \prod_{i=1}^{b} q_i(x_i | y, z)$$

   • アーキテクチャ: この手法は、Transformerの隠れ状態を$K$個のカテゴリ固有のオフセットにマッピングする軽量な「カテゴリ隠れ拡張器（category hidden expander）」と、$q(z|y)$を予測する「カテゴリ事前分布ヘッド（category-prior head）」を追加することで、DFlashスタイルのドラフターを修正している。
   • 推論: 推論時には、潜在的なカテゴリが（オプションで温度 $Z_T$ を伴って）サンプリングされ、対応するブランチが単一の並列パスで貪欲にデコードされる。これにより、ブロック・ドラフティングのワンパスの効率性を維持しつつ、モデルが整合性のある代替的なブロックレベルの継続を表現することを可能にする。

2. 受理指向の学習目的関数:
   標準的な学習は、ターゲットブロックの負の対数尤度（NLL）を最小化するが、これは提案の尤度を最適化するのみであり、逐次的な検証プロセスを無視している。DBlastは、**期待受理長（Expected Accepted Length: AL）**に基づくサロゲート損失を導入する。

   • この目的関数は、接頭辞の受理イベントに分解された、受理された接頭辞の長さの期待値から導出される。
   • ドラフターとターゲット分布のオーバーラップが低い場合（学習初期）の重要度サンプリング（importance weighting）の不安定性を扱うため、著者らは**閾値切断型ログドメイン・サロゲート（threshold-truncated log-domain surrogate）**を提案している。この損失は、重要度比が閾値 $\tau$ を超える保持された接頭辞に対して、確率の対数比と条件付き受理長の対数の和を算出する。
   • このアプローチは、検証器の逐次的な受理挙動を直接的なターゲットとしており、ブロックの尤度だけでは不十分な場合の依存関係モデリングを補完する。

主な貢献

1. ミスマッチの診断: 著者らは、独立したブロックサンプリングが、ターゲットのデコーディングが決定論的でなくなるにつれて、性能が低下することを実証的に示している。独立した提案は、逐次的検証に必要な一貫性を欠いているため、高エントロピーのレジームにおいて受理長が著しく低下することを示した。
2. 依存ブロックサンプリング: 非貪欲な投機的サンプリングのために、CPスタイルの潜在混合提案をサブワード・ブロック拡散ドラフターに初めて適用した。この手法は、並列的なトークン予測を維持しながら、ブロック内の位置間の相関を誘導する。
3. 受理指向の損失関数: 期待受理長に基づいた新しい学習サロゲートを提案した。この目的関数は、閾値切断を組み合わせることで、ドラフターを逐次的な検証プロセスに整合させる安定した学習信号を提供する。

実験結果

実験は、Qwen3-4BおよびQ問3-8Bターゲットモデルを用い、4つのベンチマーク（GSM8K（数学）、MT-Bench（チャット）、HumanEval（コーディング）、およびクリエイティブ・ライティングベンチマーク）で実施された。評価は、高決定論（Temperature 0.7）から低決定論（Temperature 1.5）までの3つのターゲットサンプリング・レジームをカバーしている。

• 性能向上: DBlastは、独立したブロックサンプリング（DFlash）およびNLLで学習されたベースラインを一貫して上回った。
  • 高エントロピーの設定（Target Temp 1.5）において、DBlastはQwen3-8Bにおいて、独立したベースラインと比較して受理長で平均12.1%のマクロ平均ゲインを達成した。
  • ゲインは低決定論領域で最も顕著であり、依存関係のモデリングが、複数の継続が妥当となる場合に最も価値があることを裏付けている。
• アブレーション研究:
  • 潜在カテゴリ数 ($K$): $K$を増やすと性能が向上したが、$K=4$以降は収穫逓減が見られた。
  • 学習目的関数: 受理指向（AL）損失は、独立した（$K=1$）ドラフターであってもNLL学習を一貫して上回ったが、依存関係とALの組み合わせが最良の結果をもたらした。
  • 推論温度: カテゴリサンプリングの温度（$Z_T$）を調整することで、ドラフターがターゲット分布の確率性をより良く一致させることができ、受理長がさらに向上した。

意義と主張

本論文は、DBlastが投機的デコーディングにおける重要なギャップ、すなわち現在のブロックドラフターが確率的な検証を効率的に扱うことができないという問題を解決すると主張している。ブロック内の依存関係モデリングと受理に整合した学習が相補的であることを示すことで、投機的デコーディングの魅力である並列効率を犠牲にすることなく、多様で非貪欲なデコーディング・シナリオにおいて高い受理率を達成できることを著者らは示した。

著者らは、提案するサロゲート損失は形式的な不偏推定量や下限ではないものの、実用的かつ効果的な学習信号として機能すると控えめに述べている。また、学習中に使用されるソフトな潜在混合と、推論時に使用される貪欲なブランチ・デコーディングとの関係については、形式的な特性付けが未開拓の領域であることを認めているが、実験結果は効果的な転移を示している。本研究は、オープンエンドな生成のための将来の投機的デコーディング・システムが、単なる独立したトークン予測を超えて、将来のシーケンスの構造的一貫性を捉える必要があることを示唆している。