Test-Time Speculation

本文介绍了测试时推测（TTS），这是一种在线蒸馏方法，它利用目标模型的验证信号在推理过程中持续调整草稿模型，从而克服现有推测器在长序列上的性能下降问题，并显著提升了接受长度。

要点

以下是用通俗易懂的语言和富有创意的类比对论文《推理时推测》（Test-Time Speculation）的解释。

核心问题：那辆“超速行驶却迷失方向”的汽车

想象一下，你正试图让一位才华横溢但思维缓慢的作者（目标模型）写一个非常长的故事（比如一本小说）。为了节省时间，你雇佣了一位快速且精力充沛的实习生（草稿模型），在作者甚至还没阅读之前，就让他猜测接下来的几句话。

在人工智能领域，这被称为推测解码（Speculative Decoding）。实习生猜测一段内容，作者快速检查。如果实习生猜对了，作者只需说声“干得好！”然后继续前进，跳过了从头撰写这些单词的繁重工作。如果实习生猜错了，作者就必须停下来，纠正错误，然后重新开始。

关键缺陷：
论文发现，这些“实习生”的训练方式存在一个重大缺陷。

• 训练阶段： 实习生是在短篇故事（如推文或简短邮件）上受训的。他们非常擅长猜测 200 字句子中的下一个词。
• 现实情况： 在现实世界中，人们要求 AI 撰写长篇报告、代码或数千字的故事。

随着故事变长，实习生开始感到困惑。因为他们只接受过短句训练，随着文本增长，他们的“思路”开始混乱。他们开始猜测与长上下文不符的单词。

• 结果： 作者不得不拒绝实习生几乎所有的猜测。这个过程非但没有节省时间，反而因为作者不断停下来纠正实习生而变慢了。论文将这种情况称为“接受长度”（Acceptance Length）降至接近 1（意味着实习生基本上毫无用处）。

解决方案：“推理时推测”（TTS）

作者提出了一种巧妙的修复方案，称为推理时推测（Test-Time Speculation, TTS）。他们不是为每项工作雇佣新的实习生，而是教导同一个实习生如何在工作中适应。

类比：现场辅导会议
想象实习生正在写故事，而作者正在检查。

1. 旧方法： 实习生猜测 10 个词。作者检查它们。如果错了，作者修正它们并继续前进。实习生从错误中学不到任何东西，因为他们从未被告知为什么错了，以及这对下一句有何帮助。
2. TTS 方法： 每当作者检查实习生的工作时，作者不仅仅说“对”或“错”。作者利用这一刻给实习生上一堂微型课程。
   • 作者说：“你猜的是‘猫’，但在这个特定的长故事中，这个词应该是‘狗’。这是我使用的确切概率分布。”
   • 实习生立即根据这堂具体的课程更新他们的大脑（内部数学计算）。
   • 现在，当实习生猜测下一组单词时，他们会变得更聪明一点，更好地与作者当前的心情和故事的漫长历史保持一致。

这有什么特别之处？
通常，你必须停下来花几天时间重新训练模型才能让它变得更好。而 TTS 是在故事撰写过程中即时完成这一点的。它利用“验证”步骤（作者无论如何都要做的步骤）作为免费的训练信号。这就像一名学生通过与老师对话来学习一门新语言，老师实时纠正他们，使他们在对话结束时变得流利。

结果：走得越远，速度越快

论文在多种不同类型的“作者”（AI 模型）和“实习生”（推测器）上测试了这种方法，涵盖了求解数学问题、编写代码和回答科学问题等困难任务。

• 改进之处： 通过使用 TTS，“实习生”在故事变长时，猜对单词的能力大大增强。
• 数据表现： 平均而言，系统接受了多 41%的实习生猜测。在某些情况下，这比之前的最佳方法提高了72%。
• 趋势： 文本越长，TTS 的效果越好。当其他方法在几千字后失效时，TTS 实际上随着生成的继续变得更准确，因为实习生在不断即时学习和适应。

总结

可以将以前的方法想象成雇佣了一名只擅长 100 米短跑的快跑者。当你要求他们跑马拉松时，他们会崩溃。

推理时推测就像是给这名跑者配备了一位教练，教练与他们并肩奔跑，在每一步都低声纠正并调整策略。跑者不那么容易疲劳，保持在正确的道路上，整个团队能更快地完成马拉松。

论文证明，通过让 AI 在生成过程中“在工作中学习”，我们可以在撰写非常长的文档时，依然保持 AI 的快速和高效。

技术摘要

技术摘要：测试时推测（TTS）

1. 问题陈述

本文指出了当前最先进的推测解码方法（如 DFlash、EAGLE-3 和 PARD）在应用于长响应任务时存在的一个关键局限性。虽然推测解码通过使用快速“草稿”模型生成令牌、再由较慢的“目标”模型进行验证来加速大语言模型（LLM）的推理，但其效率高度依赖于接受长度——即每轮中被目标模型接受的一系列连续草稿令牌的数量。

