PaperlayerView exact revision on arXiv

Source-linked AI summary

Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs

Rui Yang, Ruomeng Ding, Yong Lin, Huan Zhang, Tong Zhang

arXiv:2406.10216v2cs.CLcs.AI

TL;DR

Reward models often generalize poorly to unseen prompts and responses, contributing to reward overoptimization in RLHF. GRM regularizes shared hidden states with text-generation losses while retaining the base language-model head, and the reported results show improved OOD reward-model accuracy and reduced overoptimization. The study therefore presents GRM as a more reliable proxy reward model within the evaluated settings.

  • Problem

    Reliable and generalizable reward modeling for unseen data remains an open problem, while proxy-reward overoptimization can reduce alignment with authentic human preferences.

  • Method

    GRM retains the base model’s language-model head and regularizes shared reward-model hidden states with text-generation losses while learning a reward head.

  • Results

    GRM substantially improves reward-model accuracy across OOD evaluation tasks and consistently improves RLHF performance while alleviating overoptimization.

  • Takeaways & Limitations

    Hidden-state text-generation regularization offers a lightweight approach for making learned preference rewards more robust to unseen prompt-response pairs.

  • Takeaways & Limitations

    Robustness was evaluated with synthetic label noise rather than human-labeled noisy data, and computational limits prevented testing GRM above 10B parameters.

Abstract

from arXiv · show

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.

1 Introduction

GRM addresses reward-model generalization under distribution shifts by regularizing hidden states with text-generation losses while retaining the base model’s language-model head. Experiments report stronger OOD reward-model accuracy and reduced RLHF overoptimization, with SFT regularization the most effective and stable variant.

  • Motivation: Reward-model generalization to unseen data remains an open problem, and overoptimization can make proxy-reward gains diverge from true-reward performance.
  • Method: GRM regularizes reward-model hidden states with text-generation losses while retaining the base model’s language-model head and jointly training a reward head.The reward head minimizes reward loss, while the language-model head preserves text-generation capabilities during preference learning.
  • Results: GRM substantially improves reward-model accuracy across OOD tasks and consistently enhances RLHF performance while alleviating overoptimization.The reported gains include both reward-model evaluation and policy-optimization settings such as best-of-n sampling and PPO.
  • Results: SFT regularization is the most effective and stable among the three text-generation regularization types evaluated.

2 Background

RLHF commonly combines preference-based reward modeling with policy optimization through best-of-n sampling or PPO. Because proxy rewards can encode erroneous patterns, excessive optimization may reduce alignment with authentic human preferences.

  • Reward Modeling: Reward modeling uses the Bradley-Terry framework to score chosen and rejected responses for a prompt, training higher scores for human-preferred outputs.
  • Policy Optimization: Best-of-n sampling generates multiple policy responses and selects the one receiving the highest reward-model score.
  • Policy Optimization: PPO optimizes a language-model policy with a reward objective and a KL-divergence penalty relative to the SFT policy.The penalty limits policy deviation to help maintain reward-model reliability.
  • Overoptimization: Overoptimization occurs when policy optimization fits erroneous proxy-reward patterns, potentially diminishing alignment with authentic human preferences.

3 Method

GRM regularizes reward-model hidden states with text-generation objectives while sharing them between a reward head and language-model head. Its design targets feature distortion, overfitting, and generalization under preference-data distribution shifts.

  • Generalizable Reward Model: GRM shares hidden states between a reward head trained on preference loss and a language-model head trained to preserve text-generation ability.The two heads jointly regularize the representation during preference learning.
  • Overall Objective: The method combines reward loss with regularization through a weighted objective, where α balances the two terms.The regularization is derived from text-generation objectives and is applied during reward-model training.
  • Theoretical Motivation: GRM motivates regularization by considering reward-model learning against an adversarial policy and using a KL-regularized policy objective with an analytical inner solution.This converts the min-max formulation into a standard optimization problem under assumptions.
  • Implementations: The practical regularization suite includes DPO, DPO without a reference model, and SFT variants that differ in memory and computational requirements.DPO uses a base-model reference, DPO-noref removes that reference, and SFT maximizes chosen-response probability.
  • Advantages: The authors argue that text-generation regularization reduces feature distortion and overfitting while adapting the model to preference-learning data.The stated mechanism preserves base-model text-generation ability and discourages reliance on spurious features.

