AdaFlash: Adaptive Speculative Decoding via On-Policy Distilled Diffusion Drafters

2026-07-21

AI Summary

This work addresses the high token- and domain-level variance, unstable acceptance rates, and substantial verification overhead in speculative decoding with diffusion draft models, which stem from their bidirectional attention mechanism. To mitigate these issues, the authors propose AdaFlash, a novel framework that introduces, for the first time, an online policy distillation algorithm based on reverse KL divergence tailored for diffusion draft models, along with an adaptive length prediction head that dynamically optimizes draft sequence lengths. This approach effectively alleviates instability caused by global dependencies while preserving generation quality and significantly enhancing inference efficiency. Experimental results demonstrate that AdaFlash achieves up to a 66% higher throughput compared to the current state-of-the-art under high-concurrency settings, consistently delivering superior acceleration ratios.

Abstract

Speculative decoding, in which a lightweight draft model first generates a draft sequence that is then verified in parallel by the target model, has become a prevalent paradigm for accelerating large language model inference. Recent work such as DFlash further boosts drafting efficiency by leveraging diffusion drafters, whose parallel denoising mechanism enables draft generation in a single forward pass. In this work, we uncover a central pitfall of diffusion drafters: bidirectional attention is a double-edged sword. On one hand, it endows the model with parallel generation and global contextual modeling capabilities; on the other hand, this inherent global dependency introduces high variance at both the domain-level and the token-level: acceptance rates fluctuate substantially across different domains, and draft token quality also varies heterogeneously at different token positions. To tackle this issue, we propose AdaFlash framework, comprising two components: (i) an on-policy distillation (OPD) algorithm with reverse-KL divergence tailored for diffusion drafters, bringing stable convergence and effectively reducing domain-level variance; and (ii) an adaptive length head that dynamically adjusts the candidate sequence length on the fly, substantially lowering the verification cost of the target model and handling token-level variance. Experiments demonstrate that AdaFlash consistently improves speedup rate during deployment, with especially significant gains in high-concurrency scenarios, achieving up to approximately 66% higher throughput than previous state-of-the-art methods.

Problem

• speculative decoding
• diffusion drafters
• bidirectional attention
• domain-level variance
• token-level variance

Innovation

• speculative decoding
• diffusion drafters
• on-policy distillation
• adaptive length prediction
• large language model acceleration