作者观察到，随着生成长度的增加，现有推测器的接受长度会显著下降。在短短数千个输出令牌内，接受长度通常会降至接近 1 的值（例如 1.1），从而实际上消除了任何加速收益。这种退化发生的原因是，最先进的推测器是在离线状态下针对短序列（通常 $\le$ 2K 令牌）进行训练的，当它们被迫在推理过程中针对更长的序列（例如 20K–32K 令牌）近似目标模型时，就会产生分布不匹配。随着生成的进行，草稿模型的预测会偏离目标模型日益自信的分布，导致频繁的拒绝。

2. 方法论：测试时推测（TTS）

为了解决这种分布不匹配问题，作者提出了测试时推测（TTS），这是一种在推理过程中实时适应草稿模型的在线蒸馏方法。

核心洞察

关键的发现是，标准的推测解码验证步骤已经生成了适应所需的监督信号，且无需额外成本。在每一轮中，目标模型都会计算其在草稿令牌上的完整概率分布。TTS 利用这一点，将：

• 目标模型视为“教师”。
• 草稿模型视为“学生”。
• 已验证的草稿令牌视为蒸馏训练样本。

算法

TTS 将生成与模型更新交错进行。每轮推测的过程如下：

1. 草稿生成：当前草稿模型（$q_t$）生成一个包含 $C$ 个令牌的画布。
2. 验证：目标模型（$p$）在单次前向传递中评估该画布，通过标准拒绝采样确定接受长度（$\tau$）。
3. 蒸馏损失：在下一轮开始之前，草稿模型使用蒸馏损失函数进行单步梯度更新：
   $$L_t(q) = \tilde{KL}(p \parallel q) + \lambda \tilde{KL}(q_t \parallel q)$$
   • 第一项近似目标分布与画布上新草稿分布之间的 Kullback-Leibler (KL) 散度。
   • 第二项是正则化分量，防止草稿偏离其先前状态（$q_t$）过远。
   • 应用位置相关的权重（$w_k$），优先处理画布中较早的令牌。
4. 更新：草稿模型参数被更新（$q_{t+\tau} \leftarrow q_t - \eta \nabla L_t$）。

系统优化

为了在改善的接受长度与梯度更新带来的延迟开销之间取得平衡，TTS 采用了以下策略：

• 步进式更新：梯度更新每隔 $S$ 轮执行一次，而不是每轮都执行，从而分摊计算成本。
• 异步流水线：更新被卸载到专用的 CUDA 流中，该流与随后的 $S-1$ 轮生成并行运行，从而将延迟从关键路径中隐藏。

3. 主要贡献

1. 退化诊断：作者证明，由于草稿模型的短序列训练分布与长序列推理分布之间的不匹配，当前推测器的有效性会随着生成长度的增加而下降。
2. TTS 框架：他们提出了测试时推测，这是一种在线蒸馏方法，利用验证步骤作为监督信号，在推理时适应草稿模型，无需离线重新训练。
3. 全面评估：该方法在五个最先进的模型（Qwen-3、Qwen-3.5、Llama3.1 系列）和八个多样化的基准测试（包括 AIME、LiveCodeBench 和 GPQA）上进行了评估，显示出一致的改进。
4. 系统集成：作者在 SGLang 推理框架内实现了 TTS，解决了内核微分和 CUDA 图同步等系统级挑战。

4. 实验结果

• 接受长度提升：与 DFlash 相比，TTS 将平均接受长度提高了高达 72%，平均提高 41%；与 EAGLE-3 相比，提高了高达 67%（平均 34%）。
• 随长度扩展：TTS 的收益随生成长度而扩展。例如，在 AIME 2024 数据集上，相对于 DFlash 的改进从 0–10K 令牌范围内的 15% 增长到 20–30K 令牌范围内的 183%。
• 吞吐量：虽然频繁更新（步长 $S=1$）能最大化接受长度，但步长 $S=5$ 通过平衡适应频率与更新开销，实现了最佳的吞吐量加速（相对于 DFlash 高达 1.71$\times$）。
• 泛化性：TTS 在不同模型规模（4B 到 122B）和架构（稠密和 MoE）上均有效，特别是能够补偿那些在短上下文上训练的推测器（例如具有 2K 上下文的 EAGLE-3）在应用于具有更大上下文窗口的目标模型时的不足。

5. 意义与主张

本文声称，TTS 从根本上解决了推测解码在现实世界长响应场景中的局限性。通过在生成过程期间适应草稿模型，TTS 缩小了训练分布与推理分布之间的差距，确保即使对于跨越数万令牌的输出，推测解码依然有效。

作者强调，TTS 不需要对请求流结构做任何假设（不同于依赖特定领域缓冲区的先前在线方法），并且可以直接在现有的、公开的先进推测器之上运行。这使得 TTS 成为在生产环境中维持高推理吞吐量的实用解决方案，在这些环境中，长格式生成（例如代码、推理、内容创作）占据主导地位。这项工作被呈现为一种必要的演进，旨在随着 LLM 应用向更长上下文窗口转变，保持推测解码的可行性。