Uncertainty-Aware Variational Reward Factorization via Probabilistic Preference Bases for LLM Personalization

Gyuseok Lee; Wonbin Kweon; Zhenrui Yue; SeongKu Kang; Jiawei Han; Dong Wang

cs.CLApr 1, 2026

Uncertainty-Aware Variational Reward Factorization via Probabilistic Preference Bases for LLM Personalization

Authors: Gyuseok Lee, Wonbin Kweon, Zhenrui Yue, SeongKu Kang, Jiawei Han, Dong Wang

Organizations: University of Illinois Urbana-Champaign · Korea University

Abstract

Reward factorization personalizes large language models (LLMs) by decomposing rewards into shared basis functions and user-specific weights. Yet, existing methods estimate user weights from scarce data in isolation and as deterministic points, leading to inaccurate and unreliable inference. We introduce Variational Reward Factorization (VRF), an uncertainty-aware framework that represents each user's preferences as a variational distribution in a shared preference space. VRF infers user distributions via a variational encoder, derives weights through Wasserstein distance matching with shared probabilistic bases, and downweights uncertain estimates through a variance-attenuated loss. On three benchmarks, VRF outperforms all baselines across seen and unseen users, few-shot scenarios, and varying uncertainty levels, with gains extending to downstream alignment. Our code is available at https://github.com/Gyu-Seok-Lee/VRF_COLM26.

Figures & tables

Figure 1: Motivating example. Preferences can be consistent (certain) or diverse (uncertain).

Figure 2: Overall framework of VRF. Best viewed in color.

	PersonalLLM
Method	Very Diverse ( $\alpha{=}0.001$ )			Moderately Diverse ( $\alpha{=}0.01$ )			Near Uniform ( $\alpha{=}0.1$ )
	Seen	Unseen	Overall	Seen	Unseen	Overall	Seen	Unseen	Overall
Ref	58.6 $\pm$ 0.2	58.9 $\pm$ 0.9	58.8 $\pm$ 0.4	59.1 $\pm$ 0.3	59.0 $\pm$ 0.6	59.1 $\pm$ 0.4	59.9 $\pm$ 0.2	60.0 $\pm$ 0.2	60.0 $\pm$ 0.2
BT	90.3 $\pm$ 0.3	90.5 $\pm$ 0.4	90.4 $\pm$ 0.2	91.1 $\pm$ 0.4	91.0 $\pm$ 0.3	91.1 $\pm$ 0.2	94.8 $\pm$ 0.2	94.6 $\pm$ 0.4	94.7 $\pm$ 0.2
VPL	86.4 $\pm$ 0.4	86.9 $\pm$ 0.5	86.6 $\pm$ 0.2	87.3 $\pm$ 0.5	87.4 $\pm$ 0.3	87.3 $\pm$ 0.2	91.7 $\pm$ 0.2	91.5 $\pm$ 0.5	91.6 $\pm$ 0.3
PAL	88.8 $\pm$ 0.5	89.3 $\pm$ 0.1	89.0 $\pm$ 0.3	89.3 $\pm$ 0.5	89.2 $\pm$ 0.3	89.2 $\pm$ 0.2	93.0 $\pm$ 0.2	92.8 $\pm$ 0.4	92.9 $\pm$ 0.3

Table 1: Overall performance comparison on three datasets.

Panel of Figure 3 — Figure 3: In-depth analysis of unseen users in PRISM. VRF maintains superior accuracy across all user groups (a–c) while significantly reducing adaptation latency (d).

Dataset	Model	NLL ( $\downarrow$ )	Brier ( $\downarrow$ )	ECE ( $\downarrow$ )
PersonalLLM ( $\alpha{=}0.001$ )	BT	0.2584 $\pm$ 0.0069	0.0734 $\pm$ 0.0017	0.0162 $\pm$ 0.0034
PersonalLLM ( $\alpha{=}0.001$ )	VRF	0.1747 $\pm$ 0.0195	0.0464 $\pm$ 0.0027	0.0099 $\pm$ 0.0095
PRISM	BT	0.6862 $\pm$ 0.0071	0.2402 $\pm$ 0.0027	0.0844 $\pm$ 0.0075
PRISM	VRF	0.6348 $\pm$ 0.0049	0.2237 $\pm$ 0.0018	0.0269 $\pm$ 0.0127
TL;DR ( $n{=}100$ )	BT	0.7572 $\pm$ 0.0399	0.2600 $\pm$ 0.0086	0.1353 $\pm$ 0.0292
TL;DR ( $n{=}100$ )	VRF	0.6862 $\pm$ 0.0310	0.2394 $\pm$ 0.0099	0.0761 $\pm$ 0.0348

