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

SAI replication review · Referee report

Summary

The paper proposes the Bayesian Non-negative Reward Model (BNRM), which replaces the standard dense linear head on a Bradley–Terry (BT) preference loss with a factorized, non-negative reward $r = \theta^\top \Phi$ (plus, per Algorithm 1, a learned bias). Instance-level $\theta$ activations are meant to disentangle preference factors, while a global dictionary $\Phi$ with sparsity-inducing Gamma priors is meant to down-weight population-level spurious signals. Posteriors over both are approximated by a Weibull-parameterized amortized variational network on top of the LLM backbone, giving end-to-end training with only a small overhead over BT. The conceptual move is to import non-negative / Poisson factor-analysis structure into RLHF reward heads and read the resulting sparsity as an inductive bias against reward hacking. The problem is well motivated: dense scalar rewards are vulnerable to length and format shortcuts, and existing mitigations (ensembles, information bottlenecks, targeted debiasing) are either expensive or narrowly scoped. Empirical breadth is a genuine strength — LoRA and full fine-tune settings, ID + OOD (RewardBench, RM-Bench, HHH, MT-Bench), low-resource and noisy-label sweeps, PPO downstream on two 8B policies, BoN, Arena-Hard, and factor-level interpretability all point the same way. The main conceptual limitation is that the paper repeatedly leans on framings — a "mathematical guarantee" of identifiability, a "tightened generalization bound," a "special case of BT," "non-monotonic preference interactions" — that are softer than the language suggests, and the empirical story hides more asymmetric wins and losses (a Code regression on RM-Bench, Hellaswag/BBH/Race drops after PPO, a suspicious +39-point TriviaQA jump, a mis-bolded Table 2 Safety cell, a 25%-vs-40% noise inconsistency) than the abstract admits. An automated attempt to re-execute the code did not complete, so none of the headline numbers were independently confirmed in this review; this is a harness-side failure rather than evidence against the work. Overall the direction is promising and internally coherent, but the theoretical scaffolding and some of the empirical presentation outrun the support.

Strengths

Conceptual contribution. Injecting an NFA-style factorization at the reward head, rather than bolting uncertainty onto a dense scalar, is a cleaner way to combine sparsity and Bayesian uncertainty than the ensembling / information-bottleneck alternatives surveyed in Related Work, and it yields a reward decomposition that is at least in principle inspectable.
Empirical breadth. The paper spans direct reward-modeling accuracy (Tables 1, 2, 5, 8), length-bias diagnostics (Figures 5, 9), low-resource and noisy-label robustness (Figure 4), PPO alignment on two policy backbones (Table 3), an Arena-Hard win-rate check plus a small human study (Table 4), factor-level interpretability (Figure 6, Tables 7, 11), and BoN over-optimization curves (Figure 10). Even after discounting the presentation issues, the pattern of gains across such varied settings is real signal.
Cost profile. The reported wall-clock overhead of 1.3% on Llama-3.1-8B and 7.7% on Gemma-2B relative to BT (Table 6) is genuinely small compared with reward-model ensembles, and directly answers a fair critique of the ensemble baselines.
Mechanistic interpretability, in principle. The factor case study in Table 11 (safety refusal, terse one-liners, step-by-step instructional style) is the kind of qualitative artifact one would want from an NFA-style reward decomposition, and Table 7 attempts to show cross-run stability rather than cherry-picking one seed.
Core training path is functional. In the executed reproduction attempt, the BT baseline training launched and descended under the repository's pinned dependencies via the shipped LoRA script, a weak but real signal that the central training code runs as intended.

Weaknesses

