投機的デコードのための検証を考慮した学習

要旨

投機的デコーディングは、ドラフトモデルを用いて候補トークンを生成し、それらをターゲットモデルが単一のフォワードパスで検証することにより、大規模言語モデルの推論を高速化する。検証は逐次的に進行し、最初の棄却以降のすべての位置を破棄する。しかしながら、既存のドラフトモデルの訓練は、どちらの特性も反映していない固定の位置別重み付けによるターゲットのトークンレベル模倣に依存している。我々は、各訓練ステップで検証をシミュレートし、その結果得られる受理・棄却パターンを教師信号に変換するプラグインフレームワークである検証認識トレーニング（Verification-Aware Training, VAT）を導入する。VATは2つの構成要素からなる。(i) 検証ヘッド：各位置が逐次検証を通過するかどうかについてドラフトモデルを教師する、ドラフトモデルと同時に学習される軽量な二値分類器。(ii) 検証適応型重み付け：各サンプルの最初の棄却点まで全重みを維持し、減衰の開始をその点に再設定することで、固定の重み付けスケジュールを置き換える。VATは訓練目的関数のみを変更するため、ドラフトモデルのアーキテクチャ、ターゲットモデル、推論手順を変更することなく、既存手法の上に重ねて適用することができる。VATをEAGLE-3およびDFlashに、Qwen3-4B、Qwen3-8B、LLaMA-3.1-8B上で適用したところ、平均受理長を最大11.4%、実時間速度向上を最大8.7%改善し、数学、コード、チャットの各ベンチマークで一貫した向上が得られた。コードは https://github.com/naver-ai/vat で公開予定である。

English

Speculative decoding accelerates large language model inference by using a draft model to generate candidate tokens, which are verified by the target model in a single forward pass. Verification proceeds sequentially and discards every position from the first rejection onward, yet existing draft training relies on token-level imitation of the target with a fixed per-position weighting that reflects neither property. We introduce Verification-Aware Training (VAT), a plug-in framework that simulates verification at every training step and turns the resulting accept and reject patterns into supervision. VAT consists of two components: (i) a verification head, a lightweight jointly trained binary classifier that supervises the draft model on whether each position survives sequential verification; (ii) verification-adaptive weighting, which replaces the fixed weighting schedule by keeping full weight up to each sample's first rejection point and re-anchoring the decay to start there. VAT modifies only the training objective, so it can be layered on top of existing methods without changing the draft architecture, the target model, or the inference procedure. Applied to EAGLE-3 and DFlash on Qwen3-4B, Qwen3-8B, and LLaMA-3.1-8B, VAT improves average acceptance length by up to 11.4% and wall-clock speedup by up to 8.7%, with consistent gains across math, code, and chat benchmarks. Code will be available at https://github.com/naver-ai/vat