BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding
arXiv · HuggingFace · ▲74
摘要(原文)
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.
摘要(中译)
推测性解码通过使用轻量级草稿模型并行生成候选token,然后由目标模型进行验证,从而实现无损加速。最近,基于扩散的推测性解码通过在单次前向传播中生成多个token的块级扩散进一步提高了并行性,实现了最先进(SOTA)的性能。然而,现有方法采用固定的推理块大小,并假设所有输入都具有统一的最优解码策略。在本文中,我们表明这种假设是次优的,因为最优块大小在不同样本之间变化,并且在推测性解码性能中起着关键作用。此外,这些值表现出明显的局部结构,集中在训练块大小附近,这将问题简化为低维和结构化的决策空间。基于这些见解,我们提出了BlockPilot,这是一种样本自适应策略,可从预填充表示中预测最优块大小。具体来说,我们将块大小选择表述为一个轻量级的策略学习问题,并提出了一种实例自适应决策机制,该机制根据预填充阶段的表示来预测最优块大小。预测仅在预填充后执行一次,允许无缝集成。大量实验表明,我们的方法是即插即用的,引入的开销最小,并且始终提高效率,在Qwen3-4B上实现了接受长度为5.92和在温度T=1时的4.20倍加速。
背景剖析
背景剖析
1. 技术背景与需求
大型语言模型(LLMs)在推理任务中表现出色,但其逐token的自回归生成方式导致效率低下——每个token必须依赖前一个token的输出,使得长文本生成时延迟高、并行性差。推测解码(Speculative Decoding)通过轻量级“草稿模型”并行生成候选token,再用目标模型验证,实现了无损加速。近年来,基于扩散模型的推测解码(diffusion-based speculative decoding)进一步提升了并行性:扩散模型(dLLMs)通过块级扩散一次生成多个token,显著降低了延迟。然而,这类技术仍面临一个核心问题:如何高效平衡并行速度与生成准确性?
2. 先前方法的局限性
现有方法假设所有输入的最佳块大小(block size)相同,并直接沿用训练时的固定配置。但实际中,不同输入的语义约束、上下文确定性和预测难度差异很大:例如,结构化任务(如代码生成)可能需要更大的块以提升效率,而开放域生成(如创意写作)则需要更小的块以避免错误累积。固定块大小的策略忽略了这种输入特异性,导致性能次优。此外,尽管扩散模型的生成具有局部结构(最优块大小通常集中在训练配置附近),但先前方法未利用这一特性,导致优化效率低下。
3. 本文的解决方案
论文提出BlockPilot,一种样本自适应的块大小选择方法。其核心思路是:将块大小视为一个可学习的策略,而非静态参数。具体而言,在目标模型完成“预填充”(prefilling)阶段后,利用最终token的预测分布作为当前解码状态的表示(该分布反映了上下文不确定性和未来生成的稳定性),并训练一个轻量级预测器来推断最佳块大小。这一过程仅需一次计算,且不修改现有扩散推测解码框架,实现了即插即用。
4. 与前人工作的关键差异
与先前方法相比,BlockPilot的突破在于:
- 从工程参数到学习组件:首次将块大小选择视为一个策略学习问题,而非固定超参数。
- 利用局部结构:发现最优块大小的分布具有局部性,将问题转化为低维离散决策空间,简化了优化难度。
- 轻量级集成:预测器仅在预填充后运行一次,避免了在线搜索的高开销,同时保持与现有框架的兼容性。
通过这一设计,BlockPilot在Qwen3-4B模型上实现了4.20倍的速度提升,同时保持输出质量,证明了自适应块大小策略的有效性。

