Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Bayesian Preference Learning for Test-Time Steerable Reward Models

Jiwoo Hong, Shao Tang, Zhipeng Wang

TL;DR

The paper tackles the rigidity of static reward models by introducing Variational In-Context Reward Modeling (ICRM), a Bayesian framework that enables test-time steerability of reward models through in-context demonstrations. By casting BT preferences as a Beta posterior and applying amortized variational inference with a Beta prior, ICRM learns a context-conditioned posterior over preference probabilities, balancing reconstruction and regularization via a KL term that preserves an interior optimum. Empirically, ICRM demonstrates strong single- and multi-objective steerability, expanding Pareto frontiers with more in-context evidence and outperforming static baselines in several benchmarks, including SafeRLHF and RM-Bench, while also supporting RLRV-style verifiable rewards in math reasoning. The approach yields theoretical guarantees of a global interior optimum and shows how KL regularization mitigates reward over-optimization, offering a principled, adaptable reward modeling paradigm for scalable, preference-aware language model alignment.

Abstract

Reward models are central to aligning language models with human preferences via reinforcement learning (RL). As RL is increasingly applied to settings such as verifiable rewards and multi-objective alignment, RMs are expected to encode more complex and multifaceted preference distributions. However, classifier RMs remain static once trained, limiting their adaptability at test time. We propose Variational In-Context Reward Modeling (ICRM), a novel Bayesian reward modeling objective that enables test-time steerability via in-context preference demonstrations. ICRM casts reward modeling as amortized variational inference over a latent preference probability under the Bradley-Terry model using a conjugate Beta prior. We show that ICRM adapt to unseen preference distributions at test time for both single and multi-objective settings. With more in-context demonstrations, ICRM gains 34% accuracy on SafeRLHF and 9% accuracy on RM-Bench in the single-objective setting, while widening the Pareto frontier with a 4% gain in hypervolume on helpfulness and refusal benchmarks. We further study the practical applicability of ICRM for RL training, showing that it can effectively encode verifiable rewards by outperforming a conventional RM in math reasoning. Finally, we provide theoretical guarantees that the variational objective admits a global interior optimum with finite confidence, and we analyze how KL regularization mitigates reward over-optimization.

Bayesian Preference Learning for Test-Time Steerable Reward Models

TL;DR

The paper tackles the rigidity of static reward models by introducing Variational In-Context Reward Modeling (ICRM), a Bayesian framework that enables test-time steerability of reward models through in-context demonstrations. By casting BT preferences as a Beta posterior and applying amortized variational inference with a Beta prior, ICRM learns a context-conditioned posterior over preference probabilities, balancing reconstruction and regularization via a KL term that preserves an interior optimum. Empirically, ICRM demonstrates strong single- and multi-objective steerability, expanding Pareto frontiers with more in-context evidence and outperforming static baselines in several benchmarks, including SafeRLHF and RM-Bench, while also supporting RLRV-style verifiable rewards in math reasoning. The approach yields theoretical guarantees of a global interior optimum and shows how KL regularization mitigates reward over-optimization, offering a principled, adaptable reward modeling paradigm for scalable, preference-aware language model alignment.

Abstract

Reward models are central to aligning language models with human preferences via reinforcement learning (RL). As RL is increasingly applied to settings such as verifiable rewards and multi-objective alignment, RMs are expected to encode more complex and multifaceted preference distributions. However, classifier RMs remain static once trained, limiting their adaptability at test time. We propose Variational In-Context Reward Modeling (ICRM), a novel Bayesian reward modeling objective that enables test-time steerability via in-context preference demonstrations. ICRM casts reward modeling as amortized variational inference over a latent preference probability under the Bradley-Terry model using a conjugate Beta prior. We show that ICRM adapt to unseen preference distributions at test time for both single and multi-objective settings. With more in-context demonstrations, ICRM gains 34% accuracy on SafeRLHF and 9% accuracy on RM-Bench in the single-objective setting, while widening the Pareto frontier with a 4% gain in hypervolume on helpfulness and refusal benchmarks. We further study the practical applicability of ICRM for RL training, showing that it can effectively encode verifiable rewards by outperforming a conventional RM in math reasoning. Finally, we provide theoretical guarantees that the variational objective admits a global interior optimum with finite confidence, and we analyze how KL regularization mitigates reward over-optimization.
Paper Structure (69 sections, 2 theorems, 33 equations, 6 figures, 3 tables)

This paper contains 69 sections, 2 theorems, 33 equations, 6 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Preliminaries
  4. Theoretical Background
  5. In-context learning as implicit fine-tuning
  6. Estimating the true preference distribution in the Bradley--Terry model
  7. Bayesian treatment of the Bradley-Terry model
  8. Variational In-Context Reward Modeling
  9. Problem Setup
  10. Prior distribution
  11. Posterior distribution
  12. Reward Modeling as Variational Inference
  13. Beta prior for the Bradley-Terry model
  14. Amortized variational approximation of posterior
  15. Evidence lower bound for variational objective
  16. ...and 54 more sections

Key Result

Lemma 8.1

Let $P_\theta(y_w \succ y_l \mid x)$ denote the ICRM preference with $\mu=\sigma(\Delta u_\theta)=\sigma(u_\theta(x,y_w)-u_\theta(x,y_l))$ and $\varepsilon:=1-\mu$. For $\tau\in(0,\infty)$, as $\varepsilon\to 0^+$,

Figures (6)

  • Figure 1: Variational in-context reward modeling (ICRM) with Beta prior for the Bradley-Terry (BT) model. ICRM directly models the mean and sharpness of the Beta posterior, calibrated to how "confident" the model is for the preference triplet $(x, y_w, y_l)$ given in-context preference demonstrations. This yields multi-objective test-time steerability of the reward model for any preferences or tasks.
  • Figure 2: Ablation study. The learning curve of the preference mean $\mu$ and the concentration factor $\tau$ of the parameterized Beta posterior in the variational in-context reward modeling. Weaker KL regularization, i.e., smaller $\lambda$, leads to stronger adaptation to the training data.
  • Figure 3: Trend of the confidence factor $\tau$ as number of in-context preference demonstrations increase for Qwen3-4B-Base ICRM. $\tau$ values were collected from the SafeRLHF evaluation results.
  • Figure 4: Multi-objective steerability analysis. Pareto frontiers of ICRM trained on Llama-3.2-3B-Base (Figure \ref{['subfig:moo_l32']}) and Qwen3-4B-Base (Figure \ref{['subfig:moo_q3']}), and the Hypervolume (HV) of the Pareto frontiers plotted against the number of in-context demonstrations $N$ (Figure \ref{['subfig:hv']}).
  • Figure 5: Parameterizing verifiable rewards. Accuracy mean ("Accuracy (%)") and average rewards ("Training Reward") of eight sample responses per query.
  • ...and 1 more figures

Theorems & Definitions (4)

  • Lemma 8.1: Edge behavior at finite confidence
  • Theorem 8.2
  • proof
  • proof