Mitigating Reward Hacking in RLHF via Bayesian Non-negative Reward Modeling

AI summary

Overview

• Research area: Reward modeling for reinforcement learning from human feedback (RLHF) and large language model (LLM) alignment, combining Bayesian probabilistic modeling, non-negative factor analysis, and variational inference.
• Technical level: Advanced. The paper builds on Bradley–Terry preference modeling, Gamma/Weibull latent-variable models, amortized variational inference, and the ELBO objective.
• Scope: The paper proposes and empirically evaluates BNRM, a Bayesian non-negative reward model that replaces the standard dense linear reward head with sparse non-negative latent factors at both the instance and population levels, and tests it on in-distribution and out-of-distribution reward benchmarks, low-resource and noisy-data settings, and PPO-based RLHF fine-tuning.

What This Paper Is About

Reward models in RLHF are trained on noisy human preference labels, and they tend to latch onto superficial cues such as response length or formatting instead of genuine quality. A policy optimized against such a proxy reward will exploit those spurious cues, a failure mode called reward hacking or reward over-optimization. The paper's goal is to restructure the reward function itself so that it is uncertain-aware, sparse, and less susceptible to spurious correlations, rather than only patching the symptoms with ensembles or post-hoc penalties.

Key Contributions

1. A Bayesian non-negative reward modeling framework (BNRM). BNRM embeds non-negative factor analysis into the Bradley–Terry preference model, replacing the deterministic reward r = f(x,y)ᵀ W_bt with a reward generated from sparse, non-negative latent factors: local per-instance variables θ that induce disentangled reward representations, and a shared global dictionary Φ whose sparsity acts as a population-level debiasing mechanism.
2. A scalable amortized variational inference scheme. The LLM backbone is repurposed as an inference network (encoder) rather than part of the generative model. The variational posterior over θ is parameterized as a reparameterizable Weibull distribution, with a Softplus activation for the shape parameter k and a ReLU activation for the scale parameter λ, enabling end-to-end training via backpropagation.
3. Explicit uncertainty modeling. By marginalizing over both local latent variables θ and global factor weights Φ (Eq. 4), BNRM targets both aleatoric uncertainty from ambiguous human annotations and epistemic uncertainty about the global reward factors.
4. Empirical validation on reward modeling and RLHF. Results are reported across ID and OOD benchmarks (Unified Feedback, HHH Alignment, MT-Bench, RewardBench, RM-Bench), low-resource and label-noise settings, and a full PPO RLHF pipeline with automated and human evaluation.

Main Findings

