PaperlayerView exact revision on arXiv

Source-linked AI summary

Parametric Retrieval Augmented Generation

Weihang Su, Yichen Tang, Qingyao Ai, Junxi Yan, Changyue Wang, Hongning Wang, Ziyi Ye, Yujia Zhou, Yiqun Liu

arXiv:2501.15915v1cs.CLcs.IR

TL;DR

Existing RAG injects retrieved documents into LLM input contexts, creating computational and reasoning limitations while leaving a gap between external and parametric knowledge. Parametric RAG parameterizes documents for direct FFN integration and reports stronger, more efficient performance, including a 29% to 36% reduction in time per question versus Standard RAG.

  • Problem

    In-context knowledge injection increases context-related cost and can degrade complex reasoning, while external knowledge is not integrated into the LLM’s parameter space.

  • Method

    Parametric RAG parameterizes external documents offline and injects their representations into LLM FFN parameters through Retrieve-Update-Generate inference.

  • Results

    Parametric RAG outperforms traditional RAG methods across multiple benchmarks and reduces LLaMA3-8B inference time by 29% to 36% versus Standard RAG.

  • Takeaways & Limitations

    Parametric RAG offers a more efficient and scalable pathway for integrating external knowledge and can combine with in-context RAG for better performance.

  • Takeaways & Limitations

    Document parameterization is computationally intensive, produces larger-than-text representations, and currently ties them to specific LLMs.

Abstract

from arXiv · show

Retrieval-augmented generation (RAG) techniques have emerged as a promising solution to enhance the reliability of large language models (LLMs) by addressing issues like hallucinations, outdated knowledge, and domain adaptation. In particular, existing RAG methods append relevant documents retrieved from external corpus or databases to the input of LLMs to guide their generation process, which we refer to as the in-context knowledge injection method. While this approach is simple and often effective, it has inherent limitations. Firstly, increasing the context length and number of relevant documents can lead to higher computational overhead and degraded performance, especially in complex reasoning tasks. More importantly, in-context knowledge injection operates primarily at the input level, but LLMs store their internal knowledge in their parameters. This gap fundamentally limits the capacity of in-context methods. To this end, we introduce Parametric retrieval-augmented generation (Parametric RAG), a new RAG paradigm that integrates external knowledge directly into the parameters of feed-forward networks (FFN) of an LLM through document parameterization. This approach not only saves online computational costs by eliminating the need to inject multiple documents into the LLMs' input context, but also deepens the integration of external knowledge into the parametric knowledge space of the LLM. Experimental results demonstrate that Parametric RAG substantially enhances both the effectiveness and efficiency of knowledge augmentation in LLMs. Also, it can be combined with in-context RAG methods to achieve even better performance. We have open-sourced all the code, data, and models in the following anonymized GitHub link: https://github.com/oneal2000/PRAG

1 Introduction

Existing RAG methods inject retrieved knowledge into the input context, but longer contexts add overhead and can impair complex reasoning. Parametric RAG instead parameterizes documents and temporarily integrates them into the LLM, improving efficiency and benchmark performance.

  • Once training is complete, an LLM’s internal knowledge is effectively static, motivating RAG’s use of external corpora for adaptability.
  • In-context knowledge injection appends relevant passages or documents to the LLM’s input context.
  • Longer input contexts increase inference overhead and latency and can hurt external-knowledge use, especially in complex reasoning tasks.
  • Parametric RAG converts retrieved documents into parameter representations and integrates them into an LLM’s FFN through a Retrieve-Update-Generate workflow.
  • Parametric RAG shows superior inference efficiency and outperforms state-of-the-art in-context methods on several complex-reasoning RAG benchmarks.
  • Combining Parametric RAG with in-context RAG can produce even better benchmark performance.

2 Related Work

Prior RAG work improves retrieval timing, prompting, document selection, and knowledge organization, but retains input-context injection. Parametric RAG differs by inserting documents directly into model parameters.

  • Existing RAG extensions study retrieval timing, prompt templates, document selection, and external knowledge organization.
  • Across these frameworks, relevant passages are appended directly to the LLM’s input context for knowledge injection.
  • Parametric RAG diverges from existing RAG frameworks by injecting documents directly into the LLM’s parameters.

3 Methodology

