PaperlayerView exact revision on arXiv

Source-linked AI summary

Using Context Is Not Enough: Test-Time Training for Personalized Reward Modeling

Bohao Wang, Xiaoyan Zhao, Yang Zhang, Jinghang Guo, Chun Chen, Can Wang, Jiawei Chen

arXiv:2609.35109v1cs.AI

TL;DR

Shared reward models can overlook individual preferences, and ICL-based personalized reward models may fail to capture the preference relations in contextual pairs. P-TTT encodes those relations into user-specific fast weights using sequence-level operations and a preference-aligned objective. The paper reports improved preference-relation capture and performance over existing methods, with reduced inference latency.

  • Problem

    Shared reward models can overlook differences among users, while ICL-based personalized reward models may fail to capture preference relations conveyed by contextual pairs.

    Source: §1 · p. 1

  • Method

    P-TTT encodes contextual preference relations into user-specific fast weights through sequence-level update and apply operations and a preference-aligned objective.

    Source: §1 · p. 2

  • Results

    P-TTT achieves the highest mean accuracy across three datasets with both backbones, outperforming the strongest baseline in each setting by an average of 3.94 percentage points.

    Source: §5.2 · p. 8

  • Takeaways & Limitations

    P-TTT more effectively models user-specific preference relations while reducing inference latency compared with existing methods.

    Source: §6 · p. 10

    The theoretical analysis assumes a positive lower bound on the similarity of the matching contextual response.

    Source: §4.4 · p. 7

Abstract

from arXiv · show

Reinforcement learning from human feedback (RLHF) aligns large language models (LLMs) with human preferences, yet most pipelines learn a single reward model that overlooks individual differences in preferences. Personalized reward models (PRMs) address this by conditioning rewards on user-specific feedback, most commonly through in-context learning (ICL), where a user's historical comparisons are supplied as contextual preference pairs. However, we identify a key limitation of ICL-based PRMs: they fail to capture the preference relations conveyed by contextual pairs. To address this, we propose Preference-Aligned Test-Time Training (P-TTT), which explicitly encodes these relations into user-specific fast weights for personalized reward prediction. P-TTT introduces sequence-level update and apply operations to match the response-level granularity of preference feedback, together with a preference-aligned objective that directly uses pairwise preference relations to guide fast-weight adaptation. Notably, P-TTT is simple to implement and computationally efficient, updating fast weights within a single forward pass without inference-time backpropagation. Extensive experiments show that P-TTT more effectively captures historical preference relations and outperforms state-of-the-art methods by a large margin.

1 INTRODUCTION

Shared reward models can average conflicting preferences, while personalized reward models must infer a user’s preferences from a few historical pairs. The paper identifies that ICL may use contextual examples without capturing their preference relations, and proposes P-TTT to encode those relations in fast weights.

1 INTRODUCTION

  • 85.95% of predictions on the same target pairs remained unchanged on average after contextual preference labels were reversed, indicating that ICL-based PRMs weakly capture preference relations.Reversed labels represent a user with opposite interests, yet predictions usually did not adjust.Source: §1 · p. 2
  • P-TTT uses sequence-level update and apply operations and a preference-aligned objective to encode pairwise relations in user-specific fast weights.The sequence-level operations match response-level feedback, while the objective directly uses preference relations as the adaptation signal.Source: §1 · p. 2
  • P-TTT updates a small set of fast weights within one forward pass without inference-time backpropagation and avoids placing user history in the scoring context.The paper reports reduced self-attention overhead for lengthy contexts and faster inference than ICL.Source: §1 · p. 2
  • P-TTT consistently outperforms state-of-the-art baselines by up to 5.85 percentage points in accuracy, reusing a reward model’s MLP down-projection as fast weights for user-specific updates.Source: §1 · p. 2
  • The paper frames preference relations as central to personalized reward modeling and identifies their weak capture as a limitation of conventional ICL strategies.A shared reward model can average conflicting preferences and bias rewards toward the majority, while different users may prefer different responses to the same prompt.Source: §1 · pp. 1–2

