BlockPilot：基於擴散的推測解碼中的實例自適應策略學習

摘要

推测解码通过轻量级草稿模型并行生成候选token，再由目标模型验证，从而实现无损加速。近年来，基于扩散的推测解码进一步提升了并行性——通过块级扩散，每次前向传播可生成多个token，达到最先进性能。然而现有方法采用固定的推理块大小，并假设对所有输入应用统一的解码策略。本文表明该假设并非最优，因为最优块大小会因样本而异，且对推测解码性能至关重要。此外，这些最优块大小呈现出明显的局部结构，集中于训练块大小附近，从而将问题简化为低维且结构化的决策空间。基于这些发现，我们提出BlockPilot——一种样本自适应策略，能够从预填充表示中预测最优块大小。具体而言，我们将块大小选择问题建模为轻量级策略学习任务，并提出基于预填充阶段表示的实例自适应决策机制。预测仅在预填充完成后执行一次，因此可无缝集成。大量实验表明，我们的方法即插即用，引入极小的额外开销，并持续提升效率——在Qwen3-4B模型、温度T=1的条件下，达到5.92的接受长度与4.20倍的加速比。

English

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.20times speedup on Qwen3-4B under temperature T=1.