BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding

2026-06-30

AI Summary

This work addresses a key limitation in existing diffusion-based speculative decoding methods, which employ a fixed block size and overlook the substantial variation in optimal block sizes across different input instances, thereby constraining inference efficiency. To overcome this, the authors propose an instance-adaptive block size selection mechanism that predicts the optimal block size from the input representation after the prefill phase. This prediction is framed as a lightweight policy learning task, enabling a “predict once, apply throughout” integration with minimal overhead. The study further reveals—for the first time—that optimal block sizes exhibit significant inter-sample variation and possess local structural patterns, which are leveraged to construct a low-dimensional, structured decision space. Evaluated on Qwen3-4B, the method achieves a 4.20× speedup and an average acceptance length of 5.92, substantially enhancing inference efficiency with negligible computational cost.

Abstract

Speculative decoding accelerates inference by using a lightweight draft model to generate candidate tokens in parallel, and are then verified by the target model, enabling lossless acceleration. Recently, diffusion-based speculative decoding further improves parallelism by generating multiple tokens per forward pass via block-level diffusion, achieving state-of-the-art (SOTA) performance. However, existing methods adopt a fixed inference block size and assume a uniform optimal decoding strategy across all inputs. In this paper, we show that this assumption is suboptimal, as the optimal block size varies across samples and plays a critical role in speculative decoding performance. Moreover, these values exhibit a clear local structure, concentrating around the training block size, which reduces the problem to a low-dimensional and structured decision space. Based on these insights, we propose BlockPilot, a sample-adaptive policy that predicts the optimal block size from the prefilling representation. Specifically, we formulate block size selection as a lightweight policy learning problem and propose an instance-adaptive decision mechanism that predicts the optimal block size based on the representation of the prefilling stage. The prediction is performed only once after prefilling, allowing for seamless integration. Extensive experiments demonstrate that our method is plug-and-play, introduces minimal overhead, and consistently improves efficiency, achieving an acceptance length of 5.92 and a 4.20$\times$ speedup on Qwen3-4B under temperature $T=1$.

Problem

• speculative decoding
• diffusion-based generation
• block size adaptation
• instance-adaptive policy
• inference acceleration

Innovation

• speculative decoding
• diffusion-based generation
• adaptive block size
• policy learning
• instance-adaptive inference