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

1. Problem Statement

The paper addresses the challenge of personalizing Large Language Models (LLMs) to diverse human preferences. While Reinforcement Learning from Human Feedback (RLHF) aligns LLMs with general human values, it often collapses toward an "average" user, failing to capture individual nuances.

Existing Reward Factorization methods attempt to solve this by decomposing a user's reward function $r_u$ into shared basis functions $\phi$ and user-specific weights $w_u$ (i.e., $r_u = w_u^\top \phi$). However, the authors identify two fundamental limitations in current approaches:

1. Isolated Preference Modeling (C1): Existing methods estimate user weights in isolation. This fails to leverage shared preference patterns across the user population, leading to inaccurate estimates for users with scarce data (few-shot scenarios) and unseen users.
2. Deterministic Point Estimates (C2): Current methods treat user weights as deterministic points. This ignores the inherent uncertainty in user preferences caused by data scarcity or diverse tastes. Consequently, the model cannot distinguish between a confident preference (consistent history) and an uncertain one (sparse/diverse history), leading to unreliable inference.

2. Methodology: Variational Reward Factorization (VRF)

The authors propose VRF, a framework that represents user preferences as variational distributions within a shared probabilistic space rather than deterministic points. The framework consists of four key components:

A. Variational User Preference Encoding

Instead of a single weight vector, VRF infers a user's preference as a Gaussian distribution $q_u = \mathcal{N}(\mu_u, \text{diag}(\sigma_u^2))$ from a small reference set of pairwise feedback ($C_u$).

• Pairwise Encoding: A neural encoder processes each triplet $(x, y^+, y^-)$ to produce Gaussian parameters.
• Aggregation via Product of Gaussians: The individual Gaussians from the reference set are aggregated into a single user distribution using the Product of Gaussians. This method weights observations by their precision ($\sigma^{-2}$), allowing the model to naturally reduce uncertainty ($\sigma_u$) as more data becomes available.

B. Probabilistic Preference Bases

VRF introduces $K$ shared probabilistic preference bases $P = \{p_k\}$, where each $p_k$ is also a Gaussian distribution $\mathcal{N}(\mu_k, \text{diag}(\sigma_k^2))$.

• These bases represent distinct preference styles (e.g., "concise," "empathetic," "formal").
• They encode population-level patterns, allowing data-scarce users to benefit from collective signals.

C. Uncertainty-Aware Matching (Weight Derivation)

User weights $w_u$ are derived by matching the user distribution $q_u$ to the preference bases $p_k$ using the squared 2-Wasserstein distance ($W_2$).

• $W_2$ accounts for both the location ($\mu$) and the uncertainty ($\sigma$) of the distributions.
• The weights are computed via a softmax over the distances:
  $$w_{u,k} \propto \exp(-W_2(q_u, p_k)^2 / \tau_d)$$
• This results in a "soft assignment" that reflects uncertainty: if a user's distribution is wide (high uncertainty), the assignment to bases becomes more diffuse.

D. Training Objectives

The training loss combines three terms to ensure robust learning:

1. Variance-Attenuated Bradley-Terry Loss ($L_{VBT}$): Replaces the scalar reward difference with a Gaussian distribution. The loss function adaptively scales the gradient based on the variance of the reward estimate. High uncertainty (large variance) attenuates the gradient, preventing the model from overfitting to noisy or ambiguous signals.
2. Separation Loss ($L_{SEP}$): Maximizes the Wasserstein distance between different preference bases to ensure they capture distinct styles.
3. Regularization Loss ($L_{REG}$): Regularizes user distributions to lie within the learned preference space using a Mixture-of-Gaussians prior.

E. Inference

For unseen users, VRF performs few-shot inference in a single forward pass. It encodes the user's reference set into a distribution and derives weights via the pre-learned bases, requiring no gradient-based optimization per user.

3. Key Contributions

1. Problem Identification: Systematically identified and analyzed the limitations of isolated modeling and deterministic estimates in existing reward factorization methods.
2. Novel Framework: Proposed VRF, the first uncertainty-aware reward factorization framework that models both users and preference bases as probabilistic distributions.
3. Efficient Inference: Demonstrated a method for estimating user weights for unseen users in a single forward pass without optimization, significantly reducing latency compared to gradient-based baselines.
4. Robustness: Showed that modeling uncertainty explicitly leads to better performance in low-data regimes and diverse preference scenarios.

4. Experimental Results

The authors evaluated VRF on three benchmarks: PersonalLLM (synthetic diverse preferences), TL;DR (Reddit summaries), and PRISM (multi-turn conversations).

• Performance: VRF consistently outperformed state-of-the-art baselines (including LoRe, PReF, PAL, and VPL) across Seen, Unseen, and Overall splits.
  • On PersonalLLM with highly diverse preferences ($\alpha=0.001$), VRF achieved 95.3% overall accuracy, significantly beating the next best (PReF at 87.7%).
• Few-Shot Adaptation: VRF showed superior performance even with minimal data (e.g., $|C_u|=1$), outperforming competitors by up to 1.9% in low-data regimes.
• Uncertainty Robustness: VRF maintained high accuracy across users with low, medium, and high preference uncertainty, whereas deterministic baselines degraded significantly as uncertainty increased.
• Downstream Alignment: Improved reward estimation translated to better downstream performance in Best-of-N sampling, with VRF achieving the highest win rates against a standard Bradley-Terry baseline.
• Efficiency: VRF reduced per-user adaptation time to <1ms (single forward pass), whereas gradient-based methods required 83–114ms. For 1 million users, VRF completes adaptation in under 10 minutes, while PReF would take over 30 hours.

5. Significance

This paper provides a scalable and robust solution for LLM personalization. By shifting from deterministic point estimates to uncertainty-aware probabilistic modeling, VRF effectively handles the "cold start" problem for new users and the "sparse data" problem for existing users. The ability to quantify preference uncertainty allows the model to be more conservative with ambiguous data, leading to safer and more reliable alignment. The efficiency of the inference process makes VRF practical for large-scale deployment where millions of users require real-time personalization.