跳转至

Verification-Aware Training for Speculative Decoding

会议: NeurIPS 2026
arXiv: 2608.30135
OpenReview: https://openreview.net/forum?id=qMv8q5WApK
代码: https://github.com/naver-ai/VAT
领域: LLM 效率
关键词: 推测解码、验证感知训练、验证头、自适应权重、自回归与扩散草稿

一句话总结

VAT 在训练时实时模拟目标模型的序列验证,用轻量验证头学习 token 是否在首个拒绝前存活,并将逐位置损失衰减动态重锚定在首个拒绝位置,即插即用提升 EAGLE-3 与 DFlash 的接受长度与端到端吞吐。

研究背景与动机

推测解码已成为大语言模型无损推理加速的主流技术。其核心在于将自回归生成解耦为两阶段:低开销草稿模型先快速提出多个候选 token,目标模型再通过单次并行前向完成全部候选的序列验证。整个系统的端到端加速比由每个验证周期通过验证的平均接受长度直接决定。更为关键的是,推测解码的验证遵循严格的时序串行截断规则——一旦某一位置的候选被拒绝,其后所有草稿 token 无论本身质量多高都会被全部废弃,下一轮草稿必须从该失败点重新开始。

然而,现有的草稿模型训练目标与上述验证机制存在严重的结构性失配。无论是基于特征自回归预测的 EAGLE-3,还是基于并行块去噪的扩散模型 DFlash,本质上都沿用逐位置拟合目标模型输出的监督范式。这种训练方式带来两个本质问题:第一,模型接收到的仅是局部的逐词模仿信号,完全缺乏关于各位置能否在串行验证中实际“存活”的反馈;第二,现有方法采用先验固定的逐位置损失衰减权重(如固定指数衰减),无法感知具体样本在哪里发生首次拒绝。当某一输入上的草稿质量很高、首次拒绝发生得很晚时,靠后的有效 token 会被提前施加剧烈衰减,导致损失分配失真。

解决这一矛盾的关键在于将目标模型的验证逻辑内化到草稿训练阶段,让训练动态跟踪每个样本的真实验证截断点。核心 idea:在训练步实时模拟目标模型的前向验证,以派生的存活标签训练轻量验证头引导隐藏表征,并将逐位置损失衰减的锚点自适应平移到首个拒绝位置,在不改变推理架构的前提下对齐验证规律。

方法详解

整体框架

VAT 是一个纯训练期的即插即用框架,不修改目标模型权重,也不改变草稿模型的推理架构与无损验证流程。在训练阶段,给定上下文输入,目标模型产生预测分布 \(p_k\),草稿模型产生候选分布 \(\hat{p}_k\) 并采样出候选序列。系统在训练前向中立即模拟推测采样的验证过程,逐位置比对并确定该样本的首次拒绝位置 \(k^*\),由此推导出前缀存活标签 \(v_k\)。

基于模拟验证结果,VAT 构建两条平行的监督链路:一方面,在草稿模型顶层挂载轻量单层验证头,以二值交叉熵损失直接学习每个位置的串行存活状态 \(v_k\),使验证反馈的梯度反向塑造特征表示;另一方面,利用 \(k^*\) 对原本固定的位置衰减权重进行动态重锚定,在首个失败点之前保持全权重分配,随后才开始按原衰减曲线打折。在推理期,草稿模型保持原样工作;若启用早退出,验证头还可充当低成本的截断探针,跳过注定失败的高位生成以进一步节约推理算力。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入上下文与隐藏状态"] --> B["草稿模型生成候选 token 与隐状态"]
    B --> C["训练时模拟验证<br/>计算首个拒绝位置 k* 与存活标签 vk"]
    C --> D["验证头辅助监督<br/>预测每个 token 序列存活概率"]
    C --> E["验证自适应重加权<br/>首个拒绝前全权重,从 k* 处平移衰减"]
    D --> F["联合损失更新草稿表征<br/>联合软硬标签与二值分类目标"]
    E --> F
    F -->|无损推理| G["部署至现有验证流程<br/>支持可选的验证头早退出"]

