Energy-Based Reward Models for Robust Language Model Alignment

Anamika Lochab, Ruqi Zhang

Keywords: Reward Models, Alignment, Energy Based Models

TL;DR: We introduce Energy-Based Reward Model (EBRM), a post-hoc method to refine reward models using EBMs.

Abstract:

Reward models (RMs) are essential for aligning Large Language Models (LLMs) with human preference. However, they often struggle with capturing complex human preferences and generalizing to unseen data. To address these challenges, we introduce \emph{Energy-Based Reward Model} (EBRM), a lightweight post-hoc refinement framework that enhances RM robustness and generalization.
EBRM models the reward distribution explicitly, capturing uncertainty in human preferences and mitigating the impact of noisy or misaligned annotations. It achieves this through conflict-aware data filtering, label-noise-aware contrastive training, and hybrid initialization. Notably, EBRM enhances RMs without retraining, making it computationally efficient and adaptable across different models and tasks.
Empirical evaluations on RM benchmarks demonstrate significant improvements in both robustness and generalization, achieving up to a 5.97\% improvement in safety-critical alignment tasks compared to standard RMs. Furthermore, reinforcement learning experiments confirm that our refined rewards enhance alignment quality, effectively delaying reward hacking. These results demonstrate our approach as a scalable and effective enhancement for existing RMs and alignment pipelines.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the COLM Code of Ethics on https://colmweb.org/CoE.html

Author Guide: I certify that this submission complies with the submission instructions as described on https://colmweb.org/AuthorGuide.html

Submission Number: 579