• In-distribution and out-of-distribution reward accuracy improves. With 40K Unified Feedback training examples using LoRA, BT-BNRM reaches 74.2% on Unified Feedback, 83.6% on HHH Alignment, and 75.2% on MT-Bench — improvements of 5.4, 13.3, and 6.1 percentage points over the corresponding BT baseline. GRM-based BNRM attains 74.1%, 82.4%, and 75.1% on the same benchmarks, improving over GRM-SFT by 2.6, 3.7, and 2.1 points. Similar gains are reported for the 400K split.
• It acts as a plug-and-play module on an already strong reward model. Fully fine-tuning Skywork-Reward-Llama-3.1-8B on Skywork-Preference-v0.2 (SP) yields an overall RewardBench score of 93.6, with 89.7% on Chat-Hard and 92.6% on Safety; the base Skywork-Reward-Llama-3.1-8B scores 93.1 overall in the same table.
• Data efficiency. BNRM trained on only 1K examples matches BT trained on 20K on RewardBench, and the gap over BT widens as data volume decreases over the 1K–20K range (Gemma-2B-it backbone).
• Noise tolerance. Under label noise rates from 0.1 to 0.4 on 40K samples, BNRM improves over BT by up to 16.7% at a 40% noise rate and rivals BT trained at only 10%–20% noise.
• PPO-aligned policies score higher on standard benchmarks. Using BNBT-Reward-Llama-3.1-8B as the proxy reward, PPO on 20K alpaca-gpt4-data-en samples gives average accuracies of 74.98% and 62.25% for Llama3.1-8B-Instruct and OpenRLHF-Llama3-8B-SFT, versus 62.83% and 55.68% for their respective base policies (gains of 12.15 and 6.57 points).
• Arena-Hard and human agreement. Evaluated by GPT-4.1, the BNRM-aligned Llama-3.1-8B-Instruct achieves a 50% win rate and 28% tie rate against base Llama-3.1-8B-Instruct, versus a 22% win rate for the baseline. In a blind pairwise human evaluation of 50 randomly sampled Arena-Hard pairs by two PhD-level experts, the model averaged a 51% win rate and 18% tie rate. Against Mistral-7B-Instruct judged by GPT-4.1, win/tie/lose was 0.5850/0.2045/0.2105.
• Length bias is substantially reduced without explicit debiasing supervision. On the RM-Bench Hard subset (where rejected responses are deliberately longer and better formatted), the vanilla BT model shows a Pearson correlation of r = 0.488 between response length and reward, whereas BNRM shows 0.123, lower than all strong baselines.
• RM-Bench aggregate gains. BT-BNRM reaches 60.4 total versus 57.3 for BT (+3.1), with Chat 50.5 (+2.8), Math 55.2 (+3.0), Safety 86.5 (+11.4), and Hard 36.3 (+2.7). The one reported regression is Code, at 49.5 versus 54.0 (−4.5).
• Interpretability. The paper reports that the global factors Φ can rectify specific preference errors made by the local factors θ, providing a mechanistic decomposition that scalar reward heads do not expose. A quantitative analysis of this effect beyond the length/format case is not reported.

Methodology in Plain English

The starting point is the standard Bradley–Terry setup, where a prompt and response are encoded by a backbone into a dense vector z, and a final linear layer turns z into a single scalar reward. The paper treats that deterministic mapping as a degenerate special case of a Bayesian model: because z is a deterministic function, its "posterior" is a point mass, giving the model no way to express uncertainty. BNRM generalizes the model in two stages — first replacing the deterministic representation with a stochastic latent variable θ to capture aleatoric uncertainty, then treating the final layer weights Φ as a global stochastic variable to capture epistemic uncertainty.

The generative story is then: draw global reward factors Φ from a Gamma prior, draw an instance-specific non-negative activation vector θ from a Gamma prior for each prompt–response pair, compute the reward as the inner product θᵀΦ, and generate the observed preference through a Bradley–Terry likelihood σ(r₁ − r₂). Because Gamma priors encourage sparsity and both variables are constrained to be non-negative, θ activates only a small set of factors per example (instance-level disentanglement) while Φ retains only invariant, semantically meaningful factors across the dataset (population-level debiasing).

Since exact posteriors over θ and Φ are intractable, the authors turn to variational inference. Rather than adding a separate encoder, they reuse the LLM backbone as an amortized inference network: the dense feature z is projected through a matrix W_vi into the parameters of a Weibull variational posterior over θ, with Softplus used for the shape parameter and ReLU for the scale parameter. The global posterior q(Φ) is parameterized the same way with its own weights. Training maximizes the ELBO — a reconstruction/likelihood term plus KL divergences from the variational posteriors to the priors, balanced by a trade-off coefficient η. The whole model, including the LLM backbone, is trained end-to-end with backpropagation using the reparameterized Weibull samples.

Experiments compare against BT, BT-Margin, BT-Label Smoothing, BT-Ensemble, GRM variants (GRM-DPO, GRM-DPO-Noref, GRM-SFT), and InfoRM, using gemma-2b-it and gemma-2-2b-it backbones with LoRA for 2 epochs, and full fine-tuning of Skywork-Reward-Llama-3.1-8B for 1 epoch.

Why This Matters