Model definition drifts across sections. The main text writes $r = \theta^\top \Phi$ , but Algorithm 1 computes $r = \theta \cdot \Phi + \operatorname{ReLU}(b)$ . The bias is non-trivial — it can dominate when the bilinear term is small and is arguably what enforces the non-negativity shown in Figure 5 — yet it is absent from Section 4.2 and from the Section 5.4 "Φ suppresses erroneous θ" story. The dimensionality of $\Phi$ also shifts silently, from $V \times K$ in the PFA preliminary (Sec. 3.2) to a $K$ -vector in Sec. 4.2, without explanation.
Weibull parameterization is either mis-typeset or fragile. Algorithm 1 sets $\lambda_\theta \leftarrow z_{\text{out}} / \exp\{\Gamma(1 + 1/\kappa_\theta)\}$ , but mean-matching a Weibull requires dividing by $\Gamma(1 + 1/\kappa)$ , not $\exp\{\Gamma(\cdot)\}$ . Separately, the ReLU-parameterized scale can be exactly zero, making Weibull $(\kappa, 0)$ degenerate with an undefined KL to the Gamma(1,1) prior. Neither is addressed.
Identifiability claim is overstated. The paper cites Gillis & Rajkó's partial identifiability (selective-window plus sparsity) then upgrades to "mathematically guarantees … unique up to permutation and scaling." A Weibull prior does not enforce the selective-window condition on the learned $\Phi$ ; that is an empirical property of a specific run.
Generalization-bound claim has no bound. Appendix D reasons "as the bound indicates" and names $s$ and $\mu_s(\Phi)$ , but no inequality is stated (the display is empty). Citing Mehta & Gray is not the same as instantiating a bound for BNRM's non-negative factorized reward.
Downstream metric does not measure reward hacking. Table 3 argues that higher accuracy on capability benchmarks (GSM8K, MMLU, BBH, …) demonstrates mitigated reward hacking, but those benchmarks measure knowledge/reasoning, not the degree to which the policy exploits the proxy. The construct link is assumed rather than argued.
Table 3 headline hides regressions and a suspect outlier. The +12.15-point average on Llama-3.1-8B is largely carried by a +39.35 TriviaQA jump on a 5-shot knowledge task (with an anomalously low base of 32.64), while the OpenRLHF policy regresses on Hellaswag (−11.83), BBH (−5.49), and Race (−3.53). The prose acknowledges neither.
Best-of-N "gold" reward is itself a learned RM. The over-optimization test uses a Mistral-7B reward model trained on the same Unified-Feedback distribution as the gold signal, so "staying aligned with the gold reward" may reflect shared biases rather than alignment with human intent.
Baseline PPO protocol is under-specified. The SK/PoE/LP/ALBM/InfoRM columns of Table 3 are only comparable if their reward models were also 8B, trained on Skywork-Preference-v0.2, with matched PPO hyperparameters and seeds; Appendix C.3 does not say.
Table 2 presentation is selectively favorable. The Safety cell (92.6, ↓0.1) is bolded despite trailing Skywork's 92.7, and the narrative highlights only Chat-Hard and Safety while the model trails on Chat and Reasoning.
Arena-Hard and human-study claims are over-read. The "22% win rate" attributed to the baseline is actually BNRM's lose rate from the same comparison; and the two human raters diverge by ~22 lose-rate points, so the pooled 51/18/31 does not "consistently mirror" the LLM judge in any defended sense.
Length-correlation as debiasing evidence. A lower Pearson $r$ can also reflect a compressed reward dynamic range (visible in Figure 5) rather than better discrimination; no paired accuracy check on length-balanced cases is provided.
Monotone reward vs. "non-monotonic interactions." With $\theta \ge 0$ and $\Phi \ge 0$ , $r$ is monotone in each component; the claimed non-monotonicity lives only inside the backbone, and the restriction (a factor cannot subtract reward) is not stated.
Unclassified residual and internal inconsistencies. Signal Amplification (64.9%) + Error Rectification (25.5%) leaves ~9.6% of samples uncharacterized; the noise rate is quoted as both 25% and 40%; $\eta = 10^{-6}$ is called "weakest-performing" while "both extremes" are said to be harmful; "same compute budget" mislabels a 1.3–7.7% overhead.
Minor presentation errors. "BNRA" in Sec. 4.4 and "BNBT-Reward" in Table 2; $O(K \times 1)$ collapsing to $O(K)$ ; and a missing dominant-cost formula in Appendix A.2.

Reproducibility & code

The paper links a repository, and in this review the code was executed (full mode). The automated re-run did not complete: across two attempts the harness parked detached background training that outlived the session and ended mid-run, so 19 of 20 tracked claims are graded "not attempted" and the overall reproduction score is essentially zero. This is a replication-process failure, not evidence against the paper, but it means none of the central results were independently confirmed here.

