Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SRT: Accelerating Reinforcement Learning via Speculative Rollout with Tree-Structured Cache

Chi-Chih Chang, Siqi Zhu, Zhichen Zeng, Haibin Lin, Jiaxuan You, Mohamed S. Abdelfattah, Ziheng Jiang, Xuehai Qian

TL;DR

This work tackles the rollout bottleneck in on-policy reinforcement learning for language models by introducing Speculative Rollout with Tree-Structured Cache (SRT). SRT maintains a per-prompt tree-structured cache of past rollouts and uses it to draft speculative continuations, which are then verified against the current policy to preserve the on-policy distribution, ensuring correctness. Cache updates occur online from ongoing rollouts and via run-ahead generation during idle GPU time, providing richer drafts without affecting learning targets. Across PPO, GRPO, DAPO, and multi-turn settings, SRT achieves up to $2.08\times$ speedups in rollout and end-to-end reductions in latency and per-token cost, demonstrating a practical approach to scalable, efficient RL training for large language models.

Abstract

We present Speculative Rollout with Tree-Structured Cache (SRT), a simple, model-free approach to accelerate on-policy reinforcement learning (RL) for language models without sacrificing distributional correctness. SRT exploits the empirical similarity of rollouts for the same prompt across training steps by storing previously generated continuations in a per-prompt tree-structured cache. During generation, the current policy uses this tree as the draft model for performing speculative decoding. To keep the cache fresh and improve draft model quality, SRT updates trees online from ongoing rollouts and proactively performs run-ahead generation during idle GPU bubbles. Integrated into standard RL pipelines (\textit{e.g.}, PPO, GRPO and DAPO) and multi-turn settings, SRT consistently reduces generation and step latency and lowers per-token inference cost, achieving up to 2.08x wall-clock time speedup during rollout.

SRT: Accelerating Reinforcement Learning via Speculative Rollout with Tree-Structured Cache

TL;DR

This work tackles the rollout bottleneck in on-policy reinforcement learning for language models by introducing Speculative Rollout with Tree-Structured Cache (SRT). SRT maintains a per-prompt tree-structured cache of past rollouts and uses it to draft speculative continuations, which are then verified against the current policy to preserve the on-policy distribution, ensuring correctness. Cache updates occur online from ongoing rollouts and via run-ahead generation during idle GPU time, providing richer drafts without affecting learning targets. Across PPO, GRPO, DAPO, and multi-turn settings, SRT achieves up to speedups in rollout and end-to-end reductions in latency and per-token cost, demonstrating a practical approach to scalable, efficient RL training for large language models.

Abstract

We present Speculative Rollout with Tree-Structured Cache (SRT), a simple, model-free approach to accelerate on-policy reinforcement learning (RL) for language models without sacrificing distributional correctness. SRT exploits the empirical similarity of rollouts for the same prompt across training steps by storing previously generated continuations in a per-prompt tree-structured cache. During generation, the current policy uses this tree as the draft model for performing speculative decoding. To keep the cache fresh and improve draft model quality, SRT updates trees online from ongoing rollouts and proactively performs run-ahead generation during idle GPU bubbles. Integrated into standard RL pipelines (\textit{e.g.}, PPO, GRPO and DAPO) and multi-turn settings, SRT consistently reduces generation and step latency and lowers per-token inference cost, achieving up to 2.08x wall-clock time speedup during rollout.
Paper Structure (19 sections, 1 equation, 5 figures, 3 tables)

This paper contains 19 sections, 1 equation, 5 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Motivation
  3. Dominant rollout generation time during RL training.
  4. Imbalanced response length among questions in a batch.
  5. Similarity of rollouts across epochs.
  6. Method
  7. RL training in brief.
  8. Per‑prompt rollout cache using tree-structured cache.
  9. Speculative Rollout with Tree-Structured Cache.
  10. Cache update strategy.
  11. Experiment
  12. End-to-end Results
  13. Effectiveness of Per-prompt Rollout Cache.
  14. Effect of On-the-Fly Updates and Run-Ahead Generation
  15. Conclusion
  16. ...and 4 more sections

Figures (5)

  • Figure 1: Rollout Speedups of SRT Across Different RL Algorithms on Qwen2.5-1.5B
  • Figure 2: (a). Time breakdown across different RL algorithm. (b.) Output length distribution on DAPO-17k dataset. (c.) Example of N-gram overlap for a specific prompt. The overlap is computed by comparing rollouts from the current step against the aggregated N-grams from all prior steps.
  • Figure 3: Given the matched prefix "the cat", we choose its suffix "sit on the mat" with highest score as draft tokens. Leaf nodes show the number of times each suffix appeared in cached rollouts, while non-leaf nodes contain the sum of their children's counts.
  • Figure 4: Illustration of cache maintenance strategy in SRT.
  • Figure 5: Mean accepted tokens analysis of different cache maintenance strategy.