WARM: On the Benefits of Weight Averaged Reward Models

Alexandre Ramé; Nino Vieillard; Léonard Hussenot; Robert Dadashi; Geoffrey Cideron; Olivier Bachem; Johan Ferret

Source-linked AI summary

WARM: On the Benefits of Weight Averaged Reward Models

Alexandre Ramé, Nino Vieillard, Léonard Hussenot, Robert Dadashi, Geoffrey Cideron, Olivier Bachem, Johan Ferret

arXiv:2401.12187v1Jan 22, 2024cs.LGcs.AIcs.CL

TL;DR

Reward hacking in RLHF reflects distribution shifts and inconsistent preference labels that make reward models unreliable. WARM fine-tunes multiple reward models from shared pre-training and averages their weights, improving efficiency, reliability, and robustness; on summarization, a WARM-trained RL policy achieved a 79.4% win rate against one trained with a standard RM.

Problem
RLHF reward models can be exploited under distribution shifts and inconsistent preference labels, causing high proxy rewards without satisfying intended human preferences.
Method
WARM fine-tunes multiple diverse reward models from shared pre-training and averages their weights into one proxy reward model.
Results
WARM improves summarization performance without memory or inference overhead, and a WARM-trained RL policy achieved a 79.4% win rate against a policy trained with a standard RM.
Takeaways & Limitations
WARM provides an efficient reward-modeling strategy that improves reliability under distribution shifts and robustness to label corruption in summarization experiments.
Takeaways & Limitations
WARM cannot combine reward models from different architectures or pre-trainings and does not eliminate all spurious correlations or biases in preference data.

Abstract

from arXiv · show

Aligning large language models (LLMs) with human preferences through reinforcement learning (RLHF) can lead to reward hacking, where LLMs exploit failures in the reward model (RM) to achieve seemingly high rewards without meeting the underlying objectives. We identify two primary challenges when designing RMs to mitigate reward hacking: distribution shifts during the RL process and inconsistencies in human preferences. As a solution, we propose Weight Averaged Reward Models (WARM), first fine-tuning multiple RMs, then averaging them in the weight space. This strategy follows the observation that fine-tuned weights remain linearly mode connected when sharing the same pre-training. By averaging weights, WARM improves efficiency compared to the traditional ensembling of predictions, while improving reliability under distribution shifts and robustness to preference inconsistencies. Our experiments on summarization tasks, using best-of-N and RL methods, shows that WARM improves the overall quality and alignment of LLM predictions; for example, a policy RL fine-tuned with WARM has a 79.4% win rate against a policy RL fine-tuned with a single RM.

1. Introduction

RLHF reward models can be exploited because policy distributions shift during training and preference labels are inconsistent. WARM addresses these challenges by averaging diverse reward-model weights, improving efficiency, reliability, robustness, and downstream alignment.

Motivation: Reward hacking occurs when policies exploit proxy-RM loopholes to obtain high rewards without satisfying intended human preferences.This can produce linguistically flawed or unnecessarily verbose outputs and complicate checkpoint selection.
Challenges: The method targets two central reward-modeling difficulties: distribution shifts during RL and inconsistencies in preference labels.These challenges undermine reliable scoring as policy generations move away from offline preference data.
Approach: WARM averages the weights of multiple reward models fine-tuned from shared pre-training, using weight-space interpolation rather than prediction ensembling.The individual runs can differ through hyperparameters and other fine-tuning choices.
Benefits: WARM improves reliability under distribution shifts and robustness to label corruption while requiring only one model at inference time.This provides an efficient alternative to prediction ensembling without its memory and inference overhead.
Results: 79.4% win rate: a policy trained with WARM beat one trained with a standard RM on the preference-oracle metric.The result was observed for summarization policies trained with RL.

2. Context and challenges

Reward models guide RLHF because true environmental rewards are unavailable, but their reliability is threatened by distribution shifts and inconsistent preference labels. Existing remedies trade off reliability, robustness, efficiency, or additional data and tuning requirements.

