Energy-Based Reward Models for Robust Language Model Alignment

• URL: http://arxiv.org/abs/2504.13134v1
• Date: Thu, 17 Apr 2025 17:47:15 GMT
• Title: Energy-Based Reward Models for Robust Language Model Alignment
• Authors: Anamika Lochab, Ruqi Zhang,
• Abstract summary: We introduce Energy-Based Reward Model (EBRM), a lightweight post-hoc refinement framework for Reward Models (RMs)<n>EBRM models the reward distribution explicitly, capturing uncertainty in human preferences and mitigating the impact of noisy or misaligned annotations.<n> Empirical evaluations demonstrate significant improvements in robustness and generalization, achieving up to a 5.97% improvement in safety-critical alignment tasks.
• Score: 9.843359827321194
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Reward models (RMs) are essential for aligning Large Language Models (LLMs) with human preferences. However, they often struggle with capturing complex human preferences and generalizing to unseen data. To address these challenges, we introduce 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. The code is available at EBRM.

Related papers

• Adversarial Training of Reward Models [74.17196154247964]
  We introduce Adv-RM, a novel adversarial training framework that automatically identifies adversarial examples.<n>By leveraging reinforcement learning, Adv-RM trains a policy to expose vulnerabilities in large state-of-the-art reward models.<n>We demonstrate that Adv-RM significantly outperforms conventional reward training.
  arXiv  Detail & Related papers  (2025-04-08T15:38:25Z)
• RMB: Comprehensively Benchmarking Reward Models in LLM Alignment [44.84304822376291]
  Reward models (RMs) guide the alignment of large language models (LLMs)<n>We propose RMB, a comprehensive RM benchmark that covers over 49 real-world scenarios.<n>Based on our benchmark, we conduct extensive analysis on the state-of-the-art RMs.
  arXiv  Detail & Related papers  (2024-10-13T16:06:54Z)
• Uncertainty-aware Reward Model: Teaching Reward Models to Know What is Unknown [20.753374166695494]
  We introduce the Uncertainty-aware Reward Model (URM) and its ensemble variant, URME.<n>URM employs a probabilistic value head to capture aleatoric uncertainty by modeling the distribution of disentangled human preference attributes.<n> URME further quantifies uncertainty by examining discrepancies among individual URMs within the ensemble, enabling identification of unreliable evaluations.
  arXiv  Detail & Related papers  (2024-10-01T16:29:59Z)
• RRM: Robust Reward Model Training Mitigates Reward Hacking [51.12341734942797]
  Reward models (RMs) play a pivotal role in aligning large language models with human preferences. We introduce a causal framework that learns preferences independent of these artifacts. Experiments show that our approach successfully filters out undesirable artifacts, yielding a more robust reward model.
  arXiv  Detail & Related papers  (2024-09-20T01:46:07Z)
• Reward-Robust RLHF in LLMs [25.31456438114974]
  Large Language Models (LLMs) continue to progress toward more advanced forms of intelligence. The reliance on reward-model-based (RM-based) alignment methods introduces significant challenges. We introduce a reward-robust RLHF framework aimed at addressing these fundamental challenges.
  arXiv  Detail & Related papers  (2024-09-18T02:35:41Z)
• Semi-Supervised Reward Modeling via Iterative Self-Training [52.48668920483908]
  We propose Semi-Supervised Reward Modeling (SSRM), an approach that enhances RM training using unlabeled data. We demonstrate that SSRM significantly improves reward models without incurring additional labeling costs. Overall, SSRM substantially reduces the dependency on large volumes of human-annotated data, thereby decreasing the overall cost and time involved in training effective reward models.
  arXiv  Detail & Related papers  (2024-09-10T22:57:58Z)
• Prior Constraints-based Reward Model Training for Aligning Large Language Models [58.33118716810208]
  This paper proposes a Prior Constraints-based Reward Model (namely PCRM) training method to mitigate this problem. PCRM incorporates prior constraints, specifically, length ratio and cosine similarity between outputs of each comparison pair, during reward model training to regulate optimization magnitude and control score margins. Experimental results demonstrate that PCRM significantly improves alignment performance by effectively constraining reward score scaling.
  arXiv  Detail & Related papers  (2024-04-01T07:49:11Z)
• WARM: On the Benefits of Weight Averaged Reward Models [63.08179139233774]
  We propose Weight Averaged Reward Models (WARM) to mitigate reward hacking. Experiments on summarization tasks, using best-of-N and RL methods, shows that WARM improves the overall quality and alignment of LLM predictions.
  arXiv  Detail & Related papers  (2024-01-22T18:27:08Z)
• Confronting Reward Model Overoptimization with Constrained RLHF [114.71591361764547]
  We show that correlation between component RMs has a significant effect on the locations of these points. Our method addresses the problem of weighting component RMs by learning dynamic weights, naturally expressed by Lagrange multipliers.
  arXiv  Detail & Related papers  (2023-10-06T16:59:17Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.