Latent Abstraction for Retrieval-Augmented Generation

The gist The LAnR framework proposes a unified architecture where a single LLM performs encoding, retrieval, and generation entirely within its own latent space to enhance Retrieval-Augmented Generation (RAG) efficiency.

How it works

1. LAnR replaces explicit text-based query generation with latent query vectors derived from the hidden states of designated special tokens [PRED] (Page 3).

2. These latent query vectors are used to retrieve documents from a vector database constructed using representations from the same LLM, enabling tight integration between retrieval and generation (Page 3).

3. To support multi-hop reasoning, LAnR introduces a lightweight multi-layer perceptron (MLP) head that predicts whether additional retrieval steps are necessary (Page 3).

4. The framework operates iteratively: at each turn, the model appends [PRED] to the context and extracts the latent query vector q(r), which is used for retrieval and control decisions (Page 3).

5. A lightweight MLP head that determines whether the currently retrieved evidence is sufficient to answer the query, or whether additional retrieval is required takes q(r) as input and produces a binary decision yˆ(r) = fθ q(r) (Page 3).

Latent Query Construction

1. The LLM bypasses explicit textual query generation and instead produces a dense retrieval vector directly from its hidden representations by appending a designated token [PRED] to the input, forming x˜ = (x1,..., xT, [PRED]) (Page 3).

2. Through causal self-attention, the hidden state attends to the full preceding context, yielding a latent query vector: q = h[PRED] ∈ R d (Page 3).

3. Each representation of the document Di are obtained from the same model M via last-token pooling: di = h last(Di) ∈ R d (Page 3).

4. The [PRED] representation is trained with a standard contrastive objective LCL = − log expim(q, d+)/τ + PN− j=1 expsim(q, d−j)/τ, where d+ is a positive document and the negatives are mined from model retrieval errors (Page 3).

Multi-turn Training Objective

1. The framework extends to multi-turn training by combining next-token prediction loss (LNTP), multi-turn contrastive retrieval loss (LCL), and the control head objective (Lctrl) in the total loss L = X R r=1 L(r) CL + λLNTP + µLctrl (Page 5).

2. The control head label y(r) is derived from the model’s own retrieval state, where y(r) = 0 if P ⊆ R(r), and 1 otherwise (Page 5).

3. The contrastive loss at turn r uses a positive document sampled from P(r) = P R(r−1), steering each successive turn toward the missing evidence: L CL = − log expim(q(r), d+(r))/τ + PN− j=1 expsim(q(r), d−j)/τ (Page 5).

Retrieval Control Mechanism

1. The MLP-based retrieval control head is trained with binary cross-entropy: Lctrl = − X R r=1 h y(r) log ˆy(r) + (1 − y(r)) log(1 − yˆ(r)) (Page 5).

2. During inference, the model performs iterative retrieval beginning from the input query, and if the controller predicts that additional retrieval is needed (yˆ(r) ≥ 0.5), it retrieves the top-K documents and continues (Page 4).

3. The control head signal is motivated by empirical observations that uncertainty signals, such as entropy in the token distribution, correlate with the need for further information acquisition (Page 3).

Experimental Results

1. LAnR consistently achieves superior performance compared to state-of-the-art RAG methods across five QA benchmarks, especially on challenging multi-hop benchmarks (Page 4).

2. LAnR outperforms AutoRefine and Search-R1 in terms of both Exact Match (EM) accuracy and inference efficiency, achieving superior performance compared to existing RAG systems while significantly reducing inference latency (Page 4).

3. The retrieval control head demonstrates strong predictive power, with high AUC values on multi-hop datasets such as 2Wiki (AUC = 0.8747) and HotpotQA (AUC = 0.8300) (Page 16).

4. LAnR achieves competitive retrieval recall using substantially fewer retrieved documents than specialized retrievers, matching or exceeding Recall@5–10 performance on benchmarks like HotpotQA and 2WikiMQA (Page 7).

5. The system demonstrates reduced computational cost, generating on average only ∼5 output tokens per query, which is roughly 30× fewer than Search-R1 (∼163) or AutoRefine (∼168) (Page 4).

Ablation and Scaling

1. Component ablation studies confirm that the MLP retrieval control head is the primary driver of multi-hop gains, as removing it causes large drops on multi-hop tasks (Page 9).

2. Increasing the number of [PRED] tokens encourages more abstract latent queries, but models with multiple [PRED] tokens require substantially longer training before surpassing the performance of the single-[PRED] setting on complex reasoning tasks (Page 18).

3. Scaling from Qwen-2.5-3B to 7B backbones consistently improves performance for all methods, with LAnR achieving the strongest overall multi-hop performance at both scales (Page 17).

LAnR is a unified framework that performs retrieval and reasoning directly in the latent space of a single LLM. The paper demonstrates that replacing explicit query generation with compact latent vectors and an adaptive retrieval controller allows for fewer but more effective search calls, higher cumulative gold recall, and lower computational cost (Page 5). This approach highlights the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9).

Limitations

1. A key limitation is that LAnR encodes the entire retrieval corpus using the same LLM used for generation, which is more computationally expensive than conventional retrieval approaches (Page 9).

2. The latent query vectors are not directly interpretable, making it harder to diagnose retrieval failures or explain why specific documents were retrieved (Page 9).

