Provably Efficient Online RLHF with One-Pass Reward Modeling

TL;DR

The paper addresses the computational bottleneck of online RLHF by formulating RLHF as a contextual dueling bandit with linear rewards and introducing a one-pass reward modeling approach based on implicit online mirror descent. This method achieves constant-time, one-pass updates per iteration and eliminates the need to store historical data, while providing high-probability guarantees and improved estimation bounds compared to MLE. It applies to passive, active, and deployment-time RLHF settings, with tailored algorithms, regret/suboptimality guarantees, and practical implementations using Hessian-vector products and rejection sampling for uncertainty. Empirical results on Llama-3-8B-Instruct and Qwen2.5-7B-Instruct with Ultrafeedback and Mixture2 demonstrate stronger statistical efficiency and substantial computational savings, validating the approach for scalable, real-time alignment. Overall, the work advances the practicality of RLHF by delivering both theoretical guarantees and practical, scalable algorithms that reduce computation and storage while maintaining or improving alignment quality.

Abstract

Reinforcement Learning from Human Feedback (RLHF) has shown remarkable success in aligning Large Language Models (LLMs) with human preferences. Traditional RLHF methods rely on a fixed dataset, which often suffers from limited coverage. To this end, online RLHF has emerged as a promising direction, enabling iterative data collection and refinement. Despite its potential, this paradigm faces a key bottleneck: the requirement to continuously integrate new data into the dataset and re-optimize the model from scratch at each iteration, resulting in computational and storage costs that grow linearly with the number of iterations. In this work, we address this challenge by proposing a one-pass reward modeling method that eliminates the need to store historical data and achieves constant-time updates per iteration. Specifically, we first formalize RLHF as a contextual preference bandit and develop a new algorithm based on online mirror descent with a tailored local norm, replacing the standard maximum likelihood estimation for reward modeling. We then apply it to various online RLHF settings, including passive data collection, active data collection, and deployment-time adaptation. We provide theoretical guarantees showing that our method enhances both statistical and computational efficiency. Finally, we design practical algorithms for LLMs and conduct experiments with the Llama-3-8B-Instruct and Qwen2.5-7B-Instruct models on Ultrafeedback and Mixture2 datasets, validating the effectiveness of our approach.

Figures (13)

• Figure 1: Different settings of online RLHF. Contexts and actions selected by the environment () are shown in grey, while those selected by the algorithm () are highlighted in color.
• Figure 2: For online RLHF with passive data collection, we report the comparison of MLE and our method about (a) training loss, (b) training accuracy, (c) evaluation loss and (d) evaluation accuracy.
• Figure 3: For online RLHF with active data collection, we report the comparison of different methods about (a) training loss, (b) evaluation accuracy and (c) final evaluation accuracy and training time.
• Figure 4: For online RLHF with deployment-time adaptation, we report (a) cumulative rewards of MLE-based methods, (b) cumulative rewards of OMD-based methods, and (c) win rates.
• Figure 5: For online RLHF with passive data collection, we compare our proposed method and MLE ICML'23:Zhu-Principled in with passive data collection combined with Adam. We report the average accuracy and loss of the reward model during the training process.

Theorems & Definitions (22)

• Definition 1: Bradley-Terry Model
• Remark 1
• Definition 2
• Lemma 1
• Theorem 1
• Remark 2
• Remark 3
• Theorem 2
• Remark 4
• Theorem 3