Research map / Large language models / Alignment and preferences
Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes GRM, which regularizes a reward model's shared hidden states with text-generation losses (DPO, DPO-noref, SFT) computed through a retained LM head, improving OOD reward accuracy and mitigating over-optimization in BoN and PPO.
Problem
Reward models generalize poorly to unseen prompts and responses and distort pretrained features, leading to reward over-optimization in RLHF.
Approach
Keep the base model's LM head and add text-generation losses (DPO, DPO without reference, or SFT on chosen responses) as regularizers on the hidden states that also feed a learned reward head, derived from an adversarial policy-regularized objective; SFT regularization is the most stable variant.
Key results
- On gemma-2B-it with 400K data, GRM w/ SFT raises HHH-Alignment accuracy from 73.4 to 79.8 over the baseline classifier.
- With Llama3-8B full-parameter training on preference_700K, GRM reaches 87.0 average on RewardBench versus 84.7 for FsfairX-LLaMA3-RM-8B.
- In BoN on 7B base, GRM reaches an average gold score of about 1.5 versus 0.5 for the baseline reward model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add SFT text-generation regularization | Gemma-2B-it reward model · HHH-Alignment 400K | accuracy: 73.4 -> 79.8 | Helped |
| add SFT text-generation regularization | Gemma-2B-it reward model · RewardBench 40K | average score: 64.5 -> 69.5 | Helped |
| use SFT rather than DPO regularization (ablation) | GRM · HHH-Alignment 40K | accuracy: 71.6 -> 78.7 | Helped |
| use DPO regularization without a reference model rather than DPO regularization (ablation) | GRM · HHH-Alignment 400K | accuracy: 79.2 -> 79.7 | Helped |
| use a linear rather than nonlinear reward head (ablation) | GRM w/ SFT · RewardBench 400K | average score: 77.6 -> 79.5 | Mixed |
| replace the standard SFT objective with the log-form SFT objective, alpha=0.001 (ablation) | GRM · HHH-Alignment 400K | accuracy: 79.8 -> 80.2 | Helped |
| replace preference data with pretraining data for regularization (ablation) | GRM · HHH-Alignment 40K | accuracy: 78.7 -> 74.5 | Hurt |
| add positive reward regularization (ablation) | baseline classifier · HHH-Alignment 40K | accuracy: 70.3 -> 70.0 | No clear effect |
| vary regularization weight from the default to an extreme value (0 or 0.1) (ablation) | GRM · HHH-Alignment and MT-Bench OOD | accuracy | Hurt |
| add a second layer and ReLU activation to the reward head (ablation) | GRM w/ SFT · HHH-Alignment, MT-Bench, and RewardBench | accuracy | Mixed |
| reduce reward-model training data from 400K to 40K (ablation) | GRM w/ SFT · HHH-Alignment | accuracy: 79.8 -> 78.7 | Hurt |
| reduce reward-model training data from 400K to 40K (ablation) | GRM (linear) w/ SFT · RewardBench | average score: 79.5 -> 78.3 | Hurt |
Surprising findings
- Frozen reward-head classifier performs worst on ID and OOD tasks, and the baseline with full fine-tuning can underperform frozen when data is small (8K).
- Using a pretraining dataset (RedPajama) for regularization is less effective than using the preference data when data is small.
Limitations
- Label-noise experiments use synthetic flipped labels rather than human-labeled noise.
- Computational limits prevented testing reward models above 10B parameters.
Experimental setup
Datasets: Unified-Feedback HHH-Alignment MT-Bench Human Judgements RewardBench RedPajama-Data-1T-Sample preference_700K
Benchmarks: Unified-Feedback HHH-Alignment MT-Bench Human Judgements RewardBench
Metrics: pairwise accuracy gold reward score under best-of-n gold reward score under PPO
Baselines: Classifier (frozen) Classifier (baseline) Margin loss Label smoothing Ensemble PairRM Starling-RM-7B Starling-RM-34B UltraRM-13B
Builds on: Bradley-Terry reward modeling DPO PPO label smoothing margin loss
Closest papers by meaning
- Mitigating Reward Over-Optimization in RLHF via Behavior-Supported Regularization (ICLR 2025 · similarity 0.92)
- GRAM: A Generative Foundation Reward Model for Reward Generalization (ICML 2025 · similarity 0.92)
- On the Robustness of Reward Models for Language Model Alignment (ICML 2025 · similarity 0.91)
- Confronting Reward Model Overoptimization with Constrained RLHF (ICLR 2024 · similarity 0.91)
- Explicit Preference Optimization: No Need for an Implicit Reward Model (ICML 2025 · similarity 0.91)
- Zeroth-Order Policy Gradient for Reinforcement Learning from Human Feedback without Reward Inference (ICLR 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.