検証を予測する草稿モデルへ：VATが投機的デコードを改善

概要

投機的デコードは、大規模言語モデルの自己回帰推論を高速化する代表的な手法である。小さな草稿モデルが複数の候補トークンを先に生成し、大きなターゲットモデルがそれらを1回のフォワード計算で検証する。候補が連続して受理されれば、ターゲットモデルを1トークンずつ実行する回数を減らせる。

一方、この検証には重要な特徴がある。検証は左から順番に進み、最初の棄却が発生すると、その位置以降の候補は破棄される。従来の草稿モデル学習は、ターゲットのトークン分布をまねることを中心とし、位置ごとに固定された重みを使うことが多かった。この設定では、1つの棄却が後続位置全体に与える影響を十分に表現できない。

論文「Verification-Aware Training for Speculative Decoding」は、このずれを埋めるVerification-Aware Training（VAT）を提案する。VATは草稿モデルやターゲットモデルを作り直すのではなく、学習目的を変更するプラグイン型の枠組みである。

主な仕組み

• 検証ヘッドの追加。 学習中にターゲットモデルによる逐次検証をシミュレートし、各位置が検証を生き残るかどうかを記録する。軽量な二値分類器である検証ヘッドを草稿モデルと同時に学習し、各候補位置の生存可能性を予測させる。これにより、単なるトークン予測だけでなく、候補列をどこまで維持できるかも学習対象になる。
• 検証適応型の重み付け。 固定された位置減衰の代わりに、各サンプルで最初に棄却される位置までの損失をフルウェイトにする。その後は、その最初の棄却位置を新たな基準点として減衰を始める。検証における位置ごとの実際の価値に近い重み付けを目指す設計だ。
• 推論手順を変更しない。 ターゲットモデル、草稿モデルの構造、推測的デコードの実行方法を変える必要はなく、主な変更は学習時の教師信号と損失である。

結果と意味

VATはEAGLE-3とDFlashに適用され、Qwen3-4B、Qwen3-8B、LLaMA-3.1-8Bで評価された。報告された結果では、平均受理長が最大11.4%、実測のウォールクロック高速化が最大8.7%向上した。数学、コード、チャットのベンチマークでも一貫した改善が示されている。ただし、提供された素材には全実験条件や構成別の詳細が含まれていないため、これらはすべての環境で保証される値ではなく、論文が報告する最大改善幅として読むべきである。

この研究が示すのは、トークン単位の模倣精度と投機的デコードの効率は同じではないという点だ。草稿モデルには、局所的にもっともらしいトークンを出すだけでなく、長い受理済みプレフィックスを連続して生成する能力が求められる。受理と棄却のパターンを学習目標へ戻すVATは、既存の推論基盤を維持しながらこの目的を扱う実践的な方向性といえる。今後は公開コードを通じて、ハードウェア、バッチサイズ、タスクごとの効果を確認する必要がある。

出典：Hugging Face Daily Papers