PaperlayerView exact revision on arXiv

Source-linked AI summary

RA-DIT: Retrieval-Augmented Dual Instruction Tuning

Xi Victoria Lin, Xilun Chen, Mingda Chen, Weijia Shi, Maria Lomeli, Rich James, Pedro Rodriguez, Jacob Kahn, Gergely Szilvasy, Mike Lewis, Luke Zettlemoyer, Scott Yih

arXiv:2310.01352v4cs.CLcs.AI

TL;DR

RALMs provide external knowledge but are difficult to build because existing methods require costly retrieval-aware pre-training or suboptimal post-hoc integration. RA-DIT addresses this with two lightweight instruction-tuning stages that separately improve the LLM and retriever. RA-DIT 65B achieves state-of-the-art results and outperforms in-context RALMs by +8.9% in 0-shot and +1.4% in 5-shot settings on average.

  • Problem

    RALMs need external retrieval for long-tail and up-to-date knowledge, but existing approaches require costly retrieval-specific pre-training changes or suboptimal post-hoc integration.

  • Method

    RA-DIT retrofits pre-trained LLMs through separate instruction-tuning stages that improve retrieved-information use and train the retriever according to LLM preferences.

  • Results

    +8.9% in 0-shot and +1.4% in 5-shot average performance over MMLU, NQ, TQA, and ELI5 versus REPLUG, with state-of-the-art benchmark results from RA-DIT 65B.

  • Takeaways & Limitations

    Each fine-tuning stage improves performance, and combining the tuned LLM and retriever produces additional gains while competing effectively with extensively pretrained RALMs.

  • Takeaways & Limitations

    Direct comparison with InstructRetro 48B is challenging because the models use different base LMs, fine-tuning datasets, and inference settings.

Abstract

from arXiv · show

Retrieval-augmented language models (RALMs) improve performance by accessing long-tail and up-to-date knowledge from external data stores, but are challenging to build. Existing approaches require either expensive retrieval-specific modifications to LM pre-training or use post-hoc integration of the data store that leads to suboptimal performance. We introduce Retrieval-Augmented Dual Instruction Tuning (RA-DIT), a lightweight fine-tuning methodology that provides a third option by retrofitting any LLM with retrieval capabilities. Our approach operates in two distinct fine-tuning steps: (1) one updates a pre-trained LM to better use retrieved information, while (2) the other updates the retriever to return more relevant results, as preferred by the LM. By fine-tuning over tasks that require both knowledge utilization and contextual awareness, we demonstrate that each stage yields significant performance improvements, and using both leads to additional gains. Our best model, RA-DIT 65B, achieves state-of-the-art performance across a range of knowledge-intensive zero- and few-shot learning benchmarks, significantly outperforming existing in-context RALM approaches by up to +8.9% in 0-shot setting and +1.4% in 5-shot setting on average.

1 INTRODUCTION

RALMs address LLMs’ difficulty with long-tail and up-to-date knowledge by coupling language models with external retrieval, but existing designs require costly pre-training changes or suboptimal post-hoc integration. RA-DIT instead retrofits LLMs through lightweight dual instruction tuning of the language model and retriever.

  • Motivation: RALMs supplement parameterized LLM knowledge with external retrieval to improve access to long-tail and up-to-date information.LLMs struggle to capture long-tail knowledge and require substantial resources to remain current.
  • Existing approaches: Existing RALM approaches separately target better use of retrieved knowledge and more relevant retrieval, using training strategies that range from pre-training integration to later-stage adaptation.The cited related-work passage identifies both core challenges and differing points in the model-training process where retrieval is introduced.
  • RA-DIT: RA-DIT retrofits any LLM with retrieval capabilities through lightweight instruction tuning over tasks requiring knowledge utilization and contextual awareness.The framework is initialized with pre-trained LLAMA and the DRAGON+ dense retriever, with retrieved chunks prepended to prompts and ensembled at inference.
  • RA-DIT: The method performs separate LM-ft and R-ft stages: the LLM learns to use retrieved information and ignore distractions, while the retriever learns from LLM preferences.LM-ft uses retrieval-augmented prompts and label loss; R-ft updates the query encoder with an LM-supervised retrieval objective.

2 METHOD

