Source-linked AI summary
Verification-Aware Training for Speculative Decoding
Geonmo Gu, Byeongho Heo, HeeJae Jun, Yoohoon Kang, Sangmin Lee, Sangdoo Yun, Dongyoon Han
TL;DR
Speculative decoding depends on draft tokens surviving sequential target-model verification, but existing training uses token imitation and fixed positional weights that do not reflect this process. VAT simulates verification during training, adding a verification head and adaptive weighting; across evaluated models and benchmarks, it improves both acceptance length and wall-clock speedup.
Problem
Existing draft training imitates target outputs with fixed per-position weighting, although verification is sequential and discards positions after the first rejection.
Method
VAT simulates target verification at each training step and adds a jointly trained verification head plus weighting re-anchored at each sample’s first rejection.
Results
VAT consistently improves average acceptance length and wall-clock speedup across EAGLE-3 and DFlash, with gains up to 11.4% and 8.7%, respectively.
Takeaways & Limitations
Because VAT changes only the training objective, it can be layered onto existing speculative decoding methods without modifying the draft architecture, target model, or inference procedure.
Abstract
from arXiv · showhide
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
1 Introduction
Speculative decoding reduces inference cost by drafting candidate tokens and verifying them with the target model, but existing draft training does not reflect sequential acceptance. VAT aligns training with verification through simulated accept/reject supervision and adaptive weighting, improving acceptance length and speedup across evaluated settings.
- Motivation: Speculative decoding splits generation into low-cost drafting and target-model verification to reduce inference latency.The target model verifies multiple draft candidates in a single forward pass.
- Motivation: Existing draft training imitates target outputs with fixed per-position weighting, ignoring sequential verification and sample-specific first rejections.After the first rejection, subsequent positions are discarded regardless of their individual quality.
- Method: VAT simulates target verification during training and uses the resulting labels for a verification head and verification-adaptive weighting.The weighting keeps full weight through the first rejection and re-anchors decay afterward.
- Results: 11.4% is the maximum improvement in average acceptance length, while 8.7% is the maximum improvement in wall-clock speedup.The evaluation covers EAGLE-3 and DFlash with Qwen3-4B, Qwen3-8B, and LLaMA-3.1-8B across math, code, and chat benchmarks.
2 Related Work
Related work advances speculative decoding through improved drafting architectures and training procedures. EAGLE-3 and DFlash use different drafting paradigms but train by imitating target outputs, while prior objective-focused methods retain uniform positional weighting.
- Drafting Architectures: EAGLE and its extensions improve drafting by reusing target hidden states, while EAGLE-3 additionally trains on multi-step draft-generated contexts.These methods improve draft quality without changing the basic drafting-and-verification setup.
- Training Objectives: Objective-focused methods study distillation and train-test mismatch, but HASS and EAGLE-3 retain uniform per-position loss weighting.This leaves positional supervision separate from the verification pattern that determines accepted prefixes.
3 Preliminary
Speculative decoding drafts K candidate tokens and verifies them sequentially, so the average acceptance length determines how many tokens are produced per verification cycle. Existing draft objectives supervise each position with target outputs and fixed schedules that do not adapt to sample-specific rejection points.
- Drafting and Verification: Given a prefix, the draft model generates K candidate tokens, which the target model evaluates in one forward pass.Verification proceeds sequentially under an acceptance rule that preserves the target distribution.
- Drafting and Verification: The average acceptance length τ measures expected accepted tokens per cycle, and larger τ produces more tokens per fixed verification cost.It is the key performance metric for speculative decoding systems.
- Draft Model Training: Draft training minimizes a weighted sum of per-position cross-entropy losses against target outputs.EAGLE-3 uses soft target distributions, whereas DFlash uses sampled target tokens as hard labels.
- Draft Model Training: EAGLE-3 and DFlash use predetermined position-dependent weights that emphasize earlier positions rather than adapting to each sample’s first rejection.Their schedules are 0.8^(k−1) for EAGLE-3 and exp(−(k−1)/γ) for DFlash.
4 Method: Verification-Aware Training
VAT aligns draft-model training with sequential verification by simulating acceptance outcomes and adapting supervision to each sample’s first rejection. It adds a verification head and verification-adaptive weighting while changing only the training objective.
- VAT simulates target verification during training and uses the resulting accept/reject patterns as additional supervision.
- Verification Head: The verification head is a lightweight binary classifier jointly trained to predict whether each draft position survives sequential verification.It is attached to the draft model’s last hidden states and trained against acceptance labels.
- VAT leaves the draft architecture, target model, and inference procedure unchanged, enabling integration with existing speculative decoding methods.The verification head is used as a training-time auxiliary objective and adds no inference cost in the main setup.
- Verification-Adaptive Weighting: Verification-adaptive weighting keeps full loss weight through each sample’s first rejection, then reuses the base decay schedule from that rejection point onward.The first rejection position also receives full weight because improving it can directly extend the accepted prefix.
- Training Objective: VAT combines dynamically reweighted soft-label and hard-label cross-entropy losses with the verification-head loss.The soft label matches the target distribution, while the hard label supervises the target’s sampled token.
5 Experiments
VAT improves speculative-decoding acceptance length and wall-clock speedup across models, drafting methods, and task categories. Ablations and dynamics analyses show that its verification head and verification-adaptive weighting provide complementary benefits, while head-guided early exit reduces inference cost.
- Main Results: VAT improves average acceptance length and wall-clock speedup across EAGLE-3, DFlash, three target models, and math, code, and chat benchmarks.Average acceptance-length gains reach 11.4%, while wall-clock speedup gains reach 8.7%.
- Component Ablation: The verification head, verification-adaptive weighting, and soft-plus-hard labels each improve average acceptance length over the DFlash Qwen3-4B baseline.Average τ rises from 5.73 to 5.87 with the head, 5.91 with adaptive weighting, and 5.82 with soft-plus-hard labels.
- Component Ablation: Combining all three factors yields the best average acceptance length, 6.08, and speedup, 4.81×, across the evaluated task categories.The verification head and verification-adaptive weighting together achieve 4.76× speedup among pairs.
- Weighting Ablation: Verification-adaptive weighting performs best with either tested base-weight schedule, indicating that adaptation rather than decay form drives the gain.It reaches τ = 6.09 with the EAGLE-3-style base and τ = 6.08 with the DFlash-style base.
- Training Dynamics: Attaching the verification head shifts the first rejection later and keeps post-rejection target matches stable during training.The baseline’s post-rejection match count gradually declines, whereas the verification-head variant remains stable despite deeper post-rejection regions.
- Inference Analysis: Head-guided early exit improves speedup over drafting without early exit, recovering most of the oracle speedup with a small acceptance-length drop.The head predicts the first-rejection position to truncate the draft before target verification.
6 Conclusion
VAT aligns draft-model training with target-model verification by simulating verification during training and adapting supervision to rejection points. It can be layered onto existing methods and improves acceptance length and wall-clock speedup across diverse benchmarks.
- VAT introduces a verification head and verification-adaptive weighting to align draft training with the target’s verification process.The verification head supervises verification outcomes, while weighting adapts to each sample’s first rejection point.
- Both VAT components modify only the training objective, without changing the draft architecture, target model, or inference procedure.
- The verification head can optionally guide early-exit drafting at inference time for additional speedup.
- VAT consistently improves average acceptance length and wall-clock speedup for EAGLE-3 and DFlash across math, code, and chat benchmarks.
- Evaluation is limited to models up to 8B parameters, leaving scalability to substantially larger models as future work.
A Ablation on Weighting Schemes
The ablation shows that verification-adaptive weighting performs best among the tested schemes, retaining full weight through the first rejection and decaying afterward. Alternatives that eliminate or ignore rejection-point conditioning perform worse on averaged benchmark speedup and τ.
- The ablation keeps the rest of the VAT recipe fixed and uses a 1-epoch budget on DFlash with Qwen3-4B.
- Prefix-only weighting collapses below the DFlash baseline at 2.44× speedup / 3.11 τ because samples rejected at the first position contribute no draft loss.
- Verification-adaptive weighting achieves the best averaged result: 4.61× speedup / 6.03 τ across math, code, and chat benchmarks.
- Hard cutoff weighting recovers most of the gap but remains below VAT at 4.44× speedup / 5.75 τ.
- Unshifted decay reaches 4.46× speedup / 5.75 τ, indicating a comparable margin from re-anchoring decay at the first rejection.
- Confidence-derived alternatives remain lower, with marginal contribution at 4.38× / 5.74 and D-PACE-style weights at 4.52× / 5.89.
B Effect of the Verification Rule and Corpus Temperature
VAT’s training signal is robust to the verification rule and corpus temperature choices examined with DFlash on Qwen3-4B. Simulated accepted lengths under greedy and stochastic verification remain highly correlated throughout training.
- The comparison evaluates all combinations of corpus temperature and verification rule, at both T=0 and T=1, using speedup / τ averaged across benchmarks.
- The training corpus is generated with greedy decoding, and verification under the same scheme reduces to top-1 agreement between draft and target predictions.
- Accepted lengths under greedy and stochastic verification maintain a Pearson correlation above 0.92 throughout training.
C Training Overhead of VAT
VAT adds modest online training overhead while computing verification labels and target distributions at each step. The overhead differs between EAGLE-3 and DFlash because EAGLE-3 already computes the target LM-head distribution.
- Measurements average 100 steps on a single NVIDIA A100 80GB GPU and report per-step time and peak GPU memory.
- No verification labels or target distributions are precomputed or cached; all quantities are computed online with the frozen target model.
- VAT adds 1.2% per-step time to EAGLE-3 and 6.1% to DFlash on Qwen3-4B.
- EAGLE-3 reuses its existing target LM-head distribution for soft labels, leaving the lightweight verification head as the remaining cost.
- DFlash incurs an additional target LM-head pass per training step because its original training does not apply that head.