Verification-Aware Training for Speculative Decoding

TL;DR VAT 在训练 draft 模型时模拟 speculative decoding 的序列验证，将接受/拒绝模式作为监督信号，并动态调整损失权重，在不改变推理流程的情况下提升平均接受长度和实际加速。

问题

问题背景

在大语言模型推理加速中，推测解码 (speculative decoding) 通过小模型草稿 + 大模型并行验证来降低延迟，但其实际加速上限取决于平均接受长度 (average acceptance length)，而这又高度依赖草稿模型 (draft model) 的质量。

现有方法局限

现有草稿模型训练大多采用token 级知识蒸馏，让草稿模型模仿目标模型在每个位置上的输出分布，并配合固定位置权重（如线性或常数衰减）。这带来两个关键问题：

1. 忽略验证的顺序拒绝特性：推测解码的验证过程是从左到右依次进行，一旦某个 token 被拒绝，其后所有候选 token 都会被丢弃。但传统训练对所有位置一视同仁，导致草稿模型可能在“注定被丢弃”的位置上浪费容量，而没把重心放在更靠前、更可能被接受的 token 上。
2. 训练目标与推理过程脱节：固定权重无法感知每个样本的具体拒绝点，无法告诉草稿模型“哪些位置在真实验证中能存活”。

为什么这个问题难且重要

核心难点在于训练时无法低成本获得验证结果——验证依赖目标模型的前向计算，直接模拟代价过高，且草稿模型与目标模型分布可能存在偏差。业界对推理延迟高度关注，尤其是长文本生成、代码补全等场景，即便是几个百分点的平均接受长度提升也能转换为可观的 wall-clock 加速。因此，如何在不改变推理架构的前提下，让草稿模型训练过程与顺序验证机制对齐，成为提升 speculative decoding 效率的关键问题。

行业类比

这类似于缓存预取 (cache prefetching) 中要精准预测将被访问的内容：如果预取的一批数据中第一个就是错的，后续再多正确预取也是无效开销；同样，草稿模型应优先保证靠前 token 的正确性，而非平均优化所有 token。

核心洞察

• 将推理时的顺序验证纳入训练监督，而非仅做 token 级模仿。现有 draft 训练（如 EAGLE-3、DFlash）通常设置固定的逐位置权重，让 draft 模仿 target 分布，但完全忽略验证过程一旦拒绝便丢弃后续 token 的顺序特性。VAT 在每一步训练中显式模拟验证，生成接受/拒绝标签，使 draft 模型直接为验证服务，从而优化最终接受长度，而不只是提高单 token 预测准确率。
• 通过 verification head 和 verification-adaptive weighting 对训练目标进行细粒度修正。verification head 是一个轻量二分类器，监督每个位置是否在顺序验证中存活，赋予 draft 显式的验证意识；verification-adaptive weighting 将加权衰减锚定在首次拒绝点之前，保持到该点的完整权重，相比固定衰减更贴合推理中的实际损失分布。该设计只修改训练目标，不改动架构和推理流程，可叠加在现有方法上，实验显示平均接受长度和实际加速均有稳定提升。

方法

VAT 改造 draft model 的训练目标，不改变架构与推理流程。

输入：draft model 在训练时生成候选 token 序列，目标模型对整条序列进行单次前向验证，模拟推理时的顺序接受/拒绝过程。

关键模块：

• verification head：在 draft model 上增加一个轻量二分类头，与主模型联合训练，监督信号是每个位置是否通过顺序验证存活（accept/reject 标签）。该头在训练时辅助学习生存概率，推理时丢弃。
• verification-adaptive weighting：替代原有的固定衰减权重，对每个训练样本，保持从首 token 到首次拒绝点的完整权重，此后重新锚定衰减从拒绝点开始。这样使 draft 更关注被接受的前缀，并对拒绝后的候选恢复有效学习。

输出：训练完成后的 draft model 在推理时与目标模型配合，提升平均接受长度与 wall-clock speedup。VAT 可直接叠加于 EAGLE-3、DFlash 等现有方法，无需修改架构或推理过程。

