Bilanciamento della Traiettoria con Asincronia: Disaccoppiamento di Esplorazione e Apprendimento per un Post-Addestramento Rapido e Scalabile di LLM

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

March 24, 2025

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

Abstract

L'apprendimento per rinforzo (Reinforcement Learning, RL) è un componente cruciale del post-addestramento dei grandi modelli linguistici (Large Language Models, LLM). Tuttavia, gli algoritmi on-policy esistenti utilizzati per il post-addestramento sono intrinsecamente incompatibili con l'uso di buffer di esperienza replay, che possono essere popolati in modo scalabile da attori off-policy distribuiti per migliorare l'esplorazione all'aumentare della capacità di calcolo. Proponiamo di ottenere in modo efficiente questo vantaggio dei buffer di replay tramite Trajectory Balance with Asynchrony (TBA), un sistema RL per LLM altamente scalabile. A differenza degli approcci esistenti, TBA utilizza una frazione maggiore di risorse computazionali per la ricerca, generando costantemente dati off-policy per un buffer di replay centrale. Un nodo di addestramento campiona simultaneamente dati da questo buffer in base alla ricompensa o alla recentezza per aggiornare la politica utilizzando Trajectory Balance (TB), un obiettivo RL orientato alla diversità introdotto per le GFlowNets. TBA offre tre vantaggi chiave: (1) addestramento e ricerca disaccoppiati, accelerando il tempo di addestramento di 4 volte o più; (2) miglioramento della diversità attraverso il campionamento off-policy su larga scala; e (3) ricerca scalabile per contesti con ricompense sparse. Nel ragionamento matematico, nella sintonizzazione delle preferenze e nel red-teaming automatizzato (compiti di post-addestramento diversificati e rappresentativi), TBA produce miglioramenti in termini di velocità e prestazioni rispetto a baseline solide.

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.