4 Experimental Setup

The experiments train reward models on Unified-Feedback subsets and evaluate both in-distribution and OOD preference generalization. They use 2B and 7B-class base models and compare GRM with classifier, regularization, margin, label-smoothing, ensemble, and established reward-model baselines.

  • Datasets: Experiments train on 400K and 40K Unified-Feedback samples, use an 8K held-out evaluation set, and test OOD preference data from HHH-Alignment, MT-Bench Human Judgements, and RewardBench.The setup evaluates both data-size sensitivity and transfer to unseen preference distributions.
  • Base Models: Preference-learning experiments use gemma-2B-it and Mistral-7B-Instruct-v0.2, while RLHF experiments use gemma-2B-it for BoN and PPO.A 7B human-preference model trained on the full Unified-Feedback dataset serves as the gold reward model for RLHF experiments.
  • Baselines: Baselines include the original reward-loss classifier, frozen-feature classifier, margin and label-smoothing methods, a three-model ensemble, and open-source RewardBench models.The RewardBench references include PairRM, Starling-RM-7B/34B, and UltraRM-13B.

5 Evaluation Results

GRM improves reward-model generalization across ID and OOD evaluations, including limited-data settings, and reduces reward over-optimization in BoN and PPO. Text-generation regularization is consistently beneficial, with SFT regularization the most effective and stable variant.

  • ID and OOD Evaluation: GRM consistently outperforms baseline reward models on ID and OOD evaluation sets across both 400K and 40K training examples.The evaluations cover Unified-Feedback, HHH-Alignment, and MT-Bench; the 40K setting is reported in Table 2.
  • Comparison of Different Regularizations: SFT regularization is the most effective and stable of the three text-generation regularizers, especially when training data are limited.With 40K examples, the reported ordering is GRM w/ sft, GRM w/ dpo-noref, then GRM w/ dpo on both ID and OOD scores.
  • Results on RewardBench: On RewardBench, GRM with SFT regularization raises the 7B average score from 76.3 to 79.5, while linear reward heads outperform nonlinear heads overall.The comparison spans chat, chat-hard, safety, and reasoning task groups.
  • Comparison of Different Dataset Sizes: When training data decrease from 400K to 40K, GRM w/ sft drops only from 79.8 to 78.7 on HHH Alignment and from 73.4 to 73.0 on MT-Bench.The corresponding baseline drops are larger, from 73.4 to 70.3 and from 71.2 to 69.1.
  • Best-of-n Sampling: GRM mitigates over-optimization in BoN: its proxy and gold scores rise together, reaching an average gold score of 1.5 versus 0.5 for the baseline on the 7B model.Baseline gold scores can plateau or decline after KL > 4 while proxy scores continue increasing.
  • Proximal Policy Optimization: PPO exhibits a stronger tendency than BoN to hack learned reward models, motivating evaluation of GRM under policy optimization.The supplied PPO passage introduces this comparison and the accompanying figure reports proxy and gold scores.

6 Related Works

Related work improves reward modeling through better preference data, dense or multi-objective rewards, and methods for mitigating overoptimization through regularization, iterative labeling, and ensembles.

  • Reward-modeling studies improve performance through preference-data quality or quantity, token-wise dense rewards, and multi-objective rewards.
  • Overoptimization mitigation includes label smoothing, SFT regularization, iterative replacement of hard labels with soft labels, and ensemble reward models.

7 Conclusion

The paper presents hidden-state regularization as an efficient approach for improving reward-model generalization and robustness on unseen data while mitigating RLHF overoptimization.

  • Hidden-state regularization substantially improves reward-model generalization on unseen data and effectively mitigates overoptimization in RLHF.

Limitations

The method derives regularization from adversarial policy optimization and implements it with a reward head and language-model head sharing hidden states, using calibration and policy-based formulations.

  • A Deriving the Regularization Term: The framework begins with adversarial optimization that learns a reward model against an adversarial policy.
  • A Deriving the Regularization Term: A fixed calibration policy centers the reward function while incorporating additional policy preferences into the objective.
  • A Deriving the Regularization Term: Setting the calibration policy to chosen dataset responses turns the final regularization term into an SFT loss.
  • A Deriving the Regularization Term: The two regularizers have distinct origins: DPO regularization comes from the reward loss, whereas SFT regularization derives from the adversarial term.
  • A Deriving the Regularization Term: The final GRM design jointly learns a reward model and language-model head that share hidden states, avoiding inefficient response-distribution normalization.

