Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Uncertainty-Penalized Reinforcement Learning from Human Feedback with Diverse Reward LoRA Ensembles

Yuanzhao Zhai, Han Zhang, Yu Lei, Yue Yu, Kele Xu, Dawei Feng, Bo Ding, Huaimin Wang

TL;DR

This work tackles overoptimization in RLHF by adding uncertainty-aware regularization to the RL fine-tuning stage and by building a diverse, parameter-efficient ensemble of reward models via LoRA with nuclear-norm diversification. By calibrating reward uncertainty through diverse LoRA ensembles, the method yields better out-of-distribution detection and uncertainty quantification, enabling more robust policy optimization. The UP-RLHF framework demonstrates improved gold-reward performance across summarization and QA tasks, while highlighting a trade-off between exploration and conservatism when applying uncertainty penalties. Overall, the approach advances RLHF alignment by explicitly modeling reward uncertainty and decoupling KL regularization from the actor objective, with practical gains and scalable parameter efficiency.

Abstract

Reinforcement learning from human feedback (RLHF) emerges as a promising paradigm for aligning large language models (LLMs). However, a notable challenge in RLHF is overoptimization, where beyond a certain threshold, the pursuit of higher rewards leads to a decline in human preferences. In this paper, we observe the weakness of KL regularization which is commonly employed in existing RLHF methods to address overoptimization. To mitigate this limitation, we scrutinize the RLHF objective in the offline dataset and propose uncertainty-penalized RLHF (UP-RLHF), which incorporates uncertainty regularization during RL-finetuning. To enhance the uncertainty quantification abilities for reward models, we first propose a diverse low-rank adaptation (LoRA) ensemble by maximizing the nuclear norm of LoRA matrix concatenations. Then we optimize policy models utilizing penalized rewards, determined by both rewards and uncertainties provided by the diverse reward LoRA ensembles. Our experimental results, based on two real human preference datasets, showcase the effectiveness of diverse reward LoRA ensembles in quantifying reward uncertainty. Additionally, uncertainty regularization in UP-RLHF proves to be pivotal in mitigating overoptimization, thereby contributing to the overall performance.

Uncertainty-Penalized Reinforcement Learning from Human Feedback with Diverse Reward LoRA Ensembles

TL;DR

This work tackles overoptimization in RLHF by adding uncertainty-aware regularization to the RL fine-tuning stage and by building a diverse, parameter-efficient ensemble of reward models via LoRA with nuclear-norm diversification. By calibrating reward uncertainty through diverse LoRA ensembles, the method yields better out-of-distribution detection and uncertainty quantification, enabling more robust policy optimization. The UP-RLHF framework demonstrates improved gold-reward performance across summarization and QA tasks, while highlighting a trade-off between exploration and conservatism when applying uncertainty penalties. Overall, the approach advances RLHF alignment by explicitly modeling reward uncertainty and decoupling KL regularization from the actor objective, with practical gains and scalable parameter efficiency.

Abstract

Reinforcement learning from human feedback (RLHF) emerges as a promising paradigm for aligning large language models (LLMs). However, a notable challenge in RLHF is overoptimization, where beyond a certain threshold, the pursuit of higher rewards leads to a decline in human preferences. In this paper, we observe the weakness of KL regularization which is commonly employed in existing RLHF methods to address overoptimization. To mitigate this limitation, we scrutinize the RLHF objective in the offline dataset and propose uncertainty-penalized RLHF (UP-RLHF), which incorporates uncertainty regularization during RL-finetuning. To enhance the uncertainty quantification abilities for reward models, we first propose a diverse low-rank adaptation (LoRA) ensemble by maximizing the nuclear norm of LoRA matrix concatenations. Then we optimize policy models utilizing penalized rewards, determined by both rewards and uncertainties provided by the diverse reward LoRA ensembles. Our experimental results, based on two real human preference datasets, showcase the effectiveness of diverse reward LoRA ensembles in quantifying reward uncertainty. Additionally, uncertainty regularization in UP-RLHF proves to be pivotal in mitigating overoptimization, thereby contributing to the overall performance.
Paper Structure (18 sections, 20 equations, 5 figures, 1 table)

This paper contains 18 sections, 20 equations, 5 figures, 1 table.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Reinforcement Learning from Human feedback
  4. Low-Rank Adaptations
  5. Methods
  6. Analysis of Regularizations in RLHF
  7. Training Diverse Reward LoRA Ensembles
  8. Overall Optimization Objectives
  9. Experimental Results
  10. Datasets and Training Setups
  11. Reward Model Evaluation
  12. Effect of Uncertainty Penalty
  13. Policy Model Evaluation
  14. Related Works
  15. Reinforcement Learning from Human Feedback
  16. ...and 3 more sections

Figures (5)

  • Figure 1: Illustration of UP-RLHF. Compared to RLHF, we train diverse reward LoRA ensemble in Step 2, and add uncertainty regularization in Step 3.
  • Figure 2: Illustration of training diverse reward LoRA ensembles.
  • Figure 3: With diversity regularization, our proposed diverse reward LoRA ensemble achieves better OOD detection capabilities.
  • Figure 4: Uncertainty penalty ablation on policy model evaluation in the summarization task over 4 different seeds.
  • Figure 5: Comparison of UP-RLHF and RLHF.