Generalization through Memorization: Nearest Neighbor Language Models

By Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke S. Zettlemoyer, Michael Lewis. 2019

1,095 citations recorded by the source corpus.

We introduce $k$NN-LMs, which extend a pre-trained neural language model (LM) by linearly interpolating it with a $k$-nearest neighbors ($k$NN) model. The nearest neighbors are computed according to distance in the pre-trained LM embedding space, and can be drawn from any text collection, including the original LM training data. Applying this augmentation to a strong Wikitext-103 LM, with neighbors drawn from the original training set, our $k$NN-LM achieves a new state-of-the-art perplexity of 15.79 - a 2.9 point improvement with no additional training. We also show that this approach has implications for efficiently scaling up to larger training sets and allows for effective domain adaptation, by simply varying the nearest neighbor datastore, again without further training. Qualitatively, the model is particularly helpful in predicting rare patterns, such as factual knowledge. Together, these results strongly suggest that learning similarity between sequences of text is easier than predicting the next word, and that nearest neighbor search is an effective approach for language modeling in the long tail.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Information Retrieval & Search, Large Language Models

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Information Retrieval & Search

In Large Language Models

Citations in the corpus

Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks (2020, background): “Such tasks include open-domain question answering [5, 25], fact checking [50], fact completion [42], long-form question answering [12], Wikipedia article generation [32], dialogue [36, 59, 9, 13], translation [16], and language modeling [17, 23].”
REALM: Retrieval-Augmented Language Model Pre-Training (2020, methodology): “In the language modeling literature, the k-Nearest Neighbor Language Model (Khandelwal et al., 2019) (kNNLM) retrieves similar LM examples to improve memorization.”
Few-shot Learning with Retrieval Augmented Language Models (2022, methodology): “The related kNN-LM model (Khandelwal et al., 2020) replaced LSTMs by transformer networks, and scaled the memory to billions of tokens, leading to strong performance improvements.”
LaMDA: Language Models for Dialog Applications (2022, methodology): “Given these differences, we note that the range of existing approaches to this problem include the RNNLM [34], RAG [35], REALM [36], and FiD [37] architectures.”
Improving language models by retrieving from trillions of tokens (2021, methodology): “To validate our approach in a controlled setting, we compare ourmethodwith kNN-LM (Khandelwal et al., 2020) on the Wikitext103 dataset in Table 4.”
Augmented Language Models: a Survey (2023, methodology): “Various works augment LMs with a dense retriever by appending the retrieved documents to the current context (Chen et al., 2017; Clark and Gardner, 2017; Lee et al., 2019; Guu et al., 2020; Khandelwal et al., 2020; Lewis et al., 2020; Izacard and Grave, 2020; Zhong et al., 2022; Borgeaud et al., 202”
Language Generation Models Can Cause Harm: So What Can We Do About It? An Actionable Survey (2022, background): “Here we highlight papers in generation. augmented LMs is that, since they generate using the KBs at test time, these LMs can be dynamically updated by modifying the KBs with new information (Khandelwal et al., 2020; He et al., 2021a).”
Landmark Attention: Random-Access Infinite Context Length for Transformers (2023, background): “For example, kNN-LM [16] stores the hidden representation of tokens in memory and uses the distribution of the next token among the stored vectors that are closest to the current token to predict the next token.”
A Review of Sparse Expert Models in Deep Learning (2022, background): “Retrieval mechanisms effectively expand the capacity of models by allowing them to dynamically access information beyond the current context or what is stored in the parameters (Khandelwal et al., 2019; Guu et al., 2020; Borgeaud et al., 2022).”
Language Models that Seek for Knowledge: Modular Search & Generation for Dialogue and Prompt Completion (2022, background): “In the language modeling space, there is a large body of work on nearest neighbor and cache-based language modeling (Khandelwal et al., 2020; Grave et al., 2017; Merity et al., 2017; Khandelwal et al., 2021; Yogatama et al., 2021) for accessing a large set of documents.”
Internet-augmented language models through few-shot prompting for open-domain question answering (2022, methodology): “Semi-parametric language models Semi-parametric LMs have been recently gaining momentum [5, 6, 16, 17], extending monolithic parametric models with information from a knowledge source.”
You Only Need One Model for Open-domain Question Answering (2021, methodology): “Neural Retriever Augmented Language Modeling (NRALM): Augmenting language models with neural retrieval has been shown to be very effective, such as by retrieving nearest neighbor words for LM tasks (Khandelwal et al., 2020; Yo-gatama et al., 2021) or Machine Translation (Khan-delwal et al., 2021).”
Hindsight: Posterior-guided training of retrievers for improved open-ended generation (2021, background): “Retrieval for Language Modeling To improve the perplexity of a pre-trained language model, Khandelwal et al. (2020) retrieve similar contexts from the training set at each time-step and increase the likelihood of tokens that were predicted in similar contexts.”
End-to-End Training of Multi-Document Reader and Retriever for Open-Domain Question Answering (2021, background): “In left-to-right language modeling, retrieving similar words from an external memory has been shown to improve perplexity (Khandelwal et al., 2020, Yogatama et al., 2021).”
Adaptive Semiparametric Language Models (2021, methodology): “Models of episodic memory include kNN-LM (Khandelwal et al., 2020) and architectures that are designed for more complicated tasks such as question answering (de Masson d’Autume et al.”
Internet-Augmented Dialogue Generation (2021, methodology): “Retrieval-augmentation is also studied in the area of language modeling, where it is used for pre-training (Guu et al., 2020), and as a memory (Yogatama et al., 2021), especially using k -nearest neighbor-based cache models (Khandelwal et al., 2021, 2020; Grave et al., 2016; Merity et al., 2016).”