Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training
Authors: Brian R. Bartoldson
Siddarth Venkatraman
James Diffenderfer
Moksh Jain
Tal Ben-Nun
Seanie Lee
Minsu Kim
Johan Obando-Ceron
Yoshua Bengio
Bhavya Kailkhura
Venue: NeurIPS 2025 poster ()
Keywords: LLM
post-training
RL
async
off-policy
RLHF
reasoning
Primary area: deep_learning
TL;DR: We improve the speed and performance of LLM post-training via a new asynchronous RL approach, leveraging an off-policy objective, replay buffer, and sampling strategies.

Abstract: 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.

Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training
Brian Bartoldson'
Siddarth Venkatraman
James Diffenderfer'
Moksh Jain 2.3
Tal Ben-Nun'
Seanie Lee Yoshua Bengio 2.3,5 Bhavya Kailkhura'
Minsu Kim
Johan Obando-Ceron
Lawrence Livermore National Laboratory, Mila Quebec Al Institute
Université de Montréal 'KAIST SCIFAR Fellow
{bartoldson diffenderfer2, kaiikhuraa1))1nl gov
(siddarth. venkatraman, moksh. jainnaainnmmmmnm quebec
Abstract
Reinforcement learning (RL) is 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- off-policy actors can efficiently populate in parallel training. We propose efficiently learning on such off off- policy data via Trajectory Balance with Asynchrony (TBA), an approach to asynchronous RL for LLMs that leverages the principled off 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 (4x or more), we show its reward- and axcency-rriortizinn sampling enable further gains data generation is scaled. Our code is available at https: /ttpss//ee /ttp//ggtttt //gtttu...... com obartoldson/TB..
55%
TBA (Ours)
50x faster
GSM8K Test Pass@]
+1.2%
54%
+2.0%
53%
Online DPO
VinePPO
1.6x faster
52%
Online DPO
51%
Sync Methods
50%
GRPO
RLOO
PPO
Async Methods
49%
128
256
1024 Compute Time (Minutes)
512
2048
4096
Figure 1: TBA excels the GSM8K mathematical reasoning task. All plotted points 4xA100
GPUs (or comparable L40S GPUs). DPO and RLOO baselines taken from Noukhovitch et al. [46],
PPO and VinePPO baselines taken from Kazemnejad et al. [29]. The baseline model is the SFTed RhoMath- IB [35] model, which gets 40% accuracy after SFT and before RL. Appendix B has details.
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

Introduction
Post-training through reinforcement learning (RL) is critical for enhancing large language models (LLMs), aligning them with human preferences and improving their reasoning abilities [8, 27J. How- ever, widely used RL algorithms such as Proximal Policy Optimization (PPO) [57] and REINFORCE Leave-One-Oue Out (RLOO) [2] suffer from a fundamental limitation: they are on on-policy, meaning that data generation and policy updates occur sequentially. This dependence creates bottlenecks that reduce resource utilization. Further, benefits of scaling on on-policy data generation may be limited [22].
We introduce Trajectory Balance with Asynchrony (TBA), an asynchronous RL approach designed to efficiently and scalably leverage compute for LLM post-training. TBA uses an aff-policy training objective based on Trajectory Balance (TB) [39], which we hypothesize can improve learning when using async RL to decouple data generation from policy updates, mitigating key RL bottlenecks.
Across mathematical reasoning, preference-tuningg and automated red-teaming tasks, we find TBA broadly offers three advantages over existing LLM post-training approaches: 1) Stable off-policy RL, unlocking asynchrony for massive parallelization and reduced wall-clock times see Figures and 3. 2) Sampling from a diverse replay buffer, improving exploration and preventing mode collapse. 3) Scalable search compatibility that aids sparse reward settings like automated red-teaming.
Our key contributions are summarized as follows:
We introduce TBA, distributed asynchronous RL framework for fast, scalable LLM post-training.
We show that LLMs post-trained with TBA match or exceed performances from existing methods, illustrating TB's ' ability use off-policy data in asynchronous RL.
We demonstrate significant speedups (4x to 50x) for RL across mathematical reasoning, preference-tuning and automated red- red-teaming.
By enabling high-quality and fast off-policy post-training, TBA contributes to scalable and effective LLM alignment, ensuring that large models can be refined more efficiently for real-world deployment.
2 Related Work
RL fine-tuning of language models Reinforcement learning (RI.) has been an integral component for training LLMSs| [14, 48). particular, RL. has become the de facto approach for aligning language models with human preferences [8, 82, 62, 49]. Much of this work relies on Proximal Policy Optimization (PPO) [57], an on-policy RL algorithm that has become a default choice for fine-tuning LLMs due to its strong performance across different setups. Aside from PPO, other on-policy objectives such as REINFORCE [2] and variants like GRPO [59] and VinePPO [29] have also been studied in the context of language models. An alternative to PPO-based fine-tuning is rejection sampling fine-tuning, inspired by the best-of-n inference approach proposed by Nakano et al. [45]. Recent work by Dong et al. [10], Gulcehre et al. [15], and Wang et al. [72] extends this concept by generating candidate responses for each prompt, ranking them with learned reward function, and fine-tuning the model based on the highest-ranking responses. Separately, direct preference learning approaches [53, 3, 64] skip reward modeling entirely and train language models directly optimize responses under a preference model. Hu et al. [23] introduced GFlowNet fine-tuning, leveraging off-policy GFlowNet algorithms for fine-tuning language models, which we build upon here.
Asynchronous distributed RL Distributed RL spreads actors/searchers, learners, and environments across a collection of computing resources. Asynchronous distributed RL does this such that searcher and trainer processes do not necessarily share the same weights, which can significantly improve training speed [44, 73] and facilitate RL in complex, high-dimensional domains [19, 24, 21].
A foundational method in this area is Asynchronous Advantage Actor-Critic (A3C) [42]. In A3C, multiple parallel workers asynchronously interact with the environment and communicate gradients to a central node. Our approach to async distributed RL more closely resembles the Importance- Weighted Actor-Learner Architecture (IMPALA) method [11], which communicates experience trajectories (state, action, and reward tuples) to the central node. As opposed to standard RL environments, where value-based off- off-policy learning objectives enable principled asyne training, language models present a distinct challenge since learning value functions is challenging [16]. Recent work has also leveraged these Asyne RL frameworks within the context of LLM post-training for training on near on-policy data [46, 13]. In contrast to these recent works, our approach uses a principled off- off-policy learning objective, allowing training on any off- off-policy trajectories.
2

---

