Trajektorienbalance mit Asynchronität: Entkopplung von Exploration und Lernen für schnelles, skalierbares Post-Training von LLMs

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

March 24, 2025

Autoren: 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

Zusammenfassung

Reinforcement Learning (RL) ist eine entscheidende Komponente der Nachschulung großer Sprachmodelle (LLMs). Allerdings sind die bestehenden On-Policy-Algorithmen, die für die Nachschulung verwendet werden, inhärent inkompatibel mit der Nutzung von Experience-Replay-Buffern, die skalierbar durch verteilte Off-Policy-Akteure gefüllt werden können, um die Exploration mit steigender Rechenleistung zu verbessern. Wir schlagen vor, diesen Vorteil von Replay-Buffern effizient durch Trajectory Balance mit Asynchronität (TBA), ein massiv skalierbares RL-System für LLMs, zu nutzen. Im Gegensatz zu bestehenden Ansätzen verwendet TBA einen größeren Anteil der Rechenleistung für die Suche und erzeugt kontinuierlich Off-Policy-Daten für einen zentralen Replay-Buffer. Ein Trainingsknoten entnimmt gleichzeitig Daten aus diesem Buffer basierend auf Belohnung oder Aktualität, um die Policy mithilfe von Trajectory Balance (TB), einem diversitätsfördernden RL-Ziel, das für GFlowNets eingeführt wurde, zu aktualisieren. TBA bietet drei wesentliche Vorteile: (1) Entkopplung von Training und Suche, wodurch die Trainingszeit um das 4-fache oder mehr verkürzt wird; (2) verbesserte Diversität durch großflächiges Off-Policy-Sampling; und (3) skalierbare Suche für Umgebungen mit spärlichen Belohnungen. Bei mathematischem Denken, Präferenzabstimmung und automatisiertem Red-Teaming (vielfältige und repräsentative Nachschulungsaufgaben) erzielt TBA Geschwindigkeits- und Leistungsverbesserungen gegenüber starken Baselines.

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.