跳转至

Test-Time Speculation

会议: NeurIPS 2026
arXiv: 2605.09329
OpenReview: https://openreview.net/forum?id=elwxI8Nne1
代码: https://github.com/z-lab/dflash
领域: 其他
关键词: 投机采样、测试时自适应、长文本生成、在线蒸馏、接受长度

一句话总结

针对现有投机解码草稿模型在长文本生成中接受长度严重衰减至接近 1 的分布偏移问题，TTS 利用目标模型验证阶段免费产生的完整输出分布作为在线蒸馏信号，在测试时实时调整草稿模型，将平均接受长度提升 41%。

研究背景与动机

投机解码（Speculative Decoding）通过轻量草稿模型快速推测序列、重量级目标模型单次前向并行验证的范式，打破了自回归大语言模型每生成一个 token 都必须执行完整前向传播的延迟瓶颈。在标准评估中，这种推测-验证回环的加速比完全由每轮被目标模型接受的连续 token 数，即接受长度（Acceptance Length, AL）决定。若草稿模型预测分布与目标模型高度一致，单轮可接受多个 token，从而显著降低目标模型的调用频次；反之，若接受长度趋近于 1，投机解码将退化为带额外开销的自回归生成，加速收益彻底归零。

然而，现有最先进的投机解码系统（如 DFlash、EAGLE-3、PARD 等）在长文本生成任务（如复杂代码合成、长链条数学推理与 Agent 工作流）中遭遇了普遍而严重的性能崩塌。实验表明，在 LiveCodeBench 与 AIME 等长推理任务上，即便使用强大的基座模型，草稿模型的接受长度也会随着生成步数增加急剧下跌。例如在 Qwen3-8B 搭配 EAGLE-3 时，当生成长度跨越 2 万 token，平均接受长度甚至退化至 1.1 左右。这种退化的根源在于离线训练与长程推理之间的严重分布偏移：主流草稿模型通常在短序列（≤2K tokens）的有监督数据集上离线蒸馏完成，而在面对数万 token 的真实长程上下文时，草稿模型的隐状态表征与位置外推迅速偏离目标模型的生成轨迹，导致其 Top-1 交叉熵误差单调暴增，最终使得推测候选大面积被拒。

要解决长上下文投机失效的困境，必须在推理阶段动态修正草稿模型向目标模型分布的漂移。作者敏锐地指出，投机解码的验证环节天然具备了在线蒸馏所需的一切要素：目标模型在单次前向中对草稿画布上的所有 token 进行并行验证时，已经计算并暴露了每一步的完整条件似然分布，这构成了现成且无额外前向开销的监督信号。核心 idea：将目标模型与草稿模型视为教师-学生对，直接利用验证步骤免费输出的逐位置目标分布作为蒸馏样本，在测试时交替执行单步在线梯度更新与带正则的分布修正，让草稿模型动态适配当前生成的长程轨迹。

方法详解

整体框架

Test-Time Speculation（TTS）不改变投机解码的无损验证数学保证，而是将在线蒸馏与常规的推测-验证循环紧密交织。在传统投机解码中，草稿模型权重在部署后完全冻结，而 TTS 将其维护为一个随生成步数动态演化的模型状态 \(q_t\)。在每个生成轮次中，草稿模型根据当前历史上下文生成一段长度为 \(C\) 的候选画布；目标模型对画布执行单次前向评估，计算出所有候选位置的真实条件概率，并通过拒绝采样决定前 \(\tau\) 个被接受的 token。在轮次结算后，TTS 提取目标模型刚计算出的输出分布作为软标签，构造带有锚定正则项的自适应蒸馏损失，对草稿模型执行单步梯度反向传播，使其在新一轮推测前向目标分布完成一次微调校准。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["当前已验证序列 x1:t"] --> B["草稿推测<br/>qt 生成 C 个候选 token"]
    B --> C["目标验证<br/>教师模型单次前向输出完整分布"]
    C --> D["拒绝采样裁决<br/>确定接受长度 τ 并推进上下文"]
    C --> E["在线蒸馏与锚定正则<br/>利用目标分布计算加权 KL 散度"]
    D -.->|触发更新条件| E
    E --> F["单步在线梯度更新<br/>qt+τ ← qt - η∇Lt(q)"]
    F -->|更新草稿参数| B

