PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Reinforcement Learning from Human Feedback with Efficient Reward Model Ensemble

Shun Zhang, Zhenfang Chen, Sunli Chen, Yikang Shen, Zhiqing Sun, Chuang Gan

arXiv:2401.16635v3cs.LGcs.AIcs.CL

TL;DR

RLHF depends on preference-trained reward models, whose inaccurate predictions can contribute to reward overoptimization. This paper proposes efficient reward-model ensembles, including linear-layer and LoRA-based designs, and reports that ensembled rewards outperform standard RLHF. The authors identify their model architecture as a scope assumption: a Transformer with a linear reward layer.

  • Problem

    Reward models trained on offline preference data may predict inaccurately on out-of-distribution inputs, allowing RL optimization to produce misaligned outputs with high estimated rewards.

    Source: §1 · p. 1

  • Method

    The paper proposes efficient reward ensembles, including shared-Transformer linear-layer ensembles and LoRA-based ensembles.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    The authors report that RLHF with their ensembled reward models outperforms standard RLHF, with consistently higher win rates when using reward-model ensembles.

    Source: §1 · p. 2, §4 · p. 4

  • Takeaways & Limitations

    The authors conclude that LoRA-based reward ensembling is effective under computational constraints and can improve alignment performance on AlpacaEval and MT-Bench.

    Source: §5 · p. 5, §4 · p. 5

    The paper assumes reward models consist of a Transformer and a linear layer that maps the Transformer's last hidden layer to a reward value.

    Source: §3.1 · p. 2

Abstract

from arXiv · show

Reinforcement Learning from Human Feedback (RLHF) is a widely adopted approach for aligning large language models with human values. However, RLHF relies on a reward model that is trained with a limited amount of human preference data, which could lead to inaccurate predictions. As a result, RLHF may produce outputs that are misaligned with human values. To mitigate this issue, we contribute a reward ensemble method that allows the reward model to make more accurate predictions. As using an ensemble of large language model-based reward models can be computationally and resource-expensive, we explore efficient ensemble methods including linear-layer ensemble and LoRA-based ensemble. Empirically, we run Best-of-$n$ and Proximal Policy Optimization with our ensembled reward models, and verify that our ensemble methods help improve the alignment performance of RLHF outputs.

1 Introduction

RLHF aligns language models through preference-trained reward models, but inaccurate predictions on out-of-distribution data can lead to reward overoptimization. The paper proposes efficient reward ensembles and evaluates them with Best-of-n and PPO.

1 Introduction

  • RLHF uses preference data to train a reward model, then optimizes a supervised-finetuned model with that reward.The resulting sequences have higher rewards and are presumed to be more aligned with human values.Source: §1 · p. 1
  • Inaccurate reward predictions on out-of-distribution data can cause reinforcement learning to produce misaligned outputs with incorrectly high estimated rewards.The paper identifies this problem as reward hacking or reward overoptimization.Source: §1 · p. 1
  • The paper motivates ensembles as a way to improve reward predictions, while noting that training and loading multiple large language model-based reward models is costly.Its stated aim is to make reward ensembling more efficient.Source: §1 · pp. 1–2
  • The paper proposes linear-layer and LoRA-based reward ensembles as efficiency-oriented alternatives to independently training and loading multiple large reward models.The approaches trade off computational efficiency and alignment performance.Source: §1 · pp. 1–2
  • The authors evaluate ensembled reward models with Best-of-n and PPO on AlpacaEval and MT-Bench, reporting better performance than standard RLHF.The evaluation uses well-accepted benchmarks, as described by the authors.Source: §1 · p. 2

2 Background and Related Work

The related work situates reward ensembles within RLHF, uncertainty estimation, and offline reinforcement learning. The paper distinguishes its focus on efficient ensemble designs from concurrent and adjacent approaches.

Reinforcement learning from human feedback

  • RLHF developed from learning rewards from human feedback in TAMER and deep reinforcement learning, and was later used to finetune language models toward human values and preferences.The paper points readers to its introduction and appendix for the standard RLHF framework and algorithm details.Source: §2 · p. 2

Ensemble models and uncertainty estimation

  • Prior work uses reward-model ensembles for uncertainty-informed exploration and active learning, while broader deep-learning work uses ensembles for predictive uncertainty estimation.Lakshminarayanan et al. (2017) report better performance than Bayesian neural networks; Liang et al. (2022) and Gleave and Irving (2022) apply reward ensembles in reinforcement learning.Source: §2 · p. 2
  • Concurrent work by Coste et al. (2023) also applies reward ensembles to reward overoptimization, while this paper focuses on making ensembling efficient.Other related work studies ensembling during pretraining and finetuning, LoRA with uncertainty-penalized RL objectives, weight averaging, and shared-encoder ensembles.Source: §2 · p. 2

Offline reinforcement learning

  • The paper draws on conservative offline RL, where CQL learns a conservative Q function to mitigate overestimation for out-of-distribution state-action pairs.Inspired by CQL, the paper uses the lower confidence bound of ensembled predictions for conservative reward predictions.Source: §2 · p. 2

3 Reward Model Ensemble

The paper defines reward-model ensembles built from Transformer-based reward models, then presents architectures that trade parameter sharing against model diversity. It combines ensemble predictions using either their mean or a lower confidence bound.

