Atlas: Few-shot Learning with Retrieval Augmented Language Models

By 2022, large language models had demonstrated a remarkable ability: show them a handful of examples in the prompt — "few-shot learning" — and they could perform tasks they were never explicitly trained for. GPT-3 could answer trivia, classify sentiment, or translate, all from a few demonstrations.

But this ability came at a cost. To answer factual questions, the model needed to have the knowledge stored inside its parameters. This meant hundreds of billions of weights encoding everything from historical dates to scientific facts. PaLM used 540 billion parameters. The training cost was immense — thousands of TPUs for weeks — and the knowledge was frozen at training time.

Retrieval-augmented models offered an alternative: instead of memorizing, look things up. Models like RAG and RETRO showed that a smaller model with access to a document index could match much larger models on knowledge-intensive tasks. But there was a gap: nobody had shown that retrieval-augmented models could also be strong few-shot learners. Could a model that looks things up also learn new tasks from just a few examples?

Atlas is built from two components that work in concert, like a researcher and a writer collaborating on every answer.

The Retriever (Contriever) is a dual-encoder model based on BERT. Given a query, it encodes the query and every document in the index into dense vectors, then finds the top-$k$ documents by cosine similarity. Unlike keyword-based search (like BM25), dense retrieval captures semantic meaning — "What is the capital of France?" matches a document that says "Paris is where the French government sits" even without the word "capital."

The Reader (Fusion-in-Decoder / T5) is a sequence-to-sequence model. It receives the query concatenated with each retrieved document, processes them independently through the encoder, then fuses all encoded representations in the decoder via cross-attention. This Fusion-in-Decoder design scales linearly with the number of documents, unlike naive concatenation which would be quadratic.

The crucial innovation is joint training: the retriever does not operate in isolation. Its parameters are updated based on signals from the language model, so it learns to retrieve documents that actually help the reader produce correct answers.

The hardest challenge in retrieval-augmented models is the chicken-and-egg problem: the retriever needs to know what documents are useful, but "useful" depends on what the reader can do with them. And the reader's performance depends on getting good documents. If either starts poorly, the system stays stuck.

Atlas solves this by training both components together, with the retriever receiving training signals derived from the reader's performance. The paper explores four retriever training objectives, each providing a different kind of feedback.

The idea behind Perplexity Distillation is elegant: for each retrieved document $d_i$, ask the reader "how well can you predict the target if you only see this document?" The documents that lead to lower perplexity (better prediction) are the ones the retriever should rank higher. Formally, the retriever is trained to minimize the KL divergence between its relevance scores and a distribution derived from the reader's per-document perplexity.

A naive approach to using retrieved documents would be to concatenate them all with the query and feed them to the model as one giant input. But self-attention has quadratic complexity, so doubling the input length quadruples the cost. With 20 retrieved passages, this becomes prohibitively expensive.

Fusion-in-Decoder (FiD) solves this elegantly. Each retrieved document is concatenated with the query and processed independently through the T5 encoder. The encoder produces a set of hidden representations for each document. These representations are then simply concatenated into one long sequence and passed to the decoder. The decoder uses cross-attention to attend over all document representations simultaneously.

The key insight: the encoder's self-attention — the expensive operation — processes each document separately (linear cost in number of documents). Only the decoder's cross-attention sees all documents together, and cross-attention is much cheaper than full self-attention over a concatenated input.

Atlas is not just fine-tuned for retrieval — it is pre-trained with retrieval from the start. The retriever is initialized from the unsupervised Contriever model, and the reader from T5-lm-adapt (a T5 variant trained only on unlabeled text). Both are then jointly pre-trained on a mixture of Wikipedia and Common Crawl data.

The pre-training task is masked language modeling: take a passage, mask some tokens, retrieve related documents, and predict the masked tokens using the retrieved context. This teaches the system to seek and use external knowledge from the very beginning of training.

A critical engineering detail: during pre-training, the document index must be periodically refreshed. As the retriever's weights change, the document embeddings become stale. Atlas re-encodes and re-indexes documents every 1,000 steps — a computationally expensive but necessary step that keeps the index aligned with the evolving retriever.

The paper's headline result: Atlas-11B, with only 64 training examples, reaches 42.4% accuracy on NaturalQuestions — beating PaLM-540B (39.6%), a model with 50× more parameters that required 50× more pre-training compute.

On TriviaQA, Atlas-11B reaches 84.7% with 64 examples. On MMLU (57 diverse tasks), it achieves strong results across domains, especially those that benefit from factual retrieval. On the KILT benchmark (8 knowledge-intensive tasks), Atlas sets new state-of-the-art with full fine-tuning and remains competitive even in the few-shot regime.

During fine-tuning, Atlas uses a practical optimization: query-side fine-tuning. Instead of re-encoding the entire document index (millions of documents) after each retriever update, only the query encoder is updated. The document encoder stays frozen and the index is re-ranked using the updated query embeddings. This makes fine-tuning feasible on a single machine.

One of Atlas's most practical advantages is updateability. Because knowledge lives in the document index rather than the model weights, you can update what the model "knows" by simply replacing the index — no retraining needed.

The authors demonstrated this with a temporal experiment on NaturalQuestions: they took questions whose answers changed over time (e.g., "Who is the president of X?") and showed that swapping from a 2017 Wikipedia index to a 2021 index updated the model's answers correctly, without any fine-tuning. A parameter-only model would need expensive retraining.

They also showed that the choice of index matters enormously. Using a Wikipedia-only index gives much better performance on knowledge-intensive tasks than a Common Crawl index, likely because Wikipedia is more curated and dense with factual information.

The paper includes extensive ablation studies that reveal what makes Atlas work.

Retriever training objective: All four objectives (ADist, EMDR², PDist, LOOP) perform similarly, but PDist is the most stable and computationally efficient, making it the recommended default.

Pre-training task: Three pretext tasks were compared — prefix language modeling, masked language modeling, and title-to-section generation. Masked language modeling gives a slight edge and is adopted as the default.

Index content: Using Wikipedia as the index consistently outperforms Common Crawl for downstream tasks, even when Common Crawl was used during pre-training. This suggests that index quality matters more than index size.

Model scale: Atlas benefits from scale, but the gains from retrieval augmentation are larger than the gains from simply increasing parameters. Atlas-3B with retrieval outperforms closed-book T5-11B without retrieval on many tasks.

Atlas's deepest insight is not about any particular architecture or training trick — it is about where knowledge should live. Large language models had been treating parameters as the universal storage medium: more knowledge meant more parameters. Atlas demonstrated that a well-designed retrieval system separates the skill of reasoning from the content being reasoned about, and that this separation leads to better performance at lower cost.

The principle endures: today's most capable RAG systems — from search-augmented chatbots to enterprise question-answering pipelines — trace their lineage through Atlas's demonstration that joint training, not just bolt-on retrieval, is what unlocks few-shot knowledge-intensive performance.