Verification-Aware Training for Speculative Decoding

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 and discards every position from the first rejection onward, yet existing draft training relies on token-level cross-entropy with a fixed per-position weighting that does not reflect this process. 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-trainable binary classifier that validates per-position acceptance via an explicit prediction target; (ii) verification-adaptive weighting replaces the fixed weighting schedule with weights adapted to each sample's first rejection point during training. VAT is model-agnostic and can be layered on top of existing methods without changing the draft architecture, the target model, or the inference procedure, and trained jointly with the original loss. Applied to EAGLE-3 and DFlash on Qwen3-4B, Qwen3-8B and LLaMA3.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.