軌道バランスと非同期性：探索と学習の分離による高速でスケーラブルなLLMポストトレーニング

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

March 24, 2025

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

cs.AI

要旨

強化学習（RL）は、大規模言語モデル（LLM）のポストトレーニングにおいて重要な要素です。しかし、ポストトレーニングに使用される既存のオンポリシーアルゴリズムは、経験再生バッファの使用と本質的に互換性がありません。このバッファは、分散型オフポリシーアクターによってスケーラブルにデータを蓄積し、計算リソースが増えるにつれて探索を強化することができます。本論文では、Trajectory Balance with Asynchrony（TBA）という大規模スケーラブルなLLM RLシステムを通じて、この再生バッファの利点を効率的に得ることを提案します。既存のアプローチとは異なり、TBAは計算リソースの大部分を探索に割り当て、常にオフポリシーデータを生成して中央の再生バッファに蓄積します。トレーニングノードは、このバッファから報酬や新しさに基づいてデータをサンプリングし、Trajectory Balance（TB）を使用してポリシーを更新します。TBは、GFlowNets向けに導入された多様性を追求するRL目的関数です。TBAは以下の3つの主要な利点を提供します：（1）トレーニングと探索を分離し、トレーニングの実時間を4倍以上高速化する、（2）大規模なオフポリシーサンプリングを通じて多様性を向上させる、（3）スパースな報酬設定におけるスケーラブルな探索を実現する。数学的推論、選好チューニング、自動化されたレッドチーミング（多様で代表的なポストトレーニングタスク）において、TBAは強力なベースラインを上回る速度と性能の向上をもたらします。

English

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.