RA-DIT fine-tunes the language model and retriever separately so the model uses retrieved information effectively and the retriever returns content that improves model predictions.

  • Architecture: At inference, the dual-encoder retriever selects top-k chunks, separate augmented prompts are evaluated in parallel, and output probabilities are mixed using normalized relevance scores.The query and document embeddings are combined with a dot product to score candidate chunks.
  • Language Model Fine-Tuning: The LM fine-tuning data span dialogue, open-domain QA, reading comprehension, summarization, and chain-of-thought reasoning across 20 datasets.Reading-comprehension training also teaches the model to answer “I don’t know” when the retrieved passage is incorrect.
  • Language Model Fine-Tuning: During LM fine-tuning, retrieved chunks are prepended as background fields, and the model is trained with next-token loss on the output segment.Each retrieved chunk creates a separate fine-tuning instance, mirroring inference-time augmentation.
  • Language Model Fine-Tuning: In-context retrieval fine-tuning helps the LLM use relevant background knowledge and ignore misleading retrieved content by relying on parametric knowledge when retrieval is wrong.This twofold benefit is described as adapting knowledge use while improving robustness to inaccurate retrieval.
  • Overview: RA-DIT uses two instruction-tuning stages: one adapts the language model to retrieved information, while the other aligns the retriever with the language model.The selected tasks cultivate knowledge utilization and contextual awareness in language-model predictions.
  • Retriever Fine-Tuning: Retriever fine-tuning uses LM-supervised retrieval to assign higher scores to chunks that improve the LLM’s likelihood of generating the correct answer.The method minimizes KL-divergence between LM-supervised retrieval scores and retriever scores, updating only the query encoder in practice.

3 EXPERIMENT SETUP

The experiments evaluate RA-DIT against base and retrieval-augmented baselines on held-out knowledge-intensive tasks using zero-shot and few-shot settings.

  • Retrieval Corpus: The retrieval corpus combines Wikipedia and CommonCrawl text chunks into a data store containing 399M text chunks.The Wikipedia source is the Dec. 20, 2021 dump, supplemented with CommonCrawl dumps from 2017–2020.
  • Baselines: The main comparisons use base LLAMA models, REPLUG instantiated with LLAMA and DRAGON+, and ATLAS in a 64-shot fine-tuning setting.RA-DIT is compared with LLAMA and REPLUG in zero-shot and in-context few-shot settings.
  • Evaluation Tasks: Evaluation focuses on knowledge-intensive tasks excluded from fine-tuning, including MMLU, Natural Questions, TriviaQA, and a KILT subset.The KILT development split excluding ELI5 is used to select fine-tuning hyperparameters.

4 MAIN RESULTS

RA-DIT improves knowledge-intensive performance over prior retrieval-augmented and base-model approaches, while preserving commonsense reasoning capabilities without retrieval augmentation.

  • Knowledge-Intensive Tasks: +8.9% in 0-shot and +1.4% in 5-shot settings on average over MMLU, NQ, TQA and ELI5 versus REPLUG.RA-DIT achieves the best performance on most datasets.
  • Knowledge-Intensive Tasks: REPLUG performs better than base LLAMA 65B on knowledge-intensive tasks, confirming benefits from retrieval augmentation.RA-DIT then significantly outperforms REPLUG in both reported shot settings.
  • Knowledge-Intensive Tasks: RA-DIT outperforms ATLAS by an average of 4.1 points and achieves higher performance on 6 out of 8 datasets.The comparison uses one RA-DIT model across combined 64-shot examples, whereas ATLAS reports task-specific models.
  • Commonsense Reasoning: RA-DIT 65B improves over base LLAMA models on 7 out of 8 commonsense reasoning datasets without retrieval augmentation.This indicates that the LLM’s parametric knowledge and reasoning capabilities are generally preserved.

5 ANALYSIS

The analysis finds that language-model and retriever fine-tuning each improve retrieval-augmented performance, while their combination provides the largest gain. Retriever data mixture and encoder choices also materially affect results.

  • Evaluation caveat: 0-shot comparison caveat: LLAMA 65B scores lower on NQ and TQA than reported by Touvron et al. because this analysis uses exact match.The authors believe the prior work counted responses containing the ground-truth answer string.
  • Language model fine-tuning strategies: In-context retrieval augmentation produces substantial gains in both zero- and five-shot settings, even with only the top-1 retrieved chunk.Instruction tuning alone often gives marginal or negative five-shot changes, except on HotpotQA.
  • Retriever fine-tuning strategies: 0.6 points: “95% corpus data + 5% MTI data” achieves the best retriever accuracy across all models.Corpus-only fine-tuning improves over the base DRAGON+ model by 0.4 points on average, while MTI-only improves by 0.1 points.
  • Retriever fine-tuning strategies: Freezing the document encoder performs significantly better than jointly fine-tuning query and document encoders.RA-DIT therefore fine-tunes only the query encoder.
  • Dual instruction tuning ablation: Both LM-ft and R-ft improve over the REPLUG baseline, while combining them yields the largest gain: +0.8 points on average.The combined RA-DIT model outperforms REPLUG more than either fine-tuning stage used alone.
  • Retriever settings: DRAGON+ significantly outperforms Contriever and Contriever-MSMARCO in the REPLUG setting.All retrieval-augmented models substantially improve over the LLAMA baseline.

