Learning to Reason for Multi-Step Retrieval of Personal Context in Personalized Question Answering

1. Problem Statement

Personalized Question Answering (QA) aims to generate responses that are not only factually accurate but also aligned with a user's specific background, preferences, and historical context. While current state-of-the-art methods utilize Retrieval-Augmented Generation (RAG), they often suffer from surface-level personalization.

• Limitation of Existing Methods: Standard approaches typically use the user's query directly to retrieve documents from their profile and append them to the prompt. This static retrieval strategy fails to identify nuanced personalization aspects (e.g., a user's occupation or medical history) that might require specific, multi-step queries to uncover.
• The Gap: Effective personalization requires reasoning to identify what information is needed, followed by multi-step retrieval to fetch that specific evidence, and finally integration of that evidence into the reasoning process. Current models lack the ability to dynamically decide when to retrieve, what to retrieve, and how to use it during the generation process.

2. Methodology: PR2 Framework

The authors propose PR2 (Personalized Retrieval-Augmented Reasoning), a Reinforcement Learning (RL) framework that integrates retrieval directly into the reasoning trajectory.

Core Components

• Architecture: PR2 treats the retrieval process as an external tool within the environment. The policy model ($\pi_\theta$) can interleave three types of actions during generation:
  1. Thinking: Generating internal reasoning steps (using <thought> tags).
  2. Searching: Formulating specific queries to the user profile (using <search> tags).
  3. Answering: Generating the final response (using <answer> tags).
• Optimization Algorithm (GRPO): The framework uses Group Relative Policy Optimization (GRPO). Unlike standard PPO, GRPO does not require a separate value function (critic). Instead, it optimizes the policy based on the relative advantage of a group of sampled trajectories.
• Rollout Sampling: For each training instance, the model samples a group of $G$ personalized trajectories. Crucially, it also samples a non-personalized baseline trajectory (generated without retrieval) to serve as a control.

Reward Modeling & Advantage Computation

• Reward Function: An automatic judge (Qwen2.5-32B) evaluates the final response against a set of personalized rubrics ($E_{x_i}$) and user narratives ($r_i$). The score ($0-2$ per aspect) is normalized to $[0, 1]$.
• Group-Relative Advantage: The key innovation lies in how the advantage is calculated. To prevent the model from retrieving unnecessarily or relying on noisy signals, the advantage ($A_i$) for a personalized trajectory is computed relative to the non-personalized baseline:
  $$A_i = \frac{r_i - r_{no-per} - \mu}{\sigma}$$
  Where $r_i$ is the reward for the personalized trajectory, $r_{no-per}$ is the reward for the non-personalized baseline, and $\mu, \sigma$ are the mean and standard deviation of the group rewards.
  • Implication: The model is only incentivized to retrieve from the user profile if the retrieval leads to a response that outperforms a standard non-personalized generation. This teaches the model to be selective about when personalization is beneficial.

3. Key Contributions

1. First Reasoning-Integrated Personalization: PR2 is the first framework to explicitly model the reasoning process for retrieving personal context, moving beyond simple "query-retrieve-generate" pipelines.
2. Adaptive Retrieval Policies: The model learns to dynamically decide when to stop reasoning and initiate a search, what specific query to generate, and how to incorporate the retrieved evidence into intermediate reasoning steps.
3. Contrastive Advantage Learning: By including non-personalized baselines in the GRPO advantage calculation, the framework learns to distinguish between genuine personalization gains and generic generation improvements, avoiding over-retrieval.
4. Open-Source Implementation: The authors released the code to ensure reproducibility.

4. Experimental Results

The framework was evaluated on the LaMP-QA benchmark, which covers three domains: Art & Entertainment, Lifestyle & Personal Development, and Society & Culture. Experiments were conducted using three different Large Language Models (LLMs): Qwen2.5-3B, Qwen2.5-7B, and Gemma3-4B.

• Performance Gains: PR2 consistently outperformed strong baselines (including Search-R1, RAG-Personalization, PPlug, HYDRA, and PrLM).
  • Qwen2.5-3B: 8.8% relative improvement over the best baseline.
  • Qwen2.5-7B: 12.0% relative improvement.
  • Gemma3-4B: 9.8% relative improvement.
• Ablation Studies:
  • Non-Personalized Baseline: Removing the non-personalized baseline from the advantage calculation caused a significant performance drop (e.g., Qwen2.5-7B dropped from 0.5256 to 0.4736), confirming the necessity of the contrastive learning signal.
  • Retrieval Count ($k$): Retrieving the top 3 documents ($k=3$) yielded optimal results. $k=1$ limited recall, while $k=5$ introduced noise that interfered with reasoning.
• Training Dynamics: Analysis showed that as training progressed, models learned to retrieve more frequently and produce longer, more detailed responses, correlating with increased reward scores. Larger models (7B) demonstrated a stronger capacity to leverage retrieved evidence compared to smaller models (3B).

5. Significance

This paper addresses a critical limitation in current personalized AI systems: the inability to reason about what personal information is needed before retrieving it. By unifying reasoning and retrieval under a reinforcement learning framework with a contrastive reward signal, PR2 enables models to:

• Generate highly tailored responses that account for complex user constraints (e.g., dietary restrictions, career stage).
• Avoid the "retrieval noise" problem by learning to retrieve only when it adds value over a standard answer.
• Provide a scalable, RL-based approach for personalization that can be applied across various LLM sizes and domains.

The work establishes a new paradigm for personalized QA, shifting from static context injection to dynamic, reasoning-driven context acquisition.