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

背景

随着大语言模型（LLM）规模的不断扩大，推理速度成为制约其实际应用的关键瓶颈。推测解码（Speculative Decoding）作为一种无损加速技术，通过让轻量级的草稿模型（Draft Model）并行生成候选Token，再由目标模型（Target Model）进行一次性验证，从而在不损失生成质量的前提下显著提升推理效率。

近年来，基于扩散模型的推测解码（Diffusion-based Speculative Decoding）进一步突破了传统自回归草稿模型的并行限制。它通过块级扩散（Block-level Diffusion）机制，在每次前向传播中同时生成多个Token，实现了当前最先进的（SOTA）加速性能。然而，现有方法普遍采用固定的推理块大小（Block Size），并隐含了一个关键假设：即对于所有输入样本，最优的解码策略是一致的。

核心内容

本文指出，上述“固定块大小、统一最优策略”的假设是次优的。作者通过研究发现，最优的块大小实际上随输入样本的变化而变化，且这一参数在推测解码性能中扮演着至关重要的角色。进一步观察发现，最优块大小呈现出清晰的局部结构（Local Structure），即它们往往集中在训练时使用的块大小周围。这一发现将原本复杂的决策问题简化为一个低维且结构化的决策空间。

基于上述洞察，本文提出了 BlockPilot，一种样本自适应（Instance-Adaptive）的策略学习方法。BlockPilot 的核心思想是从预填充（Prefilling）阶段的表征中预测出当前样本的最优块大小。具体而言，作者将块大小的选择建模为一个轻量级的策略学习问题，并提出了一种实例自适应的决策机制：该机制基于 Prefilling 阶段的表征来预测最优块大小。

在实现上，最优块大小的预测仅在 Prefilling 阶段完成后执行一次，这使得 BlockPilot 能够无缝集成到现有的推理流程中，无需改变整体架构。实验结果表明，BlockPilot 具备即插即用（Plug-and-Play）的特性，引入的系统开销极小，且能持续提升推理效率。在 Qwen3-4B 模型上，当温度参数 $T=1$ 时，该方法实现了 5.92 的接受长度（Acceptance Length）和高达 4.20 倍的加速比。

关键要点

• 打破固定块大小假设：现有基于扩散的推测解码方法假设所有输入共享同一最优块大小，而本文证明最优块大小具有样本依赖性，固定策略会限制性能上限。
• 局部结构降维：最优块大小并非随机分布，而是围绕训练块大小呈现局部集中特性，这为低维策略学习提供了理论基础。
• Prefilling 表征驱动：BlockPilot 利用 Prefilling 阶段的表征作为决策依据，在推理初期即完成对最优块大小的预测。
• 单次预测与无缝集成：策略预测仅在 Prefilling 后执行一次，不增加迭代开销，实现了真正的即插即用。
• 显著的加速效果：在 Qwen3-4B 上实现了 4.20