Beat the long tail: Distribution-Aware Speculative Decoding for RL Training

强化学习（RL）在大型语言模型（LLMs）的对齐过程中已成为关键环节，但其效率日益受限于rollout阶段——在此阶段中，长序列需逐个token生成。我们发现了一个主要瓶颈：rollout长度呈现长尾分布，即少数极长的生成样本主导了实际运行时间；同时我们也注意到一个潜在机遇：训练过程中积累的历史rollout数据揭示了跨训练周期稳定的、基于提示（prompt）层面的生成模式。受此启发，我们提出了DAS——一种分布感知的投机式解码框架，可在不改变模型输出的前提下加速RL的rollout过程。DAS融合了两个核心思想：一是基于近期rollout数据构建的自适应、非参数化“草稿模型”（drafter），该模型通过持续动态维护的后缀树实现；二是长度感知的投机策略，该策略将更积极的草稿预算分配给那些主导总耗时的长序列生成任务。这一设计利用历史rollout信息，在解码过程中保持较高的采纳率，同时平衡了基础模型与token级计算开销。在数学推理和代码生成任务上的实验表明，DAS最多可减少50%的rollout时间，且训练曲线与原始结果完全一致，证明了分布感知的投机式解码能够显著加速强化学习的后期训练，同时不损害学习质量。

在大语言模型的强化学习后训练中，rollout阶段生成长轨迹时逐token解码效率低下，尤其是由于生成长度呈现长尾分布，少数长序列显著拖慢整体训练速度。该问题随着模型规模增大愈发突出，虽非全新问题，但现有方法未能有效针对长尾分布和历史rollouts中的模式进行优化。

提出DAS（Distribution Aware Speculative decoding）框架，通过构建基于近期rollouts的自适应、非参数化草稿模型（利用增量维护的后缀树），并设计长度感知的推测策略，对可能较长的序列分配更积极的草案预算，从而加速解码而不改变最终输出。其创新在于结合历史生成模式与长度感知调度，实现高效且输出一致的推测解码。

实验在数学与代码推理任务上验证，DAS可减少最多50%的rollout时间，同时保持训练曲线完全一致，证明其加速不影响学习质量。使用真实RL训练轨迹数据，具备实际部署价值。论文未明确提及开源代码，但方法设计依赖于可积累的rollout历史，为后续在线学习与自适应解码提供了新方向，值得进一步探索其在更多任务和模型上的泛化能力。

1. 'Speculative Decoding: Accelerating Large Language Model Inference with Draft Models' (2023) 2. 'Fast Inference from Transformers via Speculative Generation' (ICML 2023) 3. 'Accelerating Large Language Models with Cache-Aware Speculative Execution' (NeurIPS 2023) 4. 'Rollout Caching for Efficient Reinforcement Learning in Language Tasks' (ICLR 2024) 5. 'Prefix-Tuning Based Speculation for Efficient LLM Decoding' (ACL 2024)