Source-linked AI summary
Atlas: Few-shot Learning with Retrieval Augmented Language Models
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-Yu, Armand Joulin, Sebastian Riedel, Edouard Grave
TL;DR
Few-shot learning on knowledge-intensive tasks may require less parametric knowledge than large language models suggest, but retrieval-augmented models had not shown compelling few-shot performance. Atlas jointly pre-trains and adapts a retriever and language model, achieving strong results across benchmarks, including over 42% NaturalQuestions accuracy from 64 examples while using far fewer parameters than PaLM.
Problem
Retrieval-augmented models excel at knowledge-intensive tasks with fewer parameters, but their ability to learn such tasks in few-shot settings remained unclear.
Method
Atlas jointly pre-trains a retrieval-augmented language model and explores strategies for adapting its retriever and language model to downstream tasks.
Results
Atlas achieves strong few-shot performance across NaturalQuestions, TriviaQA, FEVER, KILT, and MMLU, including 42.4% NaturalQuestions accuracy with 64 training examples.
Takeaways & Limitations
The findings show that Atlas exhibits few-shot abilities at lower scale than standard large language models while retaining retrieval-based updateability and interpretability.
Takeaways & Limitations
MMLU questions may appear on the open web, so retrieval or web-trained models may obtain misleadingly high scores through verbatim memorisation rather than generalisation.
Abstract
from arXiv · showhide
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.
1 Introduction
Atlas is a retrieval-augmented language model designed to test whether knowledge-intensive few-shot learning can emerge with fewer parameters. It combines retrieval with extensive training and evaluation, achieving strong results across knowledge tasks while supporting index updates and memory savings.
- Motivation: Atlas investigates whether few-shot learning can decouple memorisation from generalisation by outsourcing knowledge to an external non-parametric memory.The paper studies retrieval-augmented architectures as an alternative to storing all task-relevant information in model parameters.
- Approach: Atlas jointly pre-trains a dense retriever and sequence-to-sequence language model, then explores fine-tuning strategies for adapting both components.The retriever uses a dual-encoder architecture, while the language model processes retrieved documents with the current context.
- Results: 64.0% accuracy on NaturalQuestions establishes a new state of the art in the full-dataset setting with a Wikipedia index.Atlas-11B also reaches 42.4% accuracy using 64 training examples and 45.1% with a Wikipedia-only index.
- Results: Few-shot experiments report state-of-the-art results on NaturalQuestions, TriviaQA, and FEVER, with performance on MMLU on par with or stronger than models having 15× more parameters.Reported gains include +2.8% on few-shot NaturalQuestions, +3.3% on TriviaQA, and +5.1% on FEVER.
- Analysis: Atlas experiments demonstrate that retrieval-augmented models can be updated and interpreted, while product quantisation reduces index memory by 5× with comparable performance.The paper also reports experiments on updatability and interpretability characteristics.
2 Method
Atlas frames tasks as text-to-text generation using a retriever to select documents and a Fusion-in-Decoder language model to generate outputs. It jointly trains retrieval with language-model-derived signals, while addressing the computational cost of refreshing large document indexes.
- Architecture: Atlas retrieves top-k documents for each query, then feeds them with the query to a language model that generates the output.The model combines a Contriever-based retriever with a sequence-to-sequence language model.
- Architecture: The retriever uses independently encoded query and document embeddings, scoring candidates with their dot product.Average-pooled transformer outputs provide one vector for each query or document.
- Architecture: Fusion-in-Decoder processes each document independently in the encoder, concatenates encoder outputs, and applies decoder cross-attention over the combined sequence.
- Retriever training: Retriever training uses language-model supervision, allowing adaptation from query-output pairs without document annotations.The explored objectives include attention distillation, EMDR2, perplexity distillation, and leave-one-out perplexity distillation.
- Retriever training: LOOP scores documents by how much worse language-model predictions become when each document is removed, then matches the retriever distribution using KL-divergence.LOOP conditions the language model on K−1 documents and is more expensive than PDist and EMDR.
- Efficient fine-tuning: A 37M-document index refreshed every 1000 steps adds approximately 30% overhead under the stated batch and retrieval settings.Refreshing the index is necessary because retriever updates make stored document embeddings stale.
- Efficient fine-tuning: Re-ranking 10× more documents than the language model processes yields an estimated 10% overhead, although frequent retriever updates may still require full index refreshes.
3 Related work
Prior work established retrieval for knowledge-intensive tasks and explored retrieval-augmented language models, while large pretrained models brought renewed interest in few-shot learning. Related approaches differ in retrieval mechanism, supervision, generation, and whether parameters are updated.
- Retrieval for knowledge-intensive tasks: Retrieval has improved question answering, fact checking, dialogue, and citation recommendation, using methods ranging from term matching to neural retrieval.
- Retrieval for knowledge-intensive tasks: Open-domain systems have evolved from extracting answer spans in retrieved documents to generating free-form text with sequence-to-sequence models.Retrieved context can also reduce hallucination in dialogue agents.
- Retriever training: REALM and RAG jointly train retrievers and language models by treating documents as latent variables, while REALM uses masked-language-model pretraining with an extractive BERT-style reader.
- Retrieval-augmented language models: Retrieval-augmented language models have used token caches, approximate nearest-neighbor search, and external search engines to provide additional context.
- Few-shot learning: Few-shot learning has gained renewed NLP interest because large pretrained models exhibit emergent few-shot abilities and benefit from scaling model size, data, and compute.
- Few-shot learning: Few-shot learning can use in-context instructions or parameter-updating approaches that combine textual prompts with model fine-tuning.
4 Experiments
Atlas is evaluated across few-shot and full-dataset knowledge-intensive tasks, with ablations of pre-training, index content, and retriever fine-tuning. It achieves strong results across MMLU, open-domain question answering, FEVER, and KILT, often outperforming larger or more expensive baselines.
- Experimental setup: Atlas is evaluated on KILT, MMLU, NaturalQuestions, TriviaQA, and FEVER across few-shot and full-dataset settings.The evaluation includes knowledge-intensive tasks such as question answering and fact checking, plus broad multi-domain multiple-choice understanding.
- Ablations: Jointly pre-training the retriever improves the pre-training metric, while its downstream benefit is smaller at 64 shots and nearly absent at 1024 shots.The authors interpret this pattern as evidence that pre-training primarily teaches the language model to use retrieved information.
- Ablations: Masked language modelling performs slightly better than the other two pretext tasks, so Atlas uses it for subsequent pre-training.The three pretext tasks otherwise produce similar results on NaturalQuestions, Wizard of Wikipedia, and FEVER.
- Ablations: Wikipedia-based indexes yield better downstream performance than alternative index and training-data combinations.The authors suggest adaptation to the few-shot task data or higher source quality and density as possible explanations.
- Ablations: Re-ranking matches fully updating the index while being significantly more efficient, whereas keeping the retriever fixed causes substantial performance drops.Query-encoder-only fine-tuning also produces strong results in the 64-shot setting.
- MMLU results: Atlas-11B reaches 47.1% zero-shot MMLU accuracy, exceeds GPT-3's reported 5-shot result of 43.9%, and uses 15× fewer parameters in the 5-shot comparison.It reaches 56.6% with combined 5-shot multitask training and 65.6% with auxiliary data.
- Open-domain question answering: Atlas-11B achieves state-of-the-art 64-shot results on NaturalQuestions and TriviaQA, outperforming substantially larger or more compute-intensive models.With full training data, NaturalQuestions accuracy rises from 55.9% to 60.4%; a temporally matched index adds 3.6%.
5 Analysis
Atlas’s analyses examine retrieval behavior, leakage, temporal index sensitivity, updateability, and the memory costs of dense retrieval. They show that retrieved evidence relates to MMLU accuracy, index content affects temporal performance, and compression can reduce deployment costs.
- Interpretability and Leakage: MMLU accuracy rose from 55% when the answer option was absent from retrieved passages to 77% when it appeared more than 15 times.The correct answer option appeared in the top 25 passages for 30% of test questions.
- Interpretability and Leakage: 44% of 50 correctly answered MMLU questions contained at least partially useful background information in retrieved documents.Examples included contextual clues and helpful numerical figures.
- Interpretability and Leakage: 2.8% of MMLU questions were estimated to leak into the CC-Net retrieval corpus, and filtering potentially leaked passages reduced MMLU from 56.4% to 55.8%.The authors note that detecting leakage is challenging in non-retrieval-augmented models.
- Temporal Sensitivity and Updateability: Atlas conditions strongly on its supplied temporal index, outperforming T5 on the TempLAMA-derived dataset with 57.7% versus 12.1% exact match.The evaluation uses 2017 and 2020 Wikipedia indices with answers from the corresponding years.
- Temporal Sensitivity and Updateability: The December 2018 Wikipedia dump produced the best NaturalQuestions result among the evaluated temporally different indices.NaturalQuestions contains temporal bias and questions whose answers depend on context.
- Memory Requirements: Product Quantization compressed the Wikipedia index from 49GB to 4GB and the combined index from 587GB to 50GB without serious performance degradation.The combined index could consequently fit on a single 80GB GPU.
6 Discussion
Atlas is a retrieval-augmented language model designed to learn knowledge-intensive tasks from few examples. Across broad evaluations, it achieves strong few-shot and full-dataset results while supporting updateability, interpretability, and controllability.
- Discussion: Atlas achieves strong few-shot performance across NaturalQuestions, TriviaQA, FEVER, KILT, and MMLU with jointly pre-trained retrieval and language-model components.The evaluation covers 8 KILT tasks and 57 MMLU tasks.
- Discussion: More than 42% NaturalQuestions accuracy and 84.7% TriviaQA accuracy are achieved with 64 training examples.Atlas-11B improves by almost 3 points over PaLM, a 540B-parameter model requiring 50x more pre-training compute.
- Discussion: Atlas also obtains new state-of-the-art results on NaturalQuestions, TriviaQA, FEVER, and 5 KILT tasks in full-dataset settings.The paper additionally analyzes factors important for training retrieval-augmented models.
A.1.1 Training Details
Atlas is trained and evaluated for few-shot MMLU using a masked-answer format, permutation-based debiasing, and index ablations. Hyperparameters are selected without an in-domain development set, while aggregate scores show lower variance than individual domains.
- MMLU training format: MMLU inputs list a question and four answer options, and Atlas is trained to generate a mask token followed by the correct option letter.Answer-option order is permuted during training to reduce overfitting and encourage a uniform prior over letters.
- MMLU inference: Standard MMLU inference selects the highest-softmax probability among tokens A, B, C, and D.The model returns the answer corresponding to the argmax of this four-option distribution.
- Debiased inference: 12% zero-shot and 6% 5-shot overall improvements result from full permutation debiasing, while cyclic permutations achieve nearly comparable results with 4× rather than 24× inference cost.Full debiasing requires one forward pass for each of 24 answer orderings; cyclic debiasing is used in the main paper.
- Index ablation: 3.6% is the 5-shot multitask MMLU drop without Common Crawl data, from 56.4% to 52.8%.The ablation indicates that Wikipedia data is sufficient for the task, while adding Common Crawl improves results further.
- Hyperparameter selection: Hyperparameters for the 5-shot task are selected using RACE validation data rather than an in-domain development set.The search varies batch size, language-model and retriever learning rates, retriever temperature, and training steps.
- Inter-run variance: Individual MMLU domains show variance as high as 20%, but inter-run variance is low after aggregating across all 57 domains.The reported main-paper result uses the first run, and five runs are used to assess noise and variance.
A.1.2 Full results
The full-results appendix reports domain-level MMLU scores for Atlas and closed-book T5 baselines, including a de-biased Atlas-XXL evaluation.
- Domain-level results: Tables 16 and 17 report full MMLU scores for Atlas and the closed-book T5 baseline across all 57 domains.The tables provide domain-level results for the two model families.
- Result organization: The appendix organizes the full MMLU results by model family, model size, domain, and evaluation setting.This structure supports comparison between Atlas, the closed-book baseline, and de-biased Atlas-XXL.
A.2.1 Training Details
For question answering, Atlas uses a masked-answer sequence format with retrieved passages and greedy decoding. Few-shot and full-data training use different optimization durations and index-refresh schedules.
- Question-answering format: Question-answering inputs contain a question followed by a masked answer, and the model is trained to generate the mask token followed by the answer.This format is used for question-answering training.
- Retrieval and decoding: Atlas retrieves 40 passages for both training and testing, concatenates them with the query, and truncates the result to 384 tokens.Answers are generated with greedy decoding.
- Few-shot fine-tuning: Few-shot fine-tuning uses 64 random training samples for 30 steps with query-side retriever fine-tuning and shared AdamW settings for the reader and retriever.Training uses batch size 32, learning rate 4 × 10^-5, linear decay, and five warmup iterations.
- Full-data fine-tuning: Full-dataset fine-tuning runs for 5k gradient steps, refreshes the index more frequently early in training, and selects checkpoints by validation exact match.The index is refreshed every 500 steps for the first 1,000 steps and every 2k steps thereafter.
A.2.2 Impact of scaling
Atlas is evaluated across question answering, MMLU, and KILT with task-specific retrieved inputs and few-shot training procedures. Scaling the reader improves few-shot results, while gains can saturate during full-data fine-tuning.
- Impact of scaling: Scaling the reader produces strong improvements on NaturalQuestions and TriviaQA in both few-shot learning and full fine-tuning.Table 19 reports exact-match performance as a function of reader-module parameters.
- Impact of scaling: +0.6% on NaturalQuestions and +0.5% on TriviaQA are the gains from scaling the reader from 3B to 11B parameters during full-dataset fine-tuning.These limited gains indicate saturation at larger reader sizes in the full-data setting.
- KILT evaluation: KILT experiments fine-tune Atlas separately for each dataset using a masked-answer template, 20 retrieved passages, and greedy decoding.FEVER labels are lexicalized as “true” and “false.”
- Appendix results: The appendix includes full domain-level MMLU scores for Atlas, closed-book T5, and cyclic-permutation-de-biased Atlas-XXL.It also reports question-answering scaling results and KILT development-set results in separate tables.
- KILT training: KILT few-shot fine-tuning uses 64 random samples for 30 steps with query-side retriever fine-tuning and validation-based checkpoint selection.Models are evaluated every five steps using the reported development metric.