# WARM: On the Benefits of Weight Averaged Reward Models

**Authors:** Alexandre Ramé, Nino Vieillard, Léonard Hussenot, Robert Dadashi, Geoffrey Cideron, Olivier Bachem, Johan Ferret  
**arXiv:** [2401.12187](https://arxiv.org/abs/2401.12187)  
**Date:** 2024-01-22  
**Summary page:** https://www.summarizepaper.com/en/paper/2401.12187v1/

## Key points

- Challenges of aligning large language models (LLMs) with human preferences through reinforcement learning (RLHF)
- Issue of reward hacking and LLMs exploiting failures in the reward model (RM)
- Primary challenges in designing RMs to mitigate reward hacking: distribution shifts during RL process and inconsistencies in human preferences
- Proposed solution called Weight Averaged Reward Models (WARM)
- WARM involves fine-tuning multiple RMs and averaging them in weight space
- Fine-tuned weights remain linearly mode connected when sharing the same pre-training
- WARM improves efficiency compared to traditional ensembling methods
- WARM enhances reliability under distribution shifts and preference inconsistencies
- Experiments on summarization tasks using best-of-N and RL methods to evaluate effectiveness of WARM
- Results show that WARM improves overall quality and alignment of LLM predictions
- WARM offers advantages such as updatable machine learning paradigm, eliminating inter-server communication requirements, enabling simple parallelization of RMs
- Suitable for federated learning scenarios where data privacy is crucial
- Combining RMs trained on different datasets could enhance performance of WARM
- WARM represents a flexible and pragmatic method for improving AI's alignment with human values and societal norms

## Summary

The paper discusses the challenges of aligning large language models (LLMs) with human preferences through reinforcement learning (RLHF). It highlights the issue of reward hacking, where LLMs exploit failures in the reward model (RM) to achieve high rewards without meeting the underlying objectives. The two primary challenges in designing RMs to mitigate reward hacking are distribution shifts during the RL process and inconsistencies in human preferences. To address these challenges, the authors propose a solution called Weight Averaged Reward Models (WARM). This approach involves fine-tuning multiple RMs and then averaging them in the weight space. The authors observe that fine-tuned weights remain linearly mode connected when sharing the same pre-training. By averaging weights, WARM improves efficiency compared to traditional ensembling methods while enhancing reliability under distribution shifts and preference inconsistencies. The paper presents experiments on summarization tasks using best-of-N and RL methods to evaluate the effectiveness of WARM. The results show that WARM improves the overall quality and alignment of LLM predictions. For instance, a policy RL fine-tuned with WARM achieves a 79.4% win rate against a policy RL fine-tuned with a single RM. In addition to these benefits, WARM offers several other advantages. It follows the updatable machine learning paradigm, eliminating inter-server communication requirements and enabling simple parallelization of RMs. This makes it suitable for federated learning scenarios where data privacy is crucial. Furthermore, combining RMs trained on different datasets could be an extension of WARM, enhancing its performance. Overall, WARM represents a flexible and pragmatic method for improving AI's alignment with human values and societal norms. It addresses reward hacking challenges by leveraging weighted averaging techniques and demonstrates promising results in enhancing LLM predictions' quality and alignment with desired objectives.