Reward modeling: Reward models provide scalar preference scores for generated responses and guide RL or best-of-N selection toward human preferences.They are trained on preference data because the oracle reward is not directly available.
Challenges in reward modeling: Distribution shifts arise because offline preference data differ from policy generations and intensify as RL moves the policy away from SFT initialization.Out-of-distribution inputs can reduce performance, impair model selection, and produce poorly calibrated extreme rewards.
Challenges in reward modeling: Inconsistent preferences arise from labeler fatigue, misunderstanding, imperfect incentives, heterogeneous opinions, and simplifying criteria for complex tasks.These factors lower inter-rater agreement and make preference data an imperfect representation of the underlying ground truth.
Desired properties: A good reward model should be efficient, reliable under distribution shifts, and robust to inconsistencies in binary preference labels.These properties support efficient policy optimization, exploration away from initialization, and learning from robust signals.
Existing approaches: Existing approaches include KL regularization, collecting new preference data, active learning, label smoothing or flipping, and prediction ensembling.These approaches can introduce underfitting, annotation and computation costs, or memory and inference overhead.

3. WARM

WARM combines diverse reward-model fine-tunings by averaging their weights, relying on shared pre-training to preserve linear mode connectivity. The resulting single proxy model is designed to retain efficiency while improving reliability and robustness.

3.1. Weight averaging of reward models: WARM forms a single proxy reward model by averaging the weights of multiple diverse fine-tunings initialized from shared pre-training.The averaged model then guides the RL procedure like an individual reward model.
3.2. Linear mode connectivity: Weight-space interpolation is justified by linear mode connectivity: fine-tuned weights sharing pre-training can be interpolated without losing the relevant accuracy relationship.The paper empirically validates this property on out-of-distribution test samples.
3.3. Sources of diversity: WARM requires a balance between shared pre-training, which preserves connectivity, and diverse fine-tunings, which contribute to accuracy gains.Weights that are identical lack useful diversity, while training from scratch does not preserve the same connectivity guarantee.
3.3. Sources of diversity: Diversity can come from data order, learning rates, dropout probabilities, and checkpoints collected along an SFT trajectory.Averaging checkpoints from one RM fine-tuning produced lower results because accuracy and diversity were poorly balanced.
3.3. Sources of diversity: Baklava increases diversity by initializing reward models from different SFT checkpoints while retaining the same pre-training requisite for linear mode connectivity.It avoids auxiliary tasks and adds no overhead compared with initialization from only the final SFT checkpoint.

4. On the benefits of WARM

WARM compares weight averaging with prediction ensembling for reward models, showing similar first-order behavior but stronger robustness as data moves away from training distributions. The analysis attributes these gains to reduced variance and emphasis on predictive mechanisms shared across runs.

4.1. 1st order analysis: weight averaging for reliable and more efficient ensembling: WA performs similarly to ENS while improving over individual reward models on OOD data, with a single weight set avoiding ENS’s memory and inference overhead.The experiments interpolate between reward-model weights across multiple diversity sources, including training steps, data order, learning rates, and initialization.
4.2. 2nd order analysis: weight averaging for more robust ensembling: WA’s accuracy advantage over ENS grows as data moves farther from the training distribution.Under 25% corrupted training labels, WA is worse on clean training samples but better on OOD test samples.
4.2. 2nd order analysis: weight averaging for more robust ensembling: WA reduces memorization of corrupted labels while generalizing better than prediction ensembling.The comparison reports WA ≪ ENS on corrupted training samples and WA ≥ ENS on OOD test samples.
4.3. Weight averaging enforces invariance across runs: WA’s averaged predictions tend toward expected predictions across runs, providing theoretical support for its reduced reliance on low-probability features.The simplified analysis treats WA as a regularizer toward invariant predictive mechanisms.
4.3. Weight averaging enforces invariance across runs: The analysis argues that WA favors predictive mechanisms learned consistently across independent runs and downweights run-specific features linked to memorization.This connects weight averaging with invariance-based reasoning for robustness under distribution shifts.
4. On the benefits of WARM: WARM combines efficiency, reliability under distribution shifts, and robustness to noisy preferences by returning one reward model.The paper’s setup uses the TL;DR summarization benchmark and compares WA with ENS as the central empirical analysis.

5. Experiments

Experiments on summarization evaluate WARM in best-of-N selection and RL under clean and 25% corrupted preference labels. WARM consistently improves control rewards and oracle preference outcomes over individual RMs and prediction ensembling.