Table 2: Calibration of predicted preference probabilities on the overall split. Lower is better for all metrics; ECE uses ten equal-width bins.

Model Variant	PersonalLLM ( $\alpha{=}0.001$ )			TL;DR ( $n{=}150$ )			PRISM
Model Variant	Seen	Unseen	Overall	Seen	Unseen	Overall	Seen	Unseen	Overall
(a) w/o $\mathcal{L}_{\text{VBT}}$	92.7 $\pm$ 0.3	92.5 $\pm$ 0.7	92.6 $\pm$ 0.3	61.4 $\pm$ 0.5	60.8 $\pm$ 0.3	61.1 $\pm$ 0.4	58.3 $\pm$ 0.2	58.1 $\pm$ 0.5	58.2 $\pm$ 0.4
(b) w/o $\mathcal{L}_{\text{SEP}}$	91.9 $\pm$ 0.3	91.4 $\pm$ 0.7	91.7 $\pm$ 0.3	61.5 $\pm$ 0.6	60.8 $\pm$ 0.3	61.2 $\pm$ 0.5	58.2 $\pm$ 0.1	58.2 $\pm$ 0.5	58.2 $\pm$ 0.3
(c) w/o $\mathcal{L}_{\text{REG}}$	93.2 $\pm$ 0.2	93.5 $\pm$ 0.3	93.3 $\pm$ 0.2	61.2 $\pm$ 0.6	60.3 $\pm$ 0.4	60.8 $\pm$ 0.4	60.8 $\pm$ 0.4	60.8 $\pm$ 0.9	60.8 $\pm$ 0.7
(d) w/o MoG	91.9 $\pm$ 0.3	91.4 $\pm$ 0.7	91.7 $\pm$ 0.2	61.3 $\pm$ 0.5	60.8 $\pm$ 0.3	61.1 $\pm$ 0.4	60.2 $\pm$ 0.3	60.3 $\pm$ 0.8	60.2 $\pm$ 0.6
(e) w/o PoG	94.8 $\pm$ 0.2	95.0 $\pm$ 0.2	94.9 $\pm$ 0.2	62.7 $\pm$ 0.2	61.8 $\pm$ 0.4	62.2 $\pm$ 0.3	61.7 $\pm$ 0.5	60.6 $\pm$ 0.4	61.1 $\pm$ 0.2

Table 3: Ablation study on VRF across three datasets.

Figure 4: Case studies and the learned preference space on PRISM (

K{=}8

). (a) Four users with varying data scarcity and preference diversity.

n

is the reference set size. Bars show the preference distribution classified by Claude Sonnet 4.6 (see Appendix D.1 for details). (b) Learned preference space. Stars and shaded regions denote the preference bases. For user ellipses, the center point represents

\mathbf{w}_{u}

, and the size reflects

{\sigma}_{u}

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Seen					Unseen
	#Users	Training		Evaluation		#Users	Evaluation
	#Users	$\|\mathcal{C}_{u}\|$	$\|\mathcal{D}_{u}\|$	$\|\mathcal{C}_{u}\|$	$\|\mathcal{D}_{u}\|$	#Users	$\|\mathcal{C}_{u}\|$	$\|\mathcal{D}_{u}\|$
PersonalLLM	1,000	9.0	36.0	45.0	45.0	1,000	9.0	9.0
TL;DR ( $n{=}100$ )	20	50.0	50.0	100.0	1,449.8	20	50.0	1,146.4
TL;DR ( $n{=}150$ )	20	50.0	100.0	150.0	1,449.8	20	50.0	1,146.4
PRISM	623	3.0	6.1	9.1	8.6	623	9.6	9.1

Appendix

Table 4: Average dataset statistics per user. For seen users,

C_{u}

is sampled from training data per epoch, and the full training set serves as

C_{u}

at evaluation. For unseen users, only

C_{u}

is provided at inference to estimate user weights. Both settings are evaluated on held-out

D_{u}