6 RELATED WORK

Related work combines language models with non-parametric memory through several architectural and training choices. RA-DIT builds on retrieval-augmented instruction tuning while targeting lightweight retrofitting of existing LLMs.

  • Retrieval-augmented language models: RALMs augment language models with non-parametric memory to provide external knowledge access and provenance.This addresses the limits of knowledge stored only in model parameters.
  • Retrieval-augmented language models: Prior RALM systems fuse retrieved content through separate encoders and cross-attention or by directly augmenting the LM input.RETRO and FiD use separate encoder modules, whereas other approaches prepend retrieved content to the input.
  • Training strategies: Retrieval capabilities have been introduced through end-to-end pre-training, intermediate training, or post-hoc integration of retrieval components.REALM and RETRO use end-to-end pre-training, while other approaches add retrieval later in training or inference.
  • Instruction tuning: Instruction tuning aligns pretrained LLMs to follow natural-language instructions and reduces reliance on extensive prompt engineering.RA-IT applies instruction tuning specifically to improve use of retrieved information.
  • Information retrieval: Information retrieval methods include sparse, dense, and multi-vector retrievers with different representation and search mechanisms.Dense retrievers embed queries and documents into fixed-size vectors for nearest-neighbor search.

7 CONCLUSION

The conclusion presents RA-DIT as a lightweight framework for retrofitting pretrained LLMs with retrieval capabilities. It reports strong zero- and few-shot benchmark performance while noting that direct comparisons with some models are constrained.

  • Conclusion: RA-DIT retrofits pretrained LLMs with retrieval through dual instruction tuning rather than extensive retrieval-specific pre-training.It improves both the LM’s use of retrieved knowledge and the retriever’s ability to return helpful texts.
  • Conclusion: The framework fine-tunes the retriever using supervision from the LM to retrieve texts that better support correct generation.This complements retrieval-augmented instruction tuning of the language model.
  • Conclusion: RA-DIT achieves state-of-the-art zero- and few-shot results on knowledge-intensive benchmarks and competes effectively with ATLAS.It surpasses un-tuned in-context RALM approaches such as REPLUG.
  • Evaluation scope: Direct comparison with InstructRetro 48B is challenging because the models differ in base LMs, fine-tuning datasets, and inference settings.Despite this, RA-DIT 65B compares favorably on shared zero-shot evaluation datasets.

B IMPLEMENTATION DETAILS

The implementation uses diverse instruction-tuning tasks, retrieved context during LM fine-tuning, and mixed corpus and multitask data for retriever fine-tuning. Inference ensembles predictions from multiple retrieved chunks.

  • Fine-tuning dataset selection: The LM fine-tuning data cover five categories: dialogue, open-domain QA, reading comprehension, summarization, and chain-of-thought reasoning.The collection contains 20 datasets designed to improve knowledge utilization and contextual awareness.
  • Retrieval-augmented LM fine-tuning: LM fine-tuning prepends retrieved background text to instructions so the model learns to use relevant information and ignore distractions.Ground-truth context is used for tasks such as reading comprehension and summarization.
  • Retriever fine-tuning: Retriever fine-tuning combines self-supervised corpus data with multitask instruction data.The corpus component uses 900k text chunks, while MTI data contain 286k training examples across QA and dialogue tasks.
  • Retriever fine-tuning: The retriever is trained for one epoch on 95% corpus data and 5% MTI data using KL-divergence loss over the top-10 retrieved chunks.The query encoder is fine-tuned with DRAGON+ initialization.
  • Inference: Inference normally retrieves the top-10 chunks and ensembles predictions weighted by chunk relevance scores.Multi-choice and generation tasks use weighted probabilities across augmented prompts.

C FINE-TUNING DATASET TEMPALTES

The fine-tuning templates randomize field markers to reduce overfitting while preserving explicit instruction and answer boundaries.

  • Field markers are randomized during training to avoid overfitting.Instruction starts use randomly sampled markers, while answer starts are sampled from “A:” or “Answer:”.
  • Each serialized example marks instruction and answer boundaries with dedicated fields.The instruction end marker is set to a newline.