Best-of-N experiments: WARM produced higher control rewards than individual RMs and prediction ensembling, with larger gains when averaging six weights.The BoN experiments used PaLM- and T5-based SFT policies under clean and 25% corrupted preference data.
Best-of-N experiments: WARM-selected summaries reached up to 92.5% win rate against random SFT selection in best-of-N experiments.All selection strategies beat SFT, but none beat summaries selected by WARM M=6.
RL experiments: During RL, increasing the number of averaged weights moved the control-reward Pareto front toward higher reward and lower KL.Prediction ensembling remained susceptible to early reward hacking and achieved significantly lower absolute control rewards than WARM.
RL experiments: The KL-regularization coefficient strongly affected RL outcomes: low α produced high KL, whereas high α reduced KL but risked underfitting.The reported optimal α for WARM was lower than for a single RM because WARM mitigated reward hacking.
RL experiments: 99.8% win rate against the SFT policy was reached by WARM M=6 after 3500 RL steps.No other policy beat this WARM M=6 reference in the reported oracle comparison.
RL experiments: 79.4% win rate was achieved by a WARM-trained policy against a policy RL fine-tuned with the best individual RM.The comparison uses the policy trained with WARM M=6 after 3500 steps against the policy trained with φ1 after 3000 steps.

6. Discussion

The discussion highlights WARM’s efficiency, parallelizability, and potential robustness benefits, while identifying limitations relative to prediction ensembling and persistent bias in preference data. WARM improves reward modeling but does not address every RLHF challenge.

Benefits: WARM requires one inference-time model, avoiding the memory and inference burdens of prediction ensembling.The method also supports simple parallelization and could be applied in federated-learning settings where preference data remains private.
Benefits: Weight averaging may reduce memorization, support evolving preferences, and combine models trained on different datasets or labeler groups.The discussion presents these as extensions or potential benefits rather than established experimental findings.
Limitations: Prediction ensembling can combine RMs from different architectures and pre-trainings, whereas WARM is limited to shared-pretraining weight averaging.Ensembling can also use prediction disagreement for uncertainty estimation and limiting model drift.
Limitations: WARM does not eliminate spurious correlations or biases when individual RMs share the same flawed criterion, such as summary length.The discussion points to invariance regularization or last-layer retraining as possible alternatives.
Limitations: WARM enhances reward modeling but does not solve the other challenges involved in RLHF.This places the method’s contribution specifically within reward-model design.

7. Conclusion

The paper introduces WARM, which averages weights from diverse reward-model fine-tunings to improve reliability under distribution shifts and robustness to label corruption. It applies this efficient approach to summarization and reports empirical effectiveness for mitigating reward hacking.

7. Conclusion: WARM averages weights from multiple fine-tuned reward models to address reliability under distribution shifts and robustness under label corruption.The approach builds on linear mode connectivity across fine-tuned weights and targets efficient reward modeling.
7. Conclusion: WARM is presented as an efficient solution for mitigating reward hacking in reinforcement learning from human feedback.The conclusion reports empirical effectiveness on summarization tasks and anticipates applications to aligned, transparent, and effective AI systems.
7. Conclusion: Weight averaging can flatten loss landscapes, reduce estimator variance, combine model abilities, and provide benefits under spurious correlations.These properties motivate applying weight averaging to reward models.
7. Conclusion: Reward-model design remains important because alignment lacks explicit environmental rewards and existing work has comparatively underemphasized this problem.Prior efforts address reward losses, data adaptation, active learning, tools, and label corruption, but the paper positions WARM at their intersection.

B.1. Dataset details

The experiments use Reddit TL;DR summarization data, AI-generated preference labels, and multiple reward models trained from shared SFT-derived components. Models are ranked by OOD accuracy, with WARM averaging the highest-ranked weights while ENS averages predictions.

