Research map / Reinforcement learning / RL from feedback
WARM: On the Benefits of Weight Averaged Reward Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes WARM, which averages the weights of multiple reward models fine-tuned from shared pre-trained/SFT initializations, to mitigate reward hacking in RLHF with improved reliability under distribution shift and robustness to label corruption.
Problem
Reward hacking in RLHF arises because proxy reward models are unreliable under distribution shifts from policy drift and robustness is harmed by inconsistent/noisy binary preference labels.
Approach
Average the weights (not predictions) of M reward models fine-tuned from a shared pre-trained LLM with diverse hyperparameters, data orders, and SFT checkpoints, relying on linear mode connectivity; this gives a single model with no inference overhead that approximates prediction ensembling and, theoretically, keeps only run-invariant features, reducing memorization of corrupted labels.
Key results
- WARM with M=6 reaches a 99.8% win rate against the SFT policy after 3500 RL steps, the highest among compared policies.
- A policy RL fine-tuned with WARM M=6 has a 79.4% win rate against a policy trained with the single best RM phi_1.
- WARM selected summaries reach up to 92.5% win rate against random (SFT) selection in best-of-N.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Weight-average diverse reward model weights instead of using a single reward model | WARM reward modeling · TL;DR OOD test and RL alignment | control reward | Helped |
| Use weight averaging rather than prediction ensembling | Reward model aggregation · TL;DR OOD test | accuracy | Helped |
| Use weight averaging rather than prediction ensembling | Reward model aggregation · Corrupted training labels | accuracy on corrupted training examples | Hurt |
| Use weight averaging rather than prediction ensembling | Reward model aggregation · Clean training examples | accuracy | Hurt |
| Use weight averaging rather than prediction ensembling | Reward model aggregation · ID validation | accuracy | Mixed |
| Increase the number of averaged reward models (ablation) | WARM · TL;DR RL | control reward | Helped |
| Increase the number of averaged reward models to M=10 (ablation) | WARM RL policy · TL;DR RL, clean preference dataset | peak oracle preference win rate | No clear effect |
| Average checkpoints from a single reward-model fine-tuning instead of averaging final checkpoints from separate runs (ablation) | WARM reward modeling · TL;DR reward-model accuracy | accuracy | Hurt |
| Vary the KL regularization coefficient α (ablation) | WARM RL policy · TL;DR RL | control reward | Mixed |
| Use different SFT checkpoint initializations for reward-model featurizers (BAK) (ablation) | WARM reward models · TL;DR BoN | control reward | Helped |
| Use WARM M=6 to select summaries instead of other reward-model selection strategies | TL;DR BoN · TL;DR BoN, oracle preference | win rate against WARM M=6 selections | Helped |
| Train the policy with WARM M=6 rather than a single best individual reward model | RL policy · TL;DR RL, clean preference dataset | oracle preference win rate: 79.4% | Helped |
Surprising findings
- Weight averaging is slightly better than prediction ensembling under distribution shift and much less prone to memorizing corrupted training labels.
- Using M=10 reward models delays reward hacking but does not improve peak performance, attributed to weaker individual RMs in the selection.
Limitations
- Cannot exploit diversity from different architectures or pre-trainings, unlike prediction ensembling.
- Does not remove spurious correlations or biases inherent in preference data (e.g., length bias).
- Only enhances reward modeling and does not address other RLHF challenges.
Experimental setup
Datasets: Reddit TL;DR TL;DR preference dataset from Stiennon et al. PaLM-L AI-labeled preferences
Benchmarks: TL;DR summarization OOD test set (D_ood) TL;DR ID validation
Metrics: pairwise accuracy pointwise control reward pairwise oracle preference win rate KL divergence
Baselines: prediction ensembling (ENS) individual reward models SFT policy single RM phi_1
Builds on: RLHF model soups prediction ensembling RLAIF REINFORCE
Closest papers by meaning
- RM-Bench: Benchmarking Reward Models of Language Models with Subtlety and Style (ICLR 2025 · similarity 0.92)
- RRM: Robust Reward Model Training Mitigates Reward Hacking (ICLR 2025 · similarity 0.91)
- Scaling Laws for Reward Model Overoptimization in Direct Alignment Algorithms (NeurIPS 2024 · similarity 0.91)
- Provably Efficient Online RLHF with One-Pass Reward Modeling (NeurIPS 2025 · similarity 0.90)
- Exploring Data Scaling Trends and Effects in Reinforcement Learning from Human Feedback (NeurIPS 2025 · similarity 0.90)
- Post-hoc Reward Calibration: A Case Study on Length Bias (ICLR 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.