Acceptance-Aware Draft Model Training for Speculative Decoding
作者: Tianhua Xia, Mugilan Ganesan, Yifei Feng, Haiyu Wang, Maximilian Egger, Sai Qian Zhang
分类: cs.LG
发布日期: 2026-09-21
💡 一句话要点
提出接受长度感知训练损失以优化推测解码
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 推测解码 接受长度 损失函数 大型语言模型 自然语言处理 模型训练 效率优化
📋 核心要点
- 现有的草稿模型训练方法未能直接优化推测解码中的接受长度，导致推理速度受限。
- 本文提出接受长度感知的训练损失，分别针对贪婪和采样解码设计了EAL损失和WTV损失，以直接优化接受长度。
- 实验结果显示，所提方法在不同模型和任务中均优于KL基准，特别是在采样解码场景下，WTV损失显著提升了性能。
📝 摘要（中文）
推测解码通过使用轻量级草稿模型在单次前向传播中生成多个候选标记，从而加速大型语言模型的推理。然而，现有的草稿模型训练方法主要优化交叉熵或Kullback-Leibler（KL）散度，未能直接优化接受长度。本文提出了接受长度感知训练损失，直接优化推测窗口内的预期接受标记数量。针对贪婪验证，提出了预期接受长度（EAL）损失；针对基于采样的解码，引入了窗口总变差（WTV）损失。实验结果表明，所提损失在不同模型和任务中均能显著提高接受长度，尤其在采样解码下，WTV损失表现尤为突出。
🔬 方法详解
问题定义：本文解决了现有推测解码方法中草稿模型训练未能直接优化接受长度的问题。现有方法主要依赖于交叉熵或KL散度，未能有效提升推理速度。
核心思路：提出接受长度感知的训练损失，直接优化推测窗口内的预期接受标记数量。通过设计EAL损失和WTV损失，分别针对贪婪和采样解码进行优化。
技术框架：整体框架包括草稿模型和目标模型的协同工作。草稿模型生成候选标记，目标模型进行验证，训练过程中引入新的损失函数以优化接受长度。
关键创新：最重要的创新在于提出了针对不同解码模式的损失函数，EAL损失和WTV损失直接优化接受长度，而非仅仅匹配分布。
关键设计：EAL损失显式最大化预期接受长度，WTV损失则优化温度缩放的草稿和目标分布之间的重叠，同时考虑序列接受依赖性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，所提出的损失函数在不同的目标和草稿模型、任务及解码设置下均能显著提高接受长度，相较于KL基准，WTV损失在采样解码中表现出特别强的提升，EAL损失则更适合贪婪验证。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理中的文本生成、对话系统以及机器翻译等。通过优化推测解码的效率，可以显著提升大型语言模型在实际应用中的响应速度和用户体验，具有重要的实际价值和未来影响。
📄 摘要（原文）
Speculative decoding accelerates large language model (LLM) inference by using a lightweight draft model to generate multiple candidate tokens that are verified by the target model in a single forward pass. Its speedup is largely determined by the acceptance length, yet existing draft-model training methods mainly optimize cross-entropy or Kullback-Leibler (KL) divergence as proxies. These objectives encourage distribution matching but do not directly optimize acceptance length, and the acceptance mechanism also differs between greedy and sampling-based decoding. In this work, we propose acceptance-length-aware training losses that directly optimize the expected number of accepted tokens within a speculative window. For greedy verification, we derive an expected accepted length (EAL) loss that explicitly maximizes expected acceptance length. For sampling-based decoding, we introduce a window total variation (WTV) loss that optimizes the overlap between temperature-scaled draft and target distributions while accounting for sequential acceptance dependencies. Both objectives can be further combined with a group-relative reinforcement learning stage (GRPO) using simulated acceptance length as the reward. Experiments across different target and draft models, tasks, and decoding settings show that our losses consistently improve acceptance length over KL-based training. WTV provides particularly strong gains under sampling-based decoding, while EAL better matches greedy verification. These results show that directly optimizing the acceptance objective, with losses tailored to the decoding mode, is more effective than conventional distribution-matching objectives.