NgramRL:通过轨迹感知的 N-gram 推测解码加速结构化生成的 RL 后训练
NgramRL: Accelerating RL Post-Training for Structured Generation via Rollout-Aware N-gram Speculative Decoding
精选理由
利用历史 rollout 的 n-gram 加速强化学习训练生成,免增模型负担且保持输出分布,显著提升吞吐。
AI 摘要
强化学习(RL)后训练能够提升大语言模型的推理能力,但其成本主要来自自回归轨迹生成,这一过程可能占训练时间的 70% 以上。现有推测解码方法难以适应这一场景:基于模型的草稿生成器会占用稀缺的 GPU 显存,并且可能随着策略演化而变得不再匹配;现有无模型方法则未显式利用跨训练轮次的重复内容。我们观察到,同一提示词的近期生成轨迹可提供高质量、可自我更新的草稿来源。基于这一发现,我们提出 NgramRL,这是一种针对每个提示词、利用有限的近期轨迹历史构建的 n-gram 草稿生成器。NgramRL 无需辅助模型、草稿生成器训练或持久驻留 GPU 的草稿生成器状态,同时通过目标模型验证保持轨迹分布不变。为掩蔽 CPU 缓存查找和调度开销,我们进一步引入一种自适应扇出预取器,用于预测可能的验证结果 (j, b),其中 j 为被接受的草稿长度,b 为目标模型输出的续接词元。因果预算分配器将预取集中于高概率结果,而未覆盖的结果则回退到顺序查找,且不改变最终提交的词元。在涵盖两个模型系列、三个推理与结构化生成工作负载的实验中,在缓存预热后,NgramRL 每次目标模型前向计算可提交 3.60–7.77 个词元,最高分别达到 SuffixDecoding 和提示词查找解码的 1.99 倍和 7.0 倍。在 Spider 上,当候选预算为 16 时,预取器可覆盖 Qwen3-8B 的 96.6% 验证结果,并将实测轨迹生成吞吐量提升至同步 NgramRL 的 1.38 倍,接近其通过重叠执行所能达到的 1.40 倍上限。