[Paper Review] Generalization through Memorization: Nearest Neighbor Language Models

Urvashi Khandelwal, Omer Levy+3 more|arXiv (Cornell University)|Nov 1, 2019

Topic Modeling28 references56 citations

TL;DR

KNN-LMs augment a pre-trained language model with a k-nearest neighbors datastore to interpolate predictions, achieving state-of-the-art perplexity without extra training and enabling domain adaptation and data-efficient scaling.

ABSTRACT

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.

Motivation & Objective

• Motivate the hypothesis that learning similarity between text contexts can be easier than predicting the next word.
• Propose a k-nearest neighbors augmentation to a pre-trained LM to improve next-token predictions without retraining.
• Empirically evaluate whether explicit memory of training contexts improves perplexity and enables domain adaptation and data-efficient scaling.

Proposed method

• Construct a datastore of context representations and next-word targets from a trained LM.
• Query the datastore with the test context to retrieve k nearest neighbors using an L2 distance in embedding space.
• Compute a p_kNN distribution over next words from retrieved neighbors and interpolate with the base LM distribution using a tunable lambda.
• Use FAISS for scalable nearest-neighbor search over high-dimensional keys with 64-byte quantized vectors.
• Tune the interpolation parameter lambda on validation data.
• Evaluate on WikiText-103 and Books, and test data with varying datastore sizes and domains.

Experimental results

Research questions

• RQ1 Can a pre-trained LM's context representations be leveraged via kNN retrieval to improve next-token prediction without additional training?
• RQ2 How does the size of the datastore and the interpolation weight affect perplexity and domain adaptation performance?
• RQ3 Can data from larger or different domains be used to augment a smaller LM effectively through a datastore?
• RQ4 Does explicit memory of training instances help more with long-tail patterns such as factual knowledge or proper names?

Key findings

• A kNN-LM achieved a new state-of-the-art perplexity of 15.79 on Wikitext-103 with no extra training, improving over the base model by 2.86 points.
• Using the training data as datastore yields substantial perplexity gains, and combining kNN with a continuous cache further improves results to 15.79 on Wikitext-103.
• Datastore augmentation with 100M tokens and a 3B-token datastore can outperform training the same model on 3B tokens, showing data-efficient scaling.
• Domain adaptation is effective: adding an in-domain Books datastore to a Wiki-3B model reduces Books perplexity from 34.84 to 20.47, approaching in-domain training performance.
• Retrieving from larger/datastore-backed data improves performance monotonically, and the optimal lambda increases with datastore size for domain adaptation.
• Qualitative analysis shows kNN-LM better handles long-tail patterns and factual knowledge by explicit memory, compared to relying solely on implicit parameters.