Improving language models by retrieving from trillions of tokens

Abstract

We enhance auto-regressive language models by conditioning on document chunks retrieved from a large corpus, based on local similarity with preceding tokens. With a 2 trillion token database, our Retrieval-Enhanced Transformer (R ) obtains comparable performance to GPT-3 and Jurassic-1 on the Pile, despite using 25× fewer parameters. After fine-tuning, R performance translates to downstream knowledge-intensive tasks such as question answering. R combines a frozen B retriever, a differentiable encoder and a chunked cross-attention mechanism to predict tokens based on an order of magnitude more data than what is typically consumed during training. We typically train R from scratch, yet can also rapidly R fit pre-trained transformers with retrieval and still achieve good performance. Our work opens up new avenues for improving language models through explicit memory at unprecedented scale.

LaMDA: Language Models for Dialog Applications

“…Recent work in this direction has expanded and elaborated on neural models' ability to retrieve and rank passages [40]. The RETRO architecture demonstrates that language models can be primed with results retrieved from a database as large as two trillion tokens [41]. At a broad level, our approach is also comparable to that of Byrne et al [42], which fine-tunes the model to use external APIs for movie ticketing dialog.…”

Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval

“…Granularity of Retrieval While Khandelwal et al (2020) and Yogatama et al (2021) retrieve a token at a time step, other work retrieved a sentence (Hashimoto et al, 2018;Zhang et al, 2018;Gu et al, 2018), a prototype (Guu et al, 2018;He et al, 2020), or a chunk (Guu et al, 2020;Borgeaud et al, 2021). RETOMATON implicitly generalizes these approaches by dynamically constructing the retrieved sequence, essentially being able to retrieve individual tokens as well as constructing search-free longer passages.…”

“…In our experiments in Section 5, we reported perplexity compared to FoSS (the fraction of saved searches). The other alternative of measuring wall-clock time is difficult to reproduce, is brittle to temporary hardware and system load, and affected by the specific kNN retrieval library such as FAISS as used in Khandelwal et al (2020), ScaNN (Guo et al, 2020) as used in Borgeaud et al (2021), or SPTAG (Chen et al, 2018), etc. Further, it depends on factors that are orthogonal to our contribution, such as whether the RAM is large enough to store the datastore, and the random-access reading latency of the hard-drive.…”

“…We expect that as the datastore size increases to the scales of Borgeaud et al (2021), and as the number of neighbors retrieved increases (k neigh ) -the more pressure that will be put on the kNN search, the more of a bottleneck that it will become, and the larger relative benefit that saving kNN searches will provide to wall-clock time.…”

“…In these models, the retrieval component first searches for nearest neighbor examples in an external datastore; then, the base model references these examples during the prediction. This fusion of language models (LMs) and retrieval improves the base language model from several perspectives, including higher accuracy (Xu et al, 2021), domain adaptability (Jiang et al, 2021), and reduced size (Borgeaud et al, 2021). Further, the retrieved examples provide information regarding the provenance of the model's predictions and allow for modifying the dataset without retraining the model.…”