Rethinking Personalized Generation: Test-Time Alignment via Factorized Ranking Models

AI summary

Overview

Research area: Personalization of large language models, test-time alignment, and reward modeling for LLM decoding.

Technical level: Intermediate. The paper assumes familiarity with Best-of-N sampling, reward models, and RLHF/DPO-style alignment, but its core idea (replace a large reward model with a tiny MLP over the generator's own hidden states) is explained in accessible terms.

Scope: The paper runs an oracle Best-of-N experiment to show that a frozen LLM generator already produces well-personalized candidates, then proposes a million-parameter MLP ranker that reuses the generator's hidden states to select them at negligible cost, evaluated on nine datasets across three personalization settings.

What This Paper Is About

Standard alignment methods such as RLHF and DPO optimize LLMs against a single, monolithic notion of human preference, which averages away the individual tastes that make a response read as if written for a specific person. Existing fixes treat this as a capability deficit and finetune per user or adapt the generator at training time, while reward-model-based Best-of-N selection is both poorly calibrated for personalization and expensive because a billion-parameter reward model must re-encode every candidate. The authors show the missing capability is already in the frozen generator and that the real bottleneck is selection, then build a small MLP ranker that scores candidates directly from the generator's final-layer hidden states.

Key Contributions

1. Empirical discovery of personalization headroom. An oracle Best-of-N analysis shows that a frozen base generator keeps producing better-matched candidates as the pool grows, while generalist reward models recover only a small fraction of that gain, reframing personalized alignment as a candidate selection problem rather than a generator capability bottleneck.

2. A personalized ranking model at million-parameter scale. A multi-layer perceptron of 1M to 30M parameters scores candidates by reusing the generator's own final-layer hidden states for the query, the user profile, and the response, with no separate encoder, no re-encoding of text, and no per-user learned parameters. The same ranker also extends to single-pass ranking guided generation.

3. Comprehensive experimental validation. Evaluations on nine datasets spanning three personalization settings show the ranker outperforms generalist billion-parameter reward models on every dataset, uses under 0.4% of their parameters, and has four orders of magnitude lower scoring latency.

4. Diagnostic analyses. The paper reports ranker size scaling (1M to 30M), comparison against a per-dataset finetuned 8B reward model, comparison against personalized reward models that fit per-user parameters from labeled comparisons, transfer to a held-out metric (BLEU), and stratification of results by how much user history is available.

Main Findings

• A large, untapped headroom exists. The oracle selector, which picks the candidate with the evaluation metric itself, rises steadily up to N = 64 on every task with no sign of saturation. The best of four generalist reward models plateaus by N = 16 on short-form generation and never beats its N = 1 score on long-form QA; at N = 64 it captures at most 23% of the headroom, and under 12% on eight of the nine datasets.

• Per-dataset headroom capture is low or negative. At N = 64 the strongest generalist reward model captures 4%, 2%, and 8% of the oracle headroom on News, Scholarly, and Tweet, and 23%, 6%, and 11% on Amazon, Yelp, and Google. On Art, Lifestyle, and Society it captures −5%, −5%, and 0%, meaning its Best-of-64 pick is no better than a single sample under ROUGE-L; on ROUGE-1 the shares range from −9% to 22%.

• The ranker beats reward models on every dataset. Against four generalist reward models (Skywork-Reward-V2-Llama-3.1-8B, internlm2-7b-reward, URM-LLaMa-3.1-8B, ArmoRM-Llama3-8B-v0.1), the personalized ranker is the best selector on all nine datasets and its advantage widens as the pool grows, where the generalist models stay flat or decline.

• Ranker size barely matters. Moving from 1M to 30M parameters changes the Best-of-64 score by at most 0.006 ROUGE-L on any dataset, because the semantic work is done by the generator and the ranker only learns a regression in a fixed feature space.

• A finetuned 8B model still wins, but at far higher cost. Finetuning Skywork-Reward-V2-Llama-3.1-8B on the same pools with the same objective puts it ahead on all nine datasets, by Δ = +0.019 ROUGE-L on short-form generation and below +0.01 on long-form QA and explainable recommendation. Measured against the headroom, the 30M ranker recovers 45% to 94% of the finetuned model's gain per dataset, at four orders of magnitude lower scoring latency and without finetuning a billion-parameter model per dataset.

• Gains transfer to a metric the ranker never saw. Trained only on ROUGE-L labels, the ranker selects the highest-BLEU candidate on all nine datasets under BLEU evaluation, and the four generalist reward models trail it everywhere, most clearly on explainable recommendation where the ranker improves on the strongest baseline by 20% to 64% relative.

• The two metrics disagree a lot within pools. ROUGE-L and BLEU select the same best candidate in only 12–37% of pools, and their within-pool rank correlation is 0.36–0.59, so a ranker that had memorized ROUGE-L would gain little under BLEU. Across-prompt correlation is higher (0.48–0.78), reflecting shared prompt difficulty rather than agreement on which candidate to pick.

• Guided decoding recovers a modest part of the gain. Steering uncertain tokens in a single generation reaches the level of Best-of-2 to Best-of-4 of the sampled pool on short-form generation and long-form QA, and Best-of-1 to Best-of-2 on explainable recommendation. Per dataset it exceeds Best-of-64 only on Tweet and falls below Best-of-1 on News.

• Decoding hyperparameters are not sensitive. Varying each guided-decoding hyperparameter around its default changes ROUGE-L by at most about 0.005 on every dataset; only the entropy threshold τ has a visible effect, and it controls how often the ranker intervenes rather than the final score.

• More user history generally helps, and sparsity degrades gracefully. Bucketing users by number of historical interactions, performance rises most clearly on Tweet (0.410 to 0.478) and Scholarly (0.369 to 0.410) and on all three XRec datasets; the long-form QA columns are flatter with no consistent direction, and News is the exception, dropping for users with 100 or more pairs (0.1385).

• The training loss is not the source of the gains. Pairwise objectives (Bradley-Terry, RankNet) and listwise objectives (ListNet, ListMLE, LambdaRank) perform comparably to pointwise MSE under identical architecture, data, and protocol; the gains come from embedding recycling and train-time candidate scaling.

Methodology in Plain English

The authors start with a diagnostic. They let a frozen Qwen2.5-7B-Instruct generate N candidate responses per prompt and compare two ways of picking the winner: an oracle that uses the evaluation metric itself, and the best of four off-the-shelf 8B reward models. The oracle keeps improving with N; the reward models stop improving quickly. That tells them the good responses already exist and the problem is picking them out.

Their fix is to stop re-reading the text. The generator has already processed the query, the user's profile, and each candidate, so its final-layer hidden state at the last token of each of those three texts already encodes that understanding. They take those three vectors, concatenate them, and feed them through a small MLP that outputs a single score. Because the query and profile vectors depend only on the prompt, they are computed once and shared across all N candidates, so scoring a pool costs one forward pass through a shallow network per candidate rather than a full transformer pass.

They train this ranker on pools of 64 candidates sampled from the generator itself, labeling each candidate with its ROUGE-L against the reference the user actually wrote. Because the candidates come from the same generator, the negatives are hard: structurally sound and on topic, but not in the user's voice. They standardize the targets within each pool to zero mean and unit variance so the model focuses on distinguishing candidates within a prompt rather than predicting each prompt's overall difficulty, and they adopt pointwise MSE by default while showing that pairwise and listwise alternatives do just as well.

They also use the ranker as a decoding-time signal. When the generator's next-token entropy exceeds a threshold τ, they take the gradient of the ranker's predicted score with respect to the current hidden state, project that direction onto the vocabulary through the LLM head, and add the resulting vector to the logits with an entropy-modulated step size. Confident tokens are left untouched.

Why This Matters

Impact on research. The paper shifts the framing of personalized alignment from "train the generator harder" to "select better from what the generator already produces." It argues that test-time compute is a lever for personalization only if the selector's accuracy scales with pool size, and it shows that generically trained reward models do not. It also offers an existence proof that a ranker does not need its own representation-learning stack: recycling the generator's hidden states suffices.

Real-world applications (drawn from the three benchmark settings the paper evaluates):

• Personalized writing assistance, where a model drafts in a specific user's voice rather than a generic style, matching the short-form generation tasks (news headlines, scholarly titles, tweet paraphrasing).
• Personalized question answering over long-form informational content, using an individual's past narratives and histories, matching the LaMP-QA setting.
• Personalized explainable recommendation, generating natural-language explanations of why an item suits a particular user, matching the XRec setting across Amazon, Yelp, and Google.
• Low-latency deployment of Best-of-N sampling, where the scoring step becomes nearly free relative to generation and no separate billion-parameter reward model needs to be served or finetuned per dataset.

Industry relevance. The practical claim is about cost. The ranker attaches to any generator that exposes its hidden states, requires no per-user parameters, no per-user preference labels, and no finetuning of a large model per dataset, and it scores a pool at four orders of magnitude lower latency. For a system already paying to generate N candidates, the marginal cost of picking among them drops to near zero.

Future Directions

• Human evaluation of style consistency. The authors state explicitly that their evaluation rests on lexical overlap with user-authored references and that a human study of style consistency is left for future work.
• Closing the remaining gap to finetuned 8B reward models. The finetuned Skywork-Reward-V2-Llama-3.1-8B is still ahead on all nine datasets, and the 30M ranker recovers 45% to 94% of its gain per dataset; what the missing percentage consists of is an open question.
• Removing the reliance on the evaluation metric as the training target. The ranker's supervision is ROUGE-L against user-authored references, and the paper notes semantic metrics such as BERTScore abstract away precisely the surface-level nuances where personalization manifests. Better or differently shaped supervision is an open direction.
• Extending guided decoding. Guided decoding recovers only a modest part of the headroom (Best-of-2 to Best-of-4 levels on two tasks, Best-of-1 to Best-of-2 on the third, and below Best-of-1 on News), so most of the gain still requires materializing a pool.

Target Audience

Researchers and practitioners working on LLM personalization, test-time compute scaling, and reward modeling, as well as engineers building personalized assistants, recommenders, or user-adaptive writing tools who need per-user alignment without per-user finetuning or large reward-model inference costs. Readers looking for a self-contained introduction to Best-of-N selection tradeoffs will also find the oracle analysis useful, though the paper assumes comfort with standard alignment terminology.

Authors’ abstract

Aligning large language models (LLMs) to diverse user preferences is fundamentally hindered by standard alignment paradigms that optimize for monolithic users. In this work, empirical studies are first used to reveal the existence of a massive, untapped performance headroom for personalized generation through test-time alignment. We demonstrate that personalized generation is uniquely suited for test-time scaling methods like Best-of-N (BoN) because it can be viewed primarily as a candidate matching problem rather than a generator capability bottleneck. While reward models could in principle exploit this headroom, they are poorly calibrated for personalization, and their billion-parameter scale makes scoring large candidate pools prohibitively expensive. To overcome this limitation, we propose a parameter-efficient framework utilizing million-parameter scale multi-layer perceptron (MLP) ranking models. Our personalized ranking model directly reuses the internal embeddings of the base generator with minimal overhead. By scaling train-time data to provide fine-grained personalized preferences, this million-parameter ranking model accurately scores large candidate pools and can seamlessly guide generation to reduce the cost of materializing N candidates. Extensive experiments on nine datasets spanning three personalized generation settings show that our personalized ranking model effectively exploits the discovered headroom, outperforming billion-parameter generalist reward models on every dataset, with under 0.4% of their parameters and four orders of magnitude lower scoring latency.