关键设计

1. 训练时模拟验证:构建序列级存活标签

现有训练方法仅将目标模型输出作为局部标签计算交叉熵,使草稿模型无法区分“能继续扩展接受前缀的关键 token”与“已被前序失败截断的无效 token”。VAT 在训练前向中实时模拟推测采样的验证决策,对位置 \(k\) 上草稿 token \(x\) 按照验证接受概率 \(\min(1, p_k(x)/\hat{p}_k(x))\)(贪婪解码下即判定是否命中 top-1)生成单点接受指示变量 \(m_k\)。随后定位该序列的首个拒绝点:

\[k^* = \min \{ k : m_k = 0 \}\]

当全部 \(K\) 个 token 均被接受时约定 \(k^* = K + 1\)。由此定义反映真实推测截断的累积存活标签 \(v_k = \mathbb{I}[k < k^*] = \prod_{j \le k} m_j\)。与独立判定的 \(m_k\) 不同,\(v_k\) 显式反映了推测解码“一损俱损”的序列级依赖,为后续的权重调制与辅助表征学习提供了真实依据。

2. 验证头辅助监督:用存活目标引导草稿隐藏表征

针对局部自回归/扩散损失无法感知序列级存活边界的问题,VAT 在草稿模型的最终隐状态上挂载一个极其轻量的单层密集分类头(Verification Head)。该头将各位置的隐状态映射为标量预测概率 \(\hat{v}_k \in [0, 1]\),并通过二值交叉熵损失接收来自真实存活标签 \(v_k\) 的直接监督:

\[\mathcal{L}_{\text{VH}} = -\frac{1}{K}\sum_{k=1}^K \left[ v_k \log \hat{v}_k + (1 - v_k) \log (1 - \hat{v}_k) \right]\]

由于该分类头直接依托草稿主干的隐表征进行联合训练,反向传播的梯度驱使草稿网络在其高层特征中主动编码“能否与目标模型在此达成一致”的特征,促使草稿模型在较深生成步长上依然倾向于保持与目标模型决策空间对齐。该头在基准测试中完全作为训练期辅助正则项,推理时无需额外计算;同时它也能在推理端作为轻量级置信度探针指导早退出截断。

3. 验证自适应重加权:以首次拒绝点为锚动态平移损失

常规草稿训练采用从第一位 \(k=1\) 开始递减的静态权重预设(如 EAGLE-3 的 \(w_k = 0.8^{k-1}\) 或 DFlash 的 \(w_k = \exp(-(k-1)/\gamma)\)),这一设定违背了验证动力学:当样本质量极佳、\(k^*=5\) 发生首次拒绝时,前 4 个成功扩展前缀的 token 却被过早打了很大的折扣。VAT 提出自适应重加权机制,将权重衰减曲线的基准点从固定的 \(k=1\) 动态平移至样本专属的 \(k^*\):

\[\hat{w}_k = \begin{cases} 1 & \text{if } k < k^* \\ w_{k - k^* + 1} & \text{if } k \ge k^* \end{cases}\]

在此方案下,首个拒绝点之前的所有已接受位置均享受 \(1.0\) 的全额梯度权重,保障高贡献前缀的学习强度;而在首次拒绝发生点 \(k^*\) 本身,由于其属于能够直接挽救并延长接受前缀的最邻近纠错位,该位置对应的平移权重同样为 \(w_1 = 1.0\);从 \(k^*+1\) 开始才沿原预设曲线逐渐衰减。这既避免了激进丢弃拒绝后 token(硬截断会导致早期欠拟合),又消除了静态衰减对有效前缀的惩罚。

损失函数 / 训练策略

VAT 的总体训练目标结合了动态加权的草稿预测损失与验证头辅助损失。为了最大化知识迁移效率,草稿预测部分同时引入针对目标模型预测分布的软标签交叉熵 \(\ell_k^{\text{soft}}\) 与针对目标采样 token 的硬标签交叉熵 \(\ell_k^{\text{hard}}\):

