Asynchrones RLHF: Schneller und effizienter Off-Policy-RL für Sprachmodelle

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

October 23, 2024

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

cs.AI

Zusammenfassung

Das dominante Paradigma für RLHF ist das Online- und On-Policy-RL: synchrones Generieren aus dem großen Sprachmodell (LLM), Markieren mit einem Belohnungsmodell und Lernen unter Verwendung von Feedback zu den eigenen Ausgaben des LLM. Obwohl leistungsstark, ist dieses Paradigma rechnerisch ineffizient. Inspiriert von der klassischen Literatur zum tiefen RL schlagen wir vor, Generierung und Lernen in RLHF zu trennen. Dies ermöglicht eine asynchrone Generierung neuer Proben, während gleichzeitig auf alten Proben trainiert wird, was zu schnellerem Training und einer rechnerisch optimaleren Skalierung führt. Allerdings stützt sich das asynchrone Training auf ein wenig erforschtes Regime, das Online- aber Off-Policy-RLHF: Lernen an Proben aus vorherigen Iterationen unseres Modells. Um die Herausforderungen in diesem Regime zu verstehen, untersuchen wir eine grundlegende Frage: Wie viel Abweichung von der Policy können wir tolerieren, um das Lernen durch asynchrones Training zu beschleunigen, aber die Leistung beizubehalten? Unter mehreren getesteten RLHF-Algorithmen stellen wir fest, dass Online DPO am robustesten gegenüber Off-Policy-Daten ist und die Robustheit mit dem Umfang des Richtlinienmodells zunimmt. Wir untersuchen weitere Rechenoptimierungen für asynchrones RLHF, stellen jedoch fest, dass sie mit einem Leistungskosten verbunden sind, was zu einem Kompromiss führt. Schließlich überprüfen wir die Skalierbarkeit von asynchronem RLHF, indem wir LLaMA 3.1 8B bei einer Anweisungsfolge-Aufgabe 40 % schneller trainieren als bei einem synchronen Durchlauf, während die endgültige Leistung übereinstimmt.

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.