WARM: On the Benefits of Weight Averaged Reward Models

2024/01/22 by Alexandre Ramé, Ramé, Alexandre, Nino Vieillard +11 · 34 citations
Computer Science · #Adversarial Robustness in Machine Learning #Artificial Intelligence (cs.AI) #Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Natural Language Processing Techniques #Topic Modeling

paper · pdf · doi:10.48550/arxiv.2401.12187

openalex publication_date 2024/01/22 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Cited by

• Targeting Misalignment: A Conflict-Aware Framework for Reward-Model-based LLM Alignment
• UMM-RM: An Upcycle-and-Merge MoE Reward Model for Mitigating Reward Hacking
• Bootstrapping LLMs via Preference-Based Policy Optimization
• A Systematic Study of In-the-Wild Model Merging for Large Language Models
• PIRA: Preference-Oriented Instruction-Tuned Reward Models with Dual Aggregation
• Meta-Learned Reward Shaping for Reinforcement Learning from Human Feedback
• Debiasing Reward Models by Representation Learning with Guarantees
• Direct Preference Optimization with Unobserved Preference Heterogeneity: The Necessity of Ternary Preferences
• Information-Theoretic Reward Modeling for Stable RLHF: Detecting and Mitigating Reward Hacking
• Enhancing Large Language Model Reasoning with Reward Models: An Analytical Survey
• Provably Mitigating Corruption, Overoptimization, and Verbosity Simultaneously in Offline and Online RLHF/DPO Alignment
• Noise-Guided Transport for Imitation Learning
• Diversity First, Quality Later: A Two-Stage Assumption for Language Model Alignment
• Sotopia-RL: Reward Design for Social Intelligence
• Alleviating Attention Hacking in Discriminative Reward Modeling through Interaction Distillation
• Off-Policy Corrected Reward Modeling for Reinforcement Learning from Human Feedback
• Intra-Trajectory Consistency for Reward Modeling
• Bradley-Terry and Multi-Objective Reward Modeling Are Complementary
• Reward Models Can Improve Themselves: Reward-Guided Adversarial Failure Mode Discovery for Robust Reward Modeling
• Don't Mix Rewards, Mix Policies: Policy Decomposition and Optimization for Multi-Reward RL
• Pre-Trained Policy Discriminators are General Reward Models
• Inference-Time Reward Hacking in Large Language Models
• AutoRule: Reasoning Chain-of-thought Extracted Rule-based Rewards Improve Preference Learning
• Towards Reward Fairness in RLHF: From a Resource Allocation Perspective
• Proximalized Preference Optimization for Diverse Feedback Types: A Decomposed Perspective on DPO
• Learning a Pessimistic Reward Model in RLHF
• Reward Model Overoptimisation in Iterated RLHF
• Toward Evaluative Thinking: Meta Policy Optimization with Evolving Reward Models
• Alignment Tampering: How Reinforcement Learning from Human Feedback Is Exploited to Optimize Misaligned Biases
• A Model Zoo on Phase Transitions in Neural Networks
• Establishing Reliability Metrics for Reward Models in Large Language Models
• Energy-Based Reward Models for Robust Language Model Alignment
• A Model Zoo of Vision Transformers
• Adversarial Training of Reward Models

Related

• WARP: On the Benefits of Weight Averaged Rewarded Policies
• RRM: Robust Reward Model Training Mitigates Reward Hacking
• Two Minds Better Than One: Collaborative Reward Modeling for LLM Alignment
• On the Robustness of Reward Models for Language Model Alignment
• Energy-Based Reward Models for Robust Language Model Alignment