用强化学习动态生成草稿树,让大模型推理快3倍以上
RADAR: Accelerate Large Language Model Inference With RL-Based Dynamic Draft Trees
- 将草稿树生成建模为马尔可夫决策过程,用离线强化学习训练预测模型
- 在三个大模型上实现3.17倍至4.82倍的推理加速,显著优于自回归基线
- 适合追求高效推理的部署场景,尤其对延迟敏感的应用有价值
现代大语言模型的推理成本高且速度慢,推测采样成为有效解决方案。然而,推测采样中草稿模型生成候选词的调用次数是预设超参数,缺乏灵活性。为此,我们提出RADAR,一种基于强化学习的动态草稿树推测采样方法。RADAR将草稿树生成过程建模为马尔可夫决策过程(MDP),并采用离线强化学习训练预测模型,实现实时决策是否调用草稿模型,减少冗余计算,进一步加速推理。在三个大模型和四个任务上的评估表明,RADAR相比自回归解码基线实现了3.17倍至4.82倍的加速。代码已开源:https://github.com/minaduki-sora/RADAR。
原文摘要 · Abstract (English)
Inference with modern Large Language Models (LLMs) is expensive and slow, and speculative sampling has emerged as an effective solution to this problem. However, the number of calls to the draft model for generating candidate tokens in speculative sampling is a preset hyperparameter, lacking flexibility. To generate and utilize the candidate tokens more effectively, we propose RADAR, a novel speculative sampling method with RL-based dynamic draft trees. RADAR formulates the draft tree generation process as a Markov Decision Process (MDP) and employs offline reinforcement learning to train a prediction model, which enables real-time decisions on calls to the draft model, reducing redundant computations and further accelerating inference. Evaluations across three LLMs and four tasks show that RADAR achieves a speedup of 3.17x-4.82x over the auto-regressive decoding baseline. The code is available at https://github.com/minaduki-sora/RADAR.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。