Risk-Aware Alignment Must Incorporate Reward Uncertainty

Debangshu Banerjee; Kintan Saha; Aditya Gopalan

cs.AIOct 31, 2024

Risk-Aware Alignment Must Incorporate Reward Uncertainty

Authors: Debangshu Banerjee, Kintan Saha, Aditya Gopalan

Organizations: Department of Electrical and Communication Engineering, Indian Institute of Science, India

Abstract

We consider the problem of policy optimization with noisy rewards, where the policy is constrained to be in the neighborhood of a reference policy.We study the risk of the policy optimization process, defined to be the chance of obtaining a policy worse off than the reference policy. In an idealized multivariate Gaussian model for the noisy reward function, we show that the regularizer achieving optimal risk in a suitably local neighborhood of the reference policy is the quadratic norm weighted by the reward covariance matrix. This result is used to develop a new policy optimization algorithm based on regularization via reward variance estimates. Experiments across synthetic bandit models, large language models and reward models show that this algorithm yields more robust alignment outcomes than existing approaches.

Figures & tables

Panel of Figure 1 — Figure 1: Reward scores assigned by $10$ reward models on the same prompt-response pair. The reward models are identical in that they are trained independently on the same dataset, with the same hyperparameters and number of epochs. Despite this, we see a wide variation in the score assigned by each model.

	$\V_{i}\big[r_{i}(x,y^{1})\big]$		$\V_{i}\big[r_{i}(x,y^{1})-r_{i}(x,y^{2})\big]$
dataset	mean	median	mean	median
Chat	$4.22$	$3.81$	$2.37$	$1.38$
Chat-Hard	$4.95$	$4.53$	$0.80$	$0.24$
Safety	$5.56$	$5.14$	$1.56$	$0.85$
Reasoning	$4.71$	$4.44$	$0.09$	$0.01$

Table 1: Variance of the rewards assigned by

10

independently trained reward models to the same input. For each input we take the sample variance

\V_{i}

across the ten members, of the reward given to the accepted response

y^{1}

(left) and of the accepted-minus-rejected difference (right), which would vanish were the members separated only by a constant offset; each entry is the mean or median of that variance over the inputs of a dataset. Rewards lie on a

1

–

20

scale, so these are variances of scores on that scale. Figures 5 and 6 in Appendix B plot the full distributions.

	number of arms $K$
method	$50$	$100$	$500$	$1000$
vanilla	$0.163$	$0.190$	$0.180$	$0.185$
covariance	$\bm{0.055}$	$\bm{0.067}$	$\bm{0.051}$	$\bm{0.058}$
UWO $(0.1)$	$0.169$	$0.198$	$0.181$	$0.176$
UWO $(1.0)$	$0.291$	$0.330$	$0.263$	$0.197$
KL	$0.163$	$0.190$	$0.180$	$0.185$

Table 2: Underperformance probability

\mathbb{P}\left[{\pi^{\top}r^{*}\leqslant\pi_{0}^{\top}r^{*}}\right]

for a uniform reference policy at

\kappa=50

, at each geometry’s interior radius. With

\pi_{0}

uniform, KL coincides with the unweighted geometry and variance-weighted KL with the covariance geometry.

	skew $\alpha$
method	$0$	$0.5$	$1$	$1.5$	$2$
vanilla	$0.194$	$0.194$	$0.194$	$0.194$	$0.194$
covariance	$\bm{0.058}$	$\bm{0.058}$	$\bm{0.058}$	$\bm{0.058}$	$\bm{0.058}$
KL	$0.194$	$0.216$	$0.311$	$0.382$	$0.414$
var.-wtd. KL	$0.058$	$0.087$	$0.190$	$0.302$	$0.374$

Table 3: Underperformance probability at

K=100

against the skew of the reference,

\pi_{0}\propto(i+1)^{-\alpha}

, for which

p_{\min}

falls from

10^{-2}

\alpha=0

6.1\cdot 10^{-5}

\alpha=2

	radius $\varepsilon/\varepsilon_{0}$
	$0.25$	$1$	$4$	$25$	$200$
vanilla, measured	$0.035$	$0.035$	$0.032$	$0.033$	$0.099$
vanilla, predicted	$0.035$	$0.035$	$0.035$	$0.035$	$0.035$
covariance, measured	$0.008$	$0.007$	$0.008$	$0.009$	$0.028$
covariance, predicted	$0.008$	$0.008$	$0.008$	$0.008$	$0.008$
on boundary, vanilla	$0.00$	$0.00$	$0.17$	$0.97$	$1.00$

Table 4: Behaviour across the interior threshold

\varepsilon_{0}

of Definition 4.2 . “Measured” is the simulated risk, obtained by solving the constrained problem exactly for each sampled

