Asynchrone RLHF: Sneller en Efficiënter Off-Policy RL voor Taalmodellen

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

Samenvatting

De dominante paradigm voor RLHF is online en on-policy RL: synchroon genereren vanuit het grote taalmodel (LLM) beleid, labelen met een beloningsmodel, en leren met feedback op de eigen uitvoer van het LLM. Hoewel dit efficiënt is in prestaties, is dit paradigma computationeel inefficiënt. Geïnspireerd door de klassieke diepe RL-literatuur stellen we voor om generatie en leren te scheiden in RLHF. Dit maakt asynchroon genereren van nieuwe voorbeelden mogelijk terwijl tegelijkertijd wordt getraind op oude voorbeelden, wat leidt tot snellere training en meer rekenoptimaal schalen. Echter, asynchrone training is gebaseerd op een onderbelicht regime, online maar off-policy RLHF: leren op voorbeelden van eerdere iteraties van ons model. Om de uitdagingen in dit regime te begrijpen, onderzoeken we een fundamentele vraag: hoeveel afwijking van het beleid kunnen we tolereren voor asynchrone training om het leren te versnellen maar de prestaties te behouden? Onder verschillende RLHF-algoritmen die we hebben getest, vinden we dat online DPO het meest robuust is tegen off-policy gegevens, en de robuustheid neemt toe met de schaal van het beleidsmodel. We bestuderen verdere rekenoptimalisaties voor asynchrone RLHF, maar vinden dat deze ten koste gaan van de prestaties, wat resulteert in een compromis. Tot slot verifiëren we de schaalbaarheid van asynchrone RLHF door LLaMA 3.1 8B te trainen voor een instructievolgende taak 40% sneller dan een synchrone uitvoering, terwijl de uiteindelijke prestaties overeenkomen.

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.