[Paper Review] Atlas: Few-shot Learning with Retrieval Augmented Language Models

Gautier Izacard, Patrick Lewis+8 more|arXiv (Cornell University)|Aug 5, 2022

Topic ModelingComputer Science196 citations

TL;DR

Atlas shows that retrieval-augmented language models can excel at few-shot knowledge tasks with far fewer parameters than large LLMs, by jointly training a retriever and a seq2seq reader.

ABSTRACT

Large language models have shown impressive few-shot results on a wide range of tasks. However, when knowledge is key for such results, as is the case for tasks such as question answering and fact checking, massive parameter counts to store knowledge seem to be needed. Retrieval augmented models are known to excel at knowledge intensive tasks without the need for as many parameters, but it is unclear whether they work in few-shot settings. In this work we present Atlas, a carefully designed and pre-trained retrieval augmented language model able to learn knowledge intensive tasks with very few training examples. We perform evaluations on a wide range of tasks, including MMLU, KILT and NaturalQuestions, and study the impact of the content of the document index, showing that it can easily be updated. Notably, Atlas reaches over 42% accuracy on Natural Questions using only 64 examples, outperforming a 540B parameters model by 3% despite having 50x fewer parameters.

Motivation & Objective

• In-depth study of design and training strategies for retrieval-augmented language models in few-shot settings.
• Develop and evaluate a model named Atlas that performs knowledge-intensive tasks with low parameter counts.
• Investigate how to efficiently fine-tune both retriever and language model components for few-shot tasks.
• Assess updates and interpretability of the external knowledge index used by Atlas.

Proposed method

• Use a two-component architecture: a dense dual-encoder retriever (Contriever) and a T5-based sequence-to-sequence reader with Fusion-in-Decoder.
• Train the retriever and reader jointly using multiple retrieval-based objectives that leverage LM signals (ADist, EMDR2, PDist, LOOP).
• Pre-train with unsupervised tasks (prefix LM, masked LM, title-to-section generation) to align retriever and LM capabilities.
• Fine-tune with query-side optimization and experiment with index update strategies to manage retriever index freshness.
• Frame all tasks in a text-to-text paradigm and process retrieved documents independently in the encoder with cross-attention in the decoder.
• Explore efficient retriever fine-tuning strategies to reduce index refresh overhead during training.

Experimental results

Research questions

• RQ1 Can a retrieval-augmented model learn knowledge-intensive tasks with relatively few training examples?
• RQ2 What training signals and pre-training objectives best align the retriever with the language model for few-shot learning?
• RQ3 How do index update strategies affect training efficiency and downstream performance?

Key findings

• Atlas achieves 42.4% accuracy on Natural Questions with 64 training examples using an 11B parameter model, outperforming a 540B parameter model by about 3 points.
• With a Wikipedia-only index, Atlas reaches 45.1% on Natural Questions, and achieves 64.0% in a full-dataset setting, establishing a new state-of-the-art by 8 points.
• Atlas demonstrates strong few-shot and full-dataset performance across MMLU, TriviaQA, FEVER, and other KILT tasks, often outperforming much larger models.
• Joint pre-training of retriever and LM is crucial for few-shot performance, and multiple retriever supervision losses (ADist, EMDR2, PDist, LOOP) are effective.
• An efficient retriever fine-tuning strategy and index management techniques enable practical training with large document collections.