The Year in AI Papers
The year/Together AI

Paper 2511.13841

Beat the Long Tail: Distribution-Aware Speculative Decoding for RL Training

Zelei Shao, Vikranth Srivatsa, Sanjana Srivastava, Qingyang Wu, Alpay Ariyak, Xiaoxia Wu, Ameen Patel, Jue Wang, Percy Liang, Tri Dao, Ce Zhang, Yiying Zhang, et al.

Published
Nov 2025
Research lab
Together AI
Citations
9
GitHub
Not linked

01 In brief

Summary

Reinforcement learning (RL) post-training for large language models is bottlenecked by the rollout phase, which accounts for over 70% of training time.

The authors identify a long-tail distribution of rollout lengths, where a few long generations dominate wall-clock time, and note that historical rollouts reveal stable prompt-level patterns across epochs.

They propose DAS, a distribution-aware speculative decoding framework that accelerates RL rollouts without altering model outputs.

DAS uses an adaptive, nonparametric drafter built from recent rollouts via an incrementally maintained suffix tree, and a length-aware speculation policy that allocates more aggressive draft budgets to long, high-latency trajectories.

Experiments on math and code reasoning tasks show DAS reduces rollout time by up to 50% while preserving identical training curves, demonstrating that distribution-aware speculative decoding can significantly accelerate RL post-training without compromising learning quality.

The system is implemented on VeRL and vLLM, and ablations show that per-problem suffix trees outperform global ones, and that a distribution-aware budget beats an unlimited budget by 15%.

02 From the paper

Abstract

Reinforcement learning(RL) post-training has become essential for aligning large language models (LLMs), yet its efficiency is increasingly constrained by the rollout phase, where long trajectories are generated token by token. We identify a major bottleneck:the long-tail distribution of rollout lengths, where a small fraction of long generations dominates wall clock time and a complementary opportunity; the availability of historical rollouts that reveal stable prompt level patterns across training epochs. Motivated by these observations, we propose DAS, a Distribution Aware Speculative decoding framework that accelerates RL rollouts without altering model outputs. DAS integrates two key ideas: an adaptive, nonparametric drafter built from recent rollouts using an incrementally maintained suffix tree, and a length aware speculation policy that allocates more aggressive draft budgets to long trajectories that dominate makespan. This design exploits rollout history to sustain acceptance while balancing base and token level costs during decoding. Experiments on math and code reasoning tasks show that DAS reduces rollout time up to 50% while preserving identical training curves, demonstrating that distribution-aware speculative decoding can significantly accelerate RL post training without compromising learning quality.