Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection

Standard RAG retrieves a fixed number of passages for every input, whether or not retrieval helps or the passages are relevant, which can hurt versatility and quality; and the output isn’t guaranteed to agree with the retrieved passages, since models aren’t trained to follow them.

A single language model trained to decide when to retrieve, to generate from retrieved passages, and to critique the passages and its own output by emitting special reflection tokens, which also make its behaviour controllable at inference time.

Retrieve (should I retrieve now: yes / no / continue); IsRel (is this passage relevant to the input); IsSup (is the output segment fully, partially or not supported by the passage); IsUse (overall usefulness of the response, rated 1–5). The last three are critique tokens.

Segment by segment (a sentence in the experiments): it first predicts Retrieve; if no, it continues like an ordinary LM. If yes, it retrieves several passages, processes each in parallel (predicting IsRel, a candidate segment and IsSup), then ranks the candidates by their critique tokens and finally predicts IsUse.

GPT-4 was prompted with type-specific instructions and few-shot examples to label sampled input–output pairs with reflection tokens (4K–20K examples per type, agreeing well with humans); a Llama 2-7B critic was then trained on these labels and agreed with GPT-4 over 90% of the time on most token types.

The critic and a retriever annotate training outputs offline: retrieval tokens, retrieved passages and critique tokens are inserted where they apply. The generator is trained with ordinary next-token prediction on this augmented corpus (retrieved text masked from the loss, vocabulary extended with reflection tokens), so no critic is needed at inference.

Critique is computed once, offline, and inserted into the training text; the generator then learns with a standard language-modelling objective instead of PPO, which needs separate reward models during training.

With a threshold: retrieval is triggered when the probability of Retrieve = yes exceeds it. Factual tasks can use a low threshold (retrieve often); open-ended tasks like personal essays can retrieve rarely.

Segment-level beam search scoring each candidate by its generation probability plus a weighted sum of the normalized probabilities of the desirable critique tokens (e.g. IsRel = relevant, IsSup = fully supported); the weights can be changed at test time, or undesirable critiques used as hard filters.

150K instruction–output pairs (from Open-Instruct and knowledge-intensive datasets), with Llama 2 7B and 13B as generators, Llama 2 7B as critic, and Contriever-MS MARCO retrieving up to ten documents.

Across six tasks (fact verification, multiple-choice reasoning, open-domain QA, biographies, long-form QA with citations) it beat other open models with and without retrieval, and outperformed ChatGPT on several (e.g. PubHealth, PopQA, biographies), even surpassing it in citation precision on ASQA.

On tasks where the answer can’t simply be copied from a passage, such as PubHealth fact checking and ARC-Challenge reasoning, adding retrieved documents barely improved them, and most such baselines also struggled to cite accurately.

Training without retrieved passages or without the critic’s reflection tokens lost much performance; always using the top retrieved passage regardless of relevance (as in standard RAG) hurt PopQA and ASQA; dropping the IsSup score from beam search hurt ASQA.

Raising its weight increased citation precision (claims backed by evidence) but lowered fluency (MAUVE), since longer, more fluent answers contain more unsupported claims; Self-RAG lets users pick the balance without retraining.