Equilíbrio de Trajetória com Assincronia: Desacoplando Exploração e Aprendizado para Pós-Treinamento Rápido e Escalável de LLMs

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

Resumo

O aprendizado por reforço (RL) é um componente crítico do pós-treinamento de modelos de linguagem de grande escala (LLM). No entanto, os algoritmos on-policy existentes usados para pós-treinamento são intrinsecamente incompatíveis com o uso de buffers de replay de experiência, que podem ser preenchidos de forma escalável por atores off-policy distribuídos para melhorar a exploração à medida que o poder de computação aumenta. Propomos obter eficientemente esse benefício dos buffers de replay por meio do Trajectory Balance with Asynchrony (TBA), um sistema de RL para LLM massivamente escalável. Em contraste com as abordagens existentes, o TBA usa uma fração maior de computação na busca, gerando constantemente dados off-policy para um buffer de replay central. Um nó de treinamento amostra simultaneamente dados desse buffer com base na recompensa ou na recência para atualizar a política usando o Trajectory Balance (TB), um objetivo de RL que busca diversidade introduzido para GFlowNets. O TBA oferece três vantagens principais: (1) treinamento e busca desacoplados, acelerando o tempo de treinamento em 4x ou mais; (2) diversidade aprimorada por meio de amostragem off-policy em grande escala; e (3) busca escalável para cenários de recompensa esparsa. Em tarefas de raciocínio matemático, ajuste de preferências e red-teaming automatizado (tarefas de pós-treinamento diversas e representativas), o TBA produz melhorias de velocidade e desempenho em relação a baselines robustas.

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.