2 PRELIMINARIES

The preliminaries describe reward models, personalized reward models, and test-time training as background for P-TTT. They contrast implicit personalization through contextual pairs with adapting fast weights, and note that existing TTT objectives do not directly model pairwise preferences.

2 PRELIMINARIES

  • A conventional reward model assigns a scalar score to a response given a prompt, while a personalized reward model conditions reward estimation on user-specific information.The conventional formulation learns a shared reward function across users; ICL commonly supplies a user’s historical chosen and rejected responses as contextual pairs.Source: §2 · p. 3
  • Embedding-based personalization may discard fine-grained preference information, while parameter-based methods can incur costly backpropagation and risk overfitting limited user feedback.Repeated adaptation may be costly when preferences and interactions evolve frequently.Source: §2 · p. 3
  • Test-time training adapts fast weights at inference time to retain contextual information, using update operations to write information and apply operations to retrieve it for subsequent predictions.Existing approaches commonly target context memorization or next-token prediction rather than pairwise preference relations.Source: §2 · pp. 3–4

3 EMPIRICAL ANALYSIS

The empirical analysis finds that ICL-based PRMs often retain their predictions when contextual preference labels are reversed, despite benefiting from contextual examples. This discrepancy suggests that using context does not ensure that models capture the preference relations it conveys.

3 EMPIRICAL ANALYSIS

  • 14.05% was the average flip rate for ICL-based PRMs after contextual preference labels were reversed, suggesting that those relations had limited influence on predictions.Flip rate is the fraction of initially correct predictions that change after label reversal.Source: §3 · p. 4
  • Contextual examples substantially increased training accuracy across all three datasets and both backbones, so ICL’s low flip rate was not solely due to ignoring context.On UF-P-2, accuracy rose from 48.44% to 97.06% for Qwen and from 49.55% to 99.81% for Llama.Source: §3 · p. 4
  • The analysis suggests that ICL models may rely on shortcut cues in historical prompts or responses rather than modeling preference relations, potentially limiting generalization to new target samples.The authors report that simple prompting strategies did not improve the issue.Source: §3 · pp. 4–5

4 METHOD

P-TTT adapts user-specific fast weights from contextual preference pairs to address ICL-based PRMs’ difficulty incorporating their preference relations. It uses response-level operations, a preference-aligned adaptation objective, and theoretical analysis of how adapted weights affect reward predictions.

4.1 MOTIVATION

  • P-TTT adapts fast weights from contextual pairs at the response level and uses a preference-aligned objective to encode preference relations for personalized reward prediction.It addresses ICL-based PRMs’ limited ability to capture the relations conveyed by contextual pairs.Source: §4.1 · p. 5

4.2 SEQUENCE-LEVEL UPDATE AND APPLY OPERATIONS

  • P-TTT reuses selected MLP blocks without modifying the backbone, initializing each user’s fast weights from shared parameters and adapting them using contextual preference pairs.The update uses final-token activations from chosen and rejected responses as separate keys, with cross-response attention masked; fast weights are applied at the final token for target scoring.Source: §4.2 · p. 6

4.3 PREFERENCE-ALIGNED OBJECTIVE

  • The preference-aligned objective derives values from the reward-head direction, assigning +v to chosen-response keys and −v to rejected-response keys.Its optimization encourages the chosen response’s fast-weight output to align more strongly with v than the rejected response’s output; target-pair training separately teaches the model to use the adapted weights.Source: §4.3 · pp. 6–7

4.4 THEORETICAL ANALYSIS

  • Reversing all contextual preference labels negates the adaptation-induced change in the target reward margin, showing that P-TTT’s predictions reflect preference direction.The theorem compares fast weights adapted from the same initialization before and after label reversal.Source: §4.4 · p. 7
  • When target responses match contextual keys on their corresponding preference sides, the P-TTT update increases the reward margin between chosen and rejected target responses.Retrieving a chosen-side contextual response increases target reward, whereas retrieving a rejected-side response decreases it.Source: §4.4 · p. 7

