ADAFLASH: ADAPTIVE SPECULATIVE DECODING VIA ON-POLICY DISTILLED DIFFUSION DRAFTERS

Large language model (LLM) inference is often bottlenecked by memory bandwidth. Speculative decoding addresses this by using a lightweight draft model to generate sequences that a larger target model verifies in parallel. While diffusion-based drafters offer superior parallel generation capabilities, they suffer from high variance in draft quality across different domains and token positions. This paper investigates how to mitigate this variance to improve the speedup and efficiency of speculative decoding.

The authors identify that the bidirectional attention mechanism in diffusion drafters, while powerful, introduces significant domain-level and token-level variance. To address this, they propose the AdaFlash framework, which consists of two primary components:

1. On-Policy Distillation (OPD): An online training algorithm that uses reverse-KL divergence and entry-wise gradient clipping to adapt the drafter to the target model's distribution in real-time. This reduces domain-level performance fluctuations.
2. Adaptive Length Head: A lightweight module that predicts the expected acceptance rate of a draft sequence and dynamically adjusts the candidate length. This prevents the target model from wasting computation on low-quality draft tokens.

The authors also implement an asynchronous training-inference pipeline and an adaptive request scheduler to ensure these updates do not block the main inference process, allowing for hot-reloading of model weights during deployment.

Standard speculative decoding often relies on static draft models that struggle when the input distribution shifts (e.g., moving from coding tasks to creative writing). By enabling the drafter to evolve continuously during deployment and optimizing the verification budget per request, AdaFlash maintains high acceleration rates even in high-concurrency scenarios where previous methods often see performance degradation. This makes speculative decoding more robust and efficient for production-scale LLM serving.