The methodology has offline and online stages: documents are transformed into parametric representations, then retrieved representations are merged and integrated into the LLM for generation.

  • The framework first parameterizes documents offline through Document Augmentation and Parametric Document Encoding.
  • Online inference retrieves, merges, and integrates parametric representations into the LLM to generate responses.

3.1 Problem Formulation and Overview

The problem formulation contrasts traditional context injection with Parametric RAG’s parameter-level integration. Documents are mapped to mergeable representations that update the LLM before answering the query.

  • Given query q and corpus K, traditional RAG retrieves the top k documents and appends them to the LLM’s input context.
  • Parametric RAG uses offline document parameterization followed by online Retrieve-Update-Generate inference.
  • Each document d_i is transformed into a parametric representation p_i through a mapping function f_ϕ.
  • The representations can plug into FFN weights, encode document knowledge, and merge across documents.
  • During online inference, retrieved top-k representations are merged, plugged into the LLM, and used to answer the question.

3.2 Offline Document Parameterization

Offline Document Parameterization transforms each external document into a document-specific low-rank representation through augmentation and encoding. These representations can later be loaded into the LLM without appending documents to its context.

  • Document Augmentation: Each document is augmented through multiple rewrites and generated question-answer pairs to create a more comprehensive training resource.The augmented resource preserves the document’s factual content while adding diverse linguistic forms.
  • Parametric Document Encoding: The method initializes document-specific low-rank matrices for the LLM’s FFN weights while keeping the original weights fixed.Each document receives its own trainable parameter instance Δ𝜃 = {𝐴, 𝐵}, with rank much smaller than the FFN dimensions.
  • Parametric Document Encoding: Training concatenates rewritten documents, questions, and answers into sequences and applies a sequential language-modeling objective across the augmented text.The objective is designed to internalize factual details from both documents and QA pairs.
  • Design Choice: LoRA is selected because it can be merged more easily than adapters and adds less inference overhead than prefix-tuning.The paper leaves exploration of other parameter-efficient methods for future work.
  • Parametric Document Encoding: The resulting low-rank parameters form lightweight document-specific knowledge representations that can be added to the base model during inference.The preprocessing is conducted offline, so inference loads document parameters instead of appending the original document to the context.

3.3 Online Inference

Online inference follows a Retrieve-Update-Generate pipeline. Retrieved documents’ low-rank representations are merged into a single update, applied to the FFN weights, and used for response generation.

  • Retrieve: The pipeline retrieves the top-k documents with the highest relevance scores for the user query.Each retrieved document has a corresponding precomputed parametric representation.
  • Update: Low-rank matrices from the retrieved documents are merged into one plug-in module using a scalar scaling factor α.The merged update sums the document-specific low-rank contributions.
  • Update: The merged update is added to the original FFN weights, producing updated parameters that combine base-model knowledge with retrieved-document knowledge.The updated feed-forward weights are applied across Transformer layers.
  • Generate: The temporarily updated model generates the final response with standard left-to-right decoding.Generation occurs after all feed-forward layers have received the merged update.

3.4 Discussion on Time/Space Efficiency

Parametric RAG shifts document processing offline and reduces online inference cost by excluding retrieved document tokens from the LLM context. Its storage cost is determined by the low-rank FFN representations.

  • Computation Cost: The offline preprocessing cost comes primarily from Parametric Document Encoding and adds computation relative to traditional RAG.The decoder-only LLM cost is characterized as O(|𝑑|^2ℎ + |𝑑|ℎ^2).
  • Computation Cost: O(|𝑞|^2ℎ + |𝑞|ℎ^2) is the online inference complexity of Parametric RAG when loading LoRA parameters is negligible.The method’s online cost depends mainly on the input question length.
  • Computation Cost: In-context RAG instead incurs O((𝑡|𝑑| + |𝑞|)^2ℎ + (𝑡|𝑑| + |𝑞|)ℎ^2), because retrieved documents extend the input sequence.The reported savings are O(𝑡^2|𝑑|^2ℎ + 𝑡|𝑑||𝑞|ℎ + 𝑡|𝑑|ℎ^2).
  • Practical Trade-off: When a service handles many queries, Parametric RAG can provide a more carbon-efficient solution despite its additional offline preprocessing.The paper states that it is more cost-friendly when queries exceed twice the number of documents over the service lifecycle.
  • Storage Cost: A document’s parametric representation contains 2𝑛𝑟(ℎ + 𝑙) parameters, requiring approximately 2.36M extra parameters for LLaMA3-8B with 𝑟 = 2.The example uses 32 layers, hidden size 4096, and FFN intermediate size 14336.

