BlockPilot: 实例自适应策略学习用于基于扩散的推测解码

摘要

推测解码通过使用轻量级草稿模型并行生成候选token，再由目标模型进行验证，从而实现无损加速。最近，基于扩散的推测解码进一步提升了并行性：通过块级扩散在每个前向过程中生成多个token，达到了当前最优性能。然而，现有方法采用固定的推理块大小，并假设所有输入具有相同的全局最优解码策略。本文证明这一假设并非最优，因为最佳块大小因样本而异，且对推测解码性能具有关键影响。此外，这些数值呈现明显的局部结构，集中在训练块大小附近，从而将问题简化为低维且结构化的决策空间。基于这些发现，我们提出BlockPilot——一种样本自适应策略，能够从预填充阶段的表示中预测最优块大小。具体而言，我们将块大小选择建模为一个轻量级策略学习问题，并设计了一种基于预填充阶段表示进行最优块大小预测的实例自适应决策机制。该预测仅在预填充后执行一次，便于无缝集成。大量实验表明，我们的方法具有即插即用特性，引入极低开销，并持续提升效率：在温度T=1条件下，Qwen3-4B模型的接受长度达到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.