Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection (2310.11511v1)

Abstract: Despite their remarkable capabilities, LLMs often produce responses containing factual inaccuracies due to their sole reliance on the parametric knowledge they encapsulate. Retrieval-Augmented Generation (RAG), an ad hoc approach that augments LMs with retrieval of relevant knowledge, decreases such issues. However, indiscriminately retrieving and incorporating a fixed number of retrieved passages, regardless of whether retrieval is necessary, or passages are relevant, diminishes LM versatility or can lead to unhelpful response generation. We introduce a new framework called Self-Reflective Retrieval-Augmented Generation (Self-RAG) that enhances an LM's quality and factuality through retrieval and self-reflection. Our framework trains a single arbitrary LM that adaptively retrieves passages on-demand, and generates and reflects on retrieved passages and its own generations using special tokens, called reflection tokens. Generating reflection tokens makes the LM controllable during the inference phase, enabling it to tailor its behavior to diverse task requirements. Experiments show that Self-RAG (7B and 13B parameters) significantly outperforms state-of-the-art LLMs and retrieval-augmented models on a diverse set of tasks. Specifically, Self-RAG outperforms ChatGPT and retrieval-augmented Llama2-chat on Open-domain QA, reasoning and fact verification tasks, and it shows significant gains in improving factuality and citation accuracy for long-form generations relative to these models.

This paper introduces Self-Reflective Retrieval-Augmented Generation (Self-RAG), a framework designed to enhance the factual accuracy and overall quality of LLM generations without sacrificing their versatility. The core problem addressed is that standard LLMs often produce factual errors, and traditional Retrieval-Augmented Generation (RAG) methods, while helpful, retrieve information indiscriminately, which can hinder performance on tasks not requiring factual grounding or lead to outputs inconsistent with retrieved evidence.

Self-RAG trains an LLM to adaptively retrieve relevant text passages on-demand and to self-reflect on its own generated output using special "reflection tokens". These tokens are integrated into the generation process and fall into two categories:

1. Retrieval Tokens ([Retrieval]): These tokens signal whether retrieving external information would be beneficial for generating the next segment of text. Values include Yes, No, or Continue (to reuse previously retrieved evidence).
2. Critique Tokens ([Critique]): These tokens evaluate the quality of the generation process. There are three types:
   • Is Relevant ([IsRel]): Assesses if a retrieved passage is relevant to the query (Relevant, Irrelevant).
   • Is Supported ([IsSup]): Evaluates if the generated text segment is fully supported, partially supported, or not supported by the retrieved passage (Fully supported, Partially supported, No support).
   • Is Useful ([IsUse]): Judges the overall usefulness or quality of the generated response segment on a scale (e.g., 1-5).

Inference Process:

1. Given an input prompt and preceding text, the Self-RAG model first predicts a [Retrieval] token.
2. If [Retrieval]=No, it generates the next text segment like a standard LLM.
3. If [Retrieval]=Yes, it calls a retriever (R) to fetch relevant passages ($D$).
4. It then processes multiple passages ($d \in D$) in parallel:
   • Predicts the relevance ([IsRel]) of each passage $d$.
   • Generates a candidate output segment $y_t$ based on $d$.
   • Predicts the support level ([IsSup]) of $y_t$ given $d$.
   • Predicts the overall usefulness ([IsUse]) of $y_t$.
5. A segment-level beam search ranks the generated candidates ($y_t$) based on a weighted score combining the probabilities of the desired critique tokens (Relevant, Fully supported, 5 usefulness, etc.). This allows selecting the best-supported, most relevant, and useful continuation.
6. The process repeats for subsequent segments.

Training Process:

Self-RAG involves training two main components: a Critic model and the final Generator model (M).

1. Critic Training:
   • A dataset is created by prompting a powerful LLM (like GPT-4) with specific instructions to generate the desired reflection tokens for various inputs, outputs, and retrieved passages.
   • A smaller LM (e.g., Llama2-7B) is fine-tuned on this dataset to act as the Critic model, learning to predict appropriate reflection tokens.
2. Generator Training:
   • The trained Critic model and a retriever (R) are used offline to augment a diverse instruction-following dataset. For each instance, the Critic inserts retrieval and critique tokens, along with relevant passages where needed, into the target output sequence.
   • The Generator model (M) (e.g., Llama2-7B/13B) is then trained on this augmented corpus using a standard next-token prediction objective. The vocabulary is expanded to include the reflection tokens. The loss is masked for the actual retrieved text content. This teaches the Generator model to generate both the task output and the reflection tokens itself, eliminating the need for the separate Critic model during inference.

Key Features and Benefits:

• Adaptive Retrieval: Retrieves information only when deemed necessary, preserving the LLM's abilities on tasks not requiring external knowledge.
• Self-Correction/Critique: Explicitly evaluates relevance, factual grounding (support), and usefulness during generation.
• Controllability: The inference process can be customized by adjusting weights for different critique aspects (e.g., prioritizing factuality vs. fluency) or setting thresholds for retrieval frequency, without retraining.
• Improved Factuality and Citation: Generates outputs more faithful to retrieved evidence and provides better attribution through the [IsSup] token.
• Efficiency: Training involves standard LM objectives after offline data augmentation, avoiding the complexities and costs of online reinforcement learning (like RLHF/PPO).

Experiments and Results:

Self-RAG (using Llama2 7B and 13B) was evaluated on various tasks, including open-domain QA (PopQA, TriviaQA), closed-set QA/reasoning (PubHealth, ARC-Challenge), and long-form generation with citation (Biography generation with FactScore, ALCE-ASQA).

• Self-RAG significantly outperformed standard Llama2 and Alpaca models, as well as conventional RAG approaches applied to these models.
• It outperformed ChatGPT and retrieval-augmented Llama2-chat on several tasks, particularly in factuality and citation accuracy on long-form generation.
• Ablation studies confirmed the benefits of adaptive retrieval, critique tokens, and the segment-level beam search guided by critiques.
• The framework demonstrated effective test-time customization by adjusting critique weights and retrieval thresholds.

In conclusion, Self-RAG presents a novel method for training LLMs to leverage retrieval more effectively and reflect on their outputs, leading to improved factuality, quality, and controllability compared to existing approaches.