Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models (2410.18252v3)

Abstract: The dominant paradigm for RLHF is online and on-policy RL: synchronously generating from the 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.

• The paper introduces an asynchronous RLHF framework that decouples data generation from training, significantly speeding up iterations.
• It reveals that off-policy learning gains robustness with scale, with Online DPO outperforming traditional on-policy methods like PPO.
• The research highlights a practical efficiency-performance trade-off, paving the way for scalable and resource-efficient AI training.

Asynchronous RLHF: Efficiency in Off-Policy Reinforcement Learning for LLMs

The paper "Asynchronous RLHF: Faster and More Efficient Off-Policy RL for LLMs" addresses the computational inefficiency inherent in the reinforcement learning with human feedback (RLHF) paradigm, particularly within the context of LLMs. Traditional RLHF relies heavily on synchronous, on-policy reinforcement learning, which ties up compute resources and slows training processes. This work proposes an asynchronous off-policy framework that separates generation and training, offering a more efficient alternative while maintaining or surpassing performance standards.

Key Concepts and Findings

1. Asynchronous RLHF Framework: The proposed method generates data asynchronously, allowing for simultaneous generation and training processes. This shift towards an off-policy method enables faster iterations and leverages new efficient generation technologies, reducing overall computational time.
2. Off-Policy Learning Challenges: While asynchronous setups introduce off-policy data — samples from outdated model iterations — the research investigates the tolerance levels of off-policyness on performance. It concludes that as model scale increases, robustness to off-policy data also improves, aligning with practical needs for future AI developments.
3. Robustness of Online DPO: Among various RLHF algorithms, Online Direct Preference Optimization (DPO) exhibited the strongest robustness to off-policy data. This method proved to maintain performance better than PPO and other conventional on-policy algorithms, especially as model scales increased.
4. Compute Optimization: The investigation demonstrates that asynchronous methods allow RLHF training to be completed faster compared to synchronous methods, approximately 40% faster with the LLaMA 3.1 8B training. Further optimizations in either training or generation can exacerbate these speed benefits, though with potential trade-offs in model KL divergence.
5. Scalability: Experimentation with larger models showed that asynchronous RLHF effectively scales, validating its applicability in extensive, real-world deployments and making a considerable case for transitioning away from traditional synchronous paradigms in RLHF.

Implications for AI Development

The implications of adopting asynchronous RLHF are significant for both practical deployments and theoretical exploration:

• Resource Efficiency: By decoupling training and generation, institutions can streamline hardware utilization, particularly when handling extensive models. This approach is increasingly critical as LMs continue to scale, demanding more from computational infrastructure.
• Performance vs. Efficiency Trade-offs: The framework offers insights into managing trade-offs between computational efficiency and model alignment with human feedback — a central dilemma in AI model development.
• Future Directions: As models grow in complexity and requirements, refining asynchronous processes will likely become crucial for RLHF. The paper's findings advocate for further engineering innovations to optimize asynchronous frameworks, potentially unfolding new research directions in model alignment strategies.

Conclusion

By establishing an off-policy, asynchronous RLHF framework, the paper effectively tackles computational inefficiencies while maintaining high alignment with human feedback. The results not only advance state-of-the-art RLHF methodologies but also set a foundation for more scalable and practical AI model training practices. The insights garnered from asynchronous learning paradigms anticipate future demands, aligning model development processes with emerging technological capabilities and research landscapes in AI.