Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models
作者: Michael Noukhovitch, Shengyi Huang, Sophie Xhonneux, Arian Hosseini, Rishabh Agarwal, Aaron Courville
分类: cs.LG, cs.AI, cs.CL
发布日期: 2024-10-23 (更新: 2025-04-26)
备注: accepted at ICLR 2025, code at https://github.com/mnoukhov/async_rlhf, integrated into the open-instruct library https://github.com/allenai/open-instruct
💡 一句话要点
提出异步RLHF,加速并优化语言模型离线强化学习训练。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 异步强化学习 RLHF 语言模型 离线强化学习 DPO 计算效率 指令遵循
📋 核心要点
- 现有RLHF方法计算效率低,同步生成和学习样本导致训练速度受限。
- 提出异步RLHF,分离样本生成和学习过程,允许异步训练,提高计算效率。
- 实验表明,异步RLHF在指令遵循、数学和推理任务上加速训练,且保持性能。
📝 摘要(中文)
目前RLHF的主流范式是在线且on-policy的强化学习:同步地从大型语言模型(LLM)策略中生成样本,用奖励模型进行标注,并利用LLM自身输出的反馈进行学习。虽然性能良好,但这种范式在计算上效率低下。受经典深度强化学习文献的启发,我们提出在RLHF中分离生成和学习过程。这使得可以在异步生成新样本的同时训练旧样本,从而加快训练速度并实现更优的计算扩展。然而,异步训练依赖于一个未被充分探索的领域,即在线但off-policy的RLHF:基于模型先前迭代的样本进行学习,这些样本提供的训练信号较差。我们解决了这个领域中的一个根本挑战:为了使异步训练加速学习并保持性能,我们可以容忍多少off-policyness?在我们测试的几种RLHF算法中,发现在线DPO对off-policy数据最具鲁棒性,并且鲁棒性随着策略模型规模的增大而提高。我们进一步研究了异步RLHF的计算优化,但发现它们会带来性能损失,从而产生权衡。我们通过从LLaMA 3.1 8B训练一个通用聊天机器人,在指令遵循任务上比同步运行快约40%,同时匹配最终性能,从而验证了异步RLHF的可扩展性。最后,我们将结果扩展到数学和推理,证明异步RL可以使Rho 1B在GSM8k上进行微调,速度提高约70%,同时匹配同步精度。
🔬 方法详解
问题定义:论文旨在解决现有RLHF方法中计算效率低下的问题。传统的RLHF方法采用在线、on-policy的训练方式,即同步生成样本、标注奖励并进行学习。这种方式限制了训练速度,因为模型必须等待新样本生成才能进行学习,导致计算资源的浪费。
核心思路:论文的核心思路是将样本生成和学习过程解耦,采用异步的方式进行。具体来说,模型可以在生成新样本的同时,利用之前生成的样本进行训练。这种方式可以充分利用计算资源,提高训练效率。同时,论文研究了在这种异步训练模式下,模型对off-policy数据的鲁棒性。
技术框架:异步RLHF框架主要包含两个部分:样本生成器和策略学习器。样本生成器负责从当前策略中生成新的样本,并将其存储在经验回放缓冲区中。策略学习器则从经验回放缓冲区中采样样本,并利用这些样本更新策略。这两个过程是异步进行的,互不干扰。论文主要研究了在线DPO算法在异步RLHF框架下的表现。
关键创新:论文的关键创新在于提出了异步RLHF的概念,并证明了其在加速训练方面的有效性。通过解耦样本生成和学习过程,异步RLHF可以充分利用计算资源,提高训练效率。此外,论文还研究了不同RLHF算法在异步训练模式下的鲁棒性,发现在线DPO算法表现最佳。
关键设计:论文主要关注在线DPO算法在异步RLHF框架下的应用。DPO算法是一种直接策略优化算法,它通过优化策略与奖励模型的差异来学习策略。在异步训练中,论文研究了不同off-policy程度对DPO算法性能的影响。此外,论文还探索了一些计算优化方法,例如减少样本生成频率,但发现这些方法会带来性能损失。
🖼️ 关键图片
📊 实验亮点
实验结果表明,异步RLHF在指令遵循任务上,使用LLaMA 3.1 8B模型进行训练时,比同步运行快约40%,同时保持了最终性能。在数学和推理任务上,使用Rho 1B模型在GSM8k数据集上进行微调时,异步RL可以加速训练约70%,并达到与同步训练相当的准确率。这些结果验证了异步RLHF在加速训练和保持性能方面的有效性。
🎯 应用场景
异步RLHF具有广泛的应用前景,可以应用于各种需要通过强化学习进行优化的语言模型任务,例如对话生成、文本摘要、代码生成等。该方法可以显著提高训练效率,降低计算成本,加速语言模型的开发和部署。尤其是在资源受限的情况下,异步RLHF的优势更加明显。
📄 摘要(原文)
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 which give a worse training signal. We tackle the fundamental challenge in this regime: how much off-policyness can we tolerate for asynchronous training to speed up learning but maintain performance? Among several RLHF algorithms we test, online DPO is found to be 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. We verify the scalability of asynchronous RLHF by training a general-purpose chatbot from LLaMA 3.1 8B on an instruction-following task ~40% faster than a synchronous run while matching final performance. Finally, we extend our results to math and reasoning to demonstrate asynchronous RL can finetune Rho 1B on GSM8k ~70% faster while matching synchronous accuracy.