Trajectory Balance with Asynchrony: Decoupling Exploration and Learning
  for Fast, Scalable LLM Post-Training
Reinforcement learning (RL) is a critical component of large language model
(LLM) post-training. However, existing on-policy algorithms used for
post-training are inherently incompatible with the use of experience replay
buffers, which can be populated scalably by distributed off-policy actors to
enhance exploration as compute increases. We propose efficiently obtaining this
benefit of replay buffers via Trajectory Balance with Asynchrony (TBA), a
massively scalable LLM RL system. In contrast to existing approaches, TBA uses
a larger fraction of compute on search, constantly generating off-policy data
for a central replay buffer. A training node simultaneously samples data from
this buffer based on reward or recency to update the policy using Trajectory
Balance (TB), a diversity-seeking RL objective introduced for GFlowNets. TBA
offers three key advantages: (1) decoupled training and search, speeding up
training wall-clock time by 4x or more; (2) improved diversity through
large-scale off-policy sampling; and (3) scalable search for sparse reward
settings. On mathematical reasoning, preference-tuning, and automated
red-teaming (diverse and representative post-training tasks), TBA produces
speed and performance improvements over strong baselines.
67e5691d4240c0b10d6126ea
Brian R. Bartoldson
67e5691d4240c0b10d6126eb
Siddarth Venkatraman
67e5691d4240c0b10d6126ec
James Diffenderfer
67e5691d4240c0b10d6126ed
Moksh Jain
67e5691d4240c0b10d6126ee
Tal Ben-Nun
67e5691d4240c0b10d6126ef
Seanie Lee
claimed_verified
2025-05-20T07:23:17.489Z
64ad5f59b7e4b2c1ce47eb43
/avatars/1f13ebe21a90d8c99920aa2c8cd9ac45.svg
Seanie Lee
user
Seanie-lee
67e5691d4240c0b10d6126f0
Minsu Kim
67e5691d4240c0b10d6126f1
Johan Obando-Ceron
claimed_verified
2025-03-28T08:38:10.901Z
6478f30ea68454566353ef95
https://cdn-avatars.huggingface.co/v1/production/uploads/6478f30ea68454566353ef95/hX5NiVPyQbK8TIBKnT4J1.jpeg
Johan Samir Obando Ceron
user
johanobandoc
67e5691d4240c0b10d6126f2
Yoshua Bengio
67e5691d4240c0b10d6126f3
Bhavya Kailkhura
claimed_verified
2025-06-10T09:30:26.389Z
65cb79db6427380bc21261e2
/avatars/a003eb5d0955417329c1a4170ae65879.svg
Bhavya Kailkhura
user
bhavyakailkhura
Trajectory Balance with Asynchrony: Decoupling Exploration and Learning
  for Fast, Scalable LLM Post-Training
Reinforcement learning (RL) is a critical component of large language model
(LLM) post-training. However, existing on-policy algorithms used for
post-training are inherently incompatible with the use of experience replay
buffers, which can be populated scalably by distributed off-policy actors to
enhance exploration as compute increases. We propose efficiently obtaining this
benefit of replay buffers via Trajectory Balance with Asynchrony (TBA), a
massively scalable LLM RL system. In contrast to existing approaches, TBA uses
a larger fraction of compute on search, constantly generating off-policy data
for a central replay buffer. A training node simultaneously samples data from
this buffer based on reward or recency to update the policy using Trajectory
Balance (TB), a diversity-seeking RL objective introduced for GFlowNets. TBA
offers three key advantages: (1) decoupled training and search, speeding up
training wall-clock time by 4x or more; (2) improved diversity through
large-scale off-policy sampling; and (3) scalable search for sparse reward
settings. On mathematical reasoning, preference-tuning, and automated
red-teaming (diverse and representative post-training tasks), TBA produces
speed and performance improvements over strong baselines.
Trajectory Balance with Asynchrony (TBA) decouples training and search in reinforcement learning for large language models, enhancing diversity and scalability with off-policy data.
reinforcement learning (RL) large language model (LLM) experience replay buffers Trajectory Balance with Asynchrony (TBA) off-policy sampling Trajectory Balance (TB) decoupled training sparse reward settings mathematical reasoning preference-tuning automated red-teaming