Source-linked AI summary
Optimization Methods for Personalizing Large Language Models through Retrieval Augmentation
Alireza Salemi, Surya Kallumadi, Hamed Zamani
TL;DR
Personalized LLM generation requires selecting useful personal information without costly or privacy-risking LLM fine-tuning. The paper optimizes retrieval with reinforcement learning and knowledge distillation, adds pre- and post-generation retriever selection, and improves performance across six of seven LaMP tasks.
Problem
Existing LLMs do not reliably tailor outputs to users, while it is unclear which personal documents are relevant for downstream personalized generation.
Method
The paper optimizes retrieval using generation-feedback reinforcement learning and LLM-to-retriever knowledge distillation, then selects retrievers before or after generation.
Results
The proposed methods achieve statistically significant improvements on six of seven LaMP datasets, with a 5.5% average state-of-the-art improvement across datasets.
Takeaways & Limitations
Different personalization tasks benefit from different retrieval models, motivating adaptive retriever selection for personalized LLM generation.
Takeaways & Limitations
The study uses static prompt templates and evaluates only short text generation tasks, leaving prompt optimization and personalized long text generation for future work.
Abstract
from arXiv · showhide
This paper studies retrieval-augmented approaches for personalizing large language models (LLMs), which potentially have a substantial impact on various applications and domains. We propose the first attempt to optimize the retrieval models that deliver a limited number of personal documents to large language models for the purpose of personalized generation. We develop two optimization algorithms that solicit feedback from the downstream personalized generation tasks for retrieval optimization -- one based on reinforcement learning whose reward function is defined using any arbitrary metric for personalized generation and another based on knowledge distillation from the downstream LLM to the retrieval model. This paper also introduces a pre- and post-generation retriever selection model that decides what retriever to choose for each LLM input. Extensive experiments on diverse tasks from the language model personalization (LaMP) benchmark reveal statistically significant improvements in six out of seven datasets.
1 INTRODUCTION
Existing LLMs often fail to tailor outputs to users, while retrieval augmentation offers a practical route to personalization without modifying the LLM. This paper optimizes retrieval using downstream generation feedback and selects among retrievers for different inputs.
- LLMs often fail to tailor generated outputs to users’ backgrounds and historical preferences, motivating personalization across applications.
- Retrieval augmentation can use user-specific content with off-the-shelf LLMs, avoiding the storage, computation, and cold-start costs associated with adapting model parameters.
- The paper optimizes retrieval through reinforcement learning rewards defined by personalized-generation performance and knowledge distillation from the downstream LLM.
- Retriever selection is motivated by task-specific needs, with candidate retrievers including no retrieval, recency, BM25, Contriever, and two personalized dense retrievers.
- On seven LaMP tasks, the methods achieve statistically significant improvements on six, with 5.5% average state-of-the-art improvement and 15.3% average improvement over a non-personalized LLM.
2 RELATED WORK
Prior work studies personalization in search, recommendation, and NLP, while retrieval optimization research has focused mainly on question answering. Retriever selection and rank fusion provide related foundations but introduce distinct challenges.
- Personalization research spans search, recommendation, query auto-completion, dialogue, review generation, and recipe generation.
- LaMP evaluates personalized NLP models across diverse classification and short text generation tasks, complementing work on personalized long text generation.
- Retrieval optimization research includes distilling language-model signals to retrievers and using language-model attention to estimate document importance.
- Rank fusion can be difficult when retrievers use incompatible score scales or return non-overlapping documents, motivating alternative retriever-selection approaches.
3 NOTATIONS AND TASK FORMULATION
The paper formulates personalized generation as retrieving personal documents, incorporating them into a prompt, and using a fixed LLM to generate a user-tailored output. It avoids LLM fine-tuning because of cost and privacy concerns.
- Personalized generation seeks outputs tailored to user preferences and characteristics, using training instances containing a user, input text, and personalized ground-truth output.
- A retrieval model obtains personal documents from a user profile, and a prompt-generation function combines them with the input before LLM generation.
- The objective is to minimize error between the LLM’s generated output and the ground-truth personalized output while keeping the LLM parameters fixed.
- The fixed-LLM design avoids expensive fine-tuning and reduces privacy risk from memorizing information in users’ personal data.
4 LEARNING TO RETRIEVE FOR PERSONALIZING LLMS
Standard personalized learning-to-rank feedback is not directly available for retrieval-augmented generation, so the paper uses downstream LLM performance to optimize retrieval. It develops reinforcement-learning and knowledge-distillation approaches.
- Personalized retrieval for text generation lacks clear query-document relevance labels, making standard learning-to-rank methods unsuitable.
- The proposed methods evaluate how retrieved documents affect personalized generation and use that LLM feedback to optimize the retrieval model.
- ROPG-RL: ROPG-RL treats document selection as actions in a retrieval policy and increases probabilities for documents that produce more accurate personalized generation.
Parameterized Policy (𝜋𝜃).
The retriever is parameterized as a policy over personal documents, with training using downstream LLM feedback and efficiency-oriented approximations.
- Parameterized Policy (𝜋𝜃).: The policy assigns zero probability outside the user profile and can approximate the full profile using hierarchical softmax or top-l retrieval.The experiments use the top l documents under initial retrieval weights, with l=16.
- Parameterized Policy (𝜋𝜃).: Sampling multiple documents without replacement produced little or no improvement over the simpler approach that samples one document for policy updates.The single-document approach is more efficient for precomputing document-level rewards.
- Parameterized Policy (𝜋𝜃).: ROPG-RL samples personal documents as actions and optimizes the retrieval model using rewards from downstream personalized generation.The reward can be based on an arbitrary evaluation metric for personalized text generation.
- Parameterized Policy (𝜋𝜃).: ROPG-KD instead assigns higher probability to profile items that are more useful to the LLM by distilling downstream performance into the retriever.Its target distribution is based on each retrieved document’s contribution to personalized generation and is optimized with KL divergence.
- Parameterized Policy (𝜋𝜃).: The methods are applied to dense retrieval models, with Contriever encoding queries and documents and nearest-neighbor search ranking profile documents.Inference scores documents independently and sorts them in descending score order.
5 RETRIEVAL MODEL SELECTION FOR PERSONALIZING LLMS
The paper selects among diverse retrieval models for each input because no single retriever consistently performs best across LaMP tasks. It trains pre- and post-generation selection models by distilling downstream LLM performance into retriever-selection probabilities.
- Motivation: 14.9% to 20.2% of inputs receive the highest performance from ROPG-RL or ROPG-KD, confirming that retrieval-model preferences vary across LaMP datasets.No consistent winner exists across tasks; even the proposed retrieval optimizers are best for only subsets of inputs.
- Retriever selection models: RSPG-Pre selects a retriever before generation from prompts constructed with each candidate retriever, whereas RSPG-Post uses candidate outputs and the original prompt for selection.Both approaches select among retrievers for each input, but they differ in whether selection occurs before or after LLM generation.
- Training objective: The selection model is trained by distilling a target retriever distribution derived from each retriever’s downstream personalized-generation performance.The target distribution uses an evaluation metric applied to LLM outputs generated from documents retrieved by each candidate.
- Inference pipeline: For each input, all retrieval models first retrieve relevant documents, and an encoder-only model estimates a selection score for every candidate retriever.The resulting architecture supports both pre-generation and post-generation selection scenarios.
6 EXPERIMENTS
Experiments on the seven-task LaMP benchmark evaluate personalization, retrieval optimization, and retriever selection. Retrieval selection and optimized retrievers improve performance broadly, with RSPG-Post achieving the strongest overall results while leaving room for improvement on some tasks.
- 6.2 Empirical Results: After retrieval selection, personalized models outperform the non-personalized baseline on every LaMP dataset and metric, with statistically significant gains in almost all cases.Without selection, the optimized methods improve several but not all datasets; RSPG-Pre and RSPG-Post overcome this limitation, including on LaMP-7.
- 6.2 Empirical Results: ROPG-RL improves over Contriever on LaMP-1, LaMP-3, LaMP-4, and LaMP-7, while ROPG-KD additionally improves LaMP-6 and performs better on binary-feedback tasks.Neither optimization algorithm consistently dominates the other across tasks.
- How effective are the proposed retrieval selection methods?: Both RSPG-Pre and RSPG-Post outperform baselines across nearly all datasets, while RSPG-Post improves significantly over all baselines except on LaMP-5.RSPG-Pre instead shows a significant improvement over baselines on LaMP-5 and generally performs strongly on the classification tasks.
- How effective are the proposed retrieval selection methods?: Both selection models exceed 80% success on all classification tasks and LaMP-7, whereas success rates fall to 40%–65% on LaMP-4, LaMP-5, and LaMP-6.RSPG-Post has a higher success rate than RSPG-Pre on every task except LaMP-5.
- 6.2 Empirical Results: RSPG-Post performs best on six of seven LaMP datasets, with statistically significant gains over the remaining methods; Contriever is best only on LaMP-5.The comparisons use paired t-tests for generation and ordinal classification and McNemar tests for binary and categorical classification.
- How effective are the proposed retrieval selection methods?: Retrieval selection remains below its oracle upper bound on LaMP-3 and LaMP-4, reaching 68.3% and 75.4% of oracle performance, respectively.The corresponding figures for LaMP-1, LaMP-2, LaMP-5, LaMP-6, and LaMP-7 are 84.2%, 91.8%, 88.1%, 84.7%, and 98.0%.
7 CONCLUSIONS AND FUTURE WORK
The paper personalizes LLMs by optimizing retrieval with downstream feedback and selecting retrievers before or after generation, outperforming competitive baselines on six of seven LaMP datasets. It is limited to ranking-model optimization, static prompt templates, and short-text generation, with prompt optimization and long-text personalization left for future work.
- The proposed retrieval-augmented pipeline outperforms competitive baselines on six of seven LaMP datasets with statistically significant improvements.The evaluation covers seven diverse personalization tasks, and ablation studies examine the contribution of each pipeline component.
- Two retrieval-ranking optimization methods use personalized generation feedback: reinforcement learning and knowledge distillation from the language model.The reinforcement-learning reward is based on personalized text-generation quality, while distillation transfers knowledge to the retrieval model.
- Pre-generation and post-generation retriever selection models address the observation that different personalization tasks can benefit from different retrieval models.
- The study optimizes and selects ranking models but does not optimize prompt generation, which remains a future direction.The current approach uses static prompt templates and plans to use downstream LLM feedback to optimize prompt generation.
- Because all LaMP datasets involve short text generation, personalized long-text generation remains outside the evaluated scope.The authors identify long-text personalization as a future research direction.