Learning to Draft: Adaptive Speculative Decoding with Reinforcement Learning

📄 arXiv: 2603.01639v1 📥 PDF

作者: Jiebin Zhang, Zhenghan Yu, Liang Wang, Nan Yang, Eugene J. Yu, Zheng Li, Yifan Song, Dawei Zhu, Xingxing Zhang, Furu Wei, Sujian Li

分类: cs.CL

发布日期: 2026-03-02

备注: 22pages, 7 figures

💡 一句话要点

提出Learning to Draft (LTD)，通过强化学习自适应优化推测解码吞吐量。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 推测解码 大型语言模型 强化学习 模型加速 吞吐量优化

📋 核心要点

  1. 现有推测解码方法在草稿和验证阶段的时间分配上存在静态或代理指标优化的问题，忽略了真实时间成本。
  2. LTD通过强化学习，训练两个协同自适应策略，动态协调草稿和验证阶段，直接优化每个周期的吞吐量。
  3. 实验结果表明，LTD在多个LLM和任务上实现了显著的加速，最高可达4.32倍，优于现有方法。

📝 摘要（中文）

推测解码通过使用小型草稿模型生成候选token，并由大型目标模型验证，从而加速大型语言模型(LLM)的推理。该技术的有效性取决于草稿生成和验证之间的时间分配。然而，当前最先进的方法依赖于静态时间分配，而最近的动态方法优化代理指标（如接受长度），常常忽略真实的时间成本，并将草稿和验证阶段孤立对待。为了解决这些局限性，我们引入了Learning to Draft (LTD)，一种直接优化每个草稿-验证周期吞吐量的新方法。我们将该问题建模为强化学习环境，并训练两个协同自适应策略来动态协调草稿和验证阶段。这鼓励策略相互适应，并显式地最大化解码效率。我们在五个不同的LLM和四个不同的任务上进行了广泛的评估。结果表明，LTD实现了2.24倍至4.32倍的加速比，优于最先进的方法Eagle3高达36.4%。

🔬 方法详解

问题定义：现有推测解码方法，如Eagle3，在草稿模型生成候选token以及目标模型验证这些token的过程中，时间分配策略不够灵活。静态方法无法适应不同模型的特性和任务需求，而动态方法通常优化的是接受token的长度等代理指标，而非直接优化解码速度，导致实际吞吐量提升有限。此外，草稿和验证阶段的优化是孤立进行的，没有充分考虑二者之间的相互影响。

核心思路：LTD的核心思路是将草稿和验证过程视为一个整体，通过强化学习来动态调整草稿模型的生成长度和目标模型的验证策略，从而直接优化解码的吞吐量。通过协同训练两个策略，使它们能够相互适应，从而在保证准确性的前提下，最大化解码速度。

技术框架：LTD将推测解码过程建模为一个强化学习环境。该环境包含两个智能体：草稿智能体和验证智能体。草稿智能体负责决定生成多少个候选token，验证智能体负责决定验证哪些token。环境的状态包括当前已生成的token序列、草稿模型的输出概率分布以及目标模型的输出概率分布。奖励函数被设计为每个草稿-验证周期的吞吐量，即成功验证的token数量除以所花费的时间。两个智能体通过策略梯度方法进行训练，目标是最大化累积奖励。

关键创新：LTD的关键创新在于将推测解码问题建模为一个强化学习问题，并直接优化吞吐量。与现有方法相比，LTD能够动态地调整草稿和验证策略，从而更好地适应不同模型和任务的需求。此外，LTD通过协同训练两个智能体，实现了草稿和验证阶段的联合优化，从而进一步提高了解码效率。

关键设计：草稿智能体和验证智能体都使用神经网络作为策略函数。草稿智能体的输入是当前状态，输出是生成token的数量。验证智能体的输入是当前状态和草稿模型生成的候选token序列，输出是每个token是否需要验证的概率。奖励函数被设计为：reward = (accepted_tokens / time_taken)，其中accepted_tokens是成功验证的token数量，time_taken是完成一个草稿-验证周期所花费的时间。使用Actor-Critic算法进行训练，Actor网络输出动作策略，Critic网络评估当前状态的价值。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LTD在五个不同的LLM（包括LLaMA-7B、LLaMA-13B、LLaMA-33B、LLaMA-65B和OPT-66B）和四个不同的任务上进行了评估，包括文本摘要、机器翻译、代码生成和问答。实验结果表明，LTD实现了2.24倍至4.32倍的加速比，优于最先进的方法Eagle3高达36.4%。尤其是在较大的模型上，LTD的加速效果更为显著。

🎯 应用场景

LTD可广泛应用于各种需要加速LLM推理的场景，例如在线对话系统、机器翻译、文本摘要等。通过提高LLM的推理速度，LTD可以降低部署成本，提升用户体验，并促进LLM在资源受限设备上的应用。未来，该方法可以进一步扩展到其他加速LLM推理的技术中，例如模型压缩和知识蒸馏。

📄 摘要（原文）

Speculative decoding accelerates large language model (LLM) inference by using a small draft model to generate candidate tokens for a larger target model to verify. The efficacy of this technique hinges on the trade-off between the time spent on drafting candidates and verifying them. However, current state-of-the-art methods rely on a static time allocation, while recent dynamic approaches optimize for proxy metrics like acceptance length, often neglecting the true time cost and treating the drafting and verification phases in isolation. To address these limitations, we introduce Learning to Draft (LTD), a novel method that directly optimizes for throughput of each draft-and-verify cycle. We formulate the problem as a reinforcement learning environment and train two co-adaptive policies to dynamically coordinate the draft and verification phases. This encourages the policies to adapt to each other and explicitly maximize decoding efficiency. We conducted extensive evaluations on five diverse LLMs and four distinct tasks. Our results show that LTD achieves speedup ratios ranging from 2.24x to 4.32x, outperforming the state-of-the-art method Eagle3 up to 36.4%.