R

; “predicted” is the closed form

\Phi(-z)

, which is derived assuming

\pi\geqslant 0

holds with strict inequality and is therefore constant in the radius. “On boundary” is the fraction of draws whose optimum has

\pi(x,y)=0

for at least one pair. “Reversed” counts instances in which the covariance geometry is the riskier of the two by more than three standard errors of the difference, so that sampling noise is not recorded as a violation.

	Direction Cosine			$\mathbb{P}\left[{\text{underperform}}\right]$
Geometry, Score	$\alpha=1$	$\alpha=5$	$\alpha=20$	$\alpha=1$	$\alpha=5$	$\alpha=20$
Covariance, $R$	$1.000$	$1.000$	$1.000$	$\bm{0.235}$	$\bm{0.236}$	$\bm{0.234}$
$\ell_{2}$ , $\widehat{S}$	$1.000$	$1.000$	$1.000$	$\bm{0.235}$	$\bm{0.236}$	$\bm{0.234}$
KL, $\widehat{S}/\rho\pi_{0}$	$1.000$	$1.000$	$1.000$	$\bm{0.235}$	$\bm{0.236}$	$\bm{0.234}$
$\ell_{2}$ , $R$	$0.763$	$0.764$	$0.764$	$0.292$	$0.293$	$0.291$
KL, $R$	$0.763$	$0.720$	$0.643$	$0.292$	$0.302$	$0.318$

Table 5: Update direction and risk on candidate sets, by the skewness factor

\alpha

of the reference policy. The first column group is the cosine between a geometry’s update direction and that of

\pi_{\Sigma}

, the optimum of Theorem 4.3 ; the three pairings above the rule are the optimal pairing described in Section 5 and are seen to equal the prediction. The second group is the fraction of draws on which the policy is worse than the reference under the gold reward, in bold for the three optimal pairings and for RAPO .

	trust-region radius
Objective	$0.5$	$1$	$2$	$4$	$8$
Pythia-1.4B
Vanilla	$0.155$	$0.162$	$0.202$	$0.267$	$0.378$
UWO $(0.1)$	$0.157$	$0.163$	$0.203$	$0.263$	$0.372$
RAPO	$\bm{0.148}$	$\bm{0.163}$	$\bm{0.187}$	$\bm{0.215}$	$\bm{0.247}$
Qwen2.5-7B

Table 6: Probability that the fine-tuned policy underperforms the reference policy under the gold reward, across trust-region radii and objectives. Experimental details are provided in Appendix D.2 .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Distribution of sample variances in the rewards assigned to accepted responses, each variance computed across the outputs of 10 independently trained reward models. The median sample variance exceeds $3.81$ in all datasets, with values reaching as high as $10$ , confirming substantial model-level variability in reward estimates for the same prompt–response pair. Figure 6 repeats the measurement on reward differences.

Panel of Figure 6 — Figure 6: Distribution of the sample variances of reward differences between accepted and rejected responses, computed across the 10 independently trained reward models. Were the members simple scalar translations of one another these differences would agree exactly and the distributions would collapse to a point mass at zero. The spread instead reflects genuine disagreement about preferences, driven by the statistical and stochastic nature of training.

(a) variance ratio $\kappa$ , exact $\Sigma$
$\kappa$	$1$	$3$	$10$	$50$	$200$
vanilla	$0.000$	$0.000$	$0.004$	$0.034$	$0.105$
covariance	$0.000$	$0.000$	$\bm{0.001}$	$\bm{0.008}$	$\bm{0.024}$
(b) members $m$ estimating $\Sigma$ , at $\kappa=50$
$m$	$3$	$5$	$10$	$20$	exact
covariance	$0.052$	$0.023$	$0.013$	$0.010$	$\bm{0.008}$

Appendix

Table 7: (a) Risk against the heterogeneity of the reward variances, at an interior radius, where

\kappa=\sigma^{2}_{\max}/\sigma^{2}_{\min}

is the ratio between the largest and smallest variance a pair can receive. (b) The same when

\Sigma

is estimated from

m

members rather than known; at

m=3

the estimated method is worse than ignoring uncertainty altogether.

	rank $d$ of the parameterization
metric $M$	$2$	$4$	$7$
$I$ (Euclidean)	$0.177$	$0.123$	$0.092$
$J^{\top}\Pi_{0}^{-1}J$ (Fisher)	$0.158$	$0.078$	$0.030$
$J^{\top}\Sigma J$ (covariance)	$\bm{0.150}$	$\bm{0.057}$	$\bm{0.008}$
unrestricted optimum	$0.008$	$0.008$	$0.008$

Appendix

Table 8: Risk for a policy whose updates are confined to the range of

