PaperlayerView exact revision on arXiv

Source-linked AI summary

Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training

Brian Bartoldson, Siddarth Venkatraman, James Diffenderfer, Moksh Jain, Tal Ben-Nun, Seanie Lee, Minsu Kim, Johan Obando-Ceron, Yoshua Bengio, Bhavya Kailkhura

arXiv:2503.18929v2cs.LG

TL;DR

LLM post-training RL is bottlenecked by on-policy methods that do not naturally exploit asynchronously generated replay data. TBA combines distributed asynchronous search with off-policy Trajectory Balance, and across reasoning, preference-tuning, and red-teaming tasks it delivers broad performance and speed gains. The method also uses replay sampling that prioritizes reward and recency, while its trajectory-level objective has high gradient variance and requires multiple responses per query.

  • Problem

    On-policy LLM post-training RL creates sequential data-generation and update bottlenecks, limiting resource utilization and replay-buffer scalability.

  • Method

    TBA combines an off-policy Trajectory Balance objective with distributed asynchronous search and training connected through a replay buffer.

  • Results

    Across mathematical reasoning, preference-tuning, and automated red-teaming, TBA matches or exceeds existing methods and achieves 4× to 50× RL speedups.

  • Takeaways & Limitations

    TBA supports stable asynchronous off-policy RL, diverse replay-buffer exploration, and scalable search for sparse-reward post-training settings.

  • Takeaways & Limitations

    Trajectory Balance can have high gradient variance because it operates at the trajectory level, addressed here by sampling more responses per query.

Abstract

from arXiv · show

Reinforcement learning (RL) is a critical component of large language model (LLM) post-training. However, on-policy algorithms used for post-training are not naturally robust to a diversified content of experience replay buffers, which asynchronous off-policy actors can efficiently populate in parallel to training. We propose efficiently learning on such off-policy data via Trajectory Balance with Asynchrony (TBA), an approach to asynchronous RL for LLMs that leverages the principled off-policy TB objective. On math, preference-tuning, and automated red-teaming tasks, we post-train models ranging from Pythia 410M to Qwen 2.5 7B, finding TBA offers speed and performance boosts over strong baselines like Online DPO and Dr. GRPO. Beyond TBA's performance benefits (high accuracy even as asynchrony grows) and speedups ($4\times$ or more), we show its reward- and recency-prioritizing sampling enable further gains as data generation is scaled. Our code is available at https://github.com/bbartoldson/TBA.

1 Introduction

On-policy RL for LLM post-training creates sequential data-generation and update bottlenecks, limiting resource utilization and potentially the benefits of scaling data generation. TBA addresses this with asynchronous off-policy learning and reports broad performance and speed advantages.

  • On-policy PPO and RLOO require sequential data generation and policy updates, creating bottlenecks that reduce resource utilization.Scaling on-policy data generation may also provide limited benefits.
  • TBA combines asynchronous RL with an off-policy Trajectory Balance objective to decouple data generation from policy updates.The approach is designed to mitigate RL bottlenecks while leveraging compute more efficiently.
  • Across mathematical reasoning, preference-tuning, and automated red-teaming, TBA provides stable off-policy RL, diverse replay-buffer exploration, and scalable search compatibility.These properties support reduced wall-clock time, help prevent mode collapse, and address sparse-reward settings.
  • 4× to 50× speedups are demonstrated for RL across mathematical reasoning, preference-tuning, and automated red-teaming.
  • TBA’s high-quality, fast off-policy post-training is presented as a contribution to more efficient LLM alignment.

2 Related Work

Related work covers on-policy RL for LLM alignment, asynchronous distributed RL, and automated red-teaming. These lines of work motivate separating search from training while addressing the scalability and diversity challenges of adversarial generation.

  • RL fine-tuning of language models: PPO is a default on-policy method for LLM fine-tuning, alongside REINFORCE and variants such as GRPO and VinePPO.These methods form the conventional RL backdrop for LLM alignment.
  • Asynchronous distributed RL: Asynchronous distributed RL separates searchers, learners, and environments across resources, allowing processes to operate with different weights for faster training.This paradigm has also been applied to complex, high-dimensional domains.
  • Asynchronous distributed RL: A3C communicates gradients asynchronously, whereas IMPALA communicates experience trajectories to a central learner.TBA’s distributed design more closely resembles IMPALA’s trajectory-based organization.
  • Automated red-teaming: Automated red-teaming uses adversarial interactions to uncover vulnerabilities, biases, and unintended behaviors before deployment.
  • Automated red-teaming: Standard red-teaming RL can suffer mode collapse and insufficient attack diversity, motivating curiosity bonuses and mutation-based prompt sampling.

