用强化学习动态优化大模型推理中的草案生成与验证效率。
Learning to Draft: Adaptive Speculative Decoding with Reinforcement Learning
- 通过强化学习训练双协同策略,动态协调草案生成与验证阶段。
- 在5个大模型上实现2.24倍至4.32倍的加速,最高超越当前最优方案36.4%。
- 适合追求高吞吐量推理的部署场景,尤其适用于资源受限环境。
推测解码通过小型草稿模型为大型目标模型生成候选词以加速大语言模型推理。该技术的效率取决于草稿生成与验证时间之间的权衡。然而,现有最先进方法依赖静态时间分配,而近期动态方法虽优化代理指标如接受长度,却常忽略真实耗时,并将草稿与验证阶段孤立处理。为此,我们提出学习起草(Learning to Draft, LTD),一种直接优化每个草稿-验证周期吞吐量的新方法。我们将问题建模为强化学习环境,训练两个协同自适应策略以动态协调草稿与验证阶段,促使策略相互适应并显式最大化解码效率。我们在五个不同大语言模型和四个任务上进行了广泛评估。结果表明,LTD实现了2.24x至4.32x的速度提升,相比最先进的Eagle3方法最多提升36.4%。
原文摘要 · Abstract (English)
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%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。