差异点：与现有 token-level 模仿中固定 per-position 权重的做法不同，VAT 显式建模顺序验证的拒绝模式与生存概率，使训练目标贴合推理时的验证行为。

实验

实验设计

VAT 被应用到 EAGLE-3 和 DFlash 两种 draft 训练框架中，target 模型覆盖 Qwen3-4B、Qwen3-8B、LLaMA-3.1-8B。评估基准涵盖 math、code、chat 三类任务，验证 VAT 在不同 draft 架构与 target 规模下的泛化性。训练时模拟 sequential verification 过程，利用新增的 verification head 与 verification-adaptive weighting 替换原有固定权重，但保持 draft 架构和推理流程不变。

关键发现

• 平均接受长度（average acceptance length）最高提升 11.4%。
• Wall-clock speedup 最高提升 8.7%，表明接受长度增益能转化为实际推理加速。
• 在 math、code、chat 基准上均观察到一致提升，说明 VAT 不针对特定任务过拟合。

与基线对比解读

相较于固定 per-position 权重（如 EAGLE-3/DFlash 默认策略），VAT 通过模拟验证让 draft 模型提前感知 token 在 sequential verification 中被接受或拒绝的模式。verification-adaptive weighting 将权重衰减重新锚定到每个样本的首次拒绝点之后，使训练损失更符合推理时的验证行为。VAT 只修改训练目标，因此可与现有加速方法叠加，无需重新设计 draft 架构或推理逻辑，工程落地成本低。这种训练/推理一致性是接受长度提升的关键。

行业影响

落地场景

VAT 直接适用于所有已部署或计划部署 speculative decoding 的 LLM 在线服务，尤其是对延迟敏感、并发量大的场景。典型产品包括：

• 代码补全平台（如 Copilot 类工具）：开发者输入时要求低延迟响应，草稿模型接受长度提升可显著减少等待时间。
• 电商智能客服与推荐对话：高并发场景下，推理吞吐决定单卡可服务用户数，VAT 训练草稿模型可增加每秒完成请求数。
• 企业知识库问答 / 金融投研助手：长文本生成、引用链式思维推理时，验证频繁出现 reject，VAT 针对验证模式优化可降低无效 draft。

商业价值

核心收益在降本与体验提升双线。

• 降本：wall-clock 加速最高 8.7%，意味着相同 GPU 实例可承载更多请求，或降低 SLA 所需的实例规模。
• 体验提升：平均接受长度提升最高 11.4%，减少验证轮次，首 token 延迟和生成中停顿均下降。
• 零额外推理成本：VAT 只改训练目标，不增加推理时计算或内存。

与现有产品 / 工作流集成

VAT 是训练目标插件，可直接叠加在 EAGLE-3、DFlash 等现有草稿模型训练框架上，不改变草稿架构、目标模型或推理代码。集成路径清晰：

1. 在现有草稿模型训练 pipeline 中增加 verification head；
2. 替换固定位置权重为 verification-adaptive weighting；
3. 重新训练草稿模型后，直接替换原草稿模型权重，推理侧无需任何改动。

对已使用 speculative decoding 的团队，只需重新训练草稿模型即可获得加速，迁移成本低。

局限

• VAT 需要在训练时用目标模型模拟顺序验证以获得接受/拒绝信号，这比普通 token 级蒸馏增加了额外前向计算开销；联合训练的 verification head 引入少量参数和额外训练目标。虽然作者报告训练开销可接受，但在超大目标模型或预算受限场景下的可行性仍未经过验证。
• 实验仅覆盖 Qwen3-4B/8B 和 LLaMA-3.1-8B，且仅基于 EAGLE-3 和 DFlash 两种 draft 训练框架。对于更大规模模型、不同 draft/target 架构组合或非自回归 draft 方案，VAT 的有效性尚未得到验证；评测集中在 math、code、chat 任务，长文本生成或多语言场景下的泛化表现仍不明确。
• VAT 的 verification head 和自适应权重依赖特定接受规则与采样温度，文中消融也显示不同 verification rule 和语料温度会影响效果。实际部署时若推理端解码配置（如温度、top-p）与训练不一致，可能需要重新训练或额外调参，增加了工程维护成本。