论文

BlockPilot: 面向 Diffusion-based Speculative Decoding 的实例自适应策略学习

BlockPilot: 面向 Diffusion-based Speculative Decoding 的实例自适应策略学习

投机解码 通过轻量级 草稿模型 并行生成候选 token,再由 目标模型 验证,实现无损加速。近期 基于扩散的投机解码 利用 块级扩散 在单次前向中生成多个 token,进一步提升了并行性,达到最先进性能。然而,现有方法采用固定的推理块大小,并假设所有输入适用统一的最优解码策略。 本文研究表明,这一假设是次优的:最优块大小 因样本而异,且在解码性能中起关键作用。此外,最优块大小呈现清晰的局部结构,集中在训练块大小附近,将问题降维为低维、结构化的决策空间。基于此,我们提出 BlockPilot,一种样本自适应策略,从 预填充表示 中预测最优块大小。具体地,我们将块大小选择建模为轻量级策略学习问题,并设计实例自适应决策机制,仅需在预填充后进行一次预测,即可无缝集成。 大量实验表明,该方法即插即用,引入极低开销,并持续提升效率:在 Qwen3-4B 模型、温度 T=1 下,实现了 接受长度 5.92 和 4.20 倍加速比。

论文精读

TL;DR BlockPilot 让扩散推测解码的块大小自适应每个样本,基于预填充特征低成本预测最优值,显著提升推理效率,在 Qwen3-4B 上达 5.92 接受长度与 4.20x 加速。

问题

问题背景

大语言模型(LLM)的自回归解码带来高延迟,推测解码 通过轻量 draft 模型生成候选 token,再由目标模型并行验证,实现无损加速。近期 扩散基推测解码 在单次前传中生成多个 token,达到当前最优(SOTA)性能,成为研究热点。

现有方法的局限

现有扩散基推测解码方法(如 Block Diffusion)采用固定推理块大小,假设所有输入样本共享统一的最优解码策略。然而,最优块大小高度依赖样本难度:简单样本可使用更大块获得更高并行,困难样本则需要更小块以保证接受率。固定策略导致:

  • 对简单样本,块大小偏小,未能充分利用并行潜力;
  • 对困难样本,块大小偏大,造成大量 token 被拒绝,引入额外计算浪费。 这种一刀切的方式限制了推测解码的效率上限,成为实际加速比提升的瓶颈。

技术挑战与业界关注

核心挑战 在于如何以极低开销为每个推理实例动态选择最优块大小。原因包括:

  1. 决策空间复杂:块大小与接受长度、速度提升之间呈非单调关系,难以通过简单启发式确定。
  2. 实时性要求:推理时决策必须快速,不能抵消推测解码的加速收益;任何预测模型必须轻量且与现有流水线无缝集成。
  3. 分布漂移:训练时的块大小偏好可能偏离实际应用分布,需要策略具备泛化性。 业界对更低延迟、更高吞吐的 LLM 部署需求强烈,尤其在对话、代码生成等交互场景,每一毫秒的加速都意味着成本降低和用户体验提升,因此动态自适应策略成为重要研究方向。

行业类比

自适应批处理 类似——根据请求长度和负载动态调整批次大小以平衡吞吐与延迟,而非使用全局固定批次;块大小选择同样需要实例感知的轻量决策来解锁额外加速空间。

核心洞察

  • 扩散推测解码的最优块大小并非均匀固定,而是随输入样本变化,且集中分布在训练块大小附近的局部区间内。这一发现将块大小选择从连续的全局搜索简化为低维分类问题,显著降低决策复杂度,而现有方法通常假设统一块大小或采用简单启发式规则,无法捕捉实例间的差异。
  • 基于预填充表征的一次性块大小预测仅引入极轻量开销,且天然具备即插即用特性。与需要每步动态调整或依赖额外打分模型的方案相比,BlockPilot 将决策完全限定在预填充阶段,不干扰自回归解码过程,工程集成成本极低,同时保持高加速比和可接受长度。

方法

方法概述

BlockPilot 的核心是一种实例自适应块大小选择策略,用于扩散推测解码。其工作流程遵循“输入 → 关键模块 → 输出”的管线。

输入与预处理:给定一个提示,首先通过目标模型进行预填充(prefilling),得到最后一层隐藏状态。该表示被简化为沿序列长度维度的平均池化向量,作为预测器的输入特征。

关键模块:轻量级块大小预测器:预测器是一个小型分类器(如两层 MLP),它将块大小选择形式化为 K 类分类问题。基于三个关键发现设计:

  • 实例变异性:不同样本的最优块大小差异显著,固定大小并非全局最优。
  • 局部区间性质:最优值密集集中在训练块大小附近的一个窄区间内(例如训练块大小 ±1 或 ±2)。
  • 分类形式化:因此问题退化为在候选区间 $[B_{\text{train}} - r, B_{\text{train}} + r]$ 内选出最佳类别。

训练数据构建:使用目标模型离线评估每个训练样本在不同块大小下的 接受长度(acceptance length),将获得最大接受长度的块大小作为该样本的伪标签(ground truth),限制在局部候选区间内。预测器以预填充表示作为输入,使用交叉熵损失进行监督学习,训练成本极低。

