Research map / Reinforcement learning / RL from feedback
Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Trajectory Balance with Asynchrony (TBA), an asynchronous distributed RL framework that trains LLMs with the off-policy VarGrad trajectory balance objective on replay-buffer data from parallel searcher nodes, reporting speedups and competitive or better accuracy on GSM8K, TL;DR preference tuning, and automated red-teaming.
Problem
On-policy RL used for LLM post-training is bottlenecked by sequential generation and cannot efficiently reuse asynchronously generated replay-buffer data, limiting wall-clock efficiency and scaling of data generation.
Approach
Replace the on-policy policy-gradient objective with an off-policy trajectory balance (VarGrad) loss whose KL-regularized reward and reference-policy reset make it robust to stale data, and run it with decoupled searcher and trainer nodes that sample from a global replay buffer using recency- and reward-based prioritization.
Key results
- On GSM8K, TBA trained in about 82 minutes on 4xA100 reaches 54.6% mean accuracy (700-step run), versus a 40.3% SFT starting point, and speeds up training nearly 50x relative to VinePPO's reported training time.
- On TL;DR with Pythia 410M, TBA reaches a 0.86 win rate at ~15 steps off-policy, above asynchronous Online DPO (0.82), Async PPO (0.76) and Proximal RLOO (0.77), and achieves about 5x speedup over speed-optimized async DPO.
- On GPT-2 red-teaming, TBA reaches the diversity-toxicity Pareto frontier with about 7x speedup over the synchronous GFlowNet baseline for GPT-2 and Llama 3.2 1B.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TBA instead of Online DPO at approximately 15 async steps | Pythia 410M TL;DR preference fine-tuning · TL;DR (Pythia 410M), ≈ 15 steps (async) | Win rate: 0.82 -> 0.86 | Helped |
| Use TBA instead of (Async) PPO at approximately 15 async steps | Pythia 410M TL;DR preference fine-tuning · TL;DR (Pythia 410M), ≈ 15 steps (async) | Win rate: 0.76 -> 0.86 | Helped |
| Use TBA instead of (Proximal) RLOO at approximately 15 async steps | Pythia 410M TL;DR preference fine-tuning · TL;DR (Pythia 410M), ≈ 15 steps (async) | Win rate: 0.77 -> 0.86 | Helped |
| Increase most-on-policy sampling probability from m=0.4 to m=0.5 (ablation) | TBA · TL;DR (Pythia-410M) | Win rate: 0.67 -> 0.82 | Helped |
| Increase most-on-policy sampling probability from m=0.4 to m=0.6 (ablation) | TBA · TL;DR (Pythia-410M) | Win rate: 0.67 -> 0.8 | Helped |
| Increase searcher count (ablation) | TBA · TL;DR preference fine-tuning | Win rate | No clear effect |
| Train for more steps (ablation) | TBA · TL;DR preference fine-tuning | Win rate | Helped |
| Use a less rapid beta decay schedule (ablation) | TBA 2.8B · TL;DR preference fine-tuning | Win rate | Hurt |
| Scale the number of searchers (ablation) | TBA · Automated red-teaming | Attack success rate | Helped |
| Increase maximum replay-buffer size (ablation) | TBA · Llama 3.2 1B automated red-teaming | Attack diversity | Helped |
| Add clipped importance sampling (ablation) | TBA′ · MATH-500 and Countdown | Accuracy | Mixed |
| Remove reference-policy resetting (ablation) | TBA′ · MATH-500 and Countdown | Accuracy | Hurt |
Surprising findings
- Off-policy TBA with stale data (about 15 steps off-policy) outperforms on-policy Online DPO in the PFT setting (win rate 0.86 vs 0.85 on-policy DPO baseline).
- Higher most-on-policy sampling probability m generally gave higher win rate on PFT, yet moderately off-policy data did not significantly hurt performance.
Limitations
- Trajectory balance objective can suffer from high gradient variance at the trajectory level; mitigated by sampling many responses per query (e.g., 20 vs 4 for RLOO).
- Some results are single-config or small-sample; GSM8K result varies across runs (53.9% to 55.8%) and final-step reporting is below peak performance.
- TBA implementation has buffer communication slowdowns as searcher count scales (red-teaming table); PFT runs were 32-bit without DeepSpeed.
Experimental setup
Datasets: GSM8K TL;DR summarization MATH Countdown (generated, 9000 problems) GPT-2 red-teaming setup data
Benchmarks: GSM8K TL;DR summarization win rate MATH-500 Countdown automated red-teaming attack success/diversity
Metrics: pass@1 accuracy win rate vs gold reward model perplexity/KL attack success rate diversity (pairwise cosine distance) wall-clock time
Baselines: VinePPO Online DPO PPO RLOO GRPO Dr. GRPO GFlowNet (one-actor TB) REINFORCE SFT
Builds on: trajectory balance GFlowNet fine-tuning VarGrad RLOO IMPALA
Closest papers by meaning
- DORA: A Scalable Asynchronous Reinforcement Learning System for Language Model Training (NeurIPS 2026 · similarity 0.91)
- AsyncOPD: How Stale Can On-Policy Distillation Be? (NeurIPS 2026 · similarity 0.91)
- Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models (ICLR 2025 · similarity 0.90)
- AREAL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning (NeurIPS 2025 · similarity 0.90)
- Trajectory Bellman Residual Minimization: A Simple Value-Based Method for LLM Reasoning (NeurIPS 2025 · similarity 0.89)
- LLM-Explorer: A Plug-in Reinforcement Learning Policy Exploration Enhancement Driven by Large Language Models (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.