4 Experimental Setup

The experiments evaluate Parametric RAG across four reasoning and question-answering benchmarks against several RAG baselines. They use shared retrieval and prompting conditions across models and methods.

  • Benchmarks and Metrics: The evaluation covers 2WikiMultihopQA, HotpotQA, PopQA, and ComplexWebQuestions.These datasets assess multi-hop reasoning, factual question answering, and multi-step web-based reasoning.
  • Benchmarks and Metrics: Performance is measured with F1 scores, with separate sub-task results reported for the categorized 2WikiMultihopQA and HotpotQA questions.The main table uses the first 300 questions from each sub-dataset.
  • Baselines: The comparison includes Standard RAG, DA-RAG, FLARE, DRAGIN, and P-RAG.The baselines represent direct prompt injection, data augmentation, and multi-round retrieval strategies.
  • Experimental Controls: P-RAG and all baselines use the same retriever and select the top 3 retrieved documents.They also share the same prompt template within each dataset for a fair comparison.
  • Implementation: Experiments use Qwen2.5-1.5B-Instruct, LLaMA-3.2-1B-Instruct, and Llama-3-8B-Instruct as base models.The models span different series and scales, and experiments run on NVIDIA A100 GPUs with 40GB of memory.
  • Implementation: Documents are drawn from Wikipedia dumps, and the parameterization process rewrites each document once and generates three QA pairs.LoRA is trained only in FFN matrices with learning rate 3×10^-4, one epoch, α = 32, and rank 2.
  • Implementation: BM25 with Elasticsearch serves as the retrieval backend, while generation uses official Hugging Face implementations with greedy decoding.Default model hyperparameters and chat templates are retained except for the decoding choice.

5 Experiments

Experiments show that Parametric RAG generally outperforms existing RAG frameworks, with gains from task-aware initialization, complementary document augmentation, and combined parametric plus in-context injection. It also reduces online inference time relative to Standard RAG.

  • Main Experiment: P-RAG outperforms existing RAG frameworks on most evaluated benchmarks and LLMs, especially Qwen-1.5B and LLaMA-8B.The results attribute these improvements to incorporating knowledge into model parameters for more effective recall and reasoning.
  • Main Experiment: Combining parametric and in-context RAG yields the highest overall performance across datasets and base LLMs.The result supports compatibility between the two knowledge-injection approaches.
  • LoRA Weight Initialization: Warm-up initialization consistently outperforms random initialization across model series, scales, and datasets.The paper attributes this to aligning LoRA parameters with the base LLM and incorporating task-relevant generation patterns.
  • Document Augmentation: QA generation contributes more to performance than rewriting, while using both produces the strongest overall results across most evaluated tasks and models.QA supports factual recall and application, whereas rewriting adds phrasing and structural diversity; the two components are complementary.
  • Data-augmentation Model: Performance remains consistent across different document-augmentation models, including smaller, larger, and generator-matched models.The findings indicate that document augmentation does not require a high-capacity model.
  • Runtime Analysis: 29% to 36%: P-RAG reduces time per question versus Standard RAG on 2WikiMultihopQA and ComplexWebQuestions using LLaMA3-8B.The reported comparison separates inference time from the 0.32-second merging and LoRA-loading time.

6 Conclusion and Future Directions

The paper concludes that Parametric RAG parameterizes external documents and integrates them into the model, improving efficiency and performance across LLMs. Future work must address computational and storage costs and the dependence of representations on specific LLMs.

  • Conclusion: Parametric RAG parameterizes external documents and integrates them into the model, reducing contextual overload and online computational costs while maintaining robust performance.Experiments across multiple benchmarks and LLMs report better performance than traditional RAG methods.
  • Future Directions: The parameterization process is computationally intensive, and each parametric document representation is substantially larger than plain text.The paper identifies improving computational and storage efficiency as a future research direction.
  • Future Directions: Parameterized documents are tied to specific LLMs, limiting their generalization across models.The paper proposes universal, model-agnostic representations to improve flexibility and reuse.
Loading 2501.15915v1…