J

, at

K=8

, so

d=7

is full rank, where the covariance pullback coincides with the unrestricted optimum of Theorem 4.3 .

	radius $\varepsilon$
target $\tau$	$0.5$	$1$	$2$	$4$	$8$
Pythia-1.4B
$0.25$	$15$	$15$	$15$	$15$	$15$
$0.5$	$20$	$20$	$20$	$20$	$20$
$1$	$11$	$15$	$15$	$15$	$15$
$2$	$12$	$12$	$12$	$12$	$12$

Appendix

Table 9: Number of samples

|A(\varepsilon)|

averaged at each radius, broken down by KL target. A sample is one ensemble head trained at one target, counted only when every objective has a checkpoint at or below that radius. Counts fall at the tightest radius because a run trained at a loose target may never come that close to the reference policy, and fall with the target for the same reason. The bold rows are the totals that Table 6 reports. A seventh target

\tau=6

was trained but contributes nothing, the three objectives having been run there on disjoint heads.

		trust-region radius
$m$	Objective	$0.5$	$1$	$2$	$4$	$8$
$20$	Vanilla	$0.277$	$0.363$	$0.472$	$0.553$	$0.643$
	UWO $(0.1)$	$0.267$	$0.338$	$0.462$	$0.535$	$0.623$
	Covariance	$\bm{0.070}$	$\bm{0.055}$	$\bm{0.058}$	$\bm{0.052}$	$\bm{0.047}$
$10$	Vanilla	$0.427$	$0.510$	$0.572$	$0.665$	$0.723$
	UWO $(0.1)$	$0.417$	$0.497$	$0.567$	$0.658$	$0.712$

Appendix

Table 10: Reward model scalability. The construction of Table 6 repeated on Qwen2.5-7B at three ensemble sizes of

m

. The covariance arm is the safest at every radius and every ensemble size, and the uncertainty penalty stays within Monte-Carlo error of vanilla throughout. Bold marks the lowest risk in a column. Remaining parameters as in Table 6 .

Policy Optimization

Policy optimization develops algorithms for improving decision rules in reinforcement learning, language models, and sequential decision systems.

See all 224 papers

RVPO: Risk-Sensitive Alignment via Variance Regularization
May 7, 2026Ivan Montero, Tomasz Jurczyk, Bhuwan DhingraRL for Language Models Risk-Sensitive RL
Apple
Revisiting scaling laws for reward optimization
Sep 29, 2026Ali Aouad, Aymane El Gadarri, Vivek F. FariasReward Modeling Reward Over-Optimization
MIT
CleanRM: Reward Modeling from Noisy Feedback with a Few Clean Feedback for RLHF
Oct 15, 2026Licheng Pan, Haocheng Yang, Yuzhou Chen +7RL from Human Feedback Noisy-Label Learning
Zhejiang University · Xiaohongshu Inc · National University of Singapore +3

May 7, 2026J/K move · Enter open · S save

RVPO: Risk-Sensitive Alignment via Variance Regularization

Ivan Montero, Tomasz Jurczyk, Bhuwan Dhingra

Apple

RL for Language Models Risk-Sensitive RL Multi-Objective Language Model Alignment Reward Design for RL

Open paper AskPDF arXiv

Current critic-less RLHF methods aggregate multi-objective rewards via an arithmetic mean, leaving them vulnerable to constraint neglect: high-magnitude success in one objective can numerically offset critical failures in others (e.g., safety or formatting), masking low-performing "bottleneck" rewards vital for reliable multi-objective alignment. We propose Reward-Variance Policy Optimization (RVPO), a risk-sensitive framework that penalizes inter-reward variance during advantage aggregation, shifting the objective from "maximize sum" to "maximize consistency." We show via Taylor expansion that a LogSumExp (SoftMin) operator effectively acts as a smooth variance penalty. We evaluate RVPO on rubric-based medical and scientific reasoning with up to 17 concurrent LLM-judged reward signals (Qwen2.5-3B/7B/14B) and on tool-calling with rule-based constraints (Qwen2.5-1.5B/3B). By preventing the model from neglecting difficult constraints to exploit easier objectives, RVPO improves overall scores on HealthBench (0.261 vs. 0.215 for GDPO at 14B,

p < 0.001

) and maintains competitive accuracy on GPQA-Diamond without the late-stage degradation observed in other multi-reward methods, demonstrating that variance regularization mitigates constraint neglect across model scales without sacrificing general capabilities.

Abstract

Policy Optimization

Explore similar work

RVPO: Risk-Sensitive Alignment via Variance Regularization

Revisiting scaling laws for reward optimization

CleanRM: Reward Modeling from Noisy Feedback with a Few Clean Feedback for RLHF