# Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs

arXiv: https://arxiv.org/abs/2406.10216 · HTML: https://mlsanity.com/paper/2406.10216/

## Key facts

- Cited by 142 papers in the map (data through 14 Oct 2026).
- Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 59 of its 75 references that are in the map.
- First posted on arXiv 14 Jun 2024.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)

## Builds on (most cited first, 30 of 59)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [A General Language Assistant as a Laboratory for Alignment](https://mlsanity.com/paper/2112.00861/) (2021)
- [SimPO: Simple Preference Optimization with a Reference-Free Reward](https://mlsanity.com/paper/2405.14734/) (2024)
- [Gemma: Open Models Based on Gemini Research and Technology](https://mlsanity.com/paper/2403.08295/) (2024)
- [AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback](https://mlsanity.com/paper/2305.14387/) (2023)
- [Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution](https://mlsanity.com/paper/2202.10054/) (2022)
- [Instruction Tuning with GPT-4](https://mlsanity.com/paper/2304.03277/) (2023)
- [RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment](https://mlsanity.com/paper/2304.06767/) (2023)
- [LLM-Blender: Ensembling Large Language Models with Pairwise Ranking and Generative Fusion](https://mlsanity.com/paper/2306.02561/) (2023)
- [Reinforced Self-Training (ReST) for Language Modeling](https://mlsanity.com/paper/2308.08998/) (2023)
- [RewardBench: Evaluating Reward Models for Language Modeling](https://mlsanity.com/paper/2403.13787/) (2024)
- [Is Pessimism Provably Efficient for Offline RL?](https://mlsanity.com/paper/2012.15085/) (2020)

## Built on by (most cited first, 30 of 142)

- [Generative Verifiers: Reward Modeling as Next-Token Prediction](https://mlsanity.com/paper/2408.15240/) (2024)
- [Skywork-Reward: Bag of Tricks for Reward Modeling in LLMs](https://mlsanity.com/paper/2410.18451/) (2024)
- [JudgeBench: A Benchmark for Evaluating LLM-based Judges](https://mlsanity.com/paper/2410.12784/) (2024)
- [EmbodiedBench: Comprehensive Benchmarking Multi-modal Large Language Models for Vision-Driven Embodied Agents](https://mlsanity.com/paper/2502.09560/) (2025)
- [RM-Bench: Benchmarking Reward Models of Language Models with Subtlety and Style](https://mlsanity.com/paper/2410.16184/) (2024)
- [DynaMath: A Dynamic Visual Benchmark for Evaluating Mathematical Reasoning Robustness of Vision Language Models](https://mlsanity.com/paper/2411.00836/) (2024)
- [Critique-out-Loud Reward Models](https://mlsanity.com/paper/2408.11791/) (2024)
- [Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment](https://mlsanity.com/paper/2503.21878/) (2025)
- [GenPRM: Scaling Test-Time Compute of Process Reward Models via Generative Reasoning](https://mlsanity.com/paper/2504.00891/) (2025)
- [Reward Shaping to Mitigate Reward Hacking in RLHF](https://mlsanity.com/paper/2502.18770/) (2025)
- [Uncertainty-aware Reward Model: Teaching Reward Models to Know What is Unknown](https://mlsanity.com/paper/2410.00847/) (2024)
- [InternLM-XComposer2.5-Reward: A Simple Yet Effective Multi-Modal Reward Model](https://mlsanity.com/paper/2501.12368/) (2025)
- [Reinforcement Learning Enhanced LLMs: A Survey](https://mlsanity.com/paper/2412.10400/) (2024)
- [A Comprehensive Survey of Reward Models: Taxonomy, Applications, Challenges, and Future](https://mlsanity.com/paper/2504.12328/) (2025)
- [Writing-Zero: Bridge the Gap Between Non-verifiable Tasks and Verifiable Rewards](https://mlsanity.com/paper/2506.00103/) (2025)
- [Rethinking Bradley-Terry Models in Preference-Based Reward Modeling: Foundations, Theory, and Alternatives](https://mlsanity.com/paper/2411.04991/) (2024)
- [A Survey of Test-Time Compute: From Intuitive Inference to Deliberate Reasoning](https://mlsanity.com/paper/2501.02497/) (2025)
- [Inference-Time Reward Hacking in Large Language Models](https://mlsanity.com/paper/2506.19248/) (2025)
- [Direct Judgement Preference Optimization](https://mlsanity.com/paper/2409.14664/) (2024)
- [Reward-Robust RLHF in LLMs](https://mlsanity.com/paper/2409.15360/) (2024)
- [Skywork-VL Reward: An Effective Reward Model for Multimodal Understanding and Reasoning](https://mlsanity.com/paper/2505.07263/) (2025)
- [Post-hoc Reward Calibration: A Case Study on Length Bias](https://mlsanity.com/paper/2409.17407/) (2024)
- [JuStRank: Benchmarking LLM Judges for System Ranking](https://mlsanity.com/paper/2412.09569/) (2024)
- [On the Limited Generalization Capability of the Implicit Reward Model Induced by Direct Preference Optimization](https://mlsanity.com/paper/2409.03650/) (2024)
- [A Comprehensive Survey of Direct Preference Optimization: Datasets, Theories, Variants, and Applications](https://mlsanity.com/paper/2410.15595/) (2024)
- [Open Rubric System: Scaling Reinforcement Learning with Pairwise Adaptive Rubric](https://mlsanity.com/paper/2602.14069/) (2026)
- [GRAM: A Generative Foundation Reward Model for Reward Generalization](https://mlsanity.com/paper/2506.14175/) (2025)
- [Preference Tuning with Human Feedback on Language, Speech, and Vision Tasks: A Survey](https://mlsanity.com/paper/2409.11564/) (2024)
- [RAG-RewardBench: Benchmarking Reward Models in Retrieval Augmented Generation for Preference Alignment](https://mlsanity.com/paper/2412.13746/) (2024)
- [Reasoning as an Adaptive Defense for Safety](https://mlsanity.com/paper/2507.00971/) (2025)

## Abstract (arXiv)

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.