WARM: On the Benefits of Weight Averaged Reward Models

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.

한 줄 요약

WARM은 가중치 평균을 통해 생성된 보상 모델로, 보상 해킹을 완화하고 정책 학습의 정확도를 79.4%까지 향상시킨다.

핵심 기여도

WARM은 여러 보상 모델의 가중치를 평균하는 방식으로, 보상 해킹을 완화한다.
가중치 평균은 동일한 사전 학습 기반에서 선형 모드 연결성을 유지함으로써 효율적이다.
TL;DR 요약 작업에서 WARM 기반 정책의 승률은 단일 RM 기반 정책 대비 79.4%로 개선됨.
WARM은 예측 앙상블 대비 메모리 및 추론 부담을 줄이며, 분산 학습 환경에서 활용 가능하다.

핵심 아이디어

WARM은 RLHF 과정에서 발생하는 보상 해킹 문제를 해결하기 위해, 여러 보상 모델의 가중치를 평균하는 방식을 제안한다. 기존의 예측 앙상블(ENS)은 여러 모델의 출력을 평균하는 방식으로, 메모리와 추론 효율성에 제약이 있었다. WARM은 대신, 동일한 사전 학습 모델에서 파생된 여러 보상 모델의 가중치를 평균함으로써, 선형 모드 연결성(linear mode connectivity)을 활용한다. 이는 보상 모델이 분포 이동(distribution shift)에 더 안정적으로 반응하도록 도와주며, 라벨 불일치(label inconsistency)에 대한 내성을 향상시킨다.

WARM의 핵심 통찰은, 보상 모델의 가중치가 사전 학습 기반에서 선형적으로 연결될 수 있다는 사실에 기반한다. 이는 여러 학습 조건(하이퍼파라미터, 라벨 순서 등)에서 생성된 보상 모델 간의 가중치 평균이 일반화 능력을 향상시킬 수 있음을 의미한다.

기술적 접근법

**모델**: WARM은 PaLM-XXS 기반의 보상 모델을 사용하며, 여러 하이퍼파라미터 조건에서 사전 학습된 모델을 미세 조정(fine-tuning)한다.
**가중치 평균**: 각 보상 모델의 가중치를 선형적으로 평균하여 최종 보상 모델을 생성한다.
**데이터**: TL;DR 요약 작업에서 PaLM-L 모델을 사용해 체인-오브-사고(chain-of-thought) 방식으로 라벨을 생성한 데이터셋을 사용한다.
**비교 대상**: WARM은 ENS(예측 앙상블)와 비교되며, ENS는 여러 보상 모델의 출력을 평균하는 전통적 방법이다.
**하이퍼파라미터**: 각 보상 모델은 서로 다른 하이퍼파라미터와 라벨 순서를 사용하여 학습된다.

주요 결과

**TL;DR 요약 작업에서**: WARM 기반 정책의 승률은 79.4%로, 단일 RM 기반 정책 대비 +59.4% 개선됨.
**OOD 데이터셋에서**: PaLM-XS 기반의 보상 모델이 80.1%의 정확도를 달성함.
**보상 해킹 감지**: WARM은 ENS 대비 메모리 사용량과 추론 시간을 줄이며, 라벨 오염(label corruption)에 대한 내성을 향상시킴.
**정확도 향상**: WARM은 분포 이동 상황에서도 안정적인 보상 신호를 제공함.

의의 및 한계

WARM은 보상 모델의 신뢰성과 안정성을 향상시키며, RLHF 과정에서 발생하는 보상 해킹 문제를 완화한다. 특히, 선형 모드 연결성을 활용한 가중치 평균 전략은 효율성과 확장성을 동시에 달성한다. 또한, WARM은 분산 학습 환경에서 사용 가능하며, 개인화된 정책 학습 및 다목적 최적화에도 활용 가능하다.

하지만, WARM은 ENS 대비 모델 아키텍처와 사전 학습 다양성을 활용하지 못한다는 한계가 있다. 또한, WARM은 라벨 데이터 내의 스파urious 상관관계를 완전히 제거하지 못하며, 보상 모델 외의 RLHF 과정의 문제(예: 정책 드리프트)는 해결하지 못한다. 따라서, WARM은 책임 있는 AI 전략의 일부로 사용되어야 한다.

실용적 활용

WARM은 대형 언어 모델의 정책 학습 과정에서 안정적인 보상 신호를 제공하여, 요약, 대화형 보조, 코드 생성 등 다양한 NLP 작업에 적용 가능하다. 특히, 라벨 데이터가 불확실하거나 분포 이동이 빈번한 상황에서 유용하며, 산업 현장에서의 안전한 AI 배포를 지원할 수 있다.