Verification-Aware Training for Speculative Decoding

2026-08-31Computation and Language

Computation and Language
AI summary

The authors propose a new training method called Verification-Aware Training (VAT) to make speculative decoding faster and more accurate. Speculative decoding speeds up language models by letting a smaller draft model guess tokens, which are then checked by the main model, stopping at the first wrong guess. VAT teaches the draft model to better predict which tokens will be accepted or rejected in this step-by-step check by adding a verification head and adjusting training weights based on where rejections occur. This approach improves the length of accepted guesses and overall speed without needing to change the main model or decoding process. The authors tested VAT on several models and saw consistent improvements across different tasks like math, coding, and chatting.

speculative decodinglanguage model inferencedraft modelverificationtraining objectivebinary classifieracceptance lengthweighting scheduleEAGLE-3LLaMA
Authors
Geonmo Gu, Byeongho Heo, HeeJae Jun, Yoohoon Kang, Sangmin Lee, Sangdoo Yun, Dongyoon Han
Abstract
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