5 EXPERIMENTS

Experiments compare P-TTT with existing personalized reward modeling approaches on three benchmarks, then test its accuracy, preference-relation sensitivity, ablations, and inference efficiency.

5.1 EXPERIMENTAL SETUP

  • The evaluation covers UF-P-2, UF-P-4, and PersonalLLM, comparing no-personalization, ICL-, embedding-, parameter-based, and conventional TTT baselines.Pairwise preference accuracy is averaged over three random seeds.Source: §5.1 · p. 8

5.2 MAIN RESULTS

  • P-TTT achieves the highest mean accuracy across three datasets and both backbones, exceeding the strongest baseline by 3.94 percentage points on average.The results are reported in Table 1.Source: §5.2 · p. 8, Table 1 · p. 8
  • The reported results associate P-TTT’s gains with directly aligning adaptation to preference relations, while prior approaches can lose preference detail or use mismatched objectives.The authors describe textual summaries and compact representations as potential sources of information loss, and conventional TTT as objective-mismatched.Source: §5.2 · p. 9
  • P-TTT has a higher preference-flip rate than ICL across all settings, averaging 77.81% versus 14.05%.The authors interpret this as evidence that P-TTT predictions respond more to changes in user preference relations.Source: §5.2 · p. 9

5.3 ABLATION STUDIES

  • Removing either SLUA or PAO reduces accuracy across all three benchmarks; without SLUA, UF-P-4 accuracy falls from 64.00% to 50.24%.Without PAO, UF-P-2 accuracy falls from 75.40% to 50.12%, near BTL levels.Source: §5.3 · p. 9, Table 2 · p. 9

5.4 EFFICIENCY COMPARISON

  • For batches of 32 samples, P-TTT is fastest across both backbones and all three datasets, at 2.446–4.694 seconds per batch.It achieves 1.70–1.83× speedups over ICL and 6.21–9.72× over per-user fine-tuning.Source: §5.4 · p. 9, Figure 4 · p. 9

6 CONCLUSION

The conclusion presents P-TTT as a way to encode contextual preference relations into user-specific fast weights. The related work situates this approach among personalized reward models and test-time training methods.

6 CONCLUSION

  • P-TTT uses sequence-level update and apply operations and a preference-aligned objective to encode contextual preference relations into user-specific fast weights.The paper reports that it outperforms existing methods while reducing inference latency.Source: §6 · p. 10
  • The authors identify continual personalization for evolving preferences and selective retention or updating of long-term user information as future work.Source: §6 · p. 10

A RELATED WORK

  • Prior personalized reward models use embedding-, parameter-, or ICL-based approaches, while test-time training work has often targeted memorization or next-token prediction.This paper focuses on adapting from historical chosen–rejected pairs for personalized reward modeling.Source: A RELATED WORK · p. 12

B PROOF OF THEOREM 1

The proof analyzes P-TTT in a single Transformer block and states that reversing contextual preference labels negates the change in reward margin for a fixed target pair.

B PROOF OF THEOREM 1

  • For analytical tractability, the analysis is restricted to a single Transformer block equipped with P-TTT.Source: B PROOF OF THEOREM 1 · p. 12
  • For a fixed target response pair, reversing all contextual preference labels negates the change in reward margin relative to initialization.Source: B PROOF OF THEOREM 1 · p. 12

C PROOF OF THEOREM 2

The theorem’s setup assumes an identity value transformation and conditions involving similarity to a matching contextual key and zero aggregate contribution from remaining keys.

C PROOF OF THEOREM 2

  • The reward-correction setup identifies i* as the matching contextual response and cmatch as a positive similarity lower bound.Source: C PROOF OF THEOREM 2 · p. 13
  • The stated assumptions set Wvalue to the identity transformation and require positive similarity to a matching key plus zero aggregate contribution from the remaining keys.Source: C PROOF OF THEOREM 2 · p. 13