Headline results unverified. The four Gemma-2B-it reward models were never trained to completion and no ID/OOD or RewardBench evaluations ran, so the 74.2/83.6/75.2 accuracies (and their gains over BT), the RewardBench table including the 93.6 full-parameter headline, RM-Bench, the 0.488→0.123 length correlation, the PPO downstream table, Arena-Hard, and the BoN ordering could not be regenerated.
Only "match" is a shipped asset. The single figure scored as reproduced (noise robustness, Figure 4b, +16.7 at 40% noise) is the image already committed to the repository, unchanged — not a plot regenerated from a fresh run — so it corroborates nothing.
Core path runs but full pipeline needs manual assembly. The BT baseline did train under the pinned dependencies via the shipped script; however, external RewardBench and RM-Bench harnesses must be cloned and adapted, and all base models and datasets downloaded, before end-to-end reproduction is possible.
Arena-Hard / human-eval materials absent. No prompt sampling, GPT-4.1 judge template, aggregation code, annotator rubric, blinding protocol, or per-pair labels are released, so Table 4 is neither reproducible nor auditable.
Missing baseline and analysis code. No InfoRM or BT-Ensemble trainer, no SK/PoE/LP/ALBM/InfoRM PPO paths, no RM-Bench-Hard length-correlation script, and no factor-extraction / Signal-Amplification-Error-Rectification / GPT-5 / cross-run-matching code for Figure 6 and Tables 7/11.
Hyperparameter and seed hygiene. The shipped LoRA script reportedly sets 1e-4 for BT-BNRM while Table 9 lists 5e-5, and no global seed is fixed, so sub-point deltas will fluctuate across reruns.

Recommended Changes

Essential

Reconcile the reward formula with Algorithm 1. State $r = \theta^\top \Phi + \operatorname{ReLU}(b)$ uniformly in Section 4.2, Appendix D, and the Figure 5 caption, and explain the bias's role in the non-negativity and error-rectification arguments.
Fix or explain the Weibull scale line. Confirm whether $\lambda_\theta = z_{\text{out}} / \Gamma(1 + 1/\kappa_\theta)$ was intended (the $\exp$ being a slip), and specify how ReLU-produced zero scales are handled for both sampling and the KL to Gamma(1,1).
Downgrade the identifiability claim and derive-or-reframe the bound. Replace "mathematically guarantees … unique up to permutation and scaling" with a conditional statement matching Gillis & Rajkó, and either write out the Mehta & Gray-style inequality with its assumptions or present Appendix D's bound as intuition.
Address the Table 3 anomalies. Explain and independently sanity-check the +39.35 TriviaQA gain (base 32.64 is unusually low), acknowledge the Hellaswag/BBH/Race regressions on the OpenRLHF policy, and add a direct reward-hacking metric (e.g., proxy-vs-gold divergence) rather than resting the anti-hacking claim on capability-benchmark accuracy.
Fully specify the Table 3 baseline PPO protocol — RM backbone, training dataset, PPO hyperparameters, and seed for SK/PoE/LP/ALBM/InfoRM — so RM-scale/data confounds can be ruled out.
Fix Table 2 bolding on Safety and rewrite the surrounding text so only column-leading entries are bold and the subsets where BNBT-Reward trails (Chat, Safety, Reasoning) are noted.
Rewrite the Arena-Hard sentence so "22%" is not attributed to a separate baseline evaluation (it is BNRM's lose rate), and report inter-rater agreement (e.g., Cohen's κ) for the human study or temper "consistently mirroring."
Release the missing execution materials — a single foreground train-and-eval driver, the Arena-Hard judge and human-eval materials, InfoRM/BT-Ensemble and PPO-baseline trainers, the length-correlation and factor-extraction/interpretability scripts, and regenerate Figure 4b from fresh checkpoints — so the headline numbers become independently checkable.
Reconcile the LR mismatch (script 1e-4 vs. Table 9 5e-5), fix global seeds, and report multi-seed variance on headline numbers where deltas are under one point.

Suggested

Justify the Best-of-N gold model or cross-check the ordering against a gold signal not derived from the same preference data.
Pair the length-reward Pearson result with an accuracy check on length-balanced cases (e.g., RM-Bench Easy) to rule out reduced sensitivity.
Acknowledge and, if possible, diagnose the RM-Bench Code regression.
Describe the ~9.6% Signal-Amplification / Error-Rectification residual and specify the cross-run factor-matching rule for Table 7 with an activation-based sanity check.
Reconcile the 25% vs. 40% noise references and the "weakest-performing $\eta = 10^{-6}$ " wording, and replace "same compute budget" with "same number of training epochs."
Reconcile the dimensionality of $\Phi$ between Sec. 3.2 and Sec. 4.2, state whether $K = 1024$ , insert the missing dominant-cost formula, fix $O(K \times 1) \to O(K)$ , and standardize the "BNRA"/"BNBT-Reward" naming.
Soften the "special case of BT" framing to "a stochastic-latent extension," and bound the "non-monotonic preference interactions" claim to inside the backbone.