3 Preliminaries

The paper formulates KL-regularized LLM RL as probabilistic inference and uses Trajectory Balance to learn from replay-buffer data. TB’s off-policy nature makes it suitable for asynchronous training, while coefficient choices affect stability and accuracy.

  • KL-regularized RL as probabilistic inference: The KL-regularized objective balances reward maximization against divergence from a reference policy, while DPO provides an offline alternative to online RL.PPO and REINFORCE optimize the policy online, whereas DPO trains directly on preference data.
  • KL-regularized RL as probabilistic inference: The optimal policy is proportional to the reference policy multiplied by an exponentiated, temperature-scaled reward.This posterior view connects RL fine-tuning with probabilistic inference.
  • KL-regularized RL as probabilistic inference: On-policy optimization can suffer mode collapse and cannot freely reuse stale replay data because policy updates invalidate its on-policy status.
  • Trajectory Balance: Trajectory Balance trains the language-model policy to sample responses from a reward-defined unnormalized target distribution, using a partition function Z(x).VarGrad replaces a learned Z(x) with a K-sample batch estimate, which also acts as a control variate for gradient variance.
  • Trajectory Balance: TB is off-policy: trajectories may come from any full-support distribution, enabling exploration strategies and replay buffers for asynchronous training.This off-policy property is the basis for using TB with asynchronous LLM post-training.
  • Trajectory Balance: Off-policy TB performance is sensitive to β, with larger values tending toward stability and smaller values toward accuracy improvements.Annealing β or resetting the reference policy can provide both benefits.

4 TBA: Fast, Scalable LLM Post-Training

TBA decouples continuous search and training through distributed SEARCHER and TRAINER nodes connected by a replay buffer. Its off-policy TB updates support asynchronous scaling, while recency- and reward-aware sampling balances freshness, high reward, and diversity.

  • TBA integrates off-policy TB gradients with asynchronous distributed RL so data generation and model updates run independently.This combination improves resource utilization and reduces training duration.
  • Continuous, independently running search and training are synchronized only every k steps, avoiding a bottleneck from enforcing constant off-policyness.TBA retains completions regardless of staleness and can prioritize more recent data.
  • A single TRAINER and one or more SEARCHER nodes collect off-policy trajectories into a shared global replay buffer.In the experiments, each node is one GPU, so 16 GPUs can provide 15 SEARCHER nodes and one TRAINER.
  • SEARCHER nodes generate K responses with delayed local policies, evaluate rewards, and store trajectories with synchronization timestamps for later recency prioritization.Local buffers are periodically merged into the global buffer and local policies receive trainer updates.
  • Asynchrony supports generating more than K samples per query, increasing exploration and sequence-length scalability while updating with K samples.Additional samples help mitigate limited diversity among independent rollouts.
  • Replay sampling alternates between recency prioritization and reward prioritization to balance fresh data, high-reward discovery, and policy diversity.Reward prioritization can favor high-reward sequences, but using it alone risks mode collapse and reduced diversity.

5 Empirical Evaluation

Across mathematical reasoning, preference-tuning, and automated red-teaming, TBA provides strong performance with asynchronous off-policy learning, substantial speedups, and improved scalability as search expands.

  • Overall findings: TBA produces results on or beyond the Pareto frontiers of all three tasks at multiple model scales.The evaluation covers mathematical reasoning, preference-tuning, and automated red-teaming.
  • Does off-policyness hurt performance?: Higher off-policyness generally improves PFT win rate, while reasonably high off-policy data does not significantly degrade performance when recent samples remain included.For m = 0.4, 0.5, and 0.6, win rates are 0.67, 0.82, and 0.8, respectively.
  • Overall findings: TBA matches or exceeds existing methods while using off-policy data asynchronously across the evaluated RL pipelines.The paper compares against on-policy and off-policy baselines including Online DPO, PPO, RLOO, VinePPO, and GFlowNet.
  • Discovery of high-reward samples via scaling search: Scaling searchers consistently improves attack success rate and diversity in automated red-teaming by exploring different regions of the solution space simultaneously.The authors also report evidence that scaling helps preference fine-tuning.
  • Discovery of high-reward samples via scaling search: Retaining more off-policy data can trade attack toxicity for greater attack diversity in Llama 3.2 1B red-teaming experiments.This result comes from scaling the maximum replay-buffer size.