Dataset	Epochs	Batch size	Learning rate	Dropout	Weight decay	$K$
PersonalLLM ( $\alpha{=}0.001$ )	5	8	2e-5	0.0	0.01	8
PersonalLLM ( $\alpha{=}0.01$ )	5	8	2e-5	0.0	0.001	8
PersonalLLM ( $\alpha{=}0.1$ )	5	8	2e-5	0.0	0.01	8
TL;DR ( $n{=}100$ )	3	8	6e-5	0.1	0.001	8
TL;DR ( $n{=}150$ )	3	8	6e-5	0.1	0.001	8
PRISM	5	8	6e-5	0.0	0.01	12

Appendix

Table 5: Best hyperparameter configurations for VRF across all datasets.

Backbone	Seen	Unseen	Overall	Training (min)	Inference (min)
0.6B	61.31 $\pm$ 0.57	60.58 $\pm$ 0.53	60.94 $\pm$ 0.18	80.8	20.1
1.7B	61.28 $\pm$ 0.32	60.44 $\pm$ 0.27	60.86 $\pm$ 0.21	131.4	33.6
4B	61.72 $\pm$ 0.51	61.31 $\pm$ 0.49	61.51 $\pm$ 0.43	288.1	74.6
8B	62.14 $\pm$ 0.46	61.15 $\pm$ 0.46	61.64 $\pm$ 0.32	446.8	117.3

Appendix

Table 6: Effect of backbone scale on accuracy and efficiency for the PRISM dataset.

Panel of Figure 5 — Figure 5: Hyperparameter sensitivity analysis on three datasets.

Panel of Figure 6 — Figure 6: In-depth comparison on PersonalLLM ( $\alpha{=}0.001$ ) and TL;DR ( $n{=}150$ ) unseen users.

Figure 7: Distribution of dominant preference bases for seen and unseen users on PRISM.

Personalized Language Models

Personalized language models adapt responses to individual users’ preferences, histories, writing styles, and changing needs.

See all 89 papers

PAFO: Pareto Fairness Optimization for Personalized Reward Modeling
Jun 6, 2026Xiaoyan Zhao, Haoting Ni, Yang Zhang +3Reward Modeling LLM Alignment
National University of Singapore · University of Science and Technology of China · Peking University
MiCRo: Mixture Modeling and Context-aware Routing for Personalized Preference Learning
May 30, 2025Jingyan Shen, Jiarui Yao, Rui Yang +5Pairwise Preference Learning Reward Modeling
New York University · University of Illinois Urbana-Champaign · Rice University
MoF: Preference-Aware Mixture Modeling for Black-Box LLM Personalization
Oct 6, 2026Hun ParkLLM Personalization Personalized Language Models
Department of Artificial Intelligence Korea University Seoul, Republic of Korea

Jun 6, 2026J/K move · Enter open · S save

PAFO: Pareto Fairness Optimization for Personalized Reward Modeling

Xiaoyan Zhao, Haoting Ni, Yang Zhang, Chunyuan Zheng+2

National University of Singapore · University of Science and Technology of China · Peking University

Reward Modeling LLM Alignment Personalized Reward Modeling Algorithmic Fairness User Preference Modeling

Open paper AskPDF arXiv

Large language models (LLMs) increasingly rely on reward models to align their outputs with diverse user preferences. While personalized reward models aim to capture such heterogeneity, they are often trained on imbalanced user preference data and may therefore favor users whose preferences are more common in the training population. In this paper, we identify this failure mode as personalized reward bias, where reward modeling quality varies systematically with preference support rate. We formulate its mitigation as a Pareto fairness problem over group utilities, aiming to improve under-served users without degrading other user groups. To this end, we propose PAFO, a Pareto fairness optimization framework for personalized reward modeling. PAFO first trains group-specialized reward models for majority and minority preference groups, then constructs conditional margin-level supervision to distill their heterogeneous preference boundaries into a single unified model. The resulting model uses group information only during training and requires no explicit group labels at inference time. Experiments on Personal-LLM and DSP show that PAFO improves both minority-group and majority-group accuracy while reducing user-level unfairness across multiple metrics, demonstrating its effectiveness for fairer LLM personalization.

Abstract

Personalized Language Models

Explore similar work

PAFO: Pareto Fairness Optimization for Personalized Reward Modeling

MiCRo: Mixture Modeling and Context-aware Routing for Personalized Preference Learning

MoF: Preference-Aware Mixture Modeling for Black-Box LLM Personalization