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
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
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.
reinforcement learning (RL) large language model (LLM) post-training on-policy algorithms experience replay buffers distributed off-policy actors Trajectory Balance with Asynchrony (TBA) central replay buffer Trajectory Balance (TB) GFlowNets decoupled training and search sparse reward settings mathematical reasoning preference-tuning automated red-teaming