Atlas: Few-shot Learning with Retrieval Augmented Language Models
Research code from Meta that jointly pretrains a retrieval-augmented sequence-to-sequence language model, combining a dense passage retriever with an encoder-decoder model for few-shot question answering and other tasks.
Atlas is a research codebase for pretraining, fine-tuning, retrieving and evaluating a retrieval-augmented sequence-to-sequence language model. It jointly pretrains a passage-based dense retriever together with an encoder-decoder language model, and includes pretrained models, corpora and indices used in the accompanying paper.
The code supports fusion-in-decoder sequence-to-sequence models up to 11B parameters, distilling relevance signals from those models into dense retrievers, and end-to-end retrieval-augmented training over a user-supplied corpus of passages with retrieval integrated into the training loop. It supports masked-language modeling, prefix-language modeling, Wikipedia section generation, open-domain question answering, multiple-choice question answering, fact checking and KILT tasks.
It also includes a parallel, distributed GPU-based exact and approximate maximum inner product search for dense vector retrieval, in-place index refreshes, and memory optimizations for training retrievers in the loop. The repository is no longer actively maintained.