MetaRM: Shifted Distributions Alignment via Meta-Learning

TL;DR

MetaRM employs meta-learning to adapt reward models in reinforcement learning from human feedback, effectively handling distribution shifts and improving response differentiation, especially for out-of-distribution samples.

Contribution

The paper introduces MetaRM, a novel meta-learning approach that aligns reward models with shifted environment distributions in RLHF, enhancing their generalization and differentiation capabilities.

Findings

MetaRM improves reward model performance in RLHF.

MetaRM enhances detection of out-of-distribution responses.

MetaRM outperforms existing methods in distribution shift scenarios.

Abstract

The success of Reinforcement Learning from Human Feedback (RLHF) in language model alignment is critically dependent on the capability of the reward model (RM). However, as the training process progresses, the output distribution of the policy model shifts, leading to the RM's reduced ability to distinguish between responses. This issue is further compounded when the RM, trained on a specific data distribution, struggles to generalize to examples outside of that distribution. These two issues can be united as a challenge posed by the shifted distribution of the environment. To surmount this challenge, we introduce MetaRM, a method leveraging meta-learning to align the RM with the shifted environment distribution. MetaRM is designed to train the RM by minimizing data loss, particularly for data that can improve the differentiation ability to examples of the shifted target distribution. Extensive experiments demonstrate that MetaRM significantly improves the RM's distinguishing ability in iterative RLHF optimization, and also provides the capacity to identify subtle differences in out-of-distribution samples.