推理过程:预填充后,预测器仅运行一次,输出预测的块大小 $\hat{B}$。后续所有解码步均采用此块大小进行扩散块生成与验证。由于预测器计算量可忽略,整体开销极小,且能无缝集成到现有扩散推测解码框架中。

与同类方法的差异:不同于现有扩散推测解码方法(如 DSI、DBD)采用固定全局块大小的做法,BlockPilot 首次利用预填充阶段的信息来动态调整块大小,将选择建模为局部分类问题,在几乎不增加延迟的前提下显著提升了平均接受长度和加速比。

实验

实验设计

  • 模型与基准:在 Qwen3-4B 上评估扩散推测解码效率,对比固定块大小与 BlockPilot 自适应策略。
  • 指标:采用 接受长度 (acceptance length)实际加速比 (wall-clock speedup),重点关注温度 T=1 的设置。
  • 预测器:输入为预填充阶段 (prefilling) 的表示,仅需一次额外前向,几乎不增加推理延迟。

关键发现

  • 最优块大小呈现实例依赖性局部集中性:不同样本的最佳块大小差异明显,但都聚集在训练块大小附近的窄区间内,这使块大小选择可简化为局部区间内的分类问题
  • BlockPilot 通过从预填充表示中学习一个轻量级策略,以极低开销实现了 5.92 的接受长度4.20 倍加速,验证了自适应块大小比统一固定块大小能更充分地挖掘每个样本的并行潜力。

与基线对比

  • 固定块大小对所有样本一刀切,无法兼顾简单样本(并行度不足)与困难样本(多余计算浪费)。BlockPilot 在保持无损加速的前提下,仅增加一次预填充后的预测,即换取了显著的效率增益,体现了“即插即用”的优势。

行业影响

落地场景

BlockPilot 提供的实例自适应块大小选择 机制可直接嵌入所有采用扩散推测解码 (diffusion-based speculative decoding) 的大语言模型服务。典型场景包括:

  • 实时对话系统:智能客服、虚拟助手需要低延迟生成,自适应块大小可动态平衡并行度与验证开销。
  • 批量内容生成:电商商品描述、社交平台推荐文案、代码补全等大吞吐场景,通过提升单次前向可接受的 tokens 数 (acceptance length) 增加整体吞吐。
  • 低延迟边缘部署:在资源受限设备上,轻量级的块大小预测器几乎无额外开销,使推测解码对延迟更敏感的应用更可行。

商业价值

  • 降低推理成本:实验表明在 Qwen3-4B 上可实现 4.20 倍加速 且保持无损输出。这意味着相同流量下所需 GPU 资源可减少约 75%,直接降低云服务或私有化部署的硬件与能耗成本。
  • 提升用户体验:平均接受长度达到 5.92 tokens,比固定块大小策略更高效地占用目标模型计算,缩短首 token 延迟和生成总耗时。
  • 增强竞争优势:对于 LLM API 或模型服务提供商,更低的每请求成本可转化为更灵活的定价策略,同时保持或提升服务质量。

与现有产品/工作流的接口

BlockPilot 是一个即插即用模块,集成路径如下:

  1. 推理引擎适配:在现有 vLLMTensorRT-LLM 等框架的推测解码流程中,于预填充 (prefill) 阶段后插入块大小预测器 (一个轻量 MLP),仅需前向一次得到 block_size
  2. 无侵入式调度:预测器仅依赖预填充后的隐状态 (prefilling representation),不修改草稿模型或验证模型结构,可直接复用现有 draft-verify 流水线。
  3. 训练与部署:预测器的监督数据可通过离线最优搜索构建,训练收敛快;推理时预测器计算量可忽略。

具体落地 Use Case

  • 电商平台智能客服:在促销高峰期,顾客咨询量激增,使用自适应推测解码可将单个 A100 节点支持的并发会话数提升 3–4 倍,同时保证回答质量无损,避免因延迟过高导致客户流失。
  • 代码助手服务:面向开发者的实时代码补全工具,上下文变化频繁、生成长度不一,固定块大小可能导致部分请求等待时间过长。BlockPilot 根据输入上下文自动选择最佳块大小,使短补全更快、长补全更稳,整体响应延迟分布更集中,改善开发者体验。

局限

  • **预测器依赖监督数据且与特定扩散模型绑定**。BlockPilot 需要针对每个扩散推测解码模型收集最优块大小的标签数据,并训练轻量级分类器;当扩散模型或训练块大小更新时,预测器需重新训练,增加了工程部署成本。论文未充分讨论跨模型迁移的可行性,实际应用中若频繁更换底座模型,可能削弱其即插即用的优势。
  • **对低温度或确定性解码的有效性未充分验证**。实验集中在温度 T=1 的设置,此时扩散推测解码的随机性较大,最优块大小波动明显,BlockPilot 的增益更显著;但在低温度(如 T=0.1)或贪心解码下,最优块大小的分布可能更集中,预测器的收益可能降低。论文仅在部分实验中对比了不同温度,缺乏系统分析。
  • **方法限于扩散推测解码框架,通用性有限**。BlockPilot 依赖扩散模型生成的块级候选序列,无法直接应用于传统的自回归 draft 模型或基于原始 transformer 的推测解码。若未来出现其他高效的并行解码范式,该方法难以迁移。此外,扩散推测解码本身仍在早期阶段,硬件适配和推理延时开销较大,限制了整体实用性。
论文Hao Zhang2026-06-30原文

相关内容