3 Reward Model Ensemble

  • A reward model takes an instruction and response as input and predicts a reward indicating the response's alignment performance.The paper assumes a Transformer whose last hidden layer feeds a linear reward-output layer.Source: §3 · p. 2, §3.1 · p. 2

3.1 Architecture Design of Reward Model Ensemble

  • Independent reward models require training and loading k(M+L) parameters, where M is the Transformer size and L is the linear-layer size.The models are trained with different random seeds and ensembled at inference.Source: §3.1 · p. 3
  • A linear-layer ensemble shares one Transformer across models and trains separate reward heads, requiring M+kL parameters for an ensemble of size k.The shared Transformer and all linear layers are trained.Source: §3.1 · p. 3, Figure 1 · p. 3
  • A LoRA-based ensemble adds a separate adapter and linear layer for each model, requiring M+kL+kA parameters, where A is adapter size.The paper first finetunes the Transformer and linear layers on a subset of preference data, then trains the ensemble using the remaining data.Source: §3.1 · p. 3

3.2 Predictions of Ensembled Reward Models

  • The paper combines ensemble predictions using either their mean or a lower confidence bound that subtracts β times their standard deviation.The lower confidence bound is described as conservative; its performance is comparable to mean prediction.Source: §3.2 · pp. 3–4

4 Empirical Evaluation

The evaluation compares reward ensembles with standard RLHF across Best-of-n and PPO on AlpacaEval and MT-Bench, testing whether ensemble architecture and prediction method affect alignment.

Evaluation setup

  • Best-of-n selects the highest-reward of n SFT samples, while PPO finetunes the SFT model; AlpacaEval win rate compares outputs with a GPT-3-generated baseline.Best-of-n uses n = 50, 100, and 200; PPO checkpoints are evaluated every 100 training steps.Source: §4 · p. 4
  • Models start from SFT10k; LoRA training first fully finetunes the Transformer on noisy preferences, then trains ensemble layers and adapters on human preferences.SFT10k is a Llama-7b model; the experiments also use AlpacaFarm training and evaluation datasets.Source: §4 · p. 4

Results

  • Win rates are consistently higher with reward-model ensembles; for Best-of-n, single-model and LoRA-based ensembles perform best.Figure 2 reports AlpacaEval win rates; PPO could not use the independently trained single-model ensemble because it required loading multiple reward models.Source: §4 · p. 4, Figure 2 · p. 5
  • LoRA-based ensemble has the best performance in the PPO comparison.The passage places this finding in the context of PPO training.Source: §4 · p. 5
  • MT-Bench results are consistent with AlpacaEval, and the authors report that LoRA-based ensembling improves alignment despite not fully finetuning the Transformer.The evaluation reports scores after the first turn, second turn, and their average.Source: §4 · p. 5, Table 1 · p. 5
  • Mean reward prediction and LCB have similar performance.The paper reports detailed prediction-method results in Sec. B.Source: §4 · p. 5

5 Discussion and Conclusion

The paper concludes that efficient reward-model ensembles can improve language-model alignment, with LoRA-based ensembling effective under computational constraints.

5 Discussion and Conclusion

  • LoRA-based reward-model ensembling is effective under computational constraints, and future work will extend the approach to other LLM training and inference steps.The proposed future direction includes sample-efficient reward-model training.Source: §5 · p. 5

A Preliminaries

The appendix provides background on reinforcement learning from human feedback and LoRA finetuning.

A Preliminaries

A.1 Reinforcement Learning from Human Feedback

RLHF combines supervised fine-tuning, reward modeling, and reinforcement learning. This section explains preference-based reward training and PPO finetuning.

A.1 Reinforcement Learning from Human Feedback

  • RLHF proceeds through supervised fine-tuning, reward modeling, then reinforcement learning with the learned reward model.Source: §A.1 · p. 7
  • Reward training uses context/preferred/less-preferred triples and learns scores where larger values indicate greater human preference.Source: §A.1 · p. 7
  • The sigmoid-based loss encourages higher rewards for human-preferred responses; PPO then finetunes the SFT policy by balancing policy divergence and cumulative reward.Source: §A.1 · p. 7

A.2 LoRA Finetuning

LoRA is presented as an efficient way to fine-tune pretrained language models by learning a low-rank weight change. The original weights stay frozen while the LoRA matrices are updated.

A.2 LoRA Finetuning

  • LoRA is used to fine-tune a pretrained large language model more efficiently when updating all parameters would be computationally expensive and resource-demanding.The paper describes LoRA as learning smaller matrices for each Transformer layer.Source: §A.2 · p. 7
  • The weight change ΔW is represented as A1A2, where A1 and A2 are smaller matrices with dimensions d1 × r and r × d2.The rank r is smaller than both d1 and d2.Source: §A.2 · p. 7
  • During training, only A1 and A2 are updated while the original weight matrix W remains frozen, and the modified W′ is used in the forward pass.The smaller set of trained parameters influences all model weights.Source: §A.2 · p. 7

B Additional Empirical Results

Additional AlpacaEval experiments compare mean and lower confidence bound reward predictions for the ensembles. Their differences are reported as insignificant.

B Additional Empirical Results

C Hyperparameters

The appendix identifies hyperparameter tables for reward modeling, PPO, and decoding, alongside figures depicting LCB variants for the two ensemble approaches.

C Hyperparameters

Related papers

Loading 2401.16635v3…