Trajectory Balance met Asynchronie: Ontkoppeling van Exploratie en Leren voor Snelle, Schaalbare Post-Training van LLM's

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

March 24, 2025

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

Samenvatting

Versterkend leren (Reinforcement Learning, RL) is een cruciaal onderdeel van de na-training van grote taalmodellen (Large Language Models, LLM). De bestaande on-policy algoritmen die voor na-training worden gebruikt, zijn echter inherent onverenigbaar met het gebruik van experience replay buffers, die op schaalbare wijze kunnen worden gevuld door gedistribueerde off-policy actoren om de exploratie te verbeteren naarmate de rekenkracht toeneemt. Wij stellen voor om dit voordeel van replay buffers efficiënt te benutten via Trajectory Balance with Asynchrony (TBA), een enorm schaalbaar RL-systeem voor LLM. In tegenstelling tot bestaande benaderingen gebruikt TBA een groter deel van de rekenkracht voor zoekacties, waarbij voortdurend off-policy data wordt gegenereerd voor een centrale replay buffer. Een trainingsnode bemonstert tegelijkertijd data uit deze buffer op basis van beloning of recentheid om het beleid bij te werken met behulp van Trajectory Balance (TB), een op diversiteit gericht RL-doel dat is geïntroduceerd voor GFlowNets. TBA biedt drie belangrijke voordelen: (1) ontkoppelde training en zoekacties, waardoor de trainingsduur met 4x of meer wordt verkort; (2) verbeterde diversiteit door grootschalige off-policy bemonstering; en (3) schaalbare zoekacties voor omgevingen met schaarse beloningen. Op het gebied van wiskundig redeneren, voorkeursafstemming en geautomatiseerde red-teaming (diverse en representatieve na-trainings taken), levert TBA snelheids- en prestatieverbeteringen op ten opzichte van sterke basislijnen.

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.