B Implementation Details

Implementation uses standard sequence-classification baselines and evaluates GRM across training sizes, regularization variants, and reward-model settings, with hidden-state regularization improving OOD performance.

  • Baseline Details: Baselines use randomly initialized linear classification heads, while ensemble baselines average or minimize outputs from three differently seeded reward models.
  • GRM Details: The default GRM reward head uses two linear layers with a ReLU, while text-generation regularization uses α = 0.01 and β = 0.1.
  • ID and OOD Evaluation: GRM achieves strong performance with both large and small datasets while fine-tuning all parameters through hidden-state regularization.
  • Training and Evaluation Details: Training uses two LoRA epochs after validation-loss convergence, while full-parameter experiments use one epoch because they are more prone to overfitting.

C.4 Ablation of the Regularization Weight

GRM is most sensitive to the regularization weight α, with intermediate values outperforming extreme settings. A two-layer reward head offers inconsistent gains across benchmarks, so the single-layer design remains the default.

  • C.4 Ablation of the Regularization Weight: Intermediate α values consistently achieve higher OOD scores than α=0 or α=0.1, motivating the default α=0.01.Figure 6 evaluates GRM (2B) on HHH-Alignment and MT-Bench.
  • C.4 Ablation of the Regularization Weight: The two-layer reward head slightly improves MT-Bench and RewardBench but lowers HHH-Alignment, so it is excluded from the main experiments.The comparison adds a linear layer and ReLU activation to the reward head.
  • C.4 Ablation of the Regularization Weight: Positive regularization does not improve results with 40K examples but provides slight overall enhancement with 400K examples.The comparison uses gemma-2B-it, with GRM adopting the linear variant for RewardBench.
  • C.4 Ablation of the Regularization Weight: GRM significantly improves both in-distribution and out-of-distribution accuracy, especially with limited preference data.The result is reported relative to the baseline classifier and positive-regularization variant.

C.7 Regularization with pretraining dataset

Text-generation regularization can use either preference data or an external pretraining dataset, but preference data is more effective when reward-model training data is scarce. GRM also improves PPO outcomes and resists reward hacking observed with label-smoothed models.

  • C.7 Regularization with pretraining dataset: A pretraining dataset matches GRM at 400K examples but is less effective than preference-data regularization at small dataset sizes.The pretraining dataset is matched in size to the reward-model training set.
  • C.7 Regularization with pretraining dataset: GRM achieves a significantly higher PPO win rate than the vanilla reward model across two base reward models.Win rates are evaluated with GPT-4o on 100 randomly selected Unified-Feedback test prompts.
  • C.7 Regularization with pretraining dataset: With label smoothing, proxy scores increase while gold scores rapidly decrease during BoN and PPO, indicating vulnerability to reward hacking.Figure 8 compares dashed proxy-score curves with solid gold-score curves for Mistral-7B-Instruct.

D Examples in the PPO Experiments

PPO exposes baseline reward models to exploitation that produces repetitive, refusal-based, or weaker-reasoning behaviors. GRM instead yields more robust instruction following and reasoning under identical PPO hyperparameters, with lower training cost than ensemble baselines.

  • D Examples in the PPO Experiments: Baseline reward models are exploited into repeated patterns, refusals on harmless prompts, or degraded mathematical reasoning.The cited examples include Ensemble (min), Baseline, and Margin models.
  • D Examples in the PPO Experiments: GRM produces more robust instruction-following responses and better reasoning than baseline reward models under identical PPO hyperparameters.The examples include harmless prompts and math problems.
  • D Examples in the PPO Experiments: GRM achieves these qualitative improvements with a smaller training cost than ensemble baselines.The comparison uses the same PPO hyperparameters across models.
  • D Examples in the PPO Experiments: The broader-impacts discussion identifies possible misuse if a generalizable reward model is used for harmful language-model training.The authors state that future AI-safety efforts are needed to prevent such misuse.
Loading 2406.10216v2…