\[\mathcal{L} = \sum_{k=1}^K \hat{w}_k \left( \ell_k^{\text{soft}} + \ell_k^{\text{hard}} \right) + \beta \mathcal{L}_{\text{VH}}\]

其中权重超参数 \(\beta\) 在所有实验中固定设为 \(1.0\)。训练在配对数据集(如 Perfectblend 生成数据)上进行,使用 NVIDIA A100 80GB GPU 采用 bf16 精度完成 3 个 epoch 的联合微调。

实验关键数据

主实验

主实验在 Qwen3-4B、Qwen3-8B 以及 LLaMA-3.1-8B 三个目标模型上,分别在自回归草稿基线 EAGLE-3 和扩散块草稿基线 DFlash 上评测端到端加速比与平均接受长度 \(\tau\)。评测涵盖数学(GSM8K、MATH-500、AIME25)、代码(HumanEval、MBPP、LiveCodeBench)以及对话(MT-Bench、Alpaca)三大任务大类(贪婪解码 \(T=0\) 条件下各领域综合平均结果如下表所示)。

模型与草稿基座 优化策略 数学加速 (τ) 代码加速 (τ) 对话加速 (τ) 全任务综合加速 全任务综合 τ
Qwen3-4B + EAGLE-3 基线 4.80× (7.32) 4.02× (6.25) 3.05× (4.74) 4.07× 6.28
Qwen3-4B + EAGLE-3 + VAT (本文) 5.05× (7.72) 4.50× (6.99) 3.24× (5.04) 4.39× (+7.9%) 6.78 (+8.0%)
Qwen3-4B + DFlash 基线 5.89× (7.32) 4.52× (5.55) 2.55× (3.62) 4.54× 5.73
Qwen3-4B + DFlash + VAT (本文) 6.17× (7.67) 4.83× (5.96) 2.73× (3.90) 4.81× (+5.9%) 6.08 (+6.1%)
Qwen3-8B + EAGLE-3 基线 4.81× (7.25) 3.95× (6.03) 2.99× (4.58) 4.04× 6.12
Qwen3-8B + EAGLE-3 + VAT (本文) 4.92× (7.46) 4.24× (6.49) 3.21× (4.96) 4.24× (+5.0%) 6.47 (+5.7%)
Qwen3-8B + DFlash 基线 5.71× (6.91) 4.40× (5.32) 2.72× (3.70) 4.47× 5.51
Qwen3-8B + DFlash + VAT (本文) 6.21× (7.75) 4.85× (6.06) 2.85× (3.86) 4.86× (+8.7%) 6.14 (+11.4%)
LLaMA-3.1-8B + EAGLE-3 基线 4.45× (6.45) 4.28× (6.29) 3.56× (5.23) 4.17× 6.08
LLaMA-3.1-8B + EAGLE-3 + VAT (本文) 4.60× (6.55) 4.43× (6.40) 3.79× (5.50) 4.33× (+3.8%) 6.23 (+2.5%)
LLaMA-3.1-8B + DFlash 基线 4.66× (6.26) 4.60× (6.07) 2.44× (3.80) 4.08× 5.57
LLaMA-3.1-8B + DFlash + VAT (本文) 4.84× (6.54) 4.71× (6.26) 2.53× (3.92) 4.22× (+3.4%) 5.78 (+3.8%)

消融实验

在 Qwen3-4B 目标模型搭配 DFlash 草稿模型上,系统解耦消融了验证头(VH)、自适应加权(Adaptive Weight)以及软硬标签结合(Soft+Hard)三大核心模块的独立与组合贡献。

