Optimization Methods for Personalizing Large Language Models through Retrieval Augmentation

TL;DR

The paper tackles personalizing large language models without fine-tuning the LLM by optimizing the retrieval component in a retrieval-augmented generation pipeline. It introduces two retrieval-optimization strategies—reinforcement-learning based and knowledge-distillation based—leveraging downstream LLM feedback, and couples them with pre-/post-generation retrieval-model selection to handle diverse personalization needs. Empirical evaluation on the LaMP benchmark across seven tasks shows significant improvements on six tasks, with the retrieval-selection approach achieving strong performance and approaching Oracle upper bounds on several datasets. The work advances practical, privacy-preserving LLM personalization by focusing optimization on retrieval and model selection while keeping the LLM frozen, and it points to future work in prompt optimization and long-form personalization.

Abstract

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.

Figures (4)

• Figure 1: An overview of retrieval augmentation approaches for LLM personalization. First, the query function $\phi_{q}$ produces a query from the input $x$. Relevant personal information is then retrieved and fed to the personalized prompt generation function $\phi_{p}$ for LLM consumption.
• Figure 2: Overview of training dense retrievers for personalizing LLMs using LLMs feedback with policy gradient optimization (a) and knowledge distillation (b). $\phi_q$ represents the query generation function, $\phi_p$ is the prompt generation function, and "Critic" denotes the evaluation metric employed for the personalized task.
• Figure 3: Relative winning rate for each selected retrieval model. When multiple retrieval models get the highest score, we consider all of them with the highest score as the winner.
• Figure 4: Pipeline for training the RSPG models. We minimize the KL-divergence between the scores generated for each prompt using RSPG and the performance of the prompt used for the evaluation of the LLM in a personalized task.