Recovering Off-Policy Supervision for Speculative Decoding

2026-09-29

AI Summary

This study addresses the severe loss of supervision signals caused by off-policy tokens in block draft models for speculative decoding. To this end, we propose a rollout-based training framework that introduces complementary Anchor-Label Relabelling and In-Rollout Anchors mechanisms. By integrating distribution rescaling with in-context draft block embeddings, the approach reuses features at zero additional cost to recover complete supervision signals, enabling efficient fine-tuning without modifying the original corpus. Experimental results demonstrate that our method improves the greedy acceptance length by 36.5% over DFlash. Notably, it surpasses the strongest erasure baseline within a single training round and achieves performance comparable to target regeneration data after three rounds.

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.

Problem

• speculative decoding
• off-policy supervision
• block drafter
• supervision loss

Innovation

• Speculative Decoding
• Off-Policy Supervision
• Rollout-based Training
• Anchor-Label Relabelling
• In-Rollout Anchors