验证头 (VH) 验证自适应加权 软硬联合标签 全局加速比 平均接受长度 (τ) 说明
- - - 4.54× 5.73 DFlash 基础基线(仅硬标签、静态指数衰减)
✓ - - 4.62× 5.87 单独引入验证头,表征受存活监督强化 (+0.14)
- ✓ - 4.65× 5.91 单独引入自适应权重,动态锚定首次拒绝点 (+0.18)
- - ✓ 4.58× 5.82 仅引入软标签辅助蒸馏 (+0.09)
✓ ✓ - 4.76× 5.99 验证头与自适应加权协同,捕获大部分主要收益
✓ - ✓ 4.69× 6.05 验证头结合双标签蒸馏
- ✓ ✓ 4.72× 5.99 自适应权重结合双标签蒸馏
✓ ✓ ✓ 4.81× 6.08 完整 VAT 方案,达到最佳加速比与最长接受长度

关键发现

  • 正交互补与协同增效:验证头(隐层表征对齐)与自适应权重(梯度配比调制)作用机制独立互补,二者组合在无额外软标签时便能将加速比提升至 4.76×,全模块组合则实现 4.81× 加速与 6.08 接受长度。
  • 重锚定策略不可或缺:若在首个拒绝点之后直接切断梯度(前缀截断策略),因模型训练早期阶段 \(k^*\) 普遍极小导致梯度严重匮乏,模型加速比崩塌至 2.44×(\(\tau=3.11\));而 VAT 的平移递减平滑过渡,既能让关键位置获得充足学习,又保留了后续上下文的弱正则。
  • 推理端低成本早退出潜力:以阈值对验证头的单点预测概率进行截断,仅向目标模型发送预测为存活的前缀。在 DFlash 代码任务上,早退出机制将端到端加速比从 4.83× 进一步推升至 4.97×,逼近真实 Oracle 截断的 5.20× 上界。

亮点与洞察

  • 将时序串行依赖内化为逐词存活目标:设计了以二值交叉熵监督的验证头,巧妙解决了传统序列自回归只顾“下一个词像不像”、不顾“前面的词错没错”的致命盲区。
  • 基于首次失败点的平移动态衰减设计:抛弃沿用了多年的全序列全局固定几何衰减,以样本首次拒绝位置 \(k^*\) 动态平移衰减曲线,同时兼顾了最邻近失败点 \(k^*\) 的高纠错价值与后续位置的平滑梯度供给。
  • 纯训练侧修改零运行时开销:不增加任何推理参数与前向计算开销,支持即插即用集成到任何主流推测解码范式中,具备极高的工程可落地性。

局限与展望

  • 模型规模评测局限:目前验证工作主要集中在 4B 至 8B 规模的语言模型上,更庞大的超大参数模型(如 70B+)以及 MoE 架构下的收益表现有待后续验证。
  • 早退出存在偶发误杀:验证头作为单层小分类器在用于推理端早退出时,偶尔会出现“假拒绝”,截断原本可以通过目标验证的正确候选,使平均接受长度轻微下降(如从 7.67 降至 7.35)。

相关工作与启发

  • vs EAGLE-3 / DFlash: 后两者专注于草稿网络模型结构与多层特征融合的创新,但损失函数依然采用固定位置衰减;VAT 则针对训练目标进行正交升级,使两者在各任务上的接受长度分别获得稳定跨步提升。
  • vs PARD-2 / D-PACE: PARD-2 与 D-PACE 采用基于置信度的平滑衰减代理;VAT 直接基于前向真实模拟得到的首个拒绝位置 \(k^*\) 进行绝对重锚定,直接对准推测采样的硬截断本质。
  • vs GRIFFIN: GRIFFIN 采用基于 top-\(m\) 阈值的硬掩码;VAT 表明保留首个拒绝点之后的衰减梯度比粗暴截断具有更好的泛化与深层特征稳定性。

评分

  • 新颖性: ⭐⭐⭐⭐☆ [抓住了推测解码串行截断与局部自回归训练的失配本质,构思新颖巧妙]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖自回归与扩散两大流派、3个基座模型及8个经典benchmark,消融详实到位]
  • 写作质量: ⭐⭐⭐⭐⭐ [动机叙述逻辑极其紧密,数学表达与图表清晰,论点有力]
  • 价值: ⭐⭐⭐⭐⭐ [零推理改造成本的即插即用训练方法,在系统效率落地方面实用价值极高]