RLHF Assíncrono: RL Off-Policy Mais Rápido e Eficiente para Modelos de Linguagem

Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models

October 23, 2024

Autores: Michael Noukhovitch, Shengyi Huang, Sophie Xhonneux, Arian Hosseini, Rishabh Agarwal, Aaron Courville

cs.AI

Resumo

O paradigma dominante para RLHF é o RL online e on-policy: gerando de forma síncrona a partir do grande modelo de linguagem (LLM), rotulando com um modelo de recompensa e aprendendo usando feedback sobre as próprias saídas do LLM. Embora eficaz, esse paradigma é computacionalmente ineficiente. Inspirados na literatura clássica de deep RL, propomos separar a geração e o aprendizado no RLHF. Isso possibilita a geração assíncrona de novas amostras enquanto treinamos simultaneamente em amostras antigas, resultando em um treinamento mais rápido e uma escalabilidade mais otimizada em termos de computação. No entanto, o treinamento assíncrono depende de um regime pouco explorado, o RLHF online mas off-policy: aprendendo em amostras de iterações anteriores do nosso modelo. Para compreender os desafios nesse regime, investigamos uma questão fundamental: até que ponto podemos tolerar a off-policy para acelerar o aprendizado de forma assíncrona, mantendo o desempenho? Entre vários algoritmos de RLHF que testamos, descobrimos que o DPO online é o mais robusto em relação aos dados off-policy, e a robustez aumenta com a escala do modelo de política. Estudamos mais otimizações de computação para RLHF assíncrono, mas observamos que elas implicam em um custo de desempenho, resultando em um trade-off. Por fim, verificamos a escalabilidade do RLHF assíncrono treinando o LLaMA 3.1 8B em uma tarefa de seguir instruções 40% mais rápido do que uma execução síncrona, mantendo o mesmo desempenho final.

English

The dominant paradigm for RLHF is online and on-policy RL: synchronously generating from the large language model (LLM) policy, labelling with a reward model, and learning using feedback on the LLM's own outputs. While performant, this paradigm is computationally inefficient. Inspired by classical deep RL literature, we propose separating generation and learning in RLHF. This enables asynchronous generation of new samples while simultaneously training on old samples, leading to faster training and more compute-optimal scaling. However, asynchronous training relies on an underexplored regime, online but off-policy RLHF: learning on samples from previous iterations of our model. To understand the challenges in this regime, we investigate a fundamental question: how much off-policyness can we tolerate for asynchronous training to speed up learning but maintain performance? Among several RLHF algorithms we tested, we find that online DPO is most robust to off-policy data, and robustness increases with the scale of the policy model. We study further compute optimizations for asynchronous RLHF but find that they come at a performance cost, giving rise to a trade-off. Finally, we verify the scalability of asynchronous RLHF by training LLaMA 3.1 8B on an instruction-following task 40% faster than a synchronous run while matching final performance.