6 Scaling TBA for Larger Models

For larger-model experiments, TBA′ simplifies TBA while preserving its focus on asynchronous, highly off-policy learning. Against Dr. GRPO, TBA performs well, particularly when data is highly off-policy.

  • Scaling TBA for Larger Models: TBA′ removes TBA’s β decay schedule and replay buffer, instead using a reference-policy reset interval ρ and constant β.The reported configuration uses ρ = 50 and β = 0.005.
  • Scaling TBA for Larger Models: TBA′ provides further evidence that TBA can enhance asynchronous RL training of LLMs, particularly with highly off-policy data.The comparison uses Qwen 2.5 7B on MATH and Qwen 3 4B Instruct on Countdown.

7 Discussion

TBA combines an off-policy RL objective with distributed asynchronous search to accelerate LLM post-training. The paper reports performance gains across several tasks while noting risks from scaling RL methods.

  • 7 Discussion: TBA combines an off-policy RL objective with distributed asynchronous search for efficient LLM post-training.The framework decouples searcher and trainer nodes to support distributed training.
  • 7 Discussion: TBA yields significant performance gains on mathematical reasoning, automated red-teaming, and RLHF tasks.
  • 7 Discussion: Scaling RL training strategies carries risks including misuse, reward misspecification, and unintended generalization.The paper calls for careful evaluation and responsible deployment as these methods scale.

A TBA Gradient Analysis

The gradient analysis interprets TBA's trajectory-balance objective and clarifies how its asynchronous off-policy setting changes the usual REINFORCE interpretation. TBA also introduces synchronization and sampling hyperparameters that govern training behavior.

  • A TBA Gradient Analysis: TBA's gradient analysis reproduces the variance-reduced loss and derives the corresponding maximization objective.
  • A TBA Gradient Analysis: The TBA loss is rewritten from minimization to maximization by sign changes that preserve the optimizer's minimizer.Rescaling the gradient does not change the minimizer.
  • A TBA Gradient Analysis: TBA's gradient resembles REINFORCE with a mean-reward baseline and KL-regularized reward, but off-policy sampling breaks that equivalence.The KL term is only an estimate of the on-policy KL divergence when data comes from an off-policy buffer.
  • A TBA Gradient Analysis: The sync period k sets model-buffer synchronization frequency, while m controls sampling of the most recent data.m is the probability of selecting samples added during the latest synchronization.
  • A TBA Gradient Analysis: Implementations build TBA on RLOO for mathematical reasoning and preference tuning, and on Lee et al.'s TB trainer for red-teaming.

B.1 GSM8K Mathematical Reasoning (MR)

The experiments evaluate TBA across mathematical reasoning and preference-tuning settings with controlled starting points and varied training configurations. Reported results include a 54.6% mean GSM8K accuracy, speed and variance considerations, and trade-offs involving training steps, KL, and win rate.

  • B.1 GSM8K Mathematical Reasoning (MR): The reported GSM8K comparisons start from a RhoMath-1B checkpoint with 40.3% baseline test accuracy.All methods use the same SFTed starting point.
  • B.1 GSM8K Mathematical Reasoning (MR): 54.6% mean GSM8K accuracy was obtained across three TBA runs, with individual accuracies of 55.8%, 53.9%, and 54.1%.The shortened run used 700 training steps on 4xA100 GPUs and took 82 minutes.
  • B.1 GSM8K Mathematical Reasoning (MR): TBA's 700-step GSM8K result has 0.6% standard error, exceeding the variance observed in the original 1000-step setup.The authors suggest further variance-reduction work for shorter runs and TBA/RL generally.
  • B.1 GSM8K Mathematical Reasoning (MR): Held-out checkpoint selection could improve reported results by approximately 1%, but the experiments report final-step performance without early stopping.Runs often reached higher performance earlier than at the final step.
  • B.2 TL;DR Preference Fine Tuning (PFT): For preference tuning, more training steps tended to increase win rate without significantly increasing perplexity, while slower beta decay reduced both KL/perplexity and win rate.Changing searcher count showed no significant performance pattern in the reported initial runs.
  • B.2 TL;DR Preference Fine Tuning (PFT): PFT speedups were measured under 32-bit precision without DeepSpeed, and gradient checkpointing limited the 2.8B setting to a 3.8x baseline speedup.The authors leave optimization with lower precision and acceleration packages for future work.