在系统执行层面上，为了防止反向传播与优化器更新阻塞关键推理路径，TTS 还引入了跨轮次的更新步长调度与异步 CUDA 流流水线优化，使梯度计算与后续前向推测重叠执行，从而在维持高接受长度的同时实现实际吞吐吞吐量的最大化。

关键设计

1. 验证信号重用的在线蒸馏：以零额外前向代价获取在轨监督

传统在线模型自适应通常需要维护额外的缓存队列或离线调用教师模型，在推理关键路径上造成高昂的算力开销。TTS 的关键洞察在于，投机解码在对画布 \(x_{\mathcal{C}} = (x_{t+1}, \ldots, x_{t+C})\) 进行验证时，目标模型 \(p\) 必须对整段拼接序列执行前向传播，进而产出了每个位置的序列条件概率分布 \(p(X_{t+k} \mid x_{1:t+k-1})\)。TTS 直接将目标模型视为教师网络，将草稿模型 \(q\) 视为学生网络，将本轮验证过的画布直接作为在轨（on-policy）蒸馏样本，从而完全免去了为学生自适应重新跑教师前向计算的巨大浪费。

2. 位置敏感加权与历史锚定正则化：平衡前缀敏感度与更新稳定性

在推测解码中，画布靠前位置的预测准确率具有级联决定性：一旦第 1 个 token 被拒，整段后续画布即便预测正确也会被强制丢弃。因此，TTS 在蒸馏目标中引入了单调递减的位置权重 \(w_k\)，优先保障序列头部的校准精度。与此同时，为了防止草稿模型在单一轮次的局部样本上发生灾难性参数漂移，TTS 在损失函数中加入了针对前一状态 \(q_t\) 的 KL 散度正则化约束。综合损失函数表示为：

\[L_t(q) = \sum_{k=1}^C w_k \mathrm{KL}\bigl(p(X_{t+k} \mid x_{1:t+k-1}) \,\|\, q(X_{t+k} \mid x_{1:t+k-1})\bigr) + \lambda \sum_{k=1}^C w_k \mathrm{KL}\bigl(q_t(X_{t+k} \mid x_{1:t+k-1}) \,\|\, q(X_{t+k} \mid x_{1:t+k-1})\bigr)\]

该设计确保了草稿模型在快速吸收当前前缀特征的同时，参数始终稳定在合理流形内，不至于因长文本局部的罕见分布而彻底破坏先验知识。

3. 跨轮次步长调度与异步流流水线：消除更新对主干推理延迟的阻塞

尽管在线更新能够拉高接受长度，但如果每一轮（Stride \(\mathcal{S}=1\)）都在主推理流上同步执行反向传播与 Adam 优化器更新，引入的计算延迟将抵消 token 接受数增加带来的时间收益。为此，TTS 提出了更新步长调度机制，仅在每隔 \(\mathcal{S}\) 轮时抓取最新一轮的监督信号执行更新，不作历史累加。更重要的是，TTS 将反向传播与权重更新调度到一个独立的伴随 CUDA 流（Side Stream）中，与主推理流上后续 \(\mathcal{S}-1\) 轮的常规推测-验证并发运行，并在跨轮边界设置 CPU 级同步屏障。这种异步流水线消除了大部分反向传播对关键路径的阻塞，使得吞吐量达到全局最优。

损失函数 / 训练策略

在推理自适应过程中，草稿模型初始权重设为离线预训练权重 \(q_0 = q\)。每隔 \(\mathcal{S}\) 轮更新时，固定仅执行一次单步梯度更新：

\[q_{t+\tau} \leftarrow q_t - \eta \nabla_q L_t(q)\big|_{q=q_t}\]

