Speculative decoding improved by recovering lost supervision during training
Recovering Off-Policy Supervision for Speculative Decoding
Computation and LanguageComputer Vision and Pattern RecognitionMachine Learning
Summary
When training language models to predict text in blocks, a single incorrect guess can cause the training to ignore many later parts, wasting useful information. The authors propose two new techniques—Anchor-Label Relabelling and In-Rollout Anchors—that recover this lost training information without changing the original text. These methods let models learn better from fixed datasets and improve performance without needing extra computation. The authors show their approach increases how much of the predicted text gets accepted compared to previous techniques.
What this means in practice
- •For language model trainers: Train models on fixed datasets more efficiently by recovering rejected training signals from initial model outputs.
- •For vision language system builders: Improve speculative decoding quality in systems combining vision and text without regenerating large training corpora.
Authors
Jungseob Lee, Chanjun Park, Sugyeong Eo, Hyeonseok Moon
Abstract
Block drafters for speculative decoding are commonly trained on corpora written by external models, where a single off-policy token invalidates supervision for all subsequent slots in a block. Existing approaches discard these divergent slots, resulting in severe supervision loss. To resolve this problem while preserving the training corpus, we propose a rollout-based training framework that recovers full supervision through two complementary components. The first component, Anchor-Label Relabelling (ALR), replaces corpus labels with distributions from greedy target rollouts, restoring valid supervision across all predicted slots. The second component, In-Rollout Anchors (IRA), places draft blocks directly inside these rollouts to expose the drafter to target-generated context, reusing precomputed rollout features at no additional target cost. Across fixed vision-language and text corpora, our framework increases greedy accepted length by up to 36.5% over DFlash and consistently outperforms erasing baselines. Notably, a single epoch of our method surpasses the best erase schedules. After three epochs, it matches the acceptance length of training on target-regenerated responses. These results show that our framework provides an effective and compute-efficient approach for training speculative drafters on fixed corpora without modifying the original text. Code is available at https://github.com/js-lee-AI/ALR-IRA.