RLHF Asincrónico: RL Fuera de Política Más Rápido y Eficiente para Modelos de Lenguaje

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

Resumen

El paradigma dominante para RLHF es el RL en línea y en política: generando de forma síncrona a partir del gran modelo de lenguaje (LLM), etiquetando con un modelo de recompensa y aprendiendo utilizando retroalimentación sobre las salidas del propio LLM. Aunque efectivo, este paradigma es computacionalmente ineficiente. Inspirados en la literatura clásica de RL profundo, proponemos separar la generación y el aprendizaje en RLHF. Esto permite la generación asíncrona de nuevas muestras mientras se entrena simultáneamente en muestras antiguas, lo que conduce a un entrenamiento más rápido y una escalabilidad más óptima en términos de cómputo. Sin embargo, el entrenamiento asíncrono se basa en un régimen poco explorado, el RLHF en línea pero fuera de política: aprendiendo en muestras de iteraciones anteriores de nuestro modelo. Para comprender los desafíos en este régimen, investigamos una pregunta fundamental: ¿hasta qué punto podemos tolerar la falta de política para acelerar el aprendizaje pero mantener el rendimiento? Entre varios algoritmos de RLHF que probamos, encontramos que el DPO en línea es el más robusto ante datos fuera de política, y la robustez aumenta con la escala del modelo de política. Estudiamos más optimizaciones de cómputo para RLHF asíncrono pero observamos que estas conllevan un costo en rendimiento, lo que da lugar a un compromiso. Finalmente, verificamos la escalabilidad de RLHF asíncrono entrenando LLaMA 3.1 8B en una tarea de seguimiento de instrucciones un 40% más rápido que una ejecución síncrona mientras se mantiene el rendimiento 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.