How Retrieved Context Shapes Internal Representations in RAG

TL;DR

This work systematically analyzes how different types of retrieved documents affect the hidden states of LLMs, and how these internal representation shifts relate to downstream generation behavior, revealing how context relevancy and layer-wise processing influence internal representations.

Abstract

Retrieval-augmented generation (RAG) enhances large language models (LLMs) by conditioning generation on retrieved external documents, but the effect of retrieved context is often non-trivial. In realistic retrieval settings, the retrieved document set often contains a mixture of documents that vary in relevance and usefulness. While prior work has largely examined these phenomena through output behavior, little is known about how retrieved context shapes the internal representations that mediate information integration in RAG. In this work, we study RAG through the lens of latent representations. We systematically analyze how different types of retrieved documents affect the hidden states of LLMs, and how these internal representation shifts relate to downstream generation behavior. Across four question-answering datasets and three LLMs, we analyze internal representations under controlled single- and multi-document settings. Our results reveal how context relevancy and layer-wise processing influence internal representations, providing explanations on LLMs output behaviors and insights for RAG system design.

Figures (14)

• Figure 1: The overview of our analysis framework. For each query, we retrieve and rerank a set of documents, and group them as relevant, distracting, and random. We then control the input context of RAG with different type(s) of documents and obtain the hidden representations for comparative analysis.
• Figure 2: Representations of prompts paired with semantically similar documents remain close to the no-document baseline, whereas semantically dissimilar documents cause the representations to drift away. We apply PCA on the last prompt token representations across different document types and plot them in 2D.
• Figure 3: LLMs are more likely to abstain when context induces large representation shifts. For each context type, we compute the cosine similarity between the representations of with-context prompts and query, and categorize responses as correct, incorrect, or abstain. We show the result of Gemma3-27B. For other models, see Figure \ref{['fig:sim_vs_response_full']}.
• Figure 4: Base LLMs do not have representation drifts across different context types. We apply PCA on representations of instruction-tuned models and base models. We show the result of Gemma3-27B on Trivia QA. For other models and datasets, see Figure \ref{['fig:base_vs_it_full']}.
• Figure 5: Representations remain similar when a relevant document is present, regardless of other context. We perform PCA on the last prompt token representations for multiple-document contexts that include one relevant document, and plot them alongside representations with only a relevant document in 2D. We show the result of Gemma3-27B on Trivia QA. For other models and datasets, see Figure \ref{['fig:multiple_document_full']}.

Theorems & Definitions (1)

• Definition 3.1: Retrieval-Augmented Generation.