Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs

TL;DR

This paper proposes a novel regularization method for reward models in LLMs that improves their generalization to unseen data and reduces reward over-optimization by preserving hidden state representations.

Contribution

It introduces a regularization technique that maintains language model capabilities during reward model training, enhancing out-of-distribution performance and robustness.

Findings

Improved reward model accuracy on OOD tasks

Reduced reward over-optimization in RLHF

Enhanced robustness of preference learning

Abstract

Reward models trained on human preference data have been proven to effectively align Large Language Models (LLMs) with human intent within the framework of reinforcement learning from human feedback (RLHF). However, current reward models have limited generalization capabilities to unseen prompts and responses, which can lead to an unexpected phenomenon known as reward over-optimization, resulting in a decline in actual performance due to excessive optimization of rewards. While previous research has advocated for constraining policy optimization, our study introduces a novel approach to enhance the reward model's generalization ability against distribution shifts by regularizing the hidden states. Specifically, we retain the base model's language model head and incorporate a suite of text-generation losses to preserve the hidden states' text-generation capabilities, while concurrently learning a reward head behind the same hidden states. Our experimental results demonstrate that the introduced regularization technique markedly improves the accuracy of learned reward models across a variety of out-of-distribution (OOD) tasks and effectively alleviates the over-optimization issue in RLHF, offering a more reliable and robust preference learning paradigm.