Equilibrio de Trayectoria con Asincronía: Desacoplando Exploración y Aprendizaje para un Post-Entrenamiento Rápido y Escalable de Modelos de Lenguaje Grande

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

March 24, 2025

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

Resumen

El aprendizaje por refuerzo (RL, por sus siglas en inglés) es un componente crítico del posentrenamiento de modelos de lenguaje grandes (LLM). Sin embargo, los algoritmos on-policy existentes utilizados para el posentrenamiento son inherentemente incompatibles con el uso de búferes de experiencia (replay buffers), los cuales pueden poblarse de manera escalable mediante actores off-policy distribuidos para mejorar la exploración a medida que aumenta el cómputo. Proponemos obtener eficientemente este beneficio de los búferes de experiencia mediante Trajectory Balance with Asynchrony (TBA), un sistema de RL para LLM altamente escalable. A diferencia de los enfoques existentes, TBA utiliza una mayor fracción del cómputo en la búsqueda, generando constantemente datos off-policy para un búfer de experiencia central. Un nodo de entrenamiento muestrea simultáneamente datos de este búfer basándose en la recompensa o la antigüedad para actualizar la política utilizando Trajectory Balance (TB), un objetivo de RL que busca diversidad introducido para GFlowNets. TBA ofrece tres ventajas clave: (1) entrenamiento y búsqueda desacoplados, acelerando el tiempo de entrenamiento en 4x o más; (2) mayor diversidad mediante muestreo off-policy a gran escala; y (3) búsqueda escalable en entornos de recompensas escasas. En tareas de razonamiento matemático, ajuste de preferencias y pruebas de red teaming automatizadas (tareas representativas y diversas de posentrenamiento), TBA produce mejoras en velocidad y rendimiento frente a líneas base sólidas.

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.