Reward hacking is a central obstacle to making RLHF reliable: a policy that maximizes a hackable proxy reward can score well on the reward model while behaving worse for real users. This paper's argument is that the fix should be structural — changing the reward representation from a dense, opaque scalar head into a sparse, non-negative factor decomposition — rather than adding regularizers or maintaining ensembles of models. That framing is useful for research because it connects classical Bayesian sparse factor models to modern LLM alignment pipelines and claims both robustness and interpretability from the same mechanism. The reported noise- and low-resource robustness also matters because high-quality preference labels are expensive and heterogeneous in practice.

Real-world applications implied by the work:

• LLM assistant alignment pipelines where a reward model must guide PPO or similar policy optimization without the policy drifting toward length-padding or stylistic tricks.
• Low-resource preference tuning, where an organization has only a small amount of labeled preference data (the paper reports matching a 20K-example BT baseline with 1K examples).
• Noisy or crowdsourced annotation settings, where label noise rates comparable to those tested (0.1 to 0.4) are realistic.
• Safety-sensitive reward scoring, where the largest reported RM-Bench gain is on the Safety split (75.1 to 86.5) and a reward model that is harder to game reduces the risk of policies optimizing a spurious proxy.

Industry relevance stems from the "plug-and-play" claim: the method is presented as a module that can be layered on existing reward models, including a strong open 8B model, rather than requiring a new training pipeline from scratch. The authors release code at a public GitHub repository.

Future Directions

• Generalizing the debiasing mechanism beyond length and formatting. The paper validates debiasing empirically on length bias (Pearson correlation) and RM-Bench's Hard subset, but does not report a systematic study of other bias families.
• Closing the remaining performance gaps. The reported regression on the RM-Bench Code split (54.0 to 49.5) and the mixed results on some RewardBench subsets suggest the sparse non-negative constraint may not help uniformly across domains; understanding when sparsity helps and when it hurts is an open question.
• Scaling and cost analysis. The paper reports no comparison of training or inference compute against ensembling baselines, even though reduced overhead relative to model ensembles is one of its stated motivations.
• Larger-scale human validation. The human evaluation covers 50 randomly sampled Arena-Hard pairs with two expert raters; broader human studies and additional judge models would strengthen the alignment claim.
• Sensitivity to the Bayesian hyperparameters. The priors Gamma(γ₀, δ₀) and Gamma(α₀, β₀), the number of latent factors K, and the KL trade-off η all shape the sparsity/robustness trade-off; the paper points to an η analysis in the appendix but does not report a full sensitivity study in the main text.

Target Audience

This paper is most useful to RLHF and alignment researchers who work directly on reward model training, to probabilistic machine learning researchers interested in non-negative factor analysis and amortized variational inference applied at LLM scale, and to practitioners building preference-based fine-tuning pipelines who need reward models that are robust to noisy labels and spurious correlations. Readers need comfort with Bradley–Terry modeling, Bayesian latent variable models, and variational inference; the paper is not an introductory treatment.

Authors’ abstract

Reward models learned from human preferences are central to aligning large language models (LLMs) via reinforcement learning from human feedback, yet they are often vulnerable to reward hacking due to noisy annotations and systematic biases such as response length or style. We propose Bayesian Non-Negative Reward Model (BNRM), a principled reward modeling framework that integrates non-negative factor analysis into Bradley-Terry (BT) preference model. BNRM represents rewards through a sparse, non-negative latent factor generative process that operates at two complementary levels: instance-specific latent variables induce disentangled reward representations, while sparsity over global latent factors acts as an implicit debiasing mechanism that suppresses spurious correlations. Together, this disentanglement-then-debiasing structure enables robust uncertainty-aware reward learning. To scale BNRM to modern LLMs, we develop an amortized variational inference network conditioned on deep model representations, allowing efficient end-to-end training. Extensive empirical results demonstrate that BNRM substantially mitigates reward over-optimization, improves robustness under distribution shifts, and yields more interpretable reward decompositions than strong baselines.