D PREFERENCE-FLIPPING ANALYSIS

The analysis tests whether reversing contextual preference labels changes predictions on the same target pairs. It constructs only variants that imply an unambiguous reversal of the target preference, then measures prediction flips.

Label reversal

  • Each counterfactual variant swaps chosen and rejected labels in selected context pairs while leaving prompts and response content unchanged.Not every pair must be reversed.Source: D PREFERENCE-FLIPPING ANALYSIS · p. 14
  • Variants are retained only when the original history implies an unambiguous target preference and every compatible user under the modified history prefers the opposite response.Variants with no compatible users or disagreement among compatible users are excluded.Source: D PREFERENCE-FLIPPING ANALYSIS · p. 14

Flip rate

  • For each initially correct sample with a valid variant, flip rate is the fraction of variants that change the prediction to the other candidate.Per-sample rates are averaged equally, regardless of how many valid variants each sample has; higher rates indicate greater influence of contextual preference relations.Source: D PREFERENCE-FLIPPING ANALYSIS · p. 14

E LIMITATIONS OF PROMPT-BASED STRATEGIES

The appendix evaluates whether explicit label-focused prompting improves contextual preference use, describes per-user fine-tuning, and tests whether P-TTT benefits from retaining context during prediction.

E LIMITATIONS OF PROMPT-BASED STRATEGIES

  • The prompt strategy emphasizes that chosen and rejected labels—not response order—signal user preferences, and asks the model to reward alignment with chosen responses.The evaluation prepends this instruction to contextual examples for Qwen2.5-0.5B-Instruct on UF-P-2 and UF-P-4.Source: E LIMITATIONS OF PROMPT-BASED STRATEGIES · p. 15
  • On UF-P-2, instructions raise flip rate from 6.83% to 8.98% and test accuracy from 53.33% to 53.97%, but both metrics fall on UF-P-4.On UF-P-4, flip rate decreases from 20.82% to 8.17% and test accuracy from 60.46% to 58.80%; explicit instructions alone are not reliably effective in these settings.Source: Table 3 · p. 15, E LIMITATIONS OF PROMPT-BASED STRATEGIES · p. 15

F DETAILS OF PER-USER FINE-TUNING

  • Per-user fine-tuning adapts the reward model on the historical preference pairs associated with each target sample.Source: F DETAILS OF PER-USER FINE-TUNING · p. 15
  • The target pair and label are excluded from adaptation, and the adapted model scores target responses without contextual examples.Parameters and optimizer state are reset for each target sample, with three optimization steps per sample.Source: F DETAILS OF PER-USER FINE-TUNING · p. 15

G RETAINING CONTEXT IN SELF-ATTENTION

H DATASET OVERVIEW AND STATISTICS

The experiments use three personalized preference datasets with four contextual pairs per user–target sample. The datasets represent preferences over response-quality dimensions or user-specific combinations of reward models.

Dataset overview

  • Experiments use UF-P-2, UF-P-4, and PersonalLLM, with 100 preference pairs reserved as a shared survey pool and four contextual pairs sampled per user–target sample.The reserved pool is excluded from training targets, and contextual pairs are labeled according to each user's preferences.Source: H DATASET OVERVIEW AND STATISTICS · p. 16, Table 5 · p. 16

UF-P-2 and UF-P-4

  • UF-P-2 covers users preferring helpfulness or honesty, while UF-P-4 adds instruction-following and truthfulness preferences.Each user prefers the response with the higher rating on their designated dimension; training and testing share user types but use disjoint target prompts.Source: H DATASET OVERVIEW AND STATISTICS · p. 16

PersonalLLM

  • PersonalLLM models pairwise preferences from user-specific weighted combinations of ten reward models and evaluates on 100 training users and 20 unseen users.The test set has 1,000 samples from each group, with accuracy reported over the combined 2,000 samples.Source: H DATASET OVERVIEW AND STATISTICS · p. 16
Loading 2609.35109v1…