비동기 RLHF: 언어 모델을 위한 빠르고 효율적인 오프-폴리시 강화 학습

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

October 23, 2024

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

cs.AI

초록

RLHF에 대한 우세한 패러다임은 온라인 및 온-폴리시 RL입니다: 대형 언어 모델 (LLM) 정책에서 동기적으로 생성하고 보상 모델로 라벨을 지정하며, LLM의 출력에 대한 피드백을 사용하여 학습합니다. 이 패러다임은 성능이 우수하지만 계산적으로 비효율적입니다. 고전적인 딥 RL 문헌에서 영감을 받아, RLHF에서 생성과 학습을 분리하는 것을 제안합니다. 이를 통해 새로운 샘플을 비동기적으로 생성하면서 동시에 이전 샘플에서 학습할 수 있게 되어 더 빠른 학습과 보다 계산적으로 최적화된 스케일링이 가능해집니다. 그러나 비동기적 학습은 미탐구된 온라인 그러나 오프-폴리시 RLHF에 의존합니다: 이전 모델의 반복에서 샘플을 학습합니다. 이 규칙에서의 도전을 이해하기 위해, 우리는 기본적인 질문을 조사합니다: 비동기적 학습을 가속화하기 위해 얼마나 많은 오프-폴리시를 용인할 수 있으며 성능을 유지할 수 있을까요? 우리가 테스트한 여러 RLHF 알고리즘 중에서 온라인 DPO가 오프-폴리시 데이터에 가장 견고하며, 견고성은 정책 모델의 규모와 함께 증가합니다. 비동기 RLHF에 대한 추가 계산 최적화를 연구했지만, 성능 비용이 발생하여 트레이드오프가 발생함을 발견했습니다. 마지막으로, 우리는 LLaMA 3.1 8B를 동기적 실행보다 40% 빠르게 지시 따르기 작업에 학습시킴으로써 비동기 RLHF의 확장성을 확인하며 최종 성능과 일치시킵니다.

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.