Draft-OPD：投機的ドラフトモデルのためのオン方策蒸留

要旨

投機的デコードは、ターゲットモデルと軽量なドラフトモデルを組み合わせ、ドラフトモデルが提案したトークンを並列に検証することで、大規模言語モデルの推論を高速化する。ドラフトモデルの構築方法として一般的なのは、EAGLE3やDFlashのように、ターゲットが生成した軌跡を用いた教師ありファインチューニング（SFT）である。しかし、我々はSFTが急速に頭打ちになる現象を観測した。すなわち、テストデータにおけるドラフトモデルの受理長（acceptance length）の改善が停止するのである。その原因は、オフラインと推論時のミスマッチにある。SFTではドラフトモデルが固定されたターゲット生成軌跡から学習するのに対し、投機的デコード時には自身のポリシーで提案されたブロック上で評価される。このことが、ドラフトモデルがドラフト誘発状態（draft-induced states）においてターゲットモデルから教師信号を受けるオン・ポリシー蒸留（OPD）の動機付けとなる。しかし、ドラフトモデルにとってOPDは依然として困難である。なぜなら、ドラフトモデルは独立して完全な系列を確実に生成することができず、一方でターゲット補助生成（target-assisted generation）を用いると収集される系列がターゲット分布に従ってしまい、オン・ポリシーの信号が失われるからである。そこで我々はDraft-OPDを提案する。これは、安定した継続生成のためのターゲット補助ロールアウトと、検証時に露呈した誤り位置からのドラフト再生（replay）を組み合わせる。これにより、ドラフトモデルは受理された提案と拒否された提案の両方に対してターゲットからのフィードバックを学習し、投機的受理を制限するドラフト誘発誤差に焦点を当てた訓練が可能となる。実験により、Draft-OPDは多様なタスクにおいて思考モデルに対して5倍以上のロスレス高速化を達成し、EAGLE-3およびDFlashに対してそれぞれ23%および13%の改善を示した。

English

Speculative decoding accelerates large language model inference by pairing a target model with a lightweight draft model whose proposed tokens are verified in parallel. A common way to build draft models, like EAGLE3 or DFlash is supervised fine-tuning (SFT) on target-generated trajectories. However, we observe that SFT quickly plateaus: the draft model's acceptance length on test data stops improving. The reason is an offline-to-inference mismatch: In SFT, the drafter learns from fixed target-generated trajectories, whereas during speculative decoding it is evaluated on blocks proposed under its own policy. This motivates on-policy distillation (OPD), where the target model supervises the drafter on draft-induced states. Yet OPD remains difficult for draft models, as they cannot reliably roll out complete sequences independently, whereas target-assisted generation makes the collected sequences follow the target distribution and thus eliminates the on-policy signal. We therefore propose Draft-OPD, which uses target-assisted rollout for stable continuations and replays drafting from the verification-exposed error positions. This allows the drafter to learn from target feedback on both accepted and rejected proposals, focusing training on the draft-induced errors that limit speculative acceptance. Experiments show that Draft-OPD achieves over 5times lossless acceleration for thinking models across diverse tasks, improving over EAGLE-3 and DFlash by 23\% and 13\%.