Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
作者: Yunxiao Zhao, Changxiao Cai
分类: cs.LG, cs.AI, cs.CL, stat.ML
发布日期: 2026-10-07
💡 一句话要点
通过直接最小化期望解码轮次训练并行推测草稿模型
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 推测解码 语言模型 训练优化 马尔可夫过程 深度学习 模型评估 代码生成 数学推理
📋 核心要点
- 现有的推测解码训练方法未能有效优化全局解码效率，导致解码轮次的起始位置与接受的令牌数量之间的耦合问题。
- 本文提出了一种新的理论框架，通过将推测解码视为马尔可夫奖励过程，定义了期望解码轮次（EDR）目标，以优化草稿模型的训练。
- 在对两个最先进的草稿器进行微调时，使用EDR目标显著提高了平均接受长度，并在九个基准测试中超越了现有训练目标。
📝 摘要（中文）
推测解码通过使用低成本的草稿模型并行提出令牌，从而加速大型语言模型的推理。并行和半自回归草稿器通过在单次前向传递中提出整个块来提高草稿效率，但训练它们面临新挑战：给定位置的草稿分布依赖于解码轮次的起始位置，而轮次的起始位置又依赖于早期轮次接受的令牌数量。现有训练目标通常依赖于忽略跨轮耦合的块局部代理，因此未能直接优化全局解码效率。本文提出了一种理论框架，通过将推测解码表示为马尔可夫奖励过程，开发了期望解码轮次（EDR）目标，该目标通过状态占用权重局部拒绝成本，并精确等于期望解码轮次。与先前的代理目标不同，EDR不引入辅助超参数。我们还推导出支持目标模型回放的无偏随机优化的精确时间差梯度。
🔬 方法详解
问题定义：本文旨在解决现有推测解码训练方法未能优化全局解码效率的问题，特别是草稿分布与解码轮次起始位置之间的耦合关系。
核心思路：通过将推测解码建模为马尔可夫奖励过程，提出期望解码轮次（EDR）目标，以权衡局部拒绝成本并优化训练过程。
技术框架：整体架构包括两个主要模块：草稿模型的训练和评估。训练过程中，使用EDR目标来优化草稿器的性能；评估阶段则通过精确的离线评估器来比较不同草稿器的轮次计数。
关键创新：最重要的创新在于引入EDR目标，该目标不依赖于任何辅助超参数，直接优化解码效率，解决了现有方法的局限性。
关键设计：在训练过程中，采用精确的时间差梯度来支持无偏随机优化，确保从目标模型回放中获得可靠的梯度信息。
🖼️ 关键图片
📊 实验亮点
实验结果表明，使用EDR目标微调的草稿器在平均接受长度上有显著提升，且在九个基准测试中均超越了现有的训练目标，展示了EDR方法在推测解码中的有效性和优势。
🎯 应用场景
该研究的潜在应用场景包括大型语言模型的推理加速，尤其在需要快速生成和验证文本的任务中，如代码生成、数学推理和对话系统。通过优化解码过程，能够显著提升模型的响应速度和准确性，具有重要的实际价值和未来影响。
📄 摘要（原文）
Speculative decoding accelerates large language model inference by using a low-cost draft model to propose tokens that the full-size target model verifies in parallel. Parallel and semi-autoregressive (semi- AR) drafters improve drafting efficiency by proposing an entire block in a single forward pass, but training them raises a new difficulty: the draft distribution for a given position depends on where the decoding round starts, and where rounds start depends on how many tokens earlier rounds accepted. Existing training objectives typically rely on block-local surrogates that ignore this cross-round coupling, and therefore do not directly optimize the global decoding efficiency. In this work, we develop a theoretical framework for training and evaluating these drafters by representing speculative decoding as a Markov reward process. This formulation yields the Expected Decoding Rounds (EDR) objective, which weights local rejection costs by state occupancies and exactly equals the expected number of decoding rounds. Unlike prior surrogate objectives, EDR introduces no auxiliary hyperparameters. We then derive an exact temporal-difference gradient that supports unbiased stochastic optimization from target-model rollouts. The same framework also yields an exact offline evaluator for round counts, enabling paired drafter comparisons on shared target rollouts without running speculative decoding. Finetuning two state-of-the- art drafters, DSpark and DFly, with EDR consistently improves mean accepted length and outperforms existing training objectives across nine benchmarks spanning math reasoning, code generation, and chat.