Automated red-teaming Through adversarial interactions, LLM red- red-teaming clarifies the robust- ness and risks of a target LLM. Automating the generation these adversarial scenarios, automated red-teaming frameworks can help uncover vulnerabilities, biases, and unintended behaviors in models more quickly, enabling preemptive mitigation before deployment [74, 28].
Perez al. [51] proposed training language models using RL to discover prompts that elicit harmful responses from some target LLM. Standard RL approaches, however, are susceptible to mode collapse and fail achieve the attack diversity that is critical for successful red-teaming. Hong et al. [20] introduced a curiosity bonus to encourage generation of diverse red-teaming prompts, whereas Samvelyan et al. [56] proposed sampling an attack prompt from a pool and iteratively mutating the prompt with auxiliary LLMs. Lee et al. [34] proposed using GFlowNet fine-tume ane-tuning followed by MLE smoothing generate diverse, transferable, and effective prompts. Our red-teaming experiments augment their TB objective optimization with our distributed asynchronous framework.
Preliminaries
3.1 KL regularized RL as probabilistic inference
We study the problem of fine-tuning a pretrained language model The with a reward model rig For mathematical reasoning [9] our reward model simply computes a response's 's correctness, while our preference-tuning for alignment [82] and red teaming [51] experiments use reward models optimized with human preference data. Notably, reward maximizing RL with learned reward functions is susceptible to spurious modes of the reward [61, 50, 12], which can result in poor performance and low diversity in responses. This is addressed by constraining the mee-ttned model to be close to the initial model in terms of the KL divergence:
π² = arg max Ex-D (ywna(y 3iDkI (I x) [er((- x))].
(1)
Online RL algorithms such as PPO [57, 62] and REINFORCE [75, 2] can be used to optimize the model (which is the policy), whereas offline objectives such as DPO [53] can be used to train the model directly on the preference data and achieve the same optimal policy asymptotically.
Eq. 1 can be interpreted from a probabilistic perspective as a Bayesian posterior inference prob- lem [32]. The optimal policy for Eq. 1 is given as:
'' x) T x) (xpp8-"rr((;;""
(2)
Approaches such as Gibbs sampling can be used to produce samples from the optimal policy without any fine-tuning [76]. These MCMC approaches can be very expensive at inference time and intractable for long sequences. Within the probabilistic interpretation, on-policy RL to maximize Eq. is equivalent to amortized variational inference to minimize the reverse KL with respect the posterior density [32]. However, reverse KL optimization is susceptible to mode collapse and requires on-policy samples. The reliance on on on-policy samples can limit the scalability as cannot use replay buffers that can be populated in parallel at scale, since new updates of the policy invalidate the on-policy nature of the older replay buffer examples. In practice, this means that the policy can get stuck suboptimal solutions and stop exploring, or may obtain high reward at the cost of diversity. This motivates us to consider an alternative off-policy amortized variational inference to efficiently leverage scalable computational resources for flexible exploration.
3.2 GFlowNets
Generative Flow Networks [GFlowNets; 4, 5] are a framework for off-policy training of hierarchical generative models to sample proportionally GFlowNets frame probabilistic inference as a sequential decision-making problem, learning a policy to construct the objects (e.g. sequences) by putting together building blocks (e.g. tokens) and optimizing consistency-based objectives. GFlowNet objectives have been used for fine- fine-tuning autoregressive [23, 34] well as discrete diffusion language models [68]. GFlowNets are MaxEnt RL algorithms [81], and notably equivalent to path consistency learning [43] in sequence generation problems [4].
given unnormalized density (reward) function.
To fine-tune a language model to sample from Eq. 2, we can set reward R(y:x) T
x) exp((-"rx(xxxx x)). Following Lee et al. [34], we use the trajectory balance objective [40] for training the language model policy ne, which is defined over a response y
3

---

Cy, x: x) = -log log ((x) (y x)
2
R(
(x) a positive scalar function of the query x, and the response y sequence of tokens. When LTu is minimized, (x) is the partition function of the posterior (i.e. Z(x) =y R(y: Otherwise it serves as a control variate, reducing the variance of TB gradients. Instead of training value network for 2(x), (xx with a K-sample batch estimate [55, 47, 78, 58, 68]. زク=
the VarGrad variant of trajectory balance which replaces learned
Given responses (y((5))
for a query x(1), a batch estimate of (x(1)) can be computed
K
log ((x((0) Σ jml
log ( (x()) log ne (y(()) x) + 1 3r0(x(110)xx(11)) (4)
K
The (detached) estimate Ż can be plugged into Eq. for batch B f(x(1)) y(11))) ) to get
i= B.j=K K
LTB
VarGra
(B:0) =
Σ
STOP-GRAD
(xx + log To To (4,) x)
BK
i=1.==
(5)
- llo ( (1.1)
x x()))
3 1 (7((((1)))x(0))
An important property of the trajectory balance is that it is off-policy. During training, can be sampled from any distribution with full support [4]. This enables the use of various exploration strategies [54, 30] as well as the use of replay buffers [60, 67]. In the context of fine-tuning language models, this off-policy nature of the objective makes it natural choice for asynchronous training.
3.3 Differences from other RL objectives
We compare the gradients of the VarGrad TB variant above and Proximal RLOO, an off-policy variant of RLOO with off-policy robustness stemming from (clipped) importance-sampling (IS) ratio [46].
With K samples, the gradient for Proximal RLOO for sample y(i) is
(x()())/x) A(y (1) V log To (y)))))) Tree (y6)[x)
V Jp-RLOO
(6)
with advantage A(y(0)xx) (x) = r(y(1),,)) + r(y(1) x). and IS clipping omitted for brevity.
The gradient for yτ) using the TB objective JTB corresponding to the loss in Eq. is, up to multiplicative constant, equal to the following (derivation given in Appendix A):
JJB((O) A(y (x) =
log (1( x),
where A(y(i)xx) = [r(y()) -) Blog To xx)
,x,)(y())x)) K E,(r((() x) - Bogg ( [c]
T.er (y )]. That
is, for on-policy data y, TB VarGrud is equivalent to mean- basallba REINFORCE [75] and a KL- divergence-regularizee reward [82]. However, we use TB off-policy data, i.e. not sampled from Tg. where this equivalence does not hold.
In practice, we find that performance is sensitive to the coefficient 3, with larger values tending to promote stability, and smaller values tending to promote accuracy improvements. We find it easier to obtain both benefits via coefficient annealing schedules or resetting of the reference policy [36]. Notably, Kimi k1.5 was RL-tuned with objective that nearly matches TB's 's and reset the reference policy [65]; Kimi k1.5 deviated from Eq. 7 by excluding the average log probability ratio from its control variate. Brantley et al. [7] also function that plays the role of the control variate. Future work may consider adding the (clipped) IS ratio to Eq. 7 to obtain further robustness to off off-policy data, or exploring IS clipping alternatives [80].
related regression objective, learning offline value

---

Policy Sampling
Continuoualy
Every & Iterations
h T44nh
Update searcher weights with updated policy information
Multiple off ptr-pollcyssarchher searchers populatee raplay buffer
Sample Scoring
Policy Training
Scoring of searcher
Replay Buffer
Use samples from raplay buifeer to updete polly
samples on routs
Figure 2: Fast, scalable LLM post-training with TBA. Continuously (solid lines), multiple SEARCHER nodes (left) collect trajectories, while a TRAINER node (right) samples from re- play buffer to train the policy off-policy. Periodically (dashed lines), updated policy weights are sent to SEARCHER nodes, and new trajectories added the TRAINER node's 's buffer. This avoids bottlenecks at any given node, which can be or more GPUs, keeping resource utilization high.
4 TBA: Fast, Scalable LLM Post- Post-Training
TBA integrates the TB gradient of Eq. into an asynchronous distributed RL framework for post- training language models. Asynchronous RL decouples data generation from model updates, enhanc- ing resource utilization and reducing training duration. the same time, the off-policy TB objective promotes efficient learning from the off-policy training data induced by asynchrony.
We test two TBA variants. The first implements TBA from scratch by modifying an RLOO trainer class (e.g. the Hugging Face TRL library [70]) to use Eq. 5, adding a schedule for 3, integrating a replay buffer and sampling approach, and parallelizing training and search to allow async RI. we discuss the details this variant below. The second variant TBA' uses the asynchronous PRIME-RI. codebase [26]. to which we simply add the update rule corresponding to Eq. 7 and reference- poliicy resetting. TBA' has fewer tunable knobs than TBA and supports multi-GPU training processes, allowing us to test TB for async RI. of LLMs in simpler setting and with larger models/contexts.
While TBA' prioritizes ease of use and multi-GPU training, our from-scratch TBA implementation prioritizes speed and search. For example, ensuring the data has a constant level of off-policyness (as PRIME-RL does) can reintroduce the bottleneck that async RL aims to remove, so TBA runs the training and searching processes continuously and independently. only syncing every i steps (the "sync period").. Moreover, every completion, regardless of how stale it is, is accessible in TBA 's replay buffer (e.g. via reward-based sampling). As discussed below, m is TBA's probability of sampling the data that is (relatively) the "most on-policy"; i.e., the data added at the most recent sync step.
Figure 2 visualizes our TBA implementation. It leverages overlapping training and search processes, and buffer hold trajectories. In particular, TBA SEARCHER nodes that collect off-policy data into a shared replay buffer 2 Dginbal- Here, a node is computational resource sufficient perform all the operations needed for training search can be a collection of multiple GPUs or even a subset always 1 GPU: e.g., given 16 GPUs, we 15 SEARCHER nodes and TRAINER node.
single TRAINER node and one more
GPU. In our TBA experiments, a node
Separation of the SEARCHER and TRAINER is highly desirable even in 2 node (e.g. 2 GPU) cluster because LLM policy rollouts are costly sequential decoding procedures, and training requires only parallel likelihood evaluation of entire sequence through single forward pass. Thus, given objective that can tolerate off- off-policy asynchronous online training. massive wall clock time speedups be realized by continuously running training node without pausing for rollout generation. In the following subsections, we provide more details on our from-scratch TBA implementation.
4.1 Scaling data collection with SEARCHER nodes
TBA's SEARCHER nodes each carry local delayed copy The of the TRAINER policy n.. To produce policy rollouts, queries x are sampled from a dataset, and the local policy generates batch of K responses y Te (y (x) that are evaluated with the reward model (y) Like Noukhovitch
5

---

et al. [46], we use vLLM [33] for faster generation. The (x, (x) tuples are stashed in the SEARCHER's 's local replay buffer Diccal We also add to the stored tuple the step of the trainer when syncing last occurred, giving notion of how off off-po the generated data is - this can later be used by the TRAINER to prioritize sampling from more recent generations.
Every & optimization steps, search and training pause to pull each SEARCHER node's local replay buffer Plocal into the global replay buffer Dglobal, and to update the searcher's 's local policy with the trainer's. The global buffer maintains a list of all generated responses and rewards for each query.
key motivation for asynchrony is its compatibility with rollout scaling for enhanced exploration or increased sequence lengths. For example, we generate S > K samples for given query even when only updating the model using K samples per query - mitigate the lack of diversity caused by the fact that K independent model rollouts are not guaranteed to produce K unique sequences. Relatedly, future work could apply simple off- off-policy inference techniques in the SEARCHER nodes such as randomly sampling the softmax temperature, or using alternative decoding techniques like beam search. We expect such approaches to particularly aid solution discovery in sparse reward settings.
4.2 Asynchronous updates with TRAINER
The TRAINER uses off-policy trajectory balance (Eq, 5) to train the policy on the global replay buffer D Dglohal. We sample batch of B queries, each with K responses and corresponding rewards:
{x (f.)) ro (y x x0)) D
Dytobal-
We then compute the loss in Eq. 5 and use it to update the policy. We sample with replacement if fewer than K unique samples exist for given query.
critical design choice is the strategy for sampling from the replay buffer P Pllobal. The most naive approach is uniform sampling over queries, then uniform sampling over samples associated with the selected query, which may not be optimal if high-reward samples are sparse. Reward prioritization can address this by tilting the sampling distribution toward high-reward sequences; however, focusing solely on high-reward data can lead to mode collapse and reduce policy diversity.
To balance between these concerns, alternate between two sampling strategies: one prioritizing recency -ice. iie.. whether the trajectory was added to the buffer in the most recent sync step-an and another prioritizing rewards. When prioritizing rewards, we consider both softmax of the reward value (to encourage sampling high reward responses) and a uniform distribution (to encourage sampling high and low reward responses equally). We randomly switch between prioritizing rewards and recency for each query in a batch, with the fraction of queries allocated to each strategy treated as a tunable hyperparameter m, which study in Section 5.4.
5 Empirical Evaluation
We evaluate the effectiveness of TBA in three common LLM post-training RL pipelines. Post- training RL for enhancing LLM capabilities- prraabliiess-pprtiuullff particularly for agentic and reasoning tasks-is a critical but nascent area where baseline approaches require many hours or days. These conventional methods often rely on costly-to-generate on-policy data and thus inefficiently leverage available computing resources. Notably, this inefficiency can become particularly harmful when scaling to larger distributed systems, which may be a necessity for domains with sparse rewards that demand increased sampling. Broadly, we find TBA is a highly-performant, fast, and scalable solution.
5.1 Tasks
Mathematical reasoning (MR): We study the GSM8K task which consists of grade-schoo school level math problems and a binary reward based on exact match for the correct final answer [9]. We adopt the setup from Kazemnejad al. [29], Noukhovitch et al. [46], using an SFTed RhoMath-1B [35] model as a base for RL post-training
Preference fine-tuning (PFT): We consider the task of fine-tuning a language model with reward function learned from human preference data. Specifically, we study the TL;DR summarization task where the goal is to write short summaries for reddit posts [62]. Following Noukhovitch et al. [46], we consider Pythia [6] as the base model for the policy and the reward models, using the SFTed versions used by Noukhovitch et al. [46].

---

0.975
3.8x faster
TL;DR Gold Win Win-Rate
0.950
Parameters
0.925
5.3x faster
410M
0.900
1B
2.8B
0.875
Online DPO -Sync
0.850
Online DPO Async
5.1x faster
TBA (Ours)
0.825
0
50
100
200 Compute Time (minutes)
150
250
300
TBA scales search and improves RL efficiency on the PFT summarization task. All plotted points use 4xA100 GPUs, but TBA allocates GPUs to search, and Online DPO allocates GPU to search. TBA produces large-scale off-policy data that its trajectory balance objective can leverage, creating massive efficiency benefits. Online DPO baselines taken from Noukhovitch et al. [46]. Dashed and solid lines use 256 and 425 updates, respectively, Appendix B has details.
Figure 3:
Red-teaming (RT): We investigate automated red-teaming, another critical step for LLM post- training. The goal to discover prompts that elicit harmful responses from target model, as measured by toxicity classifier. We follow the setup from Lee et al. [34], applying the same models: our smaller-scale experiments use GPT-2 [52] an attacker model, GPT-2 Instruction-tunee as victim model, and RoBERTa-baase toxicity classifier [69]; our larger-scale experiments use Llama- 3.2-1B [41] as an attacker model, Llama-3
as victim model, and LlamaGuard-334 [41] for measuring toxicity, averaging over multiple responses from the victim model.
For all tasks, study solely the RL post-training components of baselines' workflows. In particular, we start with SFTed checkpoints, then perform RL. Notably, Lee et al. [34] also include a maximum likelihood estimation training phase after RL that uses the buffer produced during RL, which boosts performance but is not investigated here. Using these baselines' codes, we follow their setup and hyperparameters, with deviations noted in Appendix B. For MR and PFT, implement TBA by augmenting the RLOO trainer of Noukhovitch et al. [46] with our distributed asynchronous RL framework and the TB objective (Equation trajectory balance trainer of Lee et al. [34] with our distributed asynchronous RL framework.
5). For RT, implement TBA by augmenting the
5.2 Metrics and baselines
MR metrics and baselines. We follow the evaluation setup of Noukhovitch et al. [46], computing the pass@ on the GSM8K [9] test dataset with greedy decoding. Prior work [29, 46] applies RL post-training with the GSM8K training set to the SFTed RhoMath- 1B [35] model, which initially obtains 40.3% accuracy on the test set. We compare TBA post-training to the methods used in these prior works: VinePPO [29], Online-DPO [17], (Async) PPO [57], and (Proximal) RLOO [2]. Additionally, we implement a GRPO [59] baseline (see Appendix D. for details).
PFT metrics and baselines. We follow the evaluation setup of Noukhovitch al. [46], using the win-rate under a 6.7B "gold" reward model [25] as the primary metric. We additionally report approximate KL distance- aprrrxemapt by perplexity to match the evaluation of Noukhovitch et al. (46]-between the learned policy and the reference policy. Following Noukhovitch et al. [46], we compare TBA with: Online-DPO [17], (Async) PPO [57], and (Proximal) RLOO [2].
RT metrics and baselines. We follow Lee al. [34], measuring the attack success rate on 1024 sampled prompts for victim model. We also measure the diversity of these test-time generated prompts by computing the average pairwise cosine distance. We compare against: SFT, PPO [57]
(with the novelty reward from Lee et al. [34]), REINFORCE [75, 63), RLOO [1], Online DPO [53],
GelowNet (one-actor TB) [34].
5.3 TBA redefines effieiency-performance
performance Pareto frontiers
We hypothesize that, by mixing asynchronous RL with the trajectory balance objective for learning from off off-policy data, TBA can (a) reduce resource waste and training time and (b) improve perfor-

---

Pythia 410M
Pythia
Pythia 2.8B
1.00
1.00
1.00
TL:DR Gold Win-Rate (1)
TL.DR Gold Win-Rate (1)
TL,DR Gold Win Ra-RRtee
0.95
performance
0.95
0.90
0.90
0.90
0.85
0.85
0.85
TBA (Ours) Online DPO
0.80
0.80
0.80
1.12
1.14
1.09
1.10
1.11
1.10
1.12
1.14
Perplexity or Approximate KL
Perplexity Approximate KL (4)
Perplexity Approximate KL
Figure 4: TBA defines a new KL vs. win-rate Pareto frontier on the PFT summarization task. The baseline "Online DPO" frontier is created by increasing the degree of off-policyness, starting from on- on-policy Online DPO, results from [46]. The TBA frontier is created by altering the training steps, searcher count, and KL annealing schedule as described in Appendix B.
Table 1: TBA outperforms off-policy and on-policy baselines the PFT task (Pythia 410M). Baselines from Noukhovitch al. [46]. The final block uses 16 steps off-policy for baselines, and sync period 10 for TBA, which corresponds to use of data that is 15 steps off- aff-policy on average.
Method
Steps/Syne
Metric
Online DPO
(Asyne) PPO
(Proximal) RLOO
TBA
Perplexity/KL Win Rate 1
1.13
1.14
1.13
(on-policy)
0.85
0.86
0.82
Perplexity/KL
1.13
1.14
1.13
2 (async)
Win Rate t
0.85
0.84
0.83
Perplexity/KL.
1.12
1.10
1.10
1.13
15 (asyne)
Win Rate t
0.82
0.76
0.77
0.86
mance via scaled generation and ingestion of diverse responses. To test this, we primarily consider compute-man matched settings where all methods have access to the same number of GPUs, though we also evaluate TBA with resource scaling. Even in our compute-matched experiments, TBA generates responses relatively rapidly by running asynchronous search on all but one of the available GPUs. Training happens quickly with TBA because it isn't bottlenecked by on-policy generation, and TBA's rapid response generation ensures training on diverse and previously unseen (though off-policy) data.
When tested on established RL problems, an alternative possibility is that TBA will underperform due its departures from conventional approaches: it is asynchronous, off-policy, reliant the less-common trajectory balance objective, and makes updates using many responses per query.' Indeed, Noukhovitch et al. [46] contemporaneously suggests potential limitations to asynchronous, off-policy RL for LLMs, finding that increasing off-policynee can harm performance metrics like win-rate exacerbate policy deviations from the reference model. Further, Hou al. [22] found limited benefits to scaling response generation from to (or 16) responses when using PPO.
We study this question by computing Pareto frontiers for MR (Figure 1), for PFT (Figures and and Table 1), and for RT (Figure 5 and Table 6). See Appendix for experimental details. Notably, TBA produces results on or beyond the Pareto frontiers of all three tasks at multiple model scales, consistent with our hypothesis that TBA can efficiently train LLMs off off-policy data.
Speed is vastly improved with TBA training, which proceeds entirely asynchronously without training- bound or generation bound processes the only non-tu con-training time occurs briefly every k steps. In the compute-mated matched MR experiments (Figure 1), TBA speeds up the training of the only method with comparable performance (VinePPO) by nearly 50x, while improving accuracy by 1.8% and speed
We compute the TBA loss with more samples per query than analogous approaches like RLOO (c.g.. 20 samples versus 4) to reduce the variance of the gradient of the TB objective estimate (Equation 5) that we optimize. We use fewer queries per batch to keep the batch size small despite this response scaling.

---

100%
98%
% Toxic
0.50
of Toxic Prompts (T)
97%
Cosine Distance
Cosine Distance (T)
80%
of Toxic Prompts
0.48
Training Approach
96%
60%
ICL
0.46
95%
x SFT
40%
PPO Novelty
0.44
94%
REINFORCE
0.42
20%
GFlowNet
93%
TBA (Ours)
0.40
0% 0.0
0.2
0.4
0.6
0.8
1.0
92%
2
2
2 Number Searchers
2
2
Cosine Distance
Figure 5: TBA reaches the RT diversity-toxicity Pareto frontier and improves as search is scaled. (Left) On the GPT-2 automated red-teaming task of Lee et al. [34]. TBA produces results on the diversity vs. toxicity Pareto frontier in less training time. Baselines taken from Lee et [34]. (Right) Each searcher uses one V100 GPU for generating attacks. We report means and standard errors from multiple runs of the automated red-teaming task with GPT-2 at each searcher/GPU count.
by 1.5x relative to the speed-optimized asynchronous DPO baseline [46]. In the compute-matched PFT experiments (Figure 3), TBA produces 5x speedups over speed-optimized asynchronous DPO baselines [46]. and even outperforms baselines that train with on-policy data (Table 1). In the non-compute-matched automated red-teaming experiments (Table 6). TBA gives 7x speedups for GPT-2 and Llama 3.2 IB compared to the non-distributed, synchronous GFlowNet baseline [34]. These results suggest that TBA is an effective, parallel, and scalable search framework for distributed learning, offering substantial speed-ups while remaining competitive with leading approaches.
5.4 Does off off-policyness hurt performance?
The results in the prior section are perhaps surprising given evidence of off-policyness's harmfulness to RL post-training [46]. Thus, we investigate the effect off off-policyness with TBA by studying its related hyperparameters, which first review here. The fraction controls how often sampling attempts to approximate an "on-policy" distribution. When m = 1, training occurs exclusively with samples added in the most recent sync step, while m 0 corresponds to selecting data without regard for how recently was generated (e.g., using reward weighting). Additionally, recall that parameters and buffer data are shared between searchers and trainers every k training steps: since TBA runs exploration and training parallel, TBA trains off- off-policy even when m and l 1. Specifically, with probability m, TBA selects a sample from the central replay buffer that is at most 2k updates off-policy. With probability point, which can be as off-po as the number of training steps.
m. TBA samples data produced by the model at any
To understand the effect of increasing off-policyness on TBA, we test the effect of modifying m for the MR (see Figure 7) and PFT tasks. For PFT, we train Pythia-410M on the TL:DR dataset with three values of m (0.4,0 0.5, 0.6) keeping all other parameters constant. We found that win rate fluctuated, with attaining higher win rates of 0.82 and 0.8, respectively. These results show that higher values of m generally lead to higher win rate, reinforeing the idea that on- nn-policy updates in fact the most effective. However, for reasonably high values of m, incorporating more off-policy data does not significantly degrade performance and, in cases, may even provides benefits. Regardless, findings for both MR and PFT further support the idea that we can perform massively distributed training that works well with off policy updates, long recent samples thrown into the mix.
0.4 corresponding
the lowest win rate (0.67), and m = 0.5 and m 0.6
Importantly, the choice of reinforcement learning (RL) algorithm is crucial in determining how effectively we can leverage off- off-policy updates. TBA is a fully off off-policy compatible algorithm, and shown in Figure 4, it significantly outperforms asynchronous Online DPO, on-policy setting. Noukhovitch et al. [46] identified Online DPO as the best best-performing off ofo-cylicy algorithm in their experiments (see Table 1), making it notable that TBA improves upon this.
in the latter's most

---

5.5 Discovery of high-reward samples via scaling search
Beyond improvements at given level of compute, we also observe improvements when we scale the amount of total compute, showing TBA's promise for large-scale distributed RL. In our asynchronous setup, we find that adding more searchers consistently improves the attack success rate and diversity for RT (see Figure 5, right).
This improvement likely stems from having more searchers exploring different regions of the solution space simultaneously. enabling more effective discovery of high-reward samples. Moreover, asyn- chronous updates introduce opportunities for beneficial randomness, and thus potential expansion of the search coverage in the combinatorial space of language. Interestingly, we also find some evidence for scaling's helpfulness in PFT (see Figure 9).
Relatedly, in Table 6, show maximum buffer size to retain more oof-po data. This RT experiment uses Llama 3.2 IB models.
trade attack toxicity for greater attack diversity by scaling the
Scaling TBA for Larger Models
As described in Section 4, TBA' allows to test a simplified version of TBA and to train with larger models/contexts. Relative TBA. TBA' has fewer knobs tune: does not have decay schedule for 3. and PRIME-RL does not sample from buffer, using constant level of off off-policyness instead. TBA' configuration entails reference policy reset interval p. and constant B3. We use p 50, g = 0.005. In Figure 6, we find further evidence that the TB objective can enhance asyne RI. training of LLMs, particularly when data is highly off-policy. See Table 4 for experimental hyperparameters.
2 Steps Off-Policy
Steps Off-Policy
0.75
WW
Math-500 Accuracy
0.70
0.60
0.55
TBA
Dr. GRPO
Dr. GRPO
50
100
150
200
250
300
350
400
50
100
150
200
250
300
350
400
Training
Training Step
Figure 6: TBA' vs. Dr. GRPO. When training Qwen 2.5 7B base [66] the MATH dataset [18], we find that TBA performs well relative to state-of- -hrt--rrt methods, especially in highly off policy settings (right).
7 Discussion
In this work, we introduced TBA, a novel post-training method for large language models (LLMs) that combined an off off-policy reinforcement learning (RL) objective with distributed asynchronous search. By decoupling searcher and trainer nodes, our approach enabled efficient distributed training and avoided bottlenecks. leading to significant performance gains in post-training tasks such as mathematical reasoning, automated red-teaming, and RLHF, We expect that highly parallelizable and performant framework for RL post-training can be extended to other valuable tasks, including seff-improvement training [77, 23, 15), search-based reasoning [71], as well the emerging paradigm of training LLMs to reason with RL [16].
Limitations
The trajectory balance objective can suffer from high gradient variance as it operates on the trajectory level. We addressed this by sampling more responses per query. Future work can leverage learning partial energy functions [38, 79] to balance bias and variance during policy updates.
Broader Impact Improving RL training strategies for LLMs enhance their usefulness across domains but also carries risks associated with misuse, reward misspecification, and unintended generalization. Careful evaluation and responsible deployment are essential as these methods scale.
10

---

Acknowledgments and Disclosure of Funding
We thank Michael Noukhovitch and Haizhong Zheng for helpful feedback.
The authors acknowledge funding from CIFAR, NSERC, IVADO, and Samsung. MJ is supported by FRQNT Doctoral Fellowship fettts:://dole (tttps://doi.oor/11... wtw.:///00060077366666
Prepared by LLNL under Contract DE-AC52-07NA22744 and supported by the LLNL-LDRD Program under Project No. 24-ERD-058 (LLNL-CONE-200022)
This manuscript has been authored by Lawrence Livermore National Security, L.L.C under Contract No. DE-AC52-07NA27344 with the U.S. Department of Energy. The United States Government retains, and the publisher, by accepting the article for publication, acknowledges that the United States Government retains non-exclusive, puid-up, irrevocable, world-wide license to publish or reproduce the published form of this manuscript, allow others do so, for United States Government purposes.
This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office of Science User Facility using NERSC awards ASCR-ERCAPO032802 and ASCR-ERCCE ASSCREERCAP0032112
References
[1] Sina Ahmadi and Aso Mahmudi. Revisiting and amending Central Kurdish data on UniMorph 4.0. In Proceedings of the 20th SIGMORPHON workshop on Computational Research in Phonetics, Phonology, and Morphology, pages 38-48, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18633/v1/0000 / aclanthon aclanthology org/2023
URL https://
sigmorphon- hrg.:grrrrr.0......rorroo....oo. 1..
[2] Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzich Fadace, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. Back basics: Revisiting reinforce style optimization for learning from human feedback in lims. arXiv preprint arXiv:2402. 14740, 2024.
[3] Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. A general theoretical paradigm to understand learning from human preferences. International Conference on Artificial Intelligence and Statistics (AISTATS), 2024.
[4] Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
[5] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, 24(210):1 2023.
[6] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' OrBiien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning (ICML), 2023.
[7] Kianté Brantley, Mingyu Chen, Zhaolin Gao, Jason D Lee, Wen Sun, Wenhao Zhan, and Xuezhou Zhang. Accelerating rl for Ilm reasoning with optimal advantage regression. arXiv preprint arXiv:2505.20686. 2025.
[8] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. @ @ariclechrissiann20777 7dep. title=Deep reinforcement learning from human preferences, thor-Christiano, Paul F and Leike, Jan and Brown, Tom and Martic, Miljan and Legg, Shane and Amodei, Dario, Advane=sAvances in neural information processing systems, volume=30. year=2017 Advances in neural information processing systems, 30, 2017.
[9] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jenyy Tworek, Jacob Hilton, Reicchro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
11

---

[10] Hanze Dong. Wei Xiong. Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint arXiv:2304.0676, 2023.
[11] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International Conference on Machine Learning (ICML), 2018.
[12] Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning (ICML), 2023.
[13] Jonas Gehring, Kunhao Zheng, Jade Copet, Vegard Mella, Quentin Carbonneaux, Taco Cohen, and Gabriel Synnaeve. Rlef: Grounding code Ilms in execution feedback with reinforcement learning. arXiv preprint arXiv:2410.0200 2024.
[14] Google Gemini Team. Gemini: A family of highly capable multimodal models, 2024. URL bttps://arxils org/ abs/2312 11805.
[15] Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, et al. Reinforced self-training (rest) for language modeling, arXiv preprint arXiv:2308.0899: 2023.
[16] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arviv preprint arXiv:2501.12944: 2025.
[17] Shangmin Guo, Biao Zhang, Tianlin Liu, Tianqi Liu, Misha Khalman, Felipe Llinares, Alexan- dre Rame, Thomas Mesnard, Yao Zhao, Bilal Piot, et al. Direct language model alignment from online AI feedback. arXiv preprint arXiv:2402.04792, 2024.
[18] Dan Hendryeks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021
[19] Matteo Hessel, Manuel Kroiss, Aidan Clark, Iurii Kemaev, John Quan, Thomas Keck, Fabio Viola, and Hado van Hasselt. Podracer architectures for scalable reinforcement learning. arXiv preprint 210.::2004.06272.. 2021.
[20] Zhang-Wei Hong. Idan Shenfeld, Tsun-Hsuua Wang. Yung-Sung Chuang, Aldo Pareja, James R. Glass, Akash Srivastava, and Pulkit Agrawal. Curiosity-driven red-teaming for large language models. In International Conference on Learning Representations (ICLR), 2024.
[21] Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado Van Hasselt, and David Silver. Distributed prioritized experience replay. arXiv preprint arXiv:1803.00933 2018
[22] Zhenyu Hou, Pengfan Du, Yilin Niu, Zhengxiao Du, Aohan Zeng, Xiao Liu, Minlie Huang, Hongning Wang. Jie Tang, and Yuxiao Dong. Does rihf scale? exploring the impacts from data, model, and method. arXiv preprint 222.60000 2024.
[23] Edward Hu, Moksh Jain, Eric Elmoznino. Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, and Nikolay Malkin. Amortizing intractable inference in large language models. In International Conference on Learning Representations (ICLR), 2024.
[24] Shengyi Huang, Jiayi Weng, Rujikorn Charakorn, Min Lin, Zhongwen Xu, and Santiago Ontañón. Cleanba: A reproducible and efficient distributed reinforcement learning platform. In International Conference on Learning Representations (ICLR), 2023.
[25] Shengyi Huang, Michael Noukhovitch, Arian Hosseini, Kashif Rasul, Weixun Wang, and Lewis Tunstall. The n+ implementation details of rlhf with ppo: A case study on tl: dr summarization. arXiv preprint arXiv:2403. 17031, 2024.
[26] Prime Intellect. Prime-rl, 2025. URL httpps///ttttuo.ooom/rrmmeeeeeee prime-rl.
12

---

[27] Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El- El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Camey, al. Openai ol system card. arXiv preprint arXiv:2412.16720) 2024.
[28] Jiaming Ji, Mickel Liu, Josef Dai, Xuehai Pan, Chi Zhang, Ce Bian, Boyuan Chen, Ruiyang Sun, Yizhou Wang, and Yaodong Yang. Beavertails: Towards improved safety alignment of Ilm via a numaan-reeerence dataset. Neural Information Processing Systems (NeurIPS), 2024.
[29] Amirhossein Kazemnejad, Milad Aghajohari. Eva Portelance, Alessandro Sordoni, Siva Reddy, Aaron Courville, and Nicolas Le Roux. Vineppo: Unlocking rl potential for Ilm reasoning through refined credit assignment, 2024. URL htttts:://aarrie org/ hts//bb/24410 01679.
[30] Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search GFlowNets. International Conference on Learning Representations (ICLR), 2024.
[31] Wouter Kool, Herke Hoof, and Max Welling. Buy 4 REINFORCE samples, get baseline for free!, 2019. URL
[32] Tomasz Korbak, Ethan Perez, and Christopher Buckley. RI with kl penalties is better viewed as bayesian inference. arXiv preprint arXiv:2205.1 1275, 2022.
[33] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large lan- guage model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
[34] Seanie Lee, Minsu Kim, Lynn Cherif, David Dobre, Juho Lee, Sung Ju Hwang, Kenji Kawaguchi, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, et al. Learning diverse attacks on large language models for robust red-teaming and safety tuning. International Conference Learning Representations (ICLR), 2025.
[35] Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, et al. Rho-1: Not all tokens are what you need. arXiv preprint arXiv:2404.078 079665, 2024.
[36] Mingjie Liu, Shizhe Diao, Ximing Lu, Jian Hu, Xin Dong, Yejin Choi, Jan Kautz, and Yi Dong. Prorl: Prolonged reinforcement learning expands reasoning boundaries in large language models. arXiv preprini arXiv:2505. 248644 2025.
[37] Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du Wee Sun Lee, and Min Lin. Understanding rl rl-zero-like training: A critical perspective. arXiv preprint arXiv:2503.2078. 2025.
[38] Kanika Madan, Jarrid Rector- -rookk, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning gflownets from partial episodes for improved convergence and stability, 2023. URL https://arxa org/ abs/ 2209. 12782.
[39] Dan Malkin, Tomasz Limisiewicz, and Gabriel Stanovsky. A balanced data approach for evaluating cross-lingual transfer: Mapping the linguistic blood bank. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1903-4915, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 186538/v12000/ 1etp53//1/2222naaccmmaaaa mainn-main.361. wwg//022..cc...oooo wrtg//02..aaclleeeeeoagm nwacl wwa.m.nn66....... geilanthologg.org/2222..
URL https://
[40] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GelowNets. Neural Information Processing Systems (NeurIPS), 2022.
[41] Meta Al Llama Team. The llama 3 herd of models. August 2024. URL httpstsss///rrx org/ abs/2407. 01:447..177 21783.
[42] Volodymyr Mnih. Asynchronous methods for deep reinforcement learning. arXiv preprint arXiv:1602.0173,, 2016.
13

---

[43] Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans, Bridging the gap between value and policy based reinforcement learning. Advances in neural information processing systems, 30, 2017.
[44] Arun Nair, Praveen Srinivasan, Sam Blackwell, Cagdas Alcicek, Rory Fearon, Alessandro De Maria, Vedavyas Panneershelvam, Mustafa Suleyman, Charles Beattie, Stig Petersen, et al. Massively parallel methods for deep reinforcement learning, arXiv preprint arXiv:1507004290 2015.
[45] Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christo- pher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, et al. Webgpt: Browser-assisted question-answering with human feedback, 2021. URL htlps://urxiv. org/abs/21 2021.
[46] Michael Noukhovitch, Shengyi Huang, Sophie Xhonneux, Arian Hosseini, Rishabh Agarwal, and Aaron Courville. Asynchronous rlhf: Faster and more efficient off-policy rl for language models. In The Thirteenth International Conference Learning Represeniations, ICLR 2025. OpenReview.net, 2025.
[47] Nikolas Nüsken and Lorenz Richter. Solving high-dimensional Hamilton-Jacohi-eelln PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Disserential Equations and Applications, 2(4):48, 2021.
[48] OpenAl. Gpt-4 technical report, 2023.
ggtt-4.
[49] Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022. URL https: arxiv. htt/abb//2000 22t.//225..........
[50] Alexander Pan, Kush Bhatia, and Jacob Steinhardt. The effects of reward misspecification: Map- ping and mitigating misaligned models. International Conference on Learning Representations (ICLR), 2022.
[51] Ethan Perez, Saffron Huang, Francis Song, Trevor Cai, Roman Ring, John Aslanides, Amelia Glaese, Nat McAleese, and Geoffrey Irving. Red teaming language models with language models. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 1444-344,,,4 Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10. 1211663/1122222.mmmmmmaaa 225. URL https:// acl aclanthology org/2022 emnlp-main 225.
[52] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
[53] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Neural Information Processing Systems (NeurIPS), 2023.
[54] Jarrid Rector-Brooks Kanika Madan, Moksh Jain, Maksym Korablyov, Cheng-Hao Liu, Sarath Chandar, Nikolay Malkin, and Yoshua Bengio. Thompson sampling for improved exploration in gflownets. arXiv preprint 2rXiv:2306.176933 2023.
[55] Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco J. R. Ruiz, and Ömer Deniz Aky- ildiz. Vargrad: A low-variance gradient estimator for variational inference. Neural Information Processing Systems (NeruIPS), 2020.
[56] Mikayel Samvelyan, Sharath Chandra Raparthy, Andrei Lupu, Eric Hambro, Aram H. Markosyan, Manish Bhatt, Yuning Mao. Minqi Jiang, Jack Parker-Holder, Jakob Foerster, Tim Rocktaschel, and Roberta Raileanu. Rainbow teaming: Open-ended generation of diverse adversarial prompts. arXiv preprint arXiv:2402.16822 2024.
[57] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. 2017.
14

---

[58] Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector- Brookks Alexandre Adam. Yoshua Bengio, and Nikolay Malkin. Improved off- aff-policy training of diffusion samplers. Neural Information Processing Systems (NewrIPS), 2024.
[59] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300 2024.
[60] Max W Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towards understanding and improving GFlowNet training. International Conference on Machine Learning (ICML), 2023.
[61] Joar Skalse, Nikolaus Howe, Dmitrii Krasheninnikov, and David Krueger. Defining and characterizing reward gaming. Neural Information Processing Systems (NeurIPS), 2022.
[62] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning summarize with human feedback. Neural Information Processing Systems (NewrIPS), 2020.
[63] Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour, Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
[64] Yunhao Tang, Zhaohan Daniel Guo, Zeyu Zheng, Daniele Calandriello, Rémi Munos, Mark Rowland, Pierre Harvey Richemond, Michal Valko, Bernardo Ávila Pires, and Bilal Piot. Generalized preference optimization: A unified approach to offline alignment. arXiv preprint 2.Xiv:2402.05749, 2024.
[65] Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao. Chenzhuang Du, Chonghua Liao, et al. Kimi kl. 5: Scaling reinforcement learning with llms. arXiv preprint arXiv:2501. 12599, 2025.
[66] Qwen Team. Qwen2.5: A party of foundation models, September 2024. URL https:// quenlm. www.gg.........oo...... httploog/wwn......oco/.// 5/.
[67] Nikhil Vemgal, Elaine Lau, and Doina Precup. An empirical study of the effectiveness of using replay buffer mode discovery in gflownets. arXiv preprint arXhv:2307.07674. 2023.
[68] Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, Alexandre Adam, Jarrid Sector-Brooks. Yoshua Bengio, Glen Berseth, and Nikolay Malkin. Amortizing intractable inference in diffusion models for vision, language, and control. Neural Information Processing Systems (NeurIPS), 2024.
[69] Bertie Vidgen, Tristan Thrush, Zeerak Waseem, and Douwe Kiela. Learning from the worst: Dynamically generated datasets improve online hate detection. In ACL, 2021.
[70] Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec. Trl: Transformer reinforce- ment learning. https:// thub. wwt. hww.hhgggggggg/ggggggggg 2020.
[71] Ziyu Wan, Xidong Feng, Muning Wen, Stephen Marcus McAleer, Ying Wen, Weinan Zhang, and Jun Wang. Alphazero-like tree-search can guide large language model decoding and training. In International Conference on Machine Learning (ICML), 2024.
[72] Haoxiang Wang, Yong Lin, Wei Xiong, Rui Yang, Shizhe Diao, Shuang Qiu, Han Zhao, and Tong Zhang. Arithmetic control of Ilms for diverse user preferences: Directional preference alignment with multi-objective rewards. arXiv preprint arXiv:2402. 18571, 2024.
[73] Taiyi Wang, Zhihao Wu, Jianheng Liu, Jianye Hao, Jun Wang, and Kun Shao. Distrl: An asynchronous distributed reinforcement learning framework for on-device control agents. In International Conference on Learning Representations (ICLR), 2025.
[74] Alexander Wei, Nika Haghtalab, and Jacob Steinhardt. Jailbroken: How does LLM safety training fail? Neural Information Processing Systems (NeurIPS), 2023.
15

---

[75] Ronald J Williams. Simple statistical cradient-following algorithms for connectionist reinforce- ment learning, Machine learning, 8:229-256. 1992.
[76] Wei Xiong, Hanze Dong, Chenlu Ye, Ziqi Wang, Han Zhong, Heng Ji, Nan Jiang, and Tong Zhang. Iterative preference learning from human feedback: Bridging theory and practice for rlhf under kl-constraint. International Conference on Machine Learning (ICML), 2024.
[77] Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. Neural Information Processing Systems (NeurIPS), 2022.
[78] David W Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust scheduling with GFlownets. International Conference on Learning Representations (ICLR), 2023.
[79] Zhenru Zhang, Chujie Zheng, Yangzhen Wu, Beichen Zhang, Runji Lin, Bowen Yu, Dayiheng Liu, Jingren Zhou, and Junyang Lin. The lessons of developing process reward models in mathematical reasoning. arXiv preprint arXiv:2501.0730) 2025.
[80] Haizhong Zheng, Jiawei Zhao, and Bedi Chen. Prosperity before collapse: How far can off-policy rl reach with stale data on llms? arXiv preprint arXiv:511000161, 2025.
[81] Brian D) Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entroyy. Carnegie Mellon University, 2010.
[82] Daniel Ziegler, Nisan Stiennon, Jeffrey Wu, Tom Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909008550 2019.
16