B.1. Dataset details: The Reddit TL;DR dataset contains 123k posts, with approximately 5% held out for in-distribution validation and a 92k-comparison OOD dataset generated from varied PaLM-XS policies.The OOD candidates include pretrained-only and supervised-fine-tuned policies with high-temperature sampling.
B.1. Dataset details: Preference labels come from a greedily decoded, chain-of-thought PaLM-L RLAIF procedure, with both response orderings used to reduce position bias.The procedure was reported to perform similarly to human labellers in inter-agreement.
B.1. Dataset details: Reward models use PaLM-XXS featurizers from shared SFT training, launched at 8k, 10k, and 12k checkpoints, with a shared linear classification layer.Additional training uses 10k steps, mild learning-rate and dropout ranges intended to preserve linear mode connectivity.
B.1. Dataset details: Ten reward models are ranked by decreasing OOD accuracy, and WARM averages the weights of the M best models while ENS averages their predictions.For M=2, the selected models are the two highest-performing reward models under distribution shifts.
B.1. Dataset details: The selection procedure aims to remove potentially poor reward models, and Figure 10 validates choosing models from best to worst as a reliable heuristic.The figure compares accuracy after averaging M weights selected by different procedures.
B.1. Dataset details: Policy and value models are PaLM-XS models initialized from the same SFT model and trained with modified REINFORCE plus a baseline value function for variance reduction.The main experiments use temperature 0.9, a 2k-step policy warmup, and KL coefficients of 0.003 without corruption and 0.01 with corruption.

C.1. 2nd order analysis: weight averaging for more robust ensembling

The second-order analysis extends earlier experiments across reward-model pairs, architectures, and initialization settings. Across the reported BoN comparisons, WARM outperforms ENS and individual networks on pointwise control reward, while increasing N improves oracle preference performance and limits reward hacking.

C.1. 2nd order analysis: weight averaging for more robust ensembling: The analysis covers train-corrupt, train-clean, validation-ID, and test-OOD settings using additional pairs of reward models.Figures 11–14 enrich the corresponding earlier analyses with different reward-model pairings.
C.1. 2nd order analysis: weight averaging for more robust ensembling: WARM performs better than ENS and individual networks in pointwise control reward for BoN experiments using both PaLM- and T5-based SFT policies.Figure 15 reports absolute control-reward values for both candidate-summary architectures.
C.1. 2nd order analysis: weight averaging for more robust ensembling: The BoN experiments compare generations from PaLM and T5 architectures, including Baklava with PaLM and Baklava with T5.These architecture-specific settings support the broader comparison across candidate-generation policies.
C.1. 2nd order analysis: weight averaging for more robust ensembling: The control-reward analysis also considers Baklava fine-tunings with different featurizer initializations taken from a shared SFT model.The two fine-tunings are collected at steps 12k and 8k.
C.1. 2nd order analysis: weight averaging for more robust ensembling: Increasing BoN sample count N improves performance and limits reward hacking relative to RL in the clean setup.The comparison includes WARM with M=6 and N=1000 against other reward models with 1 ≤ N ≤ 1000.

C.3.1. Experiments with corrupted preference dataset

The corrupted-preference experiments repeat the RL analysis with 25% corruption in the preference dataset.

C.3.1. Experiments with corrupted preference dataset: The RL experiments evaluate the same setup as Figure 1(b) after introducing 25% corruption into the preference dataset.

C.3.2. Experiments with clean preference dataset

The clean-preference RL experiments compare WARM against other policies at fixed training steps and examine its performance against a single reward model across KL-regularization strengths. They also report control rewards for two regularization settings.

C.3.2. Experiments with clean preference dataset: Figure 19 compares the fixed-step win rate of WARM with M=6 against other policies and against the best single RM across KL strengths.The single RM is φ1, selected by OOD accuracy.
C.3.2. Experiments with clean preference dataset: Figure 20 reports control rewards for the clean setup with α=0.01.
C.3.2. Experiments with clean preference dataset: Figure 21 reports control rewards for the clean setup with α=0.001.

C.4. Distillation experiments

The distillation experiments reproduce a setup where a control reward model labels data for smaller reward models, then compare WARM with individual models and prediction ensembling. The authors note that distillation changes diversity across fine-tuned reward models, motivating exploration of a more realistic RLAIF setup.

C.4. Distillation experiments: The reproduced setup uses a control PaLM-XS reward model to generate labels for PaLM-XXS reward models.
C.4. Distillation experiments: Distillation changes diversity across fine-tuned reward models, potentially altering the significance of the distillation setup and motivating RLAIF experiments.
C.4. Distillation experiments: In the distillation best-of-N experiment, WARM performs higher than individual reward models and prediction ensembling.The figure compares WARM with M weights, individual RMs, and ENS.