RLHF asynchrone : RL hors politique plus rapide et plus efficace pour les modèles de langage

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

October 23, 2024

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

cs.AI

Résumé

Le paradigme dominant pour l'apprentissage par renforcement à fonctionnement hiérarchique (RLHF) est l'apprentissage en ligne et sur politique : générer de manière synchrone à partir du grand modèle de langage (LLM), étiqueter avec un modèle de récompense, et apprendre en utilisant des retours sur les sorties du LLM lui-même. Bien que performant, ce paradigme est inefficace du point de vue computationnel. Inspirés par la littérature classique de l'apprentissage par renforcement profond, nous proposons de séparer la génération et l'apprentissage dans le RLHF. Cela permet une génération asynchrone de nouveaux échantillons tout en entraînant simultanément sur d'anciens échantillons, ce qui conduit à une formation plus rapide et à une mise à l'échelle plus optimale en termes de calcul. Cependant, l'entraînement asynchrone repose sur un régime peu exploré, en ligne mais hors politique, dans le RLHF : apprendre sur des échantillons des itérations précédentes de notre modèle. Pour comprendre les défis de ce régime, nous examinons une question fondamentale : dans quelle mesure pouvons-nous tolérer le hors-politique pour accélérer l'apprentissage tout en maintenant les performances ? Parmi plusieurs algorithmes de RLHF que nous avons testés, nous constatons que le DPO en ligne est le plus robuste aux données hors politique, et que la robustesse augmente avec l'échelle du modèle de politique. Nous étudions d'autres optimisations de calcul pour le RLHF asynchrone, mais constatons qu'elles entraînent un coût en performances, créant ainsi un compromis. Enfin, nous vérifions la scalabilité du RLHF asynchrone en formant LLaMA 3.1 8B sur une tâche de suivi d'instructions 40% plus rapidement qu'une exécution synchrone tout en conservant les performances finales.

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.