3. Performance is currently evaluated on English-language open-domain QA benchmarks with Wikipedia as the knowledge source, and results may not generalize to settings with significantly different document distributions or answer styles (Page 9).

LAnR is foundational research on latent retrieval-augmented generation efficiency and does not target any specific deployment. By reducing inference latency by up to 2.7× and token generation by ≈ 30× compared to existing RAG systems, LAnR lowers the computational cost of knowledge-grounded language model inference (Page 9). The findings highlight the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9). This approach demonstrates that replacing explicit query generation with compact latent vectors and an adaptive retrieval controller allows for fewer but more effective search calls, higher cumulative gold recall, and lower computational cost (Page 5). This approach highlights the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9). The system demonstrates reduced computational cost, generating on average only ∼5 output tokens per query, which is roughly 30× fewer than Search-R1 (∼163) or AutoRefine (∼168) (Page 4). This approach highlights the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9). The system demonstrates reduced computational cost, generating on average only ∼5 output tokens per query, which is roughly 30× fewer than Search-R1 (∼163) or AutoRefine (∼168) (Page 4). This approach highlights the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9).

The paper's findings highlight the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9). The system demonstrates reduced computational cost, generating on average only ∼5 output tokens per query, which is roughly 30× fewer than Search-R1 (∼163) or AutoRefine (∼168) (Page 4). This approach highlights the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9). The system demonstrates reduced computational cost, generating on average only ∼5 output tokens per query, which is roughly 30× fewer than Search-R1 (∼163) or AutoRefine (∼168) (Page 4). This approach highlights the potential of latent-space reasoning as a scalable alternative to conventional RAG pipelines (Page 9). The system demonstrates reduced computational cost, generating on average only ∼5 output tokens per query, which is roughly 30× fewer than Search-R1 (∼163) or AutoRefine (∼168) (Page 4).

1. Improve Retrieval Mechanism via Latent Query Construction: Instead of generating textual queries, LAnR produces dense retrieval vectors from the hidden states of a designated [PRED] token. This allows for embedding queries and documents within a shared representation space induced by the same parameters, eliminating the need for separate embedding models and preserving the model’s generative capability for downstream answer production.

2. Implement Adaptive Retrieval Control: "We design an implicit retrieval control mechanism, implemented as a lightweight auxiliary head, which use the LLM’s hidden representation to adaptively determines the necessity of additional retrieval without relying on explicit text-based reasoning. This enables the system to dynamically decide when sufficient evidence has been retrieved using a lightweight MLP control head over those same hidden states, leading to improved inference efficiency through reduced number of retrieval calls."

3. Enable Multi-Hop Reasoning via Latent Trajectories: Latent trajectories in hidden states can encode rich reasoning processes without explicit token generation. This allows the model to perform reasoning directly in the latent space, enabling multi-hop QA performance by allowing subsequent queries to be shaped by the [PRED] representation formed over the The Green Fog article rather than explicit intermediate text.

4. Enhance Query Abstraction via Token Expansion: Increasing the number of [PRED] tokens encourages the model to construct more abstract latent queries. This suggests that increasing token counts can improve abstraction capacity, allowing for better reasoning in complex tasks, although experiments show this is most beneficial on single-hop benchmarks.

5. Ensure Generative Integrity via Loss Masking: We apply loss masking to retrieved tokens and [PRED] tokens, so that the optimization objective is computed only over tokens generated by the LLM, excluding retrieved content from gradient updates. This prevents conflating retrieval query generation with next-token prediction over retrieved content, ensuring the model maintains its primary generative function.

6. Optimize Retrieval Efficiency via Entropy Signaling: "Observation 1 (Entropy-Sufficiency Correlation) Hk decreases monotonically in k; in particular, E[Hk P ⊆ Gk] < E[Hk P ≠ Gk]. This establishes that Hk is a reliable signal of evidence sufficiency, which can be used to inform stopping decisions, potentially allowing for inferred from hk alone, without requiring full text generation at each retrieval step."

7. Improve Scalability via Model Scaling: Scaling from 3B to 7B consistently improves performance for all methods, with the largest gains observed on multi-hop benchmarks that require compositional reasoning. This demonstrates that LAnR benefits more consistently from increased model capacity, suggesting it is a scalable approach for complex reasoning tasks.

Retrieval-Augmented Generation (RAG) has become a standard approach for enhancing large language models (LLMs) with external knowledge, mitigating hallucinations, and improving factuality. However, existing systems rely on generating natural language queries at each hop and maintaining a strict architectural separation between retriever and generator, preventing them from leveraging the full representational capacity of the LLM. We propose LAnR (Latent Abstraction for RAG), a unified framework in which a single LLM jointly performs encoding, retrieval, and generation entirely within its own latent space. Rather than generating textual queries, LAnR produces dense retrieval vectors from the hidden states of a designated[PRED] token and uses them to match against encoded document representations from the same model. Furthermore, LAnR adaptively decides when sufficient evidence has been retrieved using a lightweight MLP control head over those same hidden states, eliminating both the separate retriever and explicit token-level stopping reasoning. This design is motivated by our empirical observation that answer token entropy reliably signals retrieval sufficiency. Extensive experiments on six QA benchmarks spanning single-hop and multi-hop settings demonstrate that LAnR outperforms existing RAG methods, while achieving improved inference efficiency through reduced number of retrieval calls and tighter model integration.