C TBA′ Experiment Details

The TBA′ experiments compare the method with Dr. GRPO on MATH-500 and Countdown using matched hyperparameters. TBA′ reduces to Dr. GRPO under no asynchrony or when beta is zero.

  • C TBA′ Experiment Details: TBA′ is compared with Dr. GRPO on MATH-500 and held-out Countdown problems using the same hyperparameters.The experiments use Qwen 2.5 7B Base for MATH and Qwen 3 4B Instruct 2507 for Countdown.
  • C TBA′ Experiment Details: TBA′ equals Dr. GRPO when rho = 1 with no asynchrony or when beta = 0.

D TBA′ Ablations and Contextualization in RL Landscape

TBA′ ablations show that reference-policy resetting is central to off-policy robustness, while importance-sampling choices provide task-dependent benefits. Comparisons with related KL-regularization and masking approaches contextualize which components matter.

  • TBA′ Components: Reference-policy resetting is critical for TBA′ performance, with an effect on accuracy similar to clipped importance sampling.Importance sampling helps on some tasks but is not essential across all tasks.
  • TBA′ Design: TBA′ uses clipped importance sampling and periodically resets the reference policy, while its reset period and KL coefficient jointly control optimization pressure and stability.Larger reset periods and KL coefficients favor the original reference policy; smaller values permit freer reward optimization but may risk instability.
  • Contextualization: Dr. GRPO is exactly TBA′ with β = 0 and shows lower robustness to off-policy data.This comparison links the KL coefficient directly to the observed robustness difference.
  • Alternative Off-Policy Controls: IcePop masking preserves TBA′ effectiveness, but masking alone does not provide robustness to highly off-policy data.The comparison isolates TBA’s KL-regularization strategy from the gradient-masking choice.
  • Contextualization: Kimi K2’s KL-regularization approach is similar to TBA′ and delivers similar off-policy performance, using τ = 0.005 after a small hyperparameter study.The comparison concerns the K2 approach to KL regularization rather than the full training system.

E TBA GSM8K Ablation Studies

GSM8K ablations identify task-specific sensitivities to off-policyness, synchronization, KL regularization, and sampling breadth. TBA’s training and search scaling also expose speed–quality trade-offs, while its asynchronous approach outperforms the tested asynchronous GRPO setup.

  • GSM8K Hyperparameters: GSM8K performance benefits from larger completion batches and more frequent synchronization, suggesting greater sensitivity to off-policyness than preference-tuning.Smaller Beta Final Values improved performance until instability appeared around 0.003, while WSD scheduling could add stability.
  • GSM8K Hyperparameters: Doubling completions per prompt from K = 20 to K = 40 reduces variance and improves performance without increasing training time.The batch size remains constant by reducing the number of unique prompts per batch.
  • GSM8K Hyperparameters: Table 5 reports that doubling K improves TBA’s GSM8K accuracy and reduces variance.The table summarizes the completion-count ablation.
  • GRPO Baseline and Training Dynamics: Async GRPO did not achieve stable learning inside TBA’s asynchronous setup despite changes to KL coefficient, learning rate, and off-policyness.The tested modifications included raising KL, lowering the learning rate, and decreasing off-policyness.
  • GRPO Baseline and Training Dynamics: 54.51% test accuracy versus 49.05% for synchronous GRPO gives TBA a 5% improvement despite similar training performance levels.The comparison uses the final synchronous GRPO model and a random TBA replicate.
  • TL;DR Scaling Ablations: More training steps tend to improve win rate while increasing perplexity, whereas scaling searcher count has a small and inconsistent effect requiring further investigation.For preference-tuning, more searchers may reduce repeated-prompt sampling by expanding the set of prompts with generated completions; this mechanism is not confirmed.
Loading 2503.18929v2…