Verification-Aware Training for Speculative Decoding
Geonmo Gu, Byeongho Heo, HeeJae Jun, Yoohoon Kang, Sangmin Lee, Sangdoo Yun, Dongyoon Han
cs.CL
2026-08-31
VAT simulates verification during draft training with an accept head and first-rejection weights. On EAGLE-3 and DFlash, acceptance length rises up to 11.4%, speedup up to 8.7%.
Speculative decoding drafts a block of tokens with a cheap model, then verifies them in one target forward pass. Speedup tracks acceptance length. Verification walks the block from the left: after the first rejection, later drafts are thrown away even if they would have matched.
Draft training ignores that process. The objective still imitates the target token by token, with a decay schedule written before training and shared across samples. The schedule cannot follow where a given sample first fails, and the draft never sees a label for "would this prefix survive verification."
VAT changes only the training loss. Draft architecture, target model, and the inference loop stay as they are. Verification is simulated at every training step and turned into supervision.
At each step the method compares draft and target predictions, applies the speculative sampling rule to get per-position accept bits, and records the first rejection k. The label vk is 1 only if every earlier position was accepted, matching inference, where one reject voids the rest of the block.
Two losses sit on that simulation.
The verification head is a single dense layer on the draft's last hidden states, trained with binary cross-entropy on vk. It exists only at training time by default. Gradients flow into the draft so hidden states become informative about agreement with the target. That splits two cases the token loss treats as equal: a live prefix, where a correct guess extends acceptance length, and a dead prefix, where a correct guess cannot help this round.
Verification-adaptive weighting replaces the fixed schedule. Positions up to and including k get weight 1. After k, the original decay curve is shifted to start at k. EAGLE-3's base is 0.8^{k-1}; DFlash's is exp(-(k-1)/γ). The first rejection is the nearest fixable failure, so it keeps full weight. Because k varies, no sample-agnostic curve can mimic this.
The full objective sums the reweighted soft and hard cross-entropies plus the head loss with β=1.0. EAGLE-3 originally used soft labels, DFlash hard labels; VAT keeps both.
VAT is stacked on EAGLE-3 and DFlash with Qwen3-4B, Qwen3-8B, and LLaMA-3.1-8B, across eight math, code, and chat benchmarks, measured on A100 GPUs with Hugging Face Transformers. Training data pairs Perfectblend prompts with greedy target responses for 3 epochs; baselines and VAT share that corpus.
Temperature 0, eight-benchmark averages:
| Target | Method | Speedup | Acceptance length τ |
| Qwen3-4B | EAGLE-3 → +VAT | 4.07× → 4.39× | 6.28 → 6.78 |
| Qwen3-4B | DFlash → +VAT | 4.54× → 4.81× | 5.73 → 6.08 |
| Qwen3-8B | EAGLE-3 → +VAT | 4.04× → 4.24× | 6.12 → 6.47 |
| Qwen3-8B | DFlash → +VAT | 4.47× → 4.86× | 5.51 → 6.14 |
| LLaMA-3.1-8B | EAGLE-3 → +VAT | 4.17× → 4.33× | 6.08 → 6.23 |
| LLaMA-3.1-8B | DFlash → +VAT | 4.08× → 4.22× | 5.57 → 5.78 |
The largest lift is DFlash on Qwen3-8B: τ +11.4%, wall-clock speedup +8.7%. LLaMA moves less; EAGLE-3 τ only +2.5%. Temperature 1 keeps the same direction. One cell: DFlash on Qwen3-8B GSM8K goes from 6.21× / τ=7.45 to 7.03× / τ=8.71.
Ablating on DFlash + Qwen3-4B, the baseline is τ=5.73 / 4.54×. Head alone reaches 5.87, adaptive weights 5.91, soft+hard labels 5.82, all three 6.08 / 4.81×. The weighting gain comes from anchoring at k, not from the base shape: after adaptation, 0.8 decay and the DFlash exponential land at τ 6.09 vs 6.08.
During training the head pushes the mean first-rejection later, and post-rejection token matches stop drifting down. Optionally the head can early-exit at inference. Mean absolute error on k is 1.18 tokens for EAGLE-3 at threshold 0.5 and 1.76 for DFlash at 0.6. On DFlash code, speedup moves from 4.83× to 4.97× against an oracle 5.20×; τ dips slightly on false rejects.
This is a training-objective patch on methods people already run. No new draft block, no change to online verification. Teams on EAGLE-3 or block-diffusion drafters can keep the serving stack and still pick up acceptance length. Gains are larger on Qwen than on LLaMA, and only measured up to 8B.
PARD-2 and D-PACE, concurrent work, also retune positional weights, but they anchor on confidence proxies. VAT anchors on the observed first rejection, adds a head on cumulative acceptance, and covers both autoregressive and diffusion drafters.
Evaluation stops at 8B. On LLaMA-3.1-8B, τ moves only 2.5% to 3.8%; whether this holds at tens of billions is unknown. Training simulates greedy verification; eval at temperature 1 still improves, yet train-decode temperature still disagrees. Early exit with the head spends some τ, because false rejects clip prefixes that would have passed.
Weighting variants live in the appendix; the main text reports the winning slot. β is frozen at 1.0 with no sweep. The draft still needs 3 epochs on target responses, so a new target still means a new draft. VAT does not remove that cost.