Bayesian Reward Models for LLM Alignment

2024/02/20 by Adam X. Yang, Yang, Adam X., Maxime Robeyns +10 · 6 citations
Computer Science · #FOS: Computer and information sciences #Library Science and Information Systems #Machine Learning (cs.LG) #Natural Language Processing Techniques

paper · pdf · doi:10.48550/arxiv.2402.13210

openalex publication_date 2024/02/20 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

To ensure that large language model (LLM) responses are helpful and non-toxic, a reward model trained on human preference data is usually used. LLM responses with high rewards are then selected through best-of-n (BoN) sampling or the LLM is further optimized to produce responses with high rewards through reinforcement learning from human feedback (RLHF). However, these processes are susceptible to reward overoptimization or `hacking', where responses receive high rewards due to imperfections in the reward model rather than true preference, particularly as prompts or responses deviate from the training data. To address these challenges, we propose to train a Bayesian reward model, which signals higher uncertainty further from the training data distribution. We trained Bayesian reward models using Laplace approximation on LoRA weights, and found that the resulting uncertainty estimates can effectively mitigate reward overoptimization in BoN sampling.

Cited by

• Invasive Context Engineering to Control Large Language Models
• LLM Reinforcement in Context
• Inference-Time Personalized Alignment with a Few User Preference Queries
• The Alignment Auditor: A Bayesian Framework for Verifying and Refining LLM Objectives
• CodeFuse-CR-Bench: A Comprehensiveness-aware Benchmark for End-to-End Code Review Evaluation in Python Projects
• Mitigating Think-Answer Mismatch in LLM Reasoning Through Noise-Aware Advantage Reweighting

Related

• Uncertainty Quantification for Large Language Model Reward Learning under Heterogeneous Human Feedback
• Reward-Robust RLHF in LLMs
• Probabilistic Uncertain Reward Model
• Uncertainty-Penalized Reinforcement Learning from Human Feedback with Diverse Reward LoRA Ensembles
• Post-hoc Reward Calibration: A Case Study on Length Bias