arXiv:2606.31315cs.CL2026-06被引 1

根据输入自适应调整生成块大小,显著提升扩散模型推测解码效率

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

论文配图:BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding
图 1 · 摘自论文原文
  • 基于预填充表示动态预测最优块大小,实现样本自适应决策
  • 在Qwen3-4B上达5.92的接受长度和4.20倍加速,仅需一次预测
  • 无需修改模型结构,可直接插入现有推理流程

推测解码通过轻量级草稿模型并行生成候选词元,并由目标模型验证,实现无损加速。近期基于扩散的推测解码进一步提升并行性,通过块级扩散每前向传播生成多个词元,达到当前最佳(SOTA)性能。然而,现有方法采用固定推理块大小,且假设所有输入的最优解码策略一致,这一假设并不理想——最优块大小随样本变化显著,且具有明显局部结构,集中于训练块大小附近,因此可视为低维、有结构的决策空间。基于此,我们提出BlockPilot,一种从预填充表示中预测最优块大小的样本自适应策略。具体地,将块大小选择建模为轻量级策略学习问题,提出基于预填充阶段表示的实例自适应决策机制,预测仅在预填充后执行一次,支持无缝集成。大量实验表明,该方法可即插即用、开销极小,持续提升效率,在温度T=1下于Qwen3-4B上实现5.92的接受长度和4.20×速度提升。

原文摘要 · Abstract (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.20$\times$ speedup on Qwen3-4B under temperature $T=1$.

扩散模型推理加速自适应决策

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。