针对超参数，论文严格限制单轮更新步数 \(\mathcal{K}=1\)。实验发现，对同一轮样本执行多次梯度迭代（\(\mathcal{K}>1\)）会导致草稿模型迅速对当前局部 token 过拟合，反而显著降低后续轮次的泛化接受率。请求结束之后，草稿模型参数将重置回初始状态，避免不同请求之间的跨域污染与无边界漂移。

实验关键数据

主实验

评估覆盖了 AIME 2024、AIME 2025、GPQA Diamond、LiveCodeBench、MATH-500、OlympiadBench-Math、OlympiadBench-Physics 与 TheoremQA 共 8 大长文本推理与代码基准，测试模型包含 Qwen3-4B、Qwen3-8B、Qwen3.5-122B-A10B（混合 MoE）以及 Llama3.1-8B，基线投机器采用 DFlash（块扩散架构）与 EAGLE-3（自回归架构）。下表为 DFlash 架构在各尺寸基座模型上的平均接受长度对比：

数据集 Qwen3-4B (DFlash) Qwen3-4B (TTS) Qwen3-4B 相对提升 Qwen3-8B (DFlash) Qwen3-8B (TTS) Qwen3-8B 相对提升 Qwen3.5-122B (DFlash) Qwen3.5-122B (TTS) Qwen3.5-122B 相对提升
AIME 2024 3.3 4.6 +38.0% 3.3 4.6 +38.5% 2.9 4.8 +62.2%
AIME 2025 3.2 4.7 +47.0% 3.1 4.2 +33.1% 2.9 4.9 +71.8%
GPQA Diamond 3.6 5.2 +47.2% 3.2 4.7 +47.3% 2.6 3.5 +33.9%
LiveCodeBench 3.0 4.7 +55.2% 2.9 4.5 +55.2% 2.4 4.0 +63.5%
MATH-500 3.7 4.3 +13.8% 3.8 4.3 +13.3% 3.4 4.5 +32.2%
OlympiadBench-Math 3.0 4.4 +48.9% 2.7 4.1 +50.3% 3.4 4.7 +40.2%
OlympiadBench-Physics 3.1 4.5 +48.1% 2.8 4.3 +53.1% 3.0 3.9 +28.4%
TheoremQA 3.6 4.4 +23.4% 3.5 4.3 +22.8% 3.3 4.2 +25.1%

在 EAGLE-3 投机器下，TTS 在 Llama3.1-8B 上展现出更为显著的恢复能力：基线 EAGLE-3 因训练长度（2K）严重落后于目标模型上下文，接受长度普遍滑落至 1.2~1.4，而 TTS 将其拉升至 1.9~2.1，平均相对提升高达 61.0%（OlympiadBench-Math 上提升 66.8%）。

消融实验

更新步长 \(\mathcal{S}\) 直接主导了接受长度（算法上限）与实际端到端吞吐率（系统吞吐）之间的权衡。在 Qwen3.5-122B 基座模型配合 DFlash 投机器上的消融实验显示：

配置 (更新步长 \(\mathcal{S}\)) 平均接受长度 (AL) 相对基线 AL 增益 端到端吞吐量 (Tokens/s) 相对基线加速比 说明
DFlash Baseline (无更新) 3.09 1.00× 27.88 1.00× 冻结草稿模型，无反向开销
TTS (\(\mathcal{S}=20\)) 3.58 1.16× 38.68 1.39× 稀疏更新，更新开销极小
TTS (\(\mathcal{S}=10\)) 3.73 1.21× 40.59 1.46× 适中更新频率
TTS (\(\mathcal{S}=5\)) 4.01 1.30× 42.84 1.54× 最佳吞吐平衡点（最高达 1.71×）
TTS (\(\mathcal{S}=2\)) 4.36 1.41× 40.49 1.45× 接受长度显著提升，计算稍有竞争
TTS (\(\mathcal{S}=1\)) 4.43 1.43× 32.14 1.15× 接受长度最高，但反向传播拖累吞吐