---

A TBA Gradient Analysis
Below, we reproduce for convenience the TBA loss LBB the corresponding maximization objective JIB-
VarGrad from Eq. 5, then provide the gradient of
i=B.j=K Σ
VarGrad (B:0) TB
STOP-GRAD|log (xx log TO (y((1))) 1x(0)
BK
1=1,==1
(8)
log Ty (6,)
x(())
We first multiply everything inside the square operation of Eq. 5 by -1, which is equivalent to multiplying the entire equation by and gives
BBjjK Σ
VarGrad LB
(B:
log loog (er((r(( xx) - log (o((( 1x(())
BK
1=1Ø=1
(9)
log ((x(0)
3 3r
The gradient with respect to & then becomes
i=B.j=K
V VarGrad
(B;
Σ
2 log log (y((()))))) (x()) log (y(( x(())
BK
i=1ý=1
(10)
(y - log 2 v log To (fJ) |x('),
and we can turn this into maximization problem instead of a minimization problem by multiplying the right side by -1:
BBBj=KK
V JT
VanGri
(B;0)
Σ
2 log (re((y(is) |x(1) - llog no (y xx())
BK
1=1,=1
(11)
B (o log 2(x(0)) log To (f()) xx(0))
We multiply the right hand side by rescaling the gradient without changing the minimizer. Also,
define KL,00 (i) = t
Σラ= KL (K-) t Σj log (15) x()) - log (
(i,) x()) and define
7(() = 1
k Σソ (م.5) = 1 ΣØ To (b)()) (1)) to rewrite Eq. as log 2 (x(()) = 1 7(( ()
(i.)
k
t= BBjj=K
VJTa(B:0)
Σ
A (1:1) log (T( x()),
(12)
BK
1=1,=1
where (b. (f(( - 3((i. (i.)) KI. ()) In sum, TBA appears update model weights
using a REINFORCE algorithm with mean-reward baseline [75] and a -U-ddivrgeecc--rgulliizzz
reward [82]. However, note that Ki. (i.j) only an estimate of the KL divergence in the on-policy case because off off-policy data is not sampled from by definition, causing loss of equivalence to mean- besan-ba REINFORCE with KL-regularized reward in the asynchronous settings that TBA
operates in. Deviating further, TBA may also reset the reference policy similar to Liu et al. [36].
17

---

B TBA Experiment Details
As discussed Sections and 5.4, TBA introduces new hyperparameters. Most notably, these include (I) the syne period k, which is the number of training steps between two successive mrodel-buffe synchronizations; and (2) the most-on-pont on-policy probability m, which is the probability of sampling the data that is most on-policy i.e., the samples that were added to the buffer during the most recent synchronization. When listing hyperparameter values, we clarify those that are specific TBA and include references to their
the text.
For MR and PFT, we implement TBA by building on the RLOO trainer class of the Hugging Face TRL library [70]. For MR and PFT, we make modifications the baseline hyperparameters as shown in Table and Table 3, respectively.
Our RT implementation of TBA built on the TB trainer used in Lee et al. [34], and we largely utilize their code's default hyperparameters, with differences noted in Appendix B.3.
B.1 GSM8K Mathematical Reasoning (MR)
All baselines and TBA results use the same starting point, a RhoMath-IB [35] model SFTed on GSM8K training data by Kazemnejad al. [29] ""e""eearrreetunerhooo "rrrarreeuun/hh------ff-GGSS IDsfff-Gff--KK Hugging Face. The baseline model achieves 40.3% test accuracy. For the VinePPO baseline, training time is estimated using their reported 380 seconds per training step and 650 steps for GSM8K training [29]. We implement the GRPO baseline using its Hugging Face TRL. trainer, see Appendix D.1 for details.
Table TBA Training Hyperparameters for the GSM8K MR task.
Hyperparameter
Value
Reference
Model
Rho- 1B SFT on GSM8K
Learning Rate
1 x 10-5
Learning Rate Schedule
Warmup Stable Decay
Learning Rate Warmup Steps
50
Learning Rate Stable Steps
450
Learning Rate Decay Steps
500
Generation Temperature
0.7
Max Prompt Token Length Response Length Number of Prompts per Batch Number of Completions per Prompt Batch Size (effective)
512
512
7
20
K in Section 4.1
140
Number of Training Steps
1000
Total Prompts Seen Total Episodes
7000
140000
TBA specific hyperparameters
Beta (KL coefficient) Initial Value
0.012
3 in Equation
Beta Final Value
0.004
3 in Equation 5
Beta Linear Decay Schedule End Step
500
9 in Equation
Number of Samples per Prompt
24
$ in Section 4.1
Most-On-Policy Sampling Probability
0.95
in Section 5.4
Sync Period
2
1: in Section 5.4, Figure 2
Number of Searchers
Searchers Figure 2
Number of Trainers
Trainer in Figure 2
Reward-Baar Based Sampling Prioritization
Uniform
Section 4.2
Initial Completions in Buffer
500
Buffer at Step 0 in Figure 2
In Figure 1, TBA uses the settings in Table with a few modifications to shorten training time by an additional 30%. down to 82 minutes 4xA100 GPUs. In particular, (see Appendix D) that using the hyperparameters in Table 2 led to no improvement in performance for the final 300 steps. Thus, shrank the training duration from 1000 steps to 700 (98000 episodes with batch size 140). Additionally, we made the following modifications to attempt to reduce the variance of the shortened run: 350 stable learning rate steps (down from 450), 0.014 Beta Initial Value
observed in initial testing
18

---

Table 3: TBA Training Hyperparameters for the TL;DR PFT task. For the PFT task, we accelerate the decay Beta. In particular, the Beta Linear Decay Schedule End Step set to be half the number of training steps, but we abruptly end this decay and set Beta to its final value at one eighth the number of training steps (e.g., step 32 for 256 steps). This has the effect of trading off KL/perplexity, which we found to be relatively low with our TBA setup, for win rate.
Hyperparameter
Value
Reference
Model
Pythia SFTed on TL:DR
Learning Rate
3 x 10-6
Learning Rate Schedule
Linear
Generation Temperature
0.7
Max Token Length
1024
Max Prompt Token Length
512
Response Length
128
Number of Prompts per Batch
8
Number of Completions per Prompt Batch Size (effective)
20
K in Section 4.1
160
Number of Training Steps
256
Total Prompts Seen Total Episodes
2048
40960
TBA- speciifii hyperparameters
Beta (KL coefficient) Initial Value
8 in Equation 5
Beta Final Value
0.05
] Equation
Beta Linear Decay Schedule End Step
See caption
]8 in Equation 5
Number of Samples per Prompt
20
Section 4.1
Most-On-Police Sampling Probability
0.5
m in Section 5,4
Sync Period
10
4: in Section 5.4, Figure
Number of Searchers
3
Searchers in Figure 2
Number of Trainers
1
Trainer in Figure 2
Reward-Based Sampling Prioritization
Softmax of Score
Section 4.2
Initial Completions in Buffer
10000
Buffer Step 0 in Figure 2
(up from 0.012). We ran this experiment three times obtaining performances of 55.8%, 53.9%, and 54.1% and reported the mean accuracy 54.6% in Figure
Limitations and future work The 700-step result show in Figure has standard error 0.6%, which is a little more variance than what we observe in the original 1000 step setup (the blue line in the bottom left plot of Figure 7 shows the mean and standard errors for the 1000 step runs). Future work could further explore variance reduction approaches/yyeeppaamteer for shorter runs (and for TBA/RL more generally).
One way to deal with variance is to choose checkpoint based held out validation set, a strategy used to produce the VinePPO result [29]. We do not use this approach but note that our results would likely benefit significantly (<z 1%) from it. In particular, each of our runs tends to produce (at some point during training) a higher performance than the performance at the final step - this expected you consider that the model performance varies around the average value it converges to towards the end of training (e.g., see again the blue line in the bottom left plot of Figure 7). Despite its being lower than the maximum performance achieved at any point during training, we report this final step performance, which is what a user of TBA could expect to obtain without applying techniques like carly-stopping.
B.2
TL;DR Preference Fine Tuning (PFT)
For PFT Figure we use the settings in Table as well as a longer-duration run with 425 updates. For PFT in Figure 4, we create the TBA frontier by modifying the training steps, searcher count, and beta linear decay schedule shown in Table 3. We train for 256. 425, 625, 725, and 825 steps, and search with 2, 3, 4, and 7 searcher nodes. We did not notice a significant pattern in performance when changing searcher count, but we found tendency for higher KL/perplexity values and win-rates with
19

---

more training steps (an expected tradeoff from optimizing the policy further). We noticed that the 2.8B model did not create a wide range of KL/perplexity values with these initial runs, so we also performed runs at that model size (with 825 steps, and with 2 and 4 searchers) using a less rapid beta linear decay schedule (reaching the beta final value at step 140 instead of step 104). This schedule change had the effect of reducing the KL/perplexity and win-rate (another expected tradeoff).
For PFT in Table 1, we use the settings in Table 3, except we train for 625 steps (100000 episodes) because we found use of more steps tended improve win rate without a significant increase in perplexity (approximate KL). Additionally, we only use 2 searchers in this run. See Figure 9 for a depiction of the effects of step count and searcher count.
Limitations and future work All of our PFT results were run in 32-bit precision and without DeepSpeed, which was used by baselines we compared against [46]. Lower precision and such training acceleration packages could further improve the speedups we show. Relatedly, for our 2.8B runs, we used gradient checkpointing to fit more examples into a micro batch, which led to slowdowns this scale (i.e., we only have a 3.8x speedup over the baseline in this setting). We leave the optimization of our framework with appropriate packages to future work. Finally, we used large number (10000) of initial completions the buffer, and future work should confirm that a smaller number (e.g. 1000) works equally well - note that small number worked well for GSM8K.
B.3
Automated Red Teaming (RT)
Unlike our MR and PFT implementations, our RT implementation uses the code of Lee et al. [34] as a baseline and thus does not follow the Hugging Face TRL trainer style. We discuss hyperparameters
in the context of their trajectory balance trainer approach below. We adopt their code largely as is, with the exception that our TBA implementation uses larger replay buffers than Lee et al. [34] to accommodate our scaled search for trajectories. Additionally, unlike Lee al. [34], remove the oldest samples when the buffers fill up as opposed to the lowest reward samples.
For the Llama results in Table 6, our hyperparameter choices largely follow those of Lee et al. [34]. We train for 5000 steps with batch size 128. For temperature sampling, we use low and high of 0.7 and 2.0, respectively. We use a reward schedule horizon of 1000. The language model schedule end is 1.2, and its schedule's horizon is 2000. We prioritize sampling based on reward. We use Beta 0.05. We use sync period (k:) 10. We use 6 searchers. We use most-on-policy probability (m) 0.5 and 0.6 in combination with maximum buffer sizes 150000 and 130000, respectively. By having smaller maximum buffer size and larger m, the latter setting is expected to focus on more recently generated data and prioritize reward over diversity, which is what we observe in Table 6.
For the GPT2 results in Figure 5 and Table 6, our hyperparameter choices again largely follow those of Lee et al. [34]. We train for 5000 steps with batch size 128. We use Beta 0.05. We use sync period (k) 10. We use most-on-policy probability (m 0.5. We cap the maximum buffer size at 100000 in all experiments; we additionally prevent the buffer from growing past its size as of step 4000 to encourage focusing more recent data (larger most on policy probabilities may provide similar effect). We test searcher counts 2, 6, 14, 30, 62, and we use 3 runs per configuration for error bars and averages.
Limitations and future work In Table 6, there a slowdown as the number of searchers scales. We believe this is largely addressed by a newer version of our code that uses more efficient buffer communication, but we have not re-run these results yet to confirm this. In any case, developing more efficient TBA implementations is an interesting direction for future work given TBA's ability to effectively leverage large-scale data generation.
C TBA' Experiment Details
In Figure 6, we compare TBA' to the Dr. GRPO [37] implementation in PRIME-RL [26). The hyperparameters used are given in Table 4.
20

---

Table 4: TBA' Training Hyperparameters for the MATH-500 MR task.
Hyperparameter
Value
Reference
Model
Qwen 2.5 7B Base
Learning Rate
1 x 10-6
Generation Temperature
1.0
Max Prompt Sequence Length
1024
Max Model Length
3072
Number of Prompts per Batch
32
Number of Completions per Prompt
16
K in Section 4.1
Batch Size (effective)
512
Number of Training Steps
400
Total Prompts Seen
12800
Total Episodes
204800
TBA' TBA--specife hyperparameterss
Beta (KL coefficient) Reference policy reset period
0.005
in Equation
50
p in Section
D GSM8K Ablation Studies
We adopted the hyperparameters for our GSM8K result in Figure based on series of trial experiments centered around the hyperparameters shown Table In Figure show the effects of changing what we found to be key hyperparameters. We note the following observations about TBA hyperparameter effects GSM8K.
1. Unlike PFT, it was important for m to be somewhat large for the GSM8K MR task (Figure 7, top left). Similarly, syncing more frequently was beneficial (Figure top right). Together, these results suggest that GSM8K performance is more sensitive to off off-pooicyness than performance on other tasks (e.g.. PFT).
2. We found that the WSD schedule could add stability (Figure 7, bottom right).
3. Using smaller Beta Final Values tended to improve performance (Figure 7, bottom left), but training became unstable around 0.003. This suggests tradeoff between stability and accuracy for GSM8K that is mediated by the KL coefficient B.
4. A batch size 140 not provide significantly better or worse results than larger batch sizes in initial testing, but smaller batch sizes allow for faster training steps, motivating our use 140.
Subsequently, we considered increasing the number of completions per prompt K from 20 to 40. In Table 5, we see that this reduces variance and improves performance, Critically, this result does not require longer training times, simply reduce the number of unique prompts per batch to keep the batch size constant.
Table 5: On GSM8K, doubling K improves TBA's accuracy and variance.
Test Accuracy
Standard Deviation
20
54.61
0.85
40
54.89
0.49
D.1 GRPO Baseline Creation and Training Dynamics Comparison
We implement the GRPO baseline using Hugging Face TRL (version 0.17.0). Relative to the baseline only change the following: Num Train Epochs 8, Gradient Accumulation Steps = 5, and Per Device Train Batch Size 32. For training speed acceleration, we use DeepSpeed with ZeRO Stage 2. We find these settings lead to good GPU memory usage and utilization on the same 4xA100 (80 GB) node that we used for our TBA experiments. Morcover, the batch size (640), total training epochs (8), and number of optimization steps (744) are based on
GRPOConfig hyperparameters,
21

---

Effect of Most-On-Policy Probability
Effect of Synchronization Period
0.55
0.50
0.45
k
2
GSMBK Test Pass@11
0.40
0.95
4
Effect of Beta Final
Effect of WSD Schedule
0.55
0.50
0.45
Beta Final
WSD Schedule
0.004
False
0.40
0.005
True
200
400
600
800
1000
200
400
600
800
1000
Training Step
Figure 7: GSM8K ablation studies. All experiments begin with the base hyperparameters listed in Table 2 and make the depicted modifications, except when studying the synchronization period le in the top right plot (where we use Beta Final Value 0.005 because 0.004 led to instability with k = 4). We report the mean and standard error from runs of each configuration.
settings chosen for RL with VinePPO [29] - the paper that created the SFT baseline that we apply GRPO to in GSM8K experiment (see Figure 1).
To arrive the GRPO result in Figure 1, explored several implementation variations of syn- chronous GRPO. Interestingly, we also tried async GRPO by replacing the loss function used in asyne TBA code with GRPO's. We tested several hyperparameter settings with asyne GRPO inside TBA's async setup raising the KI. coefficient, lowering the learning rate, and decreasing off- uff-policyness but did not find setting that provided stable learning.
The average training data score (GSM8K correctness) by training step for final (synchronous) GRPO implementation and a random TBA replicate are shown in Figure 8. The final model from this GRPO run is used compute the test performance shown in Figure 1, 49.05%. This TBA run achieves 54.51% % test accuracy, a 5% improvement despite the similar training performance levels.
22

---

Mean CSM8K Training Data Corecciess
Melhod
GRPO
0.60
TBA
0.55
0.50
100
200
300
400
500
600
700
Training Step
Figure Training dynamics of GRPO and TBA runs. The TBA performances are heavily smoothed through an exponential moving average (without smoothing, the TBA curve oscillates such that the trend is harder see), whereas the GRPO performances are not smoothed.
Win Rate
Perplexity or Approximate (1)
1.150
825.0
825.0
0.89
1.145
1.140
0.86
Dainnng Steps
725
0.85
Training Steps
725.0
1.12
11:135
0.84
1.130
0.86
0.80
625.0
1.13
1.11
1.125
0.82
11120
425.0
425.0
0.84
0.85
1.12
1.13
-1.115
Scarcher Count
Searcher Count
Figure 9: TL;DR ablation studies. All experiments begin with the base hyperparameters listed in Table 3 then make the depicted modifications. More searcher nodes and more training steps tend to improve win rate the cost of higher perplexity.
E TL;DR Ablation Studies
With TBA, we take many more optimization steps per hour of training than competing methods take. Accordingly, we sought to understand how taking more steps or using more searchers (to reach compute-pario with competing methods) affects performance. As shown in Figure 9, win rate tends to improve with increased compute through changes to these variables, though perplexity suffers expected when over-optimize the policy (in the why scaling the searcher count seems to have an effect similar to step scaling. However, there is notable mechanism through which searcher scaling could have effect: using reduce the probability of selecting a training prompt that's 's been seen before when sampling off- off-policy (because scaling the searchers scales the number of unique prompts with completions added to the
of step scaling). It is not clear
searchers should
23

---

buffer). However, the effect size is small and inconsistent enough to suggest this searcher-scaling trend (in the context of PFT) needs further investigation before it's confirmed.
Red-Teaming Additional Results
Table 6: TBA speeds up the wall-clock time required to reach the Pareto frontier for the red- teaming task. The GFlowNet performances are taken from Lee et al. [34], while the training speeds are computed by us with their code. With the GPT-2 models, TBA performance improves with searcher count. With the Llama models. we trade attack toxicity for attack diversity by scaling the TBA buffer's maximum size from 130,000 (penultimate row) to 150,000 samples (final row), retaining more off- off-policy data.
Attacker / Victim Model
Training Method
Hardware
Time (h)
Speedup
Cosine Distance
% Toxic Prompts
GFlowNet Sync
Vx0000
11.9
1x
0.65
96.6
GPT-2/ / GPT-2 SFT
TBA (Ours)
4x V100
1.7
7x
0.42
94.5
TBA (Ours)
16x V100
2.5
4.8x
0.45
96
TBA (Ours)
64x
2.9
4.1x
0.49
97.6
Llama 3.2 IB Llama 3.1 8B - Instruct
GFlowNet Sync
1xA100
37.4
1x
0.32
100.0
TBA (Ours)
8xA100
5.7
6.6x
0.35
98.1
TBA (Ours)
8XA100
5.7
6.6x
0.37
94.8
"Notably, the effect scaling search is different in the case of RT, where there a single prompt scaling searchers makes more likely that riwher-reward samples are found and trained on. Here, with PFT, scaling search means that more prompts per second have a set of completions generated for them, but any given prompt is not expected to have higher-rewarg samples as result of scaling the searcher count this is an implementation choice and could be changed to gain the effect that scaling searchers has with RT, however.

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We present empirical results which support the claims about the task perfor- mance and computation speedups using our proposed framework Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of the method in Section and the limitations of each of the experiments in Appendix B.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
25

---

Justification:
We do not propose any novel theoretical results in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We disclose all the information required to reproduce the results presented in the paper in Section 5 and Appendix B. Additionally, we also include the code with the submission to aid reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
26

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We provide the code to reproduce the results in the paper and plan to open source it. We use publicly available datasets and models in our experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ httppsccccpppcccccceeceeecccc herbbiccgguiessssse /Caodoeoomubmiic ssi criionPooicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We discuss the experimental details for all the experiments, along with detailed information about the hyperparameters in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report mean and standard error for results in many of the main results presented in the paper (e.g. Figure 7, Figure 5). However, due the extensive computational requirements for fine-tuning language models, and the number of tasks considered, we do not report multiple seeds for every result.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
27

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper clarifies the specific GPU type and quantity used in each experiment, as well as the runtime durations.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The research conducted in this paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impact of the work in Section 7.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
28

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks. mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We focus on the methodology and do not release any trained models or new datasets.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The datasets and models are properly credited in Appendix B and Section 5,
Guidelines:
The answer NA means that the paper does
existing assets. The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
29

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-pane se-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code for experiments in the paper is submitted with basic documentation for reproduction and a valid license.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowd-sourcing experiments.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve study with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
30

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: Our paper presents novel approach for improving the post-training of LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components. pectps://neurti wttp///ees...//eeeeeeeee /
Please refer to our LLM policy (https (https://nat) for what should or should be described.
31


Reviews:
[{"review_id": "Z5x2ypJeMN", "reviewer": "Reviewer_cZos", "summary": "The paper proposes an interesting method to conduct LLM post-training without using policy-gradient style approach. The method is about 4x faster than GRPO (Figure 1's 50x faster claim is massively misleading). The domain is on safety training, which is a very important domain for LLM post-training.", "questions": "I do have a few questions about the method, since I'm only familiar with traditional RL methods.\n\nThe equation 5, if I actually expand $\\log \\hat Z$, then the equation becomes a square loss term to train the reward function: $r_{\\phi}(y, x) - r(y, x)$, but then all the policy $\\pi$ related terms will get cancelled out. \n\nI'm actually not sure how the training is done -- how is $\\theta$ learned? This is probably because I'm not familiar with the Bayesian notation, but understanding this will help me understand the main difference between this objective and GRPO style objective. I'm happy to raise my score after the discussion round!", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper is well-written. The method seems very clear.\n2. The benefits of async RL are well-documented.\n3. Using GFlowNet for RL seems novel (but frankly, I'm not familiar enough with the Bayesian RL literature)\n4. The implementation seems easy enough with existing libraries, which is a huge plus.\n\nWeaknesses:\n1. GFlowNet is already an established method. The paper does not offer innovation on methodology beyond the existing method. However, applying it for LLM post-training is interesting, and the empirical explorations seem warranted.", "was_revised": true, "final_justification": "The rebuttal is excellent and has addressed all my concerns.", "rebuttal": "We thank the reviewer for their detailed and thorough review, and we are grateful for their stated interest in raising their score after the discussion round! Their efforts have helped us improve the revision, and we would be happy to discuss any additional or remaining concerns.\n\n> The method is about 4x faster than GRPO (Figure 1's 50x faster claim is massively misleading). \n\nWe will improve the specificity of our text to avoid misleading. While our results in Figure 1 demonstrate that TBA’s speedups can be as much as 50x, it’s correct that we only find speedups this large when comparing to methods that conduct additional rollouts to estimate the value of candidate states (i.e. VinePPO). Our text description of our results will clarify that the speedups relative to common methods like GRPO are approximately 4x, while the speedup relative to more complex methods like VinePPO that get closer to TBA’s elevated accuracy are 50x.\n\n> The equation 5, if I actually expand $\\log \\hat{Z}$, then the equation becomes a square loss term to train the reward function: $r_{\\phi}(\\mathbf{y};\\mathbf{x}) - r(\\mathbf{y};\\mathbf{x})$, but then all the policy related terms will get cancelled out… I'm actually not sure how the training is done -- how is $\\theta$ learned? This is probably because I'm not familiar with the Bayesian notation, but understanding this will help me understand the main difference between this objective and GRPO style objective. I'm happy to raise my score after the discussion round!\n\nGreat question! In our objective, $\\log \\hat{Z}$ is analogous to the mean of rewards used in GRPO’s advantage computation. Specifically, given a set of model generations for a given query, Equation 4 shows that $\\log \\hat{Z}$ is an average over two quantities computed for each model generation – one quantity is the generation’s reward and the other quantity reflects the generation’s adherence to the reference policy (the latter quantity is computed as a difference of two terms). Since $\\log \\hat{Z}$ is an average over multiple generations, including $\\log \\hat{Z}$ in Equation 5 does not cause cancellation. Instead, $\\log \\hat{Z}$ is used to compute per-generation deviations from the mean, creating an advantage-like quantity for each generation. \n\nThis is easier to see if you multiply everything inside the square operation in Equation 5 by $-1$, which is equivalent to multiplying the entire equation by $1$ and gives \n\n$\\mathcal{L}\\_{\\text{TB}}^{\\text{VarGrad}}(\\mathbf{B};\\theta) = \\frac{1}{BK}\\sum_{i=1,j=1}^{i=B,j=K} \\bigg(\n\\log \\pi_{\\text{ref}}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})- \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})+ \\frac{1}{\\beta} r(\\mathbf{y}^{(i,j)};\\mathbf{x}^{(i)})- \\log \\hat{Z}(\\mathbf{x}^{(i)})  \\bigg)^2.$ \n\nThe gradient with respect to $\\theta$ then becomes \n\n$\\nabla \\mathcal{L}\\_{\\text{TB}}^{\\text{VarGrad}}(\\mathbf{B};\\theta) = \\frac{1}{BK}\\sum_{i=1,j=1}^{i=B,j=K} -2 \\bigg(\n\\log \\pi_{\\text{ref}}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}) - \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})+ \\frac{1}{\\beta} r(\\mathbf{y}^{(i,j)};\\mathbf{x}^{(i)}) - \\log \\hat{Z}(\\mathbf{x}^{(i)})  \\bigg)\\nabla \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}),$ \n\nand we can turn this into a maximization problem instead of a minimization problem by multiplying the right side by $-1$, demonstrating that $\\log \\hat{Z}$ does indeed act similarly to GRPO’s mean of rewards inside an advantage-like term that weights the gradient of the policy: \n\n$\\nabla \\mathcal{J}\\_{\\text{TB}}^{\\text{VarGrad}}(\\mathbf{B};\\theta) = \\frac{1}{BK}\\sum_{i=1,j=1}^{i=B,j=K} 2 \\bigg(\n\\log \\pi_{\\text{ref}}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}) - \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})+ \\frac{1}{\\beta} r(\\mathbf{y}^{(i,j)};\\mathbf{x}^{(i)}) - \\log \\hat{Z}(\\mathbf{x}^{(i)})  \\bigg)\\nabla \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}).$\n\nImportantly, even though $\\log \\hat{Z}$  is also a function of $\\theta$, the gradient above does not go through $\\log \\hat{Z}$ because  $\\log \\hat{Z}$ is detached from the compute graph.  Our submission had mentioned that $\\log \\hat{Z}$ is detached on Line 134, but we will emphasize this further in our revision by adding the following stop-grad notation to Equation 5: $\\text{STOP-GRAD}(\\log \\hat{Z}(\\mathbf{x}^{(i)})).$ Also, our revision will include the above derivation and discussion of TBA’s gradient to clarify its relationship to other methods. We believe this greatly augments our original discussion of $\\log \\hat{Z}$ as a control variate (see Line 131), allowing readers to see exactly how $\\log \\hat{Z}$ plays a role similar to the average of rewards that GRPO computes and includes as a control variate.\n\n\nFinally, as an aside, please note that the reward functions are the same in Equations 4 and 5 – the absence of the $\\phi$ subscript from the reward function in Equation 5 is a typo. We have corrected this in our revision.\n\n\n> GFlowNet is already an established method. The paper does not offer innovation on methodology beyond the existing method. However, applying it for LLM post-training is interesting, and the empirical explorations seem warranted.\n\n\nWe would like to emphasize that we propose an asynchronous version of TB – Trajectory Balance with Asynchrony (TBA) – that had not previously been tested or suggested (even outside LLM post-training), which is notable given the suitability of TB’s off-policy-compatibility to the off-policy data generated in asynchronous LLM RL settings. Moreover, our implementation of TBA is efficient and effectively post-trains language models much faster than existing algorithms. Finally, we note that we go beyond prior work on GFlowNets by testing TBA’s\n- scalability to 64+ GPUs in a distributed setting, \n- compatibility with a range of buffer sampling techniques (e.g., to prioritize reward or recency),\n- sensitivity to off-policy samples, and \n- hyperparameters in various ablation studies. \n\nWe will use the extra space available in the revised paper to emphasize these points, clarifying the novelty of our research."}, {"review_id": "4iri2CysNw", "reviewer": "Reviewer_C5mN", "summary": "This paper proposes a “Trajectory Balance with Asynchrony (TBA)” distributed reinforcement learning framework to improve the post-training efficiency and performance of LLMs. TBA achieves asynchronous training by decoupling data generation and policy updates. It leverages the offline policy training objective based on Trajectory Balance (TB) and utilizes large-scale computing on search, which enhances exploration and data diversity. Experimental results on multiple tasks show that this method has faster training speed and comparable performance compared to existing methods.", "questions": "1.\tAs mentioned in Weaknesses 3, does this method show performance and efficiency improvement when computing resources are limited in practice?\n\n2.\tIn writing, abbreviations that appear for the first time should be written in full, such as GFlowNets in the abstract. What exactly does MCMC refer to in line 107? Please explain and use the full name to avoid potential ambiguity.", "limitations": "yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1.\tThis work decouples data generation and strategy updates to achieve asynchronous operations, which improves training efficiency.\n\n2.\tTBA increases the diversity of data in the experience replay buffer by utilizing a large-scale offline strategy sampling.\n\n3.\tIn three common LLM post-training RL pipelines, TBA demonstrates performance comparable to or better than existing baselines, while training faster.\n\nWeaknesses:\n1.\tIn the Red-teaming, the used base model GPT-2 is a bit outdated. Many more advanced models have been released. Moreover, the base LLM sizes used for the three pipelines all include 1B, such as RhoMath-1B and Pythia 1B, with the smallest being 410M and the largest being 8B. The effectiveness and sensitivity of this method on larger sizes of LLMs have not been explored.\n\n2.\tThe motivation for this method seems insufficient. The motivation for proposing the off-policy training objective in this paper stems from the fact that existing RL algorithms, such as PPO and RLOO, suffer from on-policy constraints, but there is a lack of analysis and discussion compared with existing off-policy methods. Therefore, a more comprehensive analysis and comparison with off-policy-based methods are needed to reinforce the motivation behind the proposed method.\n\n3.\tThe proposed method enables scalable search with large-scale computing resources, which contributes to the effectiveness of the method. So, does this method still have advantages when computing resources are limited in practice?\n\n4.\tThe writing needs improvement. The section 3 preliminaries, which typically introduces prior knowledge, takes up one and a half pages, while the section 4, the writing of the core proposed method TBA, only takes one page. Figure 1 shows the experimental results, which are referred to in the experiments of Section 5, but why is this figure placed below the abstract on the first page?", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough review, we believe that our addressing their questions has notably improved the manuscript. We understand that the reviewer’s main concerns are wanting to see TBA tested with a larger model, TBA’s connections to other methods, and clarity issues. We would be happy to know that the reviewer finds their concerns addressed and supports our submission’s acceptance, and we are equally interested in discussing any remaining or new concerns.\n\n> .... GPT-2 is a bit outdated. ... base LLM sizes used for the three pipelines all include 1B, ... The effectiveness and sensitivity of this method on larger sizes of LLMs have not been explored.\n\nNotably, our submission’s red-teaming section contained experiments with GPT-2 and Llama 3 (1B and 8B). However, we agree that GPT-2 is “a bit outdated” and have added the following discussion to our revision to clarify why we consider GPT-2: We study GPT-2 for the following reasons: \n- prior work published in ICLR 2025 on automated red-teaming utilized GPT-2 as a baseline, and we adopt it here to facilitate comparisons to Lee et al., 2025;\n- it is small enough to fit onto V100 GPUs, which allows us to test the effects of scaling TBA’s number of searcher/trainer processes to 64 when using 64 V100 GPUs; and \n- models like GPT2 are particularly well-suited for automated red-teaming since they haven’t been explicitly trained on filtered datasets or safety-tuned, resulting in more effective red-teaming prompts. \n\nAlso, we would like to note that our RLHF experiments consider larger 2.8B parameter models, where we found TBA effective. However, we do agree that testing even larger models would make the results stronger. Considering the short duration of the rebuttal period, we considered a preliminary experiment on the MATH dataset with the Qwen2.5-7B-Instruct model. We use the MATH dataset for training and evaluate the performance on MATH-500. We compare asynchronous GRPO to TBA, finding the latter excels. We note that these are preliminary results, and we will discuss these in the next revision.\n\n**Table R1. RL tuning Qwen2.5-7B-Instruct on MATH**\n| Method | MATH-500 Accuracy |\n|-|-|\n| TBA | 70% |\n| Async GRPO | 67.3% |\n\n> ... there is a lack of analysis and discussion compared with existing off-policy methods. Therefore, a more comprehensive analysis and comparison with off-policy-based methods are needed to reinforce the motivation behind the proposed method.\n\nBelow, we first clarify that we provided empirical comparisons to strong off-policy approaches (and on-policy approaches). Addressing the reviewer’s concerns, we then show how we improve our revision with new empirical comparisons to an off-policy version of RLOO (Proximal RLOO) [40] as well as asynchronous PPO. Finally, we discuss how our revision provides a new comparison of the gradients of Proximal RLOO and TBA, highlighting key differences in these off-policy robust approaches.\n\nFirst, to clarify our submission, we note that our experiments extensively compared TBA to a strong off-policy method, DPO, which we gave benchmark values for in both synchronous and asynchronous (off-policy) settings. We also compared TBA to on-policy methods operating in synchronous, on-policy settings (Table 1). TBA exclusively used off-policy data in these experiments, training rapidly in our asynchronous framework. However, TBA also produced the highest win-rates/accuracies in each comparison, demonstrating comprehensive improvements.\n\nSecond, we provide the following comparison to the off-policy methods Online DPO and Proximal RLOO, as well as an async PPO baseline that is robust to moderately off-policy data [40]. Critically, TBA outperforms all of these methods, achieving both the highest win rate and the lowest perplexity (i.e., the closest adherence to the reference policy). TBA also makes training significantly faster, compared to baselines, and trains with stability on data that is relatively more off-policy (10 steps instead of 2).\n\n**Table R2. RL finetuning of Pythia 410M on TL;DR data (Compute = 4xA100s).**\n| Method | Perplexity ↓ | Win Rate ↑ | Training Steps Before Syncing |\n|--------|----------------|------------|------------------|\n| Async PPO\t     | 1.14       \t| 0.84   \t| 2          \t|\n| Proximal RLOO | 1.13       \t| 0.83   \t| 2          \t|\n| Online DPO\t     | 1.13       \t| 0.85  \t| 2          \t|\n| TBA (ours)         | 1.13       \t| 0.86   \t| 10        \t|\n\nFinally, we use gradient analysis to compare TBA to Proximal RLOO, an off-policy variant of RLOO with empirically demonstrated off-policy robustness [40]. With k=3 samples, the gradient for Proximal RLOO for a single sample $y_1$ is\n\n$\\nabla \\mathcal{J}\\_{\\text{P-RLOO}}(\\theta) = \\frac{\\pi_\\theta(y_1|x)}{\\pi_{ref}(y_1|x)}  \\hat{A}(y_1|x) \\nabla_\\theta \\log \\pi_\\theta(y_1|x),$\n\nwhere the advantage $\\hat{A}(y_1|x) = R(y_1, x) - \\frac{1}{k-1} \\sum_{i \\neq 1} R(y_i, x)$. \n\nThe gradient for TBA is\n\n $\\nabla \\mathcal{J}\\_{\\text{TB}}(\\theta) = 2 \\bigg( \\log \\pi_{\\text{ref}}(y_1 \\mid x) - \\log \\pi_{\\theta}(y_1 \\mid x) + \\frac{1}{\\beta} r(y_1; x) - \\log \\hat{Z}(x)  \\bigg)\\nabla \\log \\pi_{\\theta}(y_1 \\mid x),$\nwhere $\\log \\hat{Z}(x)= \\frac{1}{k} \\sum_{i}  \\bigg( \\log \\pi_{\\text{ref}}(y_i \\mid x) - \\log \\pi_{\\theta}(y_i \\mid x) + \\frac{1}{\\beta} r(y_i; x) \\bigg).$ \n\nWe emphasize that $\\log \\hat{Z}(x)$ plays a role similar to the mean of rewards used to compute the advantage in the Proximal RLOO gradient. The main differences in these methods are: (1) TBA's $\\log \\hat{Z}(x)$ average includes both reward and deviations from the reference policy, normalizing the latter when Proximal RLOO does not. (2) The term that precedes the gradient of the policy in TBA is increasing as the sampled sequence $y_1$ becomes more probable under the reference policy, whereas the analogous term in Proximal RLOO $\\frac{\\pi_\\theta(y_1|x)}{\\pi_{ref}(y_1|x)}$ decreases as $y_1$ becomes more probable under the reference policy. In summary, Proximal RLOO seeks to deal with off-policy data by up-weighting gradients on samples that are relatively probable under the current policy; in contrast, TBA manages off-policy data by balancing rewards with deviations from the reference policy, up-weighting gradients on samples with higher probability under the reference policy (all else equal).\n\n> The writing needs improvement. The section 3 preliminaries, which typically introduces prior knowledge, takes up one and a half pages, while the section 4, the writing of the core proposed method TBA, only takes one page. Figure 1 shows the experimental results, which are referred to in the experiments of Section 5, but why is this figure placed below the abstract on the first page?\n\n\nRegarding the Figure 1 placement, we included it below the abstract to give readers an immediate depiction of the efficiency improvements TBA creates, but we see the clarifying benefit of having it in the experiments section and will happily move it there if that would improve your rating of our submission and its clarity.\n\nRegarding the preliminary section’s length, we note that our submission’s preliminaries actually occupy fewer lines (Line 93 – Line 139) than the description of our method TBA (Line 141 – Line 191), but this was not clear in our submission because the visual illustration of TBA in Figure 2 appeared in the same section as the preliminaries – our revision will fix this Figure 2 placement such that description of our core proposed method occupies about 1.5 pages, and the preliminaries section occupies about 1 page. We would be happy to shorten this preliminaries section further if you recommend that, but we note that it does include background on and motivation for the loss function we use (Equation 5), which we believe adds clarity.\n\n\n> The proposed method enables scalable search with large-scale computing resources, which contributes to the effectiveness of the method. So, does this method still have advantages when computing resources are limited in practice?... As mentioned in Weaknesses 3, does this method show performance and efficiency improvement when computing resources are limited in practice?\n\nOur submission demonstrated that, yes, our proposed TBA approach produces improved “performance and efficiency” when computing resources are limited in practice. We have revised our submission to make this more clear. In particular, our original submission discussed both “compute-matched” experiments (e.g., Line 267) and “resource scaling” experiments (e.g., Line 308), and we have revised these discussions to specifically mention the number of GPUs involved in each type of experiment. \n\nFor example, instead of only mentioning that our compute-matched experiments use 4 GPUs in the captions of associated figures (Figure 1 and Figure 3), we now note this usage of only 4 GPUs in the text discussing these results. Use of 4 GPUs is standard in these RL experiments (e.g. it was adopted by the baseline methods we compare to). We also go beyond the baselines we compare to, using up to 64 GPUs, in our resource-scaling experiments – our submission only described these experiments’ GPU counts in Table 2 and Appendix A, and our revision will clarify these resource-scaling experiments’ GPU counts in the writing of the main text as well.\n\n> In writing, abbreviations that appear for the first time should be written in full, such as GFlowNets in the abstract. What exactly does MCMC refer to in line 107? Please explain and use the full name to avoid potential ambiguity.\n\nThank you for noting these potential ambiguities, our revision will address both of the mentioned examples as well as others – e.g., we will write the full names of all baseline methods like “Group Relative Policy Optimization” for GRPO. Regarding MCMC, our revision clarifies that this refers to “Markov chain Monte Carlo”, a type of sampling strategy that is too slow to serve as a replacement for post-training LLMs in most settings."}, {"review_id": "UnVIpzgp5i", "reviewer": "Reviewer_2z92", "summary": "This paper proposes a method called TBA (Trajectory Balance with Asynchrony) to improve the post-training process of large language models (LLMs). The approach combines the Trajectory Balance (TB) objective with an asynchronous reinforcement learning (RL) framework, aiming to enhance training efficiency, diversity, and scalability in sparse-reward scenarios. Experiments on mathematical reasoning (GSM8K), preference fine-tuning (PFT), and automated red-teaming (RT) tasks demonstrate TBA's advantages in speed and performance over baseline methods (e.g., PPO, RLOO, Online DPO). The authors emphasize the method's open-source availability and reproducibility while discussing its broader societal impact and limitations.", "questions": "Based on the concerns raised in the weaknesses section, please provide explanations and supplementary experiments addressing the following:\n\n1. Theoretical Justification of TB + Asynchronous Combination:  \n   - Can you formally derive the relationship between TB and other RL objectives (e.g., PPO)?  \n   - Why does combining TB with asynchronous search yield better performance than using TB alone in a synchronous setting?\n\n2. Variance Analysis in TB Estimation:  \n   - Can you provide theoretical bounds on the gradient variance of the TB objective under different sequence lengths?  \n   - What happens when $K$ (number of responses per prompt) is increased further—does performance continue to improve?\n\n3. Convergence Behavior Under Delayed Policies:  \n   - Does TBA converge under the assumption of a delayed policy $\\pi_{\\theta'}$?  \n   - Can you implement a theoretically grounded delay correction mechanism (e.g., importance sampling, gradient compensation) and compare it with your current heuristic approach?\n\n4. Robustness Against Adversarial or Noisy Rewards:  \n   - How does TBA perform when the reward model is perturbed or when off-policy data includes adversarial examples?  \n   - Does reward-based prioritization amplify the effect of noisy or misleading rewards?\n\n5. Fair Comparison with Strong Off-Policy Baselines:  \n   - Please include comparisons with RAFT [9], asynchronous PPO [40], and other advanced off-policy methods.  \n   - Are the claimed performance gains robust under more comprehensive evaluations?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:  \n--TBA effectively combines trajectory balance (a physics-inspired equilibrium condition) with asynchronous RL, addressing bottlenecks in training efficiency and sample diversity for LLM post-training. This integration shows clear value, particularly in sparse-reward settings.  \n--Rigorous evaluation across multiple tasks (mathematical reasoning, preference tuning, red-teaming) with comparisons against strong baselines (PPO, RLOO, Online DPO).  \n--Ablation studies analyze key hyperparameters (e.g., KL coefficient β, batch size), supported by visualizations with error bars, enhancing credibility.  \n--Code and experimental details (Appendix A) are publicly shared, including hardware specs (GPU models/counts) and training protocols, aligning with NeurIPS reproducibility standards.  \n--Clear reporting of computational costs (e.g., runtime per experiment, GPU memory usage) aids future research.  \n--Section 6 candidly discusses limitations (e.g., sensitivity to low-resolution inputs, high compute costs) and misuse risks (e.g., disinformation generation), proposing mitigation strategies (e.g., controlled model access). This aligns with academic ethics.  \n\nWeaknesses:\n-- Lack of Theoretical Innovation in Combining Existing Techniques: The core components of TBA—the trajectory balance objective and the asynchronous distributed architecture—are both existing techniques (from GFlowNets and IMPALA, respectively). While the paper is the first to combine them in the context of LLM post-training, it does not sufficiently explain why this combination can overcome the limitations of existing off-policy training approaches (e.g., [40] highlights that off-policyness harms performance). The theoretical novelty is weak. Moreover, the lack of comparison with recent high-performance off-policy methods (e.g., RAFT [9]) undermines the competitiveness of TBA’s central claim of \"efficient use of replay buffers.\"\n\n--Empirical Variance Control May Conceal TB Objective's Theoretical Deficiencies: The TB objective relies on VarGrad estimation $\\log \\hat{Z}$ (Formula 4), but the paper does not prove the bias bounds of this estimator in long-sequence generation. For instance, when sequence length = 512, the cumulative error of $\\log \\pi_{\\theta}$ may grow super-linearly. Appendix B shows that GSM8K training diverges when the KL coefficient $\\beta < 0.003$ (Figure 6, bottom left), suggesting high variance issues. However, the authors only address this problem empirically—by increasing the number of samples ($K=20$) and raising $\\beta$. Thus, an important question arises: Is TBA’s performance gain derived from engineering patches to mitigate TB’s flaws, rather than algorithmic superiority? Would these empirical strategies fail in more complex tasks (e.g., long-text reasoning)?\n\n--Does Asynchronicity Undermine TB’s Convergence Guarantees?  TB theoretically requires data distribution stationarity ([4]), but in TBA, the SEARCHER uses a delayed policy $\\pi_{\\theta'}$ to generate data, while the TRAINER updates $\\pi_{\\theta}$, leading to non-stationary data distributions. The paper employs heuristic methods (sync interval $k$, recency sampling $m$) to manage off-policyness, but experiments show that when $m=0.4$, the Win Rate on TL;DR drops significantly (0.67 vs. 0.8), indicating residual asynchronous bias. Why not introduce theoretically grounded delay correction mechanisms? Is the current approach essentially masking convergence risks through hyperparameter tuning?\n\n--Reward Model Vulnerability Threatens TB’s Robustness:The TB objective directly optimizes $\\pi_{\\theta} \\propto \\pi_{\\text{ref}} \\exp(\\beta^{-1} r_{\\phi})$, but reward models are known to be vulnerable to adversarial examples (e.g., reward hacking [55]). The paper does not test how TBA performs when off-policy data contains adversarial samples (e.g., intentionally constructed high-reward yet meaningless responses). Furthermore, TBA’s reward-based prioritization strategy (Section 4.2) may amplify the influence of such samples. In real-world scenarios where reward models are imperfect, is TB more fragile compared to more constrained RLHF methods like DPO? Why does the paper avoid validating this risk?\n\n--Experimental Design May Systematically Underestimate Baseline Methods: The paper does not compare with contemporary high-performance off-policy methods (e.g., RAFT [9] with ranking-based fine-tuning) or modified asynchronous frameworks (e.g., asynchronous PPO from [40]). Do the claimed “significant advantages” of TBA stem from insufficient baseline comparisons or biased evaluation metrics? If fairly compared, would TBA still maintain its overwhelming superiority?", "was_revised": true, "final_justification": "The author's response has addressed my concerns, and I keep my initial rating of Borderline Accept.", "rebuttal": "Thank you! This review improved our submission, and we would be happy to address any remaining or new concerns.\n\n# Theoretical Justification of TB + Asynchronous Combination\n\n> Can you formally derive the relationship between TB and other RL objectives (e.g., PPO)? \n\nOur response to reviewer C5mN derives the relationship between the gradients of TBA and the off-policy robust Proximal RLOO, and it provides intuition for the two approaches’ strategies to dealing with off-policy data, which we will provide in our revision. \n\nMore broadly, the gradient of the TB loss when evaluated with on-policy samples is exactly equal to the REINFORCE gradient with a mean baseline (equivalent to RLOO gradient within constant scaling factor), evaluated with the KL regularized reward where beta (the temperature) is the KL coefficient (Malkin et al., 2023). Our submission (Line 109) notes that TB’s gradient is equivalent to the gradient of reverse KL between the training policy and the target density $p\\_{ref}(x)exp(\\frac{1}{\\beta}r(x))$, and our revision will illustrate this mathematically as well.\n\nFinally, our revision will note that, in the off-policy setting, TB does not have any precise equivalence to other on-policy RL objectives, but for LLMs is still equivalent to PCL (path consistency learning, Nachum et al., 2017). \n\nRegarding RAFT, for binary rewards, on-policy RAFT is exactly REINFORCE without a baseline. Off-policy RAFT is REINFORCE without a baseline, and without off-policy importance weighting. Accordingly, we believe the baselines we compare against are sufficient, especially since we include off-policy REINFORCE/RLOO.\n\n> Why does combining TB with asynchronous search yield better performance than using TB alone in a synchronous setting?\n\nWhile it’s true that TBA (TB with asynchronous search) leads to better training speed than TB in a synchronous setting, our submission did not find that there were accuracy/win-rate benefits of asynchrony. Specifically, TBA leads to speed ups because the trainer process does not pause for generations as it does in synchronous settings, and we found that moving closer to a synchronous setting does provide expected accuracy benefits, especially for GSM8K with RhoMath-1B – e.g., see Figure 6’s ablation studies that reduce the sync interval $k$ (top right plot) and increase the probability of sampling recent generations $m$ (top left plot).\n\n# Variance Analysis in TB Estimation\n\n> [theoretical bounds]? \n>\n> What happens when $\\mathbf{K}$ (number of responses per prompt) is increased further—does performance continue to improve?\n\nDoubling $\\mathbf{K}$ leads to a slight performance improvement and variance reduction, as shown in the following results that we will include in a new ablation study in our revision’s appendix.\n\n**Table R1. RL finetuning of RhoMath-1B on GSM8K data.**\n| Method | Test Accuracy | Standard Deviation |\n|--------|--------|------|\n| TBA, K=20 | 54.61  | 0.85 |\n| TBA, K=40 | 54.89  | 0.49 |\n\n**Critically, addressing this limitation by increasing $\\mathbf{K}$ does not come at any compute cost.** Specifically, when increasing $\\mathbf{K}$, we do not increase the total number of rollouts, we simply reduce the number of unique queries per batch as discussed in Footnote 1 of our submission (Line 254), generating more rollouts per query and maintaining batch size. Moreover, both of these accuracies are significantly beyond the accuracies achieved by prior work that uses the same setup and significantly longer training times.\n\nOur revision will include theoretical bounds on the gradient variance as a future direction.\n\n\n# Convergence Behavior Under Delayed Policies\n\n> Does TBA converge under the assumption of a delayed policy $\\pi_{\\theta’}$? \n>\n> Can you implement a theoretically grounded delay correction mechanism (e.g., importance sampling, gradient compensation) and compare it with your current heuristic approach?\n\nYes. We would like to clarify that the only requirements for the convergence of TB is taking the loss for all complete trajectories to 0 [4,34]. Crucially, this guarantee does not depend on the policy generating the trajectories, which is what makes the objective suitable for off-policy training, as noted in [4]. \n\nNote that this convergence does not rely on any external mechanism like importance sampling, etc. Moreover, we wish to clarify that our experience is that TBA does not require hyperparameter tuning beyond what is required for other RL methods like PPO, GRPO, etc. Relatedly, the reviewer correctly states that RLHF ablation studies in Section 5.4 show that win-rates are finally harmed when $m$ falls to $0.4$, and the GSM8K ablation studies similarly show how accuracy can decline with $m$ (Figure 6). However, we emphasize that these results are not necessarily due to any flaw in TB or its off-policy robustness; rather, these results may relate to the fact that increasing the off-policyness can nearly eliminate the amount of feedback the rollout policy receives when training LLMs for a fixed step count to facilitate rapid training (about 1 hour on 4 GPUs).\n\n# Robustness Against Adversarial or Noisy Rewards\n\n> How does TBA perform when the reward model is perturbed or when off-policy data includes adversarial examples? \n>\n> Does reward-based prioritization amplify the effect of noisy or misleading rewards?\n\nWe plan to incorporate the following discussion in the revised manuscript to address this important matter.\n\nIn our RLHF experiments, we perform test-data evaluations with a large ‘gold-standard’ reward model, and we perform training with a smaller proxy reward model that was trained on rankings from the gold-standard model (i.e., we follow the setup of [40]). Thus, our setup provides an opportunity to determine the potential for overoptimizing feedback from an imperfect reward model. This is important to test because TBA’s reward-based prioritization could, through identifying and selecting higher-reward samples, accelerate reward hacking. Potentially offsetting this, however, TBA operates asynchronously on off-policy data, which may slow reward hacking by increasing the time it takes for the searcher policy to be updated with information about high-reward samples. Additionally, like other methods, TBA may reduce reward hacking by penalizing deviation from the reference policy via $\\beta$.\n\nWe test whether overoptimization of the imperfect proxy model occurs by simultaneously considering win-rate according to the ‘gold-standard’ reward model and the divergence from the reference policy via evaluation of the reference policy’s perplexity on the RLHF policy’s generations. Again, this setup follows [40], which notes that the latter metric can be thought of as a proxy for the KL divergence from the reference policy. If the latter metric is high (i.e., the reference policy is surprised by the generations) or the former metric is low (i.e., the gold-standard reward model does not grade the generations highly), then there is evidence for overoptimization of the proxy model. Critically, Figure 4 shows that TBA establishes a new Pareto frontier for these two metrics, surpassing that of DPO’s, which had previously been shown to outperform PPO and RLOO on these metrics [40]. \n\nTherefore, our findings are consistent with TBA discouraging reward hacking relative to other post-training methods. Like many methods, TBA may discourage reward hacking through the $\\beta$ parameter, but it also may avoid reward hacking via its asynchronous training on off-policy data, which increases the number of optimization steps between the searcher policy’s generation of a rollout and the searcher policy’s parameter update based on that rollout. In other words, TBA prevents the searcher policy from receiving frequent feedback about the value of the samples it generates, which may explain why (relative to competing methods) it did not overoptimize imperfect proxy models in our experiments.\n\n# Fair Comparison with Strong Off-Policy Baselines\n\n> Please include comparisons with RAFT [9], asynchronous PPO [40], and other advanced off-policy methods. Are the claimed performance gains robust under more comprehensive evaluations?\n\nAddressing this concern, the revision will provide the following comparison to the advanced off-policy methods Online DPO and Proximal RLOO, as well as an async PPO baseline that is robust to moderately off-policy data [40]. The revision will also include a theoretical comparison of the gradients of TBA and the off-policy robust Proximal RLOO, which we provide in our response to Reviewer C5mN. \n\n**Table R2. RL finetuning of Pythia 410M on TL;DR data (Compute = 4xA100s).**\n| Method | Perplexity ↓ | Win Rate ↑ | Training Steps Before Syncing |\n|--------|----------------|------------|------------------|\n| Async PPO\t     | 1.14       \t| 0.84   \t| 2          \t|\n| Proximal RLOO | 1.13       \t| 0.83   \t| 2          \t|\n| Online DPO\t     | 1.13       \t| 0.85  \t| 2          \t|\n| TBA (ours)         | 1.13       \t| 0.86   \t| 10        \t|\n\nCritically, TBA outperforms three advanced off-policy methods, achieving both the highest win rate and the lowest perplexity – i.e., the closest adherence to the reference policy, consistent with low reward-hacking. TBA also makes training significantly faster, compared to baselines, and trains with stability on data that is relatively more off-policy (10 steps instead of 2).\n\nBeyond off-policy method comparisons, we emphasize that TBA can outperform state-of-the-art on-policy methods learning synchronously: Figures 1 and 3 show that TBA outperforms *synchronous on-policy* Online DPO, while Table 1 shows that TBA outperforms *synchronous on-policy* RLOO and PPO. Moreover, in the context of Table R2 above, Synchronous PPO with on-policy data obtains the same win-rate as TBA (0.86). Thus, TBA leverages off-policy learning to obtain the speed benefits of asynchrony while matching or surpassing the accuracy/win-rate of on-policy learning."}, {"review_id": "y3WMDnFQsS", "reviewer": "Reviewer_ZkfV", "summary": "I summarize the paper as follows. This paper proposes Trajectory Balance with Asynchrony (TBA), a reinforcement learning (RL) framework for post-training large language models. In detail, TBA decouples data generation and policy updates using asynchronous, off-policy training via the trajectory balance objective, kind of diversity-seeking RL objective. It achieves 4x to 50x speedups and improves performance on tasks like mathematical reasoning, preference tuning, and red-teaming compared to existing methods.", "questions": "See the weakness part:\n\n**(A) There requires an efficiency discussion in sequence selection.** As mentioned in **Line169-172**, the proposed method generates more samples for each query and selects the subset of generated sequences based on recency and rewards to update the model, this inevitably brings additional memory and inference cost. When the base model is large enough, e.g., >7B params, the efficiency can be degraded. \n\nBased on the above analysis, reinforcement finetuning LLMs is expensive in computations, memories and reward annotations. As far as I know, there exist some amortization methods [2-4] to approximately score the sequence reward or predict the preference of tasks in RL. Hence, *I suggest to add a few discussion sentences regarding the chance of marrying amortized evaluation [3,4] with TBA to improve efficiency and robustness in the future* in **Limitation part Line 325-Line 327**.\n\n**(B) I want to see the discussion between GRPO [5] in DeepSeek-R1 and asynchronous off-policy methods such as TBA.** The analysis is not from quantitative results but some technical discussions. Will removing out the critic be better? I guess this is up to memory constraints.", "limitations": "yes, also please include the extra discussion about model predictive task sampling as suggested in the comment.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The writing is overall well. I will summarize the merits and flaws of this work based on my understanding. \n\n---\n\nIn terms of merits, I think this work is advantageous in four points: \n\n(1) The decoupling desgin is novel. TBA separates data generation from policy updates, leading to efficient resource uses. \n\n(2) The method is scalable and fast.Experiments demonstrated significant wall-clock time improvements, with up to 50× faster training. I think it is extremely beneficial for post-training large models.\n\n(3) TBA can achieve off-policy robustness.The use of trajectory balance enable diverse and scalable off-policy learning. \n\n(4) The experiments are extensive and convincing. It shows strong experimental results across three key LLM tasks with rigorous baseline comparisons.\n\n---\n\nThough there exist no major concerns, I have a few comments as well as suggestions to improve this manuscript. These include:\n\n(1) More discussions on implementation details can be added. Note that the reward model and policy optimization rely on GFlowNets [1] and DPO, and these involve several hyperparameters to tune during post-training. \n\n(2) There requires an efficiency discussion in sequence selection. As mentioned in **Line169-172**, the proposed method generates more samples for each query and selects the subset of generated sequences based on recency and rewards to update the model, this inevitably brings additional memory and inference cost. *When the base model is large enough, e.g., >7B params, the efficiency can be degraded.*\n\nBased on the above analysis, reinforcement finetuning LLMs is expensive in computations, memories and reward annotations. As far as I know, there exist some amortization methods [2-4] to approximately score the sequence reward or predict the preference of tasks in RL. Hence, *I suggest to add a few discussion sentences, i.e., combining amortized evaluation methods in [3,4] with TBA to improve efficiency and robustness in the future* in **Limitation part Line 325-Line 327**.\n\n(3) I want to see the discussion between GRPO [5] in DeepSeek-R1 and asynchronous off-policy methods such as TBA.The analysis is not from quantitative results but some technical discussions. Will removing out the critic be better? I guess this is up to memory constraints.\n\n(4) Some typos, such as Line 277, “these results…” should be “These results”.\n\n**Reference:**\n\n[1] Bengio, Yoshua, et al. \"Gflownet foundations.\" Journal of Machine Learning Research 24.210 (2023): 1-55.\n\n[2] Greenberg, Ido, et al. \"Train hard, fight easy: Robust meta reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2023): 68276-68299.\n\n[3] Wang, Qi Cheems, et al. \"Model Predictive Task Sampling for Efficient and Robust Adaptation.\" (2025).\n\n[4] Qu, Yun, et al. \"Fast and Robust: Task Sampling with Posterior and Diversity Synergies for Adaptive Decision-Makers in Randomized Environments.\" arXiv preprint arXiv:2504.19139 (2025).\n\n[5] Guo, Daya, et al. \"Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.\" arXiv preprint arXiv:2501.12948 (2025).\n\n---\n\n*I would further update my rating of some items if my concerns are well addressed during the rebuttal.*", "was_revised": true, "final_justification": "After reading the rebuttal, I appreciate the author's response and solid work, and all my concerns are well addressed. This work proposed a fast and scalable optimization method for post-training LLMs and significantly boosted the performance and efficiency. As far as I know, it is non-trivial to develop off-policy methods for RL finetuning reasoning models. This work is solid in both algorithm design and experiments. Given these, I increased some sub-scores.\n\nFinally, the studied topic in this paper is really crucial for both academia and industry, considering the huge energy consumption of developing reasoning models.", "rebuttal": "We thank the reviewer for their helpful review, we believe it has improved our submission. Also, we are grateful for the reviewer’s stated interest in updating their ratings based on our responses. We would be happy to address any remaining or new concerns during the discussion period.\n\n> More discussions on implementation details can be added. Note that the reward model and policy optimization rely on GFlowNets [1] and DPO, and these involve several hyperparameters to tune during post-training.\n\nThis is a great point! As noted in our response to Reviewer 2z92, our revision will include more details about the reward models used in our experiments, which are from prior work that we emulate the experimental settings of – Noukhovitch et al. (2025) and Lee et al. (2025) for RLHF and red-teaming, respectively. In particular, we will include in our revision the details from their papers on training the reward models that we use. We will also clarify that our GSM8K experiments have programmatically verifiable rewards and thus no reward model.\n\nRegarding policy optimization, many hyperparameter details are in Table 3, Table 4, and Appendix A.3, but our revision will clarify the following missing details. We use AdamW in all experiments with weight decay = 0, except in red-teaming where it is 0.1. For red-teaming, we use learning rate 1e-4 (other experiments have their learning rates provided, e.g. see Table 3). Please let us know if you have any other related concerns that our revision can address!\n\n\n> There requires an efficiency discussion in sequence selection. As mentioned in Line169-172, the proposed method generates more samples for each query and selects the subset of generated sequences based on recency and rewards to update the model, this inevitably brings additional memory and inference cost. When the base model is large enough, e.g., >7B params, the efficiency can be degraded.\n>\n>Based on the above analysis, reinforcement finetuning LLMs is expensive in computations, memories and reward annotations. As far as I know, there exist some amortization methods [2-4] to approximately score the sequence reward or predict the preference of tasks in RL. Hence, I suggest to add a few discussion sentences, i.e., combining amortized evaluation methods in [3,4] with TBA to improve efficiency and robustness in the future in Limitation part Line 325-Line 327… please include the extra discussion about model predictive task sampling as suggested in the comment.\n\nOur revision will address this limitation by including the following discussion, please let us know if you would add anything else! Our TBA implementation samples off-policy data from a buffer without notable overhead by utilizing a maximum buffer size, which we enforce by dropping the oldest rollout data, in order to reduce the time needed for buffer operations. However, future work could explore ways to make TBA’s buffer operations like data selection more efficient for settings that (e.g.) have significantly longer rollouts, must retain all rollout data, need more complex sampling strategies, or have a diversity of tasks [2,3,4]. \n\n> I want to see the discussion between GRPO [5] in DeepSeek-R1 and asynchronous off-policy methods such as TBA.The analysis is not from quantitative results but some technical discussions. Will removing out the critic be better? I guess this is up to memory constraints.\n\nOur revision will include the following technical discussion that clarifies the connection between TBA and GRPO. The high-level summary is that these methods have similar update rules and can be similarly performant when data is on-policy, but we emphasize that (as discussed in the DeepSeekMath paper) GRPO is designed for online training, while TBA is designed to work offline or online (consistent with its outperforming GRPO in off-policy settings). Finally, we note that TBA can free up memory on the trainer process by not loading the reference policy onto the trainer process’s GPUs – deviations from the reference policy can instead be computed on the inference process’s GPUs due to TBA’s separation of these processes. Please let us know if you have any questions not answered by the following, more detailed discussion.\n\n\nIn the TBA objective, $\\log \\hat{Z}$ is analogous to the mean of rewards used in GRPO’s advantage computation. Specifically, given a set of model generations for a given query, Equation 4 shows that $\\log \\hat{Z}$ is an average over two quantities computed for each model generation – one quantity is the generation’s reward and the other quantity reflects the generation’s adherence to the reference policy (the latter quantity is computed as a difference of two terms). In other words, $\\log \\hat{Z}$ is used to create an advantage-like quantity for each generation that incorporates the deviation from the reference policy, while the deviation from the reference policy (the KL divergence) is separate from the advantage calculation in GRPO.\n\nThis is easier to see if you multiply everything inside the square operation in Equation 5 by $-1$, which is equivalent to multiplying the entire equation by $1$ and gives \n\n$\\mathcal{L}\\_{\\text{TB}}^{\\text{VarGrad}}(\\mathbf{B};\\theta) = \\frac{1}{BK}\\sum_{i=1,j=1}^{i=B,j=K} \\bigg(\n\\log \\pi_{\\text{ref}}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})- \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})\n\\\\+ \\frac{1}{\\beta} r(\\mathbf{y}^{(i,j)};\\mathbf{x}^{(i)})- \\log \\hat{Z}(\\mathbf{x}^{(i)})  \\bigg)^2.$ \n\nThe gradient with respect to $\\theta$ then becomes \n\n$\\nabla \\mathcal{L}\\_{\\text{TB}}^{\\text{VarGrad}}(\\mathbf{B};\\theta) = \\frac{1}{BK}\\sum_{i=1,j=1}^{i=B,j=K} -2 \\bigg(\n\\log \\pi_{\\text{ref}}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}) - \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})\n\\\\+ \\frac{1}{\\beta} r(\\mathbf{y}^{(i,j)};\\mathbf{x}^{(i)}) - \\log \\hat{Z}(\\mathbf{x}^{(i)})  \\bigg)\\nabla \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}),$ \n\nand we can turn this into a maximization problem instead of a minimization problem by multiplying the right side by $-1$, demonstrating that $\\log \\hat{Z}$ does indeed act similarly to GRPO’s mean of rewards inside an advantage-like term that weights the gradient of the policy: \n\n$\\nabla \\mathcal{J}\\_{\\text{TB}}^{\\text{VarGrad}}(\\mathbf{B};\\theta) = \\frac{1}{BK}\\sum_{i=1,j=1}^{i=B,j=K} 2 \\bigg(\n\\log \\pi_{\\text{ref}}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}) - \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)})\n\\\\+ \\frac{1}{\\beta} r(\\mathbf{y}^{(i,j)};\\mathbf{x}^{(i)}) - \\log \\hat{Z}(\\mathbf{x}^{(i)})  \\bigg)\\nabla \\log \\pi_{\\theta}(\\mathbf{y}^{(i,j)} \\mid \\mathbf{x}^{(i)}).$\n\n> Some typos, such as Line 277, “these results…” should be “These results”.\n\nThank you for noting this! Our revised manuscript is being thoroughly proofread to address this typo and a few others that we found (e.g., the subscript $\\phi$ was previously missing from the reward function in Equation 5 and has been added to the revision).\n\n[1] Bengio, Yoshua, et al. \"Gflownet foundations.\" Journal of Machine Learning Research 24.210 (2023): 1-55.\n\n[2] Greenberg, Ido, et al. \"Train hard, fight easy: Robust meta reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2023): 68276-68299.\n\n[3] Wang, Qi Cheems, et al. \"Model Predictive Task Sampling for Efficient and Robust Adaptation.\" (2025).\n\n[4] Qu, Yun, et al. \"Fast and Robust: Task Sampling with Posterior and Diversity Synergies for Adaptive Decision-Makers in Randomized Environments.\" arXiv preprint arXiv:2504.19139 (2025).\n\n[5] Guo, Daya, et al. \"Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.\" arXiv preprint arXiv:2501.12948 (2025)."}]

Author rebuttal:

Meta-review: The work makes solid technical contributions by decoupling data generation from policy updates and demonstrates effectiveness across mathematical reasoning, preference tuning, and red-teaming tasks. Two reviewers were strongly supportive: cZos explicitly stated "I do think this paper should be accepted" and raised their score to 5, while ZkfV increased to Strong Accept (6) after discussion.

Authors provided comprehensive rebuttals addressing theoretical concerns, adding 7B model experiments, and comparing against stronger off-policy baselines. The main issue was reviewer C5mN's non-participation in required discussions despite thorough responses to their concerns.

The asynchronous framework addresses real scalability challenges in LLM post-training with convincing experimental validation. This represents meaningful progress in making RL-based LLM training more efficient and practical.
Decision: Accept (poster)