Technical Summary: Test-Time Personalization: A Diagnostic Framework and Probabilistic Fix for Scaling Failures

1. Problem Formulation

Current approaches to Large Language Model (LLM) personalization primarily focus on constructing better personalized models or inputs (via prompting, fine-tuning, or alignment) while treating inference as a single-shot generation process. This work introduces Test-Time Personalization (TTP), a paradigm that scales inference-time computation for personalized text generation.

In TTP, a personalized policy model $\pi_u$ generates $N$ candidate responses for a given query. A personalized reward model (RM) $\hat{r}_u$ then selects the best candidate $x^* = \arg\max_{x_i} \hat{r}_u(q, x_i)$. The goal is to maximize the expected true utility $U(N) = \mathbb{E}[r^*_u(q, x^*)]$, where $r^*_u$ is the user's true preference function.

While theoretical analysis suggests that an "oracle" reward model (one with perfect knowledge of $r^*_u$) can achieve expected utility growing logarithmically with $N$ ($O(\sqrt{\ln N})$), the authors observe that standard learned reward models fail to realize this potential. In practice, standard reward models often perform near random selection or worse, leaving the "oracle gap" wide open.

2. Methodology

2.1 Diagnostic Framework and Scaling Law

The authors develop an analytical framework to diagnose why standard reward models fail to scale. They derive a Unified Scaling Law that decomposes the Best-of-$N$ utility curve of any reward model into four measurable quantities:

1. Collapse Rate ($\alpha$): The fraction of users for whom the reward model degenerates to near-constant predictions (user-level collapse).
2. Reward Hacking Rate ($\beta$): The fraction of queries where the reward model's predictions negatively correlate with true quality (query-level reward hacking).
3. Mean Positive Correlation ($\bar{\rho}_+$): The average correlation on non-hacked queries.
4. Mean Negative Correlation ($\bar{\rho}_-$): The average correlation on hacked queries.

The population-level utility is approximated as:
$$\bar{U}(N) \approx \bar{\mu} + (1-\alpha) \cdot [(1-\beta)\bar{\rho}_+ - \beta|\bar{\rho}_-|] \cdot \bar{\sigma} \cdot c\sqrt{\ln N}$$

This law reveals two distinct failure modes:

• User-Level Collapse: Occurs when a reward model is trained on users with low-variance ground-truth labels, causing the model to predict constant values (near-zero correlation).
• Query-Level Reward Hacking: Occurs when the model overfits to spurious patterns on specific queries, resulting in negative correlation with true quality.

2.2 Probabilistic Personalized Reward Model

To mitigate these failure modes, the authors propose a Probabilistic Personalized Reward Model. Instead of predicting a single deterministic score, the model predicts a mean $\mu(x)$ and a variance $\sigma^2(x)$ for each input.

The model is trained using a Gaussian Negative Log-Likelihood (NLL) loss combined with a contrastive loss for high-score discrimination:
$$\mathcal{L} = \mathcal{L}_{NLL} + \lambda \mathcal{L}_{contrast}$$
$$\mathcal{L}_{NLL} = \frac{1}{2} \log \sigma^2(x) + \frac{(y - \mu(x))^2}{2\sigma^2(x)}$$

Theoretical analysis (Lemmas 5.1 and 5.2) demonstrates that the learned variance $\sigma^2$ acts as:

1. Gradient Buffering: For inputs with weak signals (prone to collapse), the model increases $\sigma^2$, attenuating the gradient on $\mu$. This prevents the model from collapsing into constant predictions.
2. Implicit Regularization: For inputs with inconsistent signals (prone to hacking), the model increases $\sigma^2$, down-weighting their contribution to the gradient. This prevents overfitting to spurious patterns.

At inference time, candidates are selected based solely on the predicted mean $\mu(x)$; the variance is used only during training to stabilize learning.

3. Key Contributions

1. Test-Time Personalization (TTP): Introduction of a new paradigm focused on scaling policy-model compute at inference time for personalized generation.
2. Diagnostic Scaling Law: Derivation of a predictive formula that forecasts the Best-of-$N$ scaling curve of any reward model based on four measurable quantities ($\alpha, \beta, \bar{\rho}_+, \bar{\rho}_-$). This law identifies user-level collapse and query-level reward hacking as the primary barriers to scaling.
3. Probabilistic Reward Modeling: Proposal of a probabilistic reward model that effectively mitigates both failure modes via learned variance, enabling stable and monotonic scaling across diverse tasks and policy types.

4. Experimental Results

The framework was evaluated on five personalized text generation tasks (News Headlines, Scholarly Titles, Abstracts, Product Reviews, Topic Writing) using the LaMP and LongLaMP benchmarks.

• Scaling Performance: The Probabilistic User RM is the only method that scales reliably across all tasks. At $N=30$, it recovers between 51% and 84% of the oracle gap, consistently outperforming training-based baselines (per-user LoRA fine-tuning) and standard Global/Deterministic User RMs.
• Theoretical Validation: The Unified Scaling Law accurately predicts observed scaling curves with a relative Mean Absolute Error (relMAE) of less than 3% across all tasks and reward model variants.
• Failure Mode Analysis:
  • Collapse: Deterministic models exhibited collapse rates ($\alpha$) up to 25% on certain tasks, whereas the Probabilistic model reduced this to 0%.
  • Hacking: Deterministic models showed reward hacking rates ($\beta$) up to 33%, while the Probabilistic model reduced this to near zero (e.g., <1% on Product Review).
• Policy Orthogonality: The benefits of TTP were consistent across different policy families (RAG, Persona prompting, and Persona+RAG), confirming that the scaling effect is driven by the reward model rather than the generator.
• Efficiency: The computational bottleneck remains candidate generation. Reward model scoring is negligible (2.3 ms per candidate vs. ~1500 ms for generation), making TTP highly efficient for high-throughput scenarios.

5. Significance and Claims

The paper claims that Test-Time Personalization offers a promising pathway to performance beyond training-based methods, provided effective reward models are constructed. The authors emphasize that:

• Standard reward models fail to scale due to specific, diagnosable failure modes (collapse and hacking) rather than a lack of personalization capacity.
• Probabilistic modeling provides a unique solution by leveraging learned variance as a training-time regularizer, rather than an inference-time selection mechanism.
• The diagnostic framework serves as a predictive tool, allowing practitioners to estimate TTP performance and identify high-risk users or queries before deployment.

The work is modest in its scope, noting that it focuses on text generation with ROUGE-derived rewards and assumes sufficient per-user data for training. It acknowledges that cold-start scenarios and generalization to non-text modalities (e.g., dialogue, agentic tasks) remain open challenges. The authors suggest that combining their probabilistic per-user RMs with other test-time axes (such as scaling user interactions or reward-model reasoning) represents a natural direction for future work.