Équilibrage de Trajectoire avec Asynchronie : Découplage de l'Exploration et de l'Apprentissage pour un Post-Entraînement Rapide et Évolutif des LLM

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

Résumé

L'apprentissage par renforcement (RL) est un élément crucial du post-entraînement des grands modèles de langage (LLM). Cependant, les algorithmes on-policy existants utilisés pour le post-entraînement sont intrinsèquement incompatibles avec l'utilisation de tampons de relecture d'expérience, qui peuvent être remplis de manière scalable par des acteurs off-policy distribués pour améliorer l'exploration à mesure que la puissance de calcul augmente. Nous proposons d'obtenir efficacement cet avantage des tampons de relecture via le Trajectory Balance with Asynchrony (TBA), un système RL massivement scalable pour les LLM. Contrairement aux approches existantes, TBA utilise une plus grande fraction de la puissance de calcul pour la recherche, générant constamment des données off-policy pour un tampon de relecture central. Un nœud d'entraînement échantillonne simultanément des données de ce tampon en fonction de la récompense ou de la récence pour mettre à jour la politique en utilisant le Trajectory Balance (TB), un objectif RL favorisant la diversité introduit pour les GFlowNets. TBA offre trois avantages clés : (1) un entraînement et une recherche découplés, accélérant le temps d'entraînement d'un facteur 4 ou plus ; (2) une diversité améliorée grâce à un échantillonnage off-policy à grande échelle ; et (3) une recherche scalable pour des environnements à récompenses rares. Sur des tâches de raisonnement mathématique, de réglage des préférences et de red-teaming automatisé (tâches de post-entraînement diversifiées et représentatives), TBA produit des améliorations en termes de vitesse et de performance par rapport à des bases de référence solides.

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.