D EVALUATION DATASETS AND TEMPLATES

The evaluation uses standardized datasets, prompts, retrieval settings, and scaling comparisons to assess retrieval-augmented instruction tuning. Results show broad retrieval benefits, with task-dependent effects from model size and augmentation format.

  • D EVALUATION DATASETS AND TEMPLATES: Dev evaluation uses up to 2500 randomly sampled examples from official development sets to reduce computational cost.Test evaluation uses the full set for fair comparison with previous work.
  • D EVALUATION DATASETS AND TEMPLATES: The evaluation includes standardized language-model prompts, retriever queries, and few-shot examples drawn from official training splits or released 64-shot sets.Commonsense reasoning evaluations do not use retrieval.
  • E.1 SCALING LAWS OF RETRIEVAL AUGMENTED LANGUAGE MODEL FINE-TUNING: The scaling experiment combines fine-tuned models of 7B, 13B, and 65B with the base DRAGON+ retriever, distinguishing 0-shot and 5-shot performance.Figure 2 uses dashed lines for 0-shot results and solid lines for 5-shot results.
  • E.1 SCALING LAWS OF RETRIEVAL AUGMENTED LANGUAGE MODEL FINE-TUNING: Retrieval augmentation substantially benefits all model sizes, with smaller models often receiving larger improvements.A 7B model using more than one retrieved chunk surpasses a vanilla 65B model on several tasks.
  • E.1 SCALING LAWS OF RETRIEVAL AUGMENTED LANGUAGE MODEL FINE-TUNING: One-hop fact lookup tasks gain across model sizes, whereas larger language models retain a prominent advantage on complex tasks.The contrast is illustrated with Zero-Shot RE and T-REx versus HotpotQA and Wizard of Wikipedia.
  • E.2 COMPARE PARALLEL RETRIEVAL AUGMENTATION TO CHUNK CONCATENATION: Parallel retrieval augmentation and chunk concatenation perform closely on average, with chunk concatenation showing a small benefit.The approaches differ more on CoNLL-YAGO and T-REx: concatenation performs much better on the former but worse on the latter.
  • E.2 COMPARE PARALLEL RETRIEVAL AUGMENTATION TO CHUNK CONCATENATION: Parallel retrieval augmentation may benefit RA-IT 65B because that model was fine-tuned with the same configuration used at inference.The effect of fine-tuning with chunk concatenation is left for future studies.

E.3 RETRIEVAL CORPORA ABLATION

The retrieval-corpus ablation compares retriever settings and corpus subsets, including CommonCrawl and multiple Wikipedia snapshots, in a 5-shot REPLUG evaluation.

  • Retriever settings are evaluated with LLAMA 65B in the REPLUG setting using 5-shot development performance.The comparison varies the retriever while holding the language model fixed.
  • The corpus ablation compares the 399M retrieval corpus with CommonCrawl-only, Wikipedia-only, and Wiki 2018 alternatives.Wikipedia-only variants are evaluated with and without infoboxes.
  • Wikipedia-only retrieval benefits several KILT tasks because Wikipedia was the intended corpus for those tasks.AIDA and zsRE are cited as examples.
  • The section presents selected task prompts, retrieved passages, and predictions from retrieval-augmented and conventionally instruction-tuned LLAMA 65B models.The examples compare RA-IT 65B with IT 65B.

F.1 HOTPOTQA

HotpotQA examples illustrate how RA-IT 65B handles noisy retrieval in zero-shot settings compared with conventionally instruction-tuned LLAMA 65B. The analysis focuses on distractor resistance and answer confidence.

  • In zero-shot HotpotQA, RA-IT 65B outperforms IT 65B by a large margin on the development set.Table 15 shows examples where RA-IT 65B answers correctly while IT 65B answers incorrectly.
  • The dense retriever often returns unhelpful chunks for HotpotQA’s multi-hop questions, creating distractors for the language model.Many retrieved chunks contain no information useful for prediction.
  • RA-IT 65B can ignore noisy retrieved passages and answer from parametric knowledge, while IT 65B is more easily misled.IT 65B also predicts “I don’t know” more frequently in these examples.
  • The examples ensemble 10 retrieved text chunks and display the top three chunks with corresponding predictions.The comparison is between RA-IT 65B and IT 65B on HotpotQA development examples.
Loading 2310.01352v4…