WARM: On the Benefits of Weight Averaged Reward Models

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Ramé, Alexandre, Vieillard, Nino, Hussenot, Léonard, Dadashi, Robert, Cideron, Geoffrey, Bachem, Olivier, Ferret, Johan
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866916101986713600
author Ramé, Alexandre
Vieillard, Nino
Hussenot, Léonard
Dadashi, Robert
Cideron, Geoffrey
Bachem, Olivier
Ferret, Johan
author_facet Ramé, Alexandre
Vieillard, Nino
Hussenot, Léonard
Dadashi, Robert
Cideron, Geoffrey
Bachem, Olivier
Ferret, Johan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2401_12187
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle WARM: On the Benefits of Weight Averaged Reward Models
Ramé, Alexandre
Vieillard, Nino
Hussenot, Léonard
Dadashi, Robert
Cideron, Geoffrey
Bachem, Olivier
Ferret, Johan
Machine Learning
Artificial Intelligence
Computation and Language
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.
title WARM: On the Benefits of Weight Averaged Reward Models
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2401.12187