关键发现

  • 长文本增益复利放大：TTS 的改善幅度随生成序列长度增加而急剧扩大。在 Qwen3-8B 跑 AIME 2024 时，前 0~10K token 范围内 TTS 相比 DFlash 仅提升 15%，而在 20K~30K token 极长区间内，接受长度增益扩大至 183%；在 LiveCodeBench 上，后程增益更暴增至 245%。这印证了持续在线蒸馏对长程分布漂移的累积抵抗效果。
  • 单步更新优于多次迭代：单轮更新步数 \(\mathcal{K}=1\) 效果最好。当强制在单轮执行 \(\mathcal{K}=2,3,4\) 步迭代时，AIME 2024 上的接受长度由 4.6 持续下跌至 4.1。这是因为草稿模型容量较小，过度拟合当前画布这一单一局部样本破坏了其在下一未知步的泛化能力。
  • 短上下文草稿模型的救赎：对于 EAGLE-3 + Llama3.1-8B 这种草稿模型训练窗口仅有 2K、而目标模型支持 128K 的极端失配场景，TTS 取得了高达 61% 的平均接受长度飞跃，大幅弥补了离线预训练数据长度受限带来的外推缺陷。

亮点与洞察

  • 在轨监督的“免费午餐”挖掘：巧妙地将投机解码固有的验证步骤重释为在线蒸馏的数据生成器。目标模型为了拒绝采样必须计算全分布，TTS 直接将这一中间计算废料变废为宝，实现了无需额外前向传播的测试时自适应。
  • 动态闭环抵抗长文本熵减与漂移：论文深入揭示了长文本生成中“目标模型熵降低（更确定）、而草稿模型交叉熵升高（更发散）”的悖论现象，通过测试时蒸馏构建了动态误差纠偏反馈环，突破了静态推测网络在长序列推理下的效能衰减瓶颈。
  • 理论收益与工程实效的优雅妥协：没有盲目追求理论最优的每轮更新（\(\mathcal{S}=1\)），而是系统性分析了反向传播在硬件执行层上的延迟损耗，设计了基于步长跳跃与异步伴随 CUDA 流的工程架构，真正使理论上的接受长度跃升转化为实际的墙上时钟加速。

局限与展望

  • 显存与硬件调度开销：在线反向传播需要临时保留草稿模型前向计算的激活张量，并为优化器分配动量缓存，这使得推理阶段的显存开销有所上升；此外伴随 CUDA 流在同 GPU 上与主推理流存在硬件计算单元竞争，在小型基座模型上加速比受到一定压制。
  • 固定步长调度的次优性：当前系统的更新步长 \(\mathcal{S}\) 是静态预设的超参数，无法根据当前上下文动态辨识是否发生分布漂移；未来可设计基于即时接受率跌幅的自适应触发更新机制（Adaptive Stride）。
  • 优化器算子未深度融合：目前草稿模型的在线微调依然依赖标准 PyTorch 算子链条，未与 SGLang 等现代推理系统深度融合为单内核算子，限制了反向更新的极限执行速度。

相关工作与启发

  • vs OSD (Online Speculative Decoding)：OSD 依赖跨用户、跨请求的全局样本缓冲池来周期性微调多个专用草稿模型，要求请求流具有明确的业务领域隔离性；而 TTS 专注于单次推理长响应内部的分布纠偏，在单个请求内按轮次自适应，对跨领域和任意请求流天然免疫。
  • vs ATLAS：ATLAS 通过在生产流量中对轻量草稿器进行跨请求更新并配合门控路由，但面对单请求万级别 token 的长链条推理时，跨请求更新无法解决当次请求中程接受率崩塌的问题；TTS 直接在请求生命周期内完成闭环更新，且无缝适配开源主流的 EAGLE-3 与 DFlash 架构。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将投机解码验证过程抽象为自监督在线蒸馏闭环，思路极其优美深刻。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 8 大主流长程评测基准、4 款参数规模横跨 4B 至 122B 的基座模型以及两种核心投机架构，消融与系统吞吐分析非常扎实。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机阐述环环相扣，从实证崩塌现象到分布偏移剖析，再到算法与工程解决，逻辑极具说服力。
  • 价值: ⭐⭐⭐⭐⭐ 直击当今长思维链推理与 Agent 生成中投机解码失效的工业界痛点，具备极高的部署落地与二次演进价值。

相关论文