Technical Summary: Bayesian Preference Learning for Test-Time Steerable Reward Models

Problem Statement

Reward models (RMs) are central to aligning language models with human preferences via Reinforcement Learning from Human Feedback (RLHF). Traditionally, classifier RMs are trained as static estimators of the Bradley-Terry (BT) model parameters using a fixed dataset. Once trained, these models lack adaptability; they cannot adjust to unseen preference distributions or dynamic trade-offs (e.g., balancing helpfulness vs. safety) without retraining. Furthermore, standard classifier RMs are prone to "over-optimization," where preference probabilities converge to 1, leading to local optima that fail to capture the true complexity of human preferences. While "LLM-as-a-Judge" approaches offer flexibility, they often rely on proprietary models with high computational costs. The paper addresses the need for an efficient, classifier-based RM that remains adaptable at test time and avoids over-optimization.

Methodology: Variational In-Context Reward Modeling (ICRM)

The authors propose Variational In-Context Reward Modeling (ICRM), a Bayesian framework that treats reward modeling as amortized variational inference over a latent preference probability.

Core Formulation

1. Latent Variable Definition: The model introduces a latent random variable $z \in [0, 1]$ representing the probability that a chosen response $y_w$ is preferred over a rejected response $y_l$ given a prompt $x$ and a set of in-context demonstrations $C$.
2. Beta Prior and Posterior:
   • The BT model is specialized for $N=2$ (one pair), where the likelihood of the preference outcome follows a Bernoulli distribution parameterized by $z$.
   • A Beta prior $p(z) = \text{Beta}(\alpha_0, \beta_0)$ is placed on $z$. The authors typically use a uniform prior $(\alpha_0, \beta_0) = (1, 1)$ to represent an uninformative starting point.
   • The approximate posterior $q_\theta(z | x, y_w, y_l, C)$ is parameterized by a neural network with a two-dimensional projection head, outputting a utility score $u_\theta$ and a confidence (evidence) score $s_\theta$.
   • The posterior is reparameterized as $\text{Beta}(\alpha_q, \beta_q)$ where:
     • $\mu = \sigma(u_w - u_l)$ represents the posterior mean (preference probability).
     • $\tau = \text{Softplus}(s_w) + \text{Softplus}(s_l) + 1$ represents the concentration (confidence).
     • $\alpha_q = \mu\tau$ and $\beta_q = (1-\mu)\tau$.

Learning Objective

The model is trained by maximizing the Evidence Lower Bound (ELBO), minimizing the negative ELBO loss:
$$\mathcal{L}_{ICRM} = -\mathbb{E}_{q_\theta}[\log z] + \lambda(N) D_{KL}(q_\theta(z|\cdot) \parallel p(z))$$

• Reconstruction Term: The first term measures how well the posterior explains the observed preference. For a Beta distribution, this has a closed form involving the digamma function $\psi$. Minimizing this pushes $\mu$ toward 1 for the chosen response.
• Regularization Term: The second term is the KL divergence between the learned posterior and the fixed prior. The hyperparameter $\lambda(N)$ is a monotonically decreasing function of the number of in-context demonstrations $N$.
  • When $N$ is small, $\lambda(N)$ is large, forcing the posterior to remain close to the prior (high uncertainty).
  • As $N$ increases, $\lambda(N)$ decreases, allowing the reconstruction term to dominate and the model to form a confident, data-driven posterior.

Theoretical Guarantees

The authors prove that the KL-regularized objective prevents the "over-optimization" failure mode common in standard BT models.

• Lemma 3.1: As the preference probability $\mu \to 1$, the gradient of the loss with respect to the utility score is penalized by a factor proportional to $1/\epsilon$ (where $\epsilon = 1-\mu$), preventing uncontrolled maximization.
• Theorem 3.2: The global minimizer of the loss function is guaranteed to lie in the interior of the domain ($0 < \mu^* < 1$ and $0 < \tau^* < \infty$), ensuring the model does not collapse to boundary solutions.

Key Contributions

1. Principled Variational Preference Learning: ICRM enables RMs to infer test-time preferences from few-shot in-context samples. The KL-regularized objective guarantees a global interior optimum, tempering excessive maximization of preference means.
2. Single and Multi-Objective Steerability: The variational design allows for both directional steerability (adapting to specific preference directions) and calibrated steerability (encoding the degree of preference shift). It extends to multi-objective settings, allowing the model to navigate trade-offs between conflicting objectives (e.g., safety vs. helpfulness) by mixing in-context demonstrations.
3. Parameterizing Verifiable Rewards: The framework is applied to Reinforcement Learning with Verifiable Rewards (RLVR). The model interprets the utility and confidence scores to generate a reward signal that is both directional and calibrated to the reliability of the context, outperforming standard verifiers and BT models in math reasoning tasks.

Experimental Results

The authors evaluated ICRM using Qwen3-4B and Llama-3.2-3B backbones against baselines including standard classifier RMs (BTRM, ArmoRM, GRM) and generative judges (LLM-as-a-Judge, GRAM-r2).

• Calibrated Steerability: On moral dilemma preference pairs, ICRM demonstrated the ability to follow the degree of preference shift specified by the composition of in-context demonstrations. As the number of demonstrations ($N$) increased from 8 to 64, the Expected Calibration Error (ECE) decreased from 0.099 to 0.076, outperforming the LLM-as-a-Judge baseline.
• Directional Steerability: On benchmarks like SafeRLHF and HHH-Alignment, ICRM adapted to "reversed" preference labels (where harmful responses were labeled as chosen) with high accuracy (up to 89%), whereas static baselines failed to adapt.
• RM-Bench Performance: Increasing the number of in-context demonstrations improved RM-Bench accuracy significantly. With $N=64$, ICRM achieved an average accuracy of 70.8%, up from 60.5% at $N=1$, surpassing the LLM-as-a-Judge baseline (70.3%).
• Multi-Objective Trade-offs: In the RM-Bench Safety domain, ICRM produced smooth Pareto frontiers between "Should Respond" and "Should Refuse" behaviors. The hypervolume (HV) of these frontiers increased with $N$, demonstrating that ICRM can expand the attainable trade-off region, a capability static RMs lack.
• RLVR Application: In math reasoning tasks using the INTELLECT-MATH dataset, a policy trained with ICRM rewards achieved 53.8% accuracy on MATH-500, outperforming policies trained with a gold verifier (50.4%) and a standard BT reward model (51.4%).

Significance and Claims

The paper claims that ICRM provides a theoretically grounded and practically effective route toward reward models that can adapt to unseen preference standards after training. By treating preferences as latent probabilities with Beta posteriors conditioned on in-context demonstrations, ICRM unifies Bayesian regularization with in-context steerability.

The authors emphasize that this approach:

• Solves the static nature of traditional classifier RMs without the computational overhead of proprietary generative judges.
• Provides a principled mechanism to mitigate reward over-optimization through KL regularization.
• Enables dynamic, test-time control over multi-objective alignment, allowing for continuous navigation of preference spaces rather than fixed operating points.

The work suggests that future reward modeling can move beyond fixed datasets toward dynamic, demonstration-driven alignment, though the authors note limitations regarding context length constraints and the need for safeguards against steering models toward harmful behaviors.