验证感知训练:对齐投机解码验证过程的草稿模型训练框架 Verification-Aware Training for Speculative Decoding
训练时模拟目标模型逐位验证,以验证头与自适应加权对齐草稿训练,加速最高提升8.7%
前置知识
投机解码(Speculative Decoding)
无损的 LLM 推理加速技术:轻量草稿模型连续起草 $K$ 个候选 token,目标模型用一次前向传播并行验证所有候选;验证从第一位起逐位进行,一旦某个位置被拒绝,其后候选(无论质量如何)全部丢弃,并从最后接受的位置重新起草。因最终输出严格服从目标模型分布,加速不损失质量,已成为最主流的推理加速手段。
本文的全部动机与实验都围绕投机解码展开:接受长度决定加速比,作者指出现有草稿模型训练目标与这一逐位顺序验证机制脱节,改进训练目标就是改进整个推理系统的速度。
投机采样接受规则
草稿 token $x$ 在位置 $k$ 被目标接受的概率为 $\min(1, p_k(x)/\hat{p}_k(x))$,其中 $p_k$ 为目标分布、$\hat{p}_k$ 为草稿分布;被拒绝时按目标分布重新采样。该规则在数学上保证最终输出与目标分布完全一致,是无损加速的理论基石。
VAT 在每个训练步用同一规则模拟验证,得到逐位接受指示 $m_k$ 与首个拒绝点 $k^*$,它们是验证头标签与自适应权重的数据来源,不理解该规则就无法理解标签如何生成。
平均接受长度 τ
每个验证循环中被目标模型接受的期望 token 数。每次验证循环开销固定,$\tau$ 越大每轮产出越多、加速比越高,因此它是衡量投机解码系统性能的核心指标。论文中 Qwen3-4B 上 EAGLE-3 基线 $\tau=6.28$、DFlash 基线 $\tau=5.73$,DFlash 每轮并行起草 16 个 token。
论文所有实验报告的核心指标就是 $\tau$ 与实际加速比,VAT 的设计目标正是把训练信号集中在决定 $\tau$ 的被接受前缀上。
软标签与硬标签蒸馏
训练草稿模型时,硬标签用目标模型实际采样出的 token 做 one-hot 交叉熵 $\ell_k^{hard}$,软标签用目标的完整输出分布做交叉熵 $\ell_k^{soft}$。软标签提供更密集的分布级信号,硬标签提供直接监督;标准知识蒸馏常将二者结合,本文消融显示联合使用优于任一单独使用。
VAT 的总损失 $\mathcal{L}=\sum_k\hat{w}_k(\ell_k^{soft}+\ell_k^{hard})+\beta\mathcal{L}_{VH}$ 建立在两种损失之上,Table 2 显示加入软标签使 $\tau$ 从 5.73 升至 5.82。
EAGLE-3 与 DFlash
两种 SOTA 草稿范式:EAGLE-3 是自回归草稿器,复用目标模型隐藏状态、配合动态草稿树与训练时测试(让草稿器见到自己的 rollout);DFlash 用块扩散模型一次前向并行生成全部 16 个草稿 token。两者训练目标形式相同,均为 $\mathcal{L}_{draft}=\sum_{k=1}^{K}w_k\ell_k$,只是权重计划不同:EAGLE-3 用 $w_k=0.8^{k-1}$,DFlash 用 $w_k=\exp(-(k-1)/\gamma)$。
VAT 作为插件同时叠加在这两种方法上验证,证明收益不依赖具体草稿架构,这是论文说服力的关键实验设计。
研究动机
投机解码的加速比由接受长度决定:验证逐位进行,第一个被拒绝位置之后的所有草稿 token 都会被丢弃,因此每个位置只有在全部前缀都被接受时才对接受长度有贡献。然而现有草稿模型训练与这一验证机制严重脱节。以 EAGLE-3 和 DFlash 为代表的方法只用逐位交叉熵让草稿模仿目标输出($w_k=0.8^{k-1}$ 或 $w_k=\exp(-(k-1)/\gamma)$ 的固定权重),草稿模型完全不知道自己的 token 能否在实际的顺序验证中存活;同时逐位权重计划预先固定、全样本共享,无法适应每个样本首个拒绝点出现的位置——当拒绝点很靠后时,本应获得满权重的靠前位置仍被持续衰减,训练信号被错配到对接受长度无贡献的区域。作者指出这两个错配在 Qwen3-4B、Qwen3-8B、LLaMA-3.1-8B 等主流模型上普遍存在,是进一步提升加速比的直接障碍。
本文的目标是本文的目标是让草稿模型的训练目标与目标模型的验证过程对齐,遵循两条原则:第一,训练不应止步于逐位模仿目标输出,而要引入来自目标模型的验证信号,让草稿模型显式学习'我的每个位置能否在顺序验证中存活';第二,逐位训练信号应随每个样本自身的验证结果自适应调整,即以该样本首个拒绝点 $k^*$ 为锚点重新分配权重。同时整个方案必须以插件形式存在:只修改训练目标,不改变草稿架构、目标模型和推理流程,可直接叠加到 EAGLE-3、DFlash 等现有方法上。量化目标是在数学、代码、对话基准上同时提升平均接受长度 $\tau$ 与实际墙钟加速比——论文最终实现 $\tau$ 最高提升 11.4%、加速最高提升 8.7%。
与已有工作不同的是,已有工作的切入角度各不相同:Medusa、Hydra、EAGLE 系列聚焦草稿架构设计,训练目标长期停留在'逐位交叉熵 + 样本无关位置权重'的标准范式;HASS 与 EAGLE-3 的训练时测试虽让草稿器见到自己的 rollout,但仍用均匀逐位损失;近期一些工作把接受/拒绝信息引入草稿器,或作为推理期自适应草稿长度的信号,或作为训练期在拒绝后直接掩码、清零损失的信号。同期工作 PARD-2 用目标对前缀的累计置信度重加权,D-PACE 用草稿自身置信度构造期望接受长度的可微代理。VAT 的独特之处在于:权重以训练时观测到的首个拒绝位置 $k^*$ 为条件(而非置信度代理),并与直接监督累计接受结果的验证头相耦合,同时适用于自回归与扩散两种草稿范式,且在拒绝后保留衰减信号而非完全掩码。
核心方法
VAT 的直觉很直接:推理时的加速比由'多少个草稿 token 通过目标验证'决定,训练时就应在每个训练步模拟这一验证,把接受/拒绝模式变成监督信号。技术路线分三步:首先,训练时草稿与目标在每个位置 $k$ 都产出分布 $\hat{p}_k$ 与 $p_k$,用投机采样规则 $\min(1, p_k(x)/\hat{p}_k(x))$ 模拟接受,得到逐位指示 $m_k$ 和首个拒绝点 $k^*$(全接受时 $k^*=K+1$);其次,挂在草稿最后隐状态上的单层二分类验证头学习预测逐位接受标签 $v_k=\mathbb{1}[k<k^*]$,梯度回传塑造草稿表征;最后,把固定权重 $w_k$ 替换为以 $k^*$ 为锚的自适应权重 $\hat{w}_k$,$k^*$ 前满权重、之后沿用基础衰减曲线从 $k^*$ 重新起算。总损失为 $\mathcal{L}=\sum_k\hat{w}_k(\ell_k^{soft}+\ell_k^{hard})+\beta\mathcal{L}_{VH}$($\beta=1.0$)。推理流程完全不变,应用于 EAGLE-3 与 DFlash 后接受长度最高提升 11.4%、加速最高提升 8.7%。
核心创新是把'顺序验证的因果结构'显式写进训练目标。关键标签是 $v_k=\mathbb{1}[k<k^*]=\prod_{j\le k}m_j$:它不衡量单个位置预测得对不对,而衡量该位置在所有前缀都被接受的前提下能否存活——这正是接受长度的真实构成。围绕 $v_k$ 有两个配合的机制:其一,验证头通过二元交叉熵把'与目标是否一致'这一隐式性质变成显式梯度,训练动力学实验(Figure 2)显示加验证头后首个拒绝位置持续后移、且拒绝后仍与目标匹配的 token 数不再随训练衰退;其二,验证自适应加权把衰减锚点从固定的 $k=1$ 移到每个样本自己的 $k^*$:$k<k^*$ 时 $\hat{w}_k=1$,$k\ge k^*$ 时 $\hat{w}_k=w_{k-k^*+1}$,而 $k^*$ 本身也拿满权重,因为它是'最近的可修正失败',改善此处能直接延长接受前缀。与固定计划的本质区别在于:$k^*$ 逐样本变化,任何样本无关的调度都无法逼近这种行为。
方法步骤详情
第一步,取训练样本(Perfectblend 提示 + 目标模型贪心生成的回复),目标模型提取隐藏状态作草稿器条件输入,草稿器在每个位置 $k\in\{1,\dots,K\}$ 输出 $\hat{p}_k$,目标输出 $p_k$。第二步,按 $\min(1,p_k(x)/\hat{p}_k(x))$ 采样接受指示 $m_k$,得首个拒绝点 $k^*=\min\{k:m_k=0\}$,全接受时 $k^*=K+1$。第三步,由 $v_k=\prod_{j\le k}m_j$ 得逐位接受标签。第四步,验证头(单个全连接层,输入为草稿最后隐状态)输出 $\hat{v}_k$,用 $\mathcal{L}_{VH}=-\frac{1}{K}\sum_k[v_k\log\hat{v}_k+(1-v_k)\log(1-\hat{v}_k)]$ 训练,梯度回传至草稿模型。第五步,计算自适应权重:$k<k^*$ 时 $\hat{w}_k=1$,$k\ge k^*$ 时 $\hat{w}_k=w_{k-k^*+1}$。第六步,计算总损失 $\mathcal{L}=\sum_k\hat{w}_k(\ell_k^{soft}+\ell_k^{hard})+\beta\mathcal{L}_{VH}$($\beta=1.0$),同时更新草稿器与验证头。第七步,推理流程不变;可选地用 $\hat{v}_k$ 做早退草稿,预测到首个拒绝即截断。训练共 3 个 epoch,A100 80GB bf16。
技术新颖性
技术新颖性体现在四点。第一,验证模拟进入损失:之前接受/拒绝信息或用于推理期自适应草稿长度,或在训练期简单掩码拒绝后位置,VAT 把'模拟验证的累计结果'作为显式预测目标(验证头),并同时驱动逐位权重分配。第二,权重锚定在观测到的拒绝点:PARD-2 依赖目标对前缀的累计置信度、D-PACE 依赖草稿自身置信度的可微代理,VAT 直接条件于模拟验证观测到的 $k^*$,避免代理偏差。第三,拒绝后保留衰减而非清零:既有做法把拒绝后损失置零,VAT 认为这些位置虽不直接影响当前样本的接受长度、但仍有泛化价值,Figure 2(b) 显示验证头使拒绝后匹配 token 数在训练后期保持稳定。第四,方法完全正交于草稿架构:同一插件在自回归 EAGLE-3 与扩散式 DFlash 上均有效;消融还表明无论基础衰减取 $0.8^{k-1}$ 还是 $\exp(-(k-1)/\gamma)$,自适应化后 $\tau$ 几乎相同(6.09 对 6.08),说明收益来自自适应机制本身而非特定函数形式。
实验结果
主实验覆盖 3 个目标模型、2 种基线、8 个基准,温度 0 与 1 两种设置。接受长度 $\tau$:EAGLE-3+VAT 在 Qwen3-4B 提升 8.0%(6.28→6.78)、Qwen3-8B 提升 5.7%(6.12→6.47)、LLaMA-3.1-8B 提升 2.5%;DFlash+VAT 分别提升 6.1%(5.73→6.08)、11.4%(5.51→6.14)、3.8%。加速比:Qwen3-4B 上 EAGLE-3 从 4.07×升至 4.39×(+7.9%),DFlash 从 4.54×升至 4.81×(+5.9%);Qwen3-8B 上 DFlash 加速提升最大达 8.7%(4.47×→4.86×)。逐基准看,Qwen3-8B 的 DFlash+VAT 在 GSM8K 达 7.03×($\tau$=8.71),Qwen3-4B 的 EAGLE-3+VAT 在 HumanEval 从 4.33×跃升至 4.83×。组件消融(Table 2,DFlash+Qwen3-4B):单独验证头使 $\tau$ 5.73→5.87,单独自适应加权→5.91,软+硬标签→5.82,三者合计达 6.08、加速 4.81×,三类任务均最优。权重消融(Table 3):固定计划下均匀权重 $\tau$=5.72、EAGLE 式衰减 5.73、DFlash 式衰减 5.99,自适应化后分别达 6.09 与 6.08,证明收益来自自适应机制而非函数形式。验证头推理用途:预测首个拒绝位置的 MAE 最优阈值下仅 1.18(EAGLE-3,p=0.5)与 1.76(DFlash,p=0.6)个 token;早退草稿在 DFlash 代码任务把加速从 4.83×提至 4.97×(oracle 上限 5.20×)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GSM8K(数学,Qwen3-4B,温度0) | 加速比 / 平均接受长度 τ | DFlash+VAT:7.01× / 8.66;EAGLE-3+VAT:5.31× / 8.25 | DFlash:6.69× / 8.28;EAGLE-3:4.94× / 7.62 | DFlash +4.8% 加速、+4.6% τ;EAGLE-3 +7.5% 加速、+8.3% τ |
| HumanEval(代码,Qwen3-4B,温度0) | 加速比 / 平均接受长度 τ | EAGLE-3+VAT:4.83× / 7.37;DFlash+VAT:4.89× / 5.97 | EAGLE-3:4.33× / 6.67;DFlash:4.58× / 5.57 | EAGLE-3 +11.6% 加速、+10.5% τ;DFlash +6.8% 加速、+7.2% τ |
| GSM8K(数学,Qwen3-8B,温度0) | 加速比 / 平均接受长度 τ | DFlash+VAT:7.03× / 8.71 | DFlash:6.21× / 7.45 | +13.2% 加速、+16.9% τ(该表最大单项增益之一) |
| 八基准平均(Qwen3-8B,DFlash) | 平均加速比 / 平均 τ | 4.86× / 6.14 | 4.47× / 5.51 | +8.7% 加速、+11.4% τ(论文报告的最大平均增益) |
| 八基准平均(Qwen3-4B,EAGLE-3) | 平均加速比 / 平均 τ | 4.39× / 6.78 | 4.07× / 6.28 | +7.9% 加速、+8.0% τ |
| 八基准平均(LLaMA-3.1-8B) | 平均加速比 / 平均 τ | EAGLE-3+VAT:4.33× / 6.23;DFlash+VAT:4.22× / 5.78 | EAGLE-3:4.17× / 6.08;DFlash:4.08× / 5.57 | EAGLE-3 +3.8% 加速、+2.5% τ;DFlash +3.4% 加速、+3.8% τ(最弱增益组合) |
局限与改进
作者承认评估局限于最多 8B 参数的模型(Qwen3-4B/8B、LLaMA-3.1-8B),更大模型上的可扩展性留作未来工作。我的补充观察:第一,收益不均匀——LLaMA-3.1-8B 上 $\tau$ 提升只有 2.5%(EAGLE-3)和 3.8%(DFlash),温度 1 时 LLaMA-3.1-8B 的 DFlash 在 GSM8K 上甚至从 4.87 降到 4.64、加速从 3.32×降到 3.30×,说明强基线或高熵解码下增益可能收窄甚至波动;第二,训练数据完全来自目标模型贪心解码的自生成回复,与实际高温度解码的接受模式存在分布偏移(附录 B 补充了温度 1 数据的结果,但系统研究不足);第三,模拟接受 $m_k$ 是单次随机采样,$k^*$ 有噪声,权重锚点在样本间可能不稳定;第四,早退草稿以牺牲 $\tau$ 为代价(DFlash 数学任务 7.67→7.35),且验证头与特定目标-草稿对绑定,换目标模型需重训;第五,加速比仅在 A100 + HuggingFace Transformers 环境测得,未涉及 vLLM 等生产级推理栈的端到端收益。
独立分析的弱点
独立分析几个弱点。其一,对 $k^*$ 的估计是单次采样:$m_k$ 按 $\min(1,p_k/\hat{p}_k)$ 只采样一次,$k^*$ 方差较大,靠近 $k^*$ 的位置标签在'接受'与'拒绝'之间剧烈抖动,权重锚点随之不稳,可用多次采样平均或期望形式稳定标签。其二,验证头标签 $v_k$ 严重不平衡:多数样本 $k^*$ 靠后,标签几乎全为 1,二元交叉熵易被多数类主导,可引入 focal loss 或按 $k^*$ 分层采样。其三,训练语料完全由贪心解码的目标自生成回复构成,温度 1 部署时的接受模式属于分布外,论文虽在附录 B 给了初步结果,但未系统研究训练温度与部署温度的匹配。其四,验证头必须与草稿器联合训练,无法直接嫁接到已训好的草稿模型上,削弱了'即插即用'的程度,可探索冻结草稿器只训头或对已有草稿器做短程微调的方案。其五,早退只用固定阈值 p 扫描(0.1–0.9),未见动态或校准后的阈值策略,且 DFlash 上 $\tau$ 损失(数学任务 7.67→7.35)表明误拒仍有优化空间。
未来方向
作者明确提出的方向是探索验证感知训练在显著更大模型上的可扩展性。基于本文成果还可延伸:第一,把验证头的二分类 $v_k$ 升级为对接受概率 $\min(1,p_k/\hat{p}_k)$ 的回归或分布预测,直接建模期望、消除单次采样噪声;第二,将验证头用于推理期自适应草稿长度(当前只做固定阈值早退),与 EAGLE-2 式动态草稿树结合,用 $\hat{v}_k$ 剪枝树分支;第三,系统研究训练语料解码温度与部署温度的匹配,例如混合温度语料或在线收集真实拒绝日志做持续学习;第四,把'锚定首个失败点'的自适应加权迁移到 Medusa/Hydra 等多头并行草稿与树形验证中,检验原则的普适性;第五,与 PARD-2 的置信度加权、D-PACE 的可微代理做受控对比,厘清观测锚点与置信度代理各自的适用边界;第六,把验证头用作草稿质量的在线监控信号,在服务端动态调度草稿长度与回退策略。
复现评估
复现条件较好。代码承诺开源于 https://github.com/naver-ai/VAT(论文标注 will be available,动手前需确认仓库已就绪);两个基线 EAGLE-3 与 DFlash 均有公开代码(SafeAILab/EAGLE 与 z-lab/dflash),评测使用 Hugging Face Transformers 的公开评测代码。训练配置交代完整:Perfectblend 用户提示配目标模型贪心生成的回复,训练 3 个 epoch,DFlash 权重 $\gamma=7$,软硬标签混合、$\beta=1.0$,其余超参沿用原论文设置,无需额外搜索;训练与评测均在 NVIDIA A100 80GB、bf16 下进行。复现难点有三:需要为每个目标模型生成回复语料(贪心解码的推理开销不小);模拟验证的采样细节(如是否随温度调整)需读代码确认;全套实验含 3 个目标模型 × 2 个基线 × 2 种温度 × 8 个基准,完整复现所有消融成本较高,但只验证主结果(单模型单基线)可在数个 GPU 日内完成。
论文图表
主结果表:三个目标模型(Qwen3-4B/8B、LLaMA-3.1-8B)× 两种基线(EAGLE-3、DFlash)及其 VAT 版本 × 八个基准(GSM8K、MATH-500、AIME25、HumanEval、MBPP、LCB、MT-Bench、Alpaca),温度 0 与 1 两种设置,报告加速比与 $\tau$。所有组合上 VAT 均为最优,如 Qwen3-4B DFlash 4.54×/5.73→4.81×/6.08,Qwen3-8B DFlash 4.47×/5.51→4.86×/6.14。
论文所有核心声明($\tau$ 最高 +11.4%、加速最高 +8.7%、跨模型与跨范式的一致性)都来自这张表,是评估方法有效性的第一证据。