Generalization through Memorization: Nearest Neighbor Language Models

1 INTRODUCTION . Neural language models ( LMs ) typically solve two subproblems : ( 1 ) mapping sentence prefixes to fixed-sized representations , and ( 2 ) using these representations to predict the next word in the text ( Bengio et al. , 2003 ; Mikolov et al. , 2010 ) . We present a new language modeling approach that is based on the hypothesis that the representation learning problem may be easier than the prediction problem . For example , any English speaker knows that Dickens is the author of and Dickens wrote will have essentially the same distribution over the next word , even if they do not know what that distribution is . We provide strong evidence that existing language models , similarly , are much better at the first problem , by using their prefix embeddings in a simple nearest neighbor scheme that significantly improves overall performance . We introduce kNN-LM , an approach that extends a pre-trained LM by linearly interpolating its next word distribution with a k-nearest neighbors ( kNN ) model . The nearest neighbors are computed according to distance in the pre-trained embedding space and can be drawn from any text collection , including the original LM training data . This approach allows rare patterns to be memorized explicitly , rather than implicitly in model parameters . It also improves performance when the same training data is used for learning the prefix representations and the kNN model , strongly suggesting that the prediction problem is more challenging than previously appreciated . To better measure these effects , we conduct an extensive empirical evaluation . Applying our kNN augmentation to a strong WIKITEXT-103 LM using only the original dataset achieves a new stateof-the-art perplexity of 15.79 – a 2.86 point improvement over the base model ( Baevski & Auli , 2019 ) – with no additional training . We also show that the approach has implications for efficiently scaling up to larger training sets and allows for effective domain adaptation , by simply varying the nearest neighbor datastore . Training a model on 100-million tokens and using kNN search over a 3-billion token dataset can outperform training the same model on all 3-billion tokens , opening a ∗Work done while the first author was interning at Facebook AI Research . new path for efficiently using large datasets in language models . Similarly , adding out-of-domain data to the datastore makes a single LM useful across multiple domains , again without further training . Qualitatively , we find the model is particularly helpful for long-tail patterns , such as factual knowledge , which might be easier to access via explicit memory . 2 NEAREST NEIGHBOR LANGUAGE MODELING . Language models ( LMs ) assign probabilities to sequences . Given a context sequence of tokens ct = ( w1 , . . . wt−1 ) , autoregressive LMs estimate p ( wt|ct ) , the distribution over the target token wt . The kNN-LM involves augmenting such a pre-trained LM with a nearest neighbors retrieval mechanism , without any additional training ( the representations learned by the LM remain unchanged ) . This can be done with a single forward pass over a text collection ( potentially including the original LM training set ) , where the resulting context-target pairs are stored in a key-value datastore that is queried during inference , as illustrated in Figure 1 . Datastore Let f ( · ) be the function that maps a context c to a fixed-length vector representation computed by the pre-trained LM . For instance , in a Transformer LM , f ( c ) could map c to an intermediate representation that is output by an arbitrary self-attention layer . Then , given the i-th training example ( ci , wi ) ∈ D , we define the key-value pair ( ki , vi ) , where the key ki is the vector representation of the context f ( ci ) and the value vi is the target word wi . The datastore ( K , V ) is thus the set of all key-value pairs constructed from all the training examples in D : ( K , V ) = { ( f ( ci ) , wi ) | ( ci , wi ) ∈ D } ( 1 ) Inference At test time , given the input context x the model generates the output distribution over next words pLM ( y|x ) and the context representation f ( x ) . The model queries the datastore with f ( x ) to retrieve its k-nearest neighbors N according to a distance function d ( · , · ) ( squared L2 distance in our experiments , making the similarity function an RBF kernel ) .Then , it computes a distribution over neighbors based on a softmax of their negative distances , while aggregating probability mass for each vocabulary item across all its occurrences in the retrieved targets ( items that do not appear in the retrieved targets have zero probability ) : pkNN ( y|x ) ∝ ∑ ( ki , vi ) ∈N 1y=vi exp ( −d ( ki , f ( x ) ) ) ( 2 ) Finally , we follow Grave et al . ( 2017a ) and interpolate the nearest neighbor distribution pkNN with the model distribution pLM using a tuned parameter λ to produce the final kNN-LM distribution : p ( y|x ) = λ pkNN ( y|x ) + ( 1− λ ) pLM ( y|x ) ( 3 ) Implementation The datastore contains an entry for each target in the training set , which for LMs can be up to billions of examples . To search over this large datastore , we use FAISS ( Johnson et al. , 2017 ) , an open source library for fast nearest neighbor retrieval in high dimensional spaces . FAISS speeds up search by clustering the keys and looking up neighbors based on the cluster centroids , while reducing memory usage by storing compressed versions of the vectors . We found in preliminary experiments that using L2 distance for FAISS retrieval results in better performance for kNN-LM , compared to inner product distance . Related Cache Models Prior work ( Grave et al. , 2017c ; Merity et al. , 2017 ) used a similar approach to compute similarity to the previous hidden states of test documents , making it easier to copy rare vocabulary items from the recent past . Such techniques have been less popular since the development of Transformers ( Vaswani et al. , 2017 ) , which can learn to copy recent words using self-attention ; in Section 4.1 , we observe relatively small gains from caching recent items in the same test document à la Grave et al . ( 2017c ) . Most relatedly , Grave et al . ( 2017a ) describe an online language model using nearest neighbor search over all previous hidden states , to improve domain adaptation . In our work , we only save training data , with the goal of explicitly memorizing training examples to better generalize to similar cases at test time . 3 EXPERIMENTAL SETUP . Data Experiments in this paper use the following English corpora : WIKITEXT-103 is a standard benchmark by Merity et al . ( 2017 ) for autoregressive language modeling with a 250K word-level vocabulary . It consists of 103M tokens of Wikipedia in the training set and 250K tokens in each of the development and test sets . BOOKS is the Toronto Books Corpus ( Zhu et al. , 2015 ) , containing 0.7B . Complete books are held out for validation/test . WIKI-3B is English Wikipedia , containing about 2.87B tokens . Whole articles are held out for validation/test . WIKI-100M is a random 100M token subset of WIKI-3B , consisting of complete articles . Except for WIKITEXT-103 , text is tokenized using the byte-pair encoding ( Sennrich et al. , 2015 ) with the 29K subword vocabulary from BERT ( Devlin et al. , 2019 ) . Model Architecture kNN-LM is compatible with any model that produces fixed size context representations . We use decoder-only Transformers ( Vaswani et al. , 2017 ) for language modeling , which are the current state of the art . Since the kNN-LM makes no changes to the underlying LM , we take the exact architecture and optimization described by Baevski & Auli ( 2019 ) and use it to create a kNN-LM for inference . This model consists of 16 layers , each with 16 self-attention heads , 1024 dimensional hidden states , and 4096 dimensional feedforward layers , amounting to 247M trainable parameters . It processes 3072 tokens of context per example for WIKITEXT-103 and 1024 tokens for the rest of the corpora . Following Baevski & Auli ( 2019 ) , we use adaptive inputs and an adaptive softmax ( Grave et al. , 2017b ) with tied weights ( Press & Wolf , 2017 ) for the WIKITEXT-103 experiments . On other datasets we do not use adaptive inputs or an adaptive softmax . Evaluation LMs are trained to minimize the negative log-likelihood of the training corpus , and evaluated by perplexity ( exponentiated negative log-likelihood ) on held out data . Following Baevski & Auli ( 2019 ) , 512 tokens are scored per test example , but up to 2560 tokens of extra prior context is provided for WIKITEXT-103 and up to 512 tokens of extra prior context is provided for the rest of the corpora . kNN-LM The keys used for kNN-LM are the 1024-dimensional representations fed to the feedforward network in the final layer of the Transformer LM ( after self-attention and layernorm ; see Section 5 for further explanation ) . We perform a single forward pass over the training set with the trained model , in order to save the keys and values . During this forward pass , each target token is provided a minimum of 1536 tokens of prior context for WIKITEXT-103 and a minimum of 512 tokens for the rest of the corpora . A FAISS index is then created using 1M randomly sampled keys to learn 4096 cluster centroids . For efficiency , keys are quantized to 64-bytes . During inference , we retrieve k = 1024 neighbors , and the index looks up 32 cluster centroids while searching for the nearest neighbors . For WIKITEXT-103 experiments , we compute squared L2 distances with full precision keys , but for the other datasets we use the FAISS L2 distances ( not squared ) between quantized keys directly , for faster evaluation . We tune the interpolation parameter λ on the validation set.1 Computational Cost Although the kNN-LM requires no training given an existing LM , it does add some other computational overheads . Storing the keys and values requires a single forward pass over the training set , which amounts to a fraction of the cost of training for one epoch on the same examples . Once the keys are saved , for WIKITEXT-103 building the cache with 103M entries takes roughly two hours on a single CPU . Finally , running on the validation set took approximately 25 minutes when retrieving 1024 keys . While the cost of building a large cache grows linearly in the number of entries , it is trivial to parallelize and requires no GPU-based training . 4 EXPERIMENTS . 4.1 USING THE TRAINING DATA AS THE DATASTORE . We first experiment with creating a datastore from the same data used to train the LM . Table 1 shows that kNN-LM improves perplexity on WIKITEXT-103 from 18.65 ( Baevski & Auli , 2019 ) to a new state-of-the-art of 16.12 . We also provide reported perplexities from two other recent models that also build upon Baevski and Auli ’ s , suggesting that further improvements may be possible by augmenting the kNN-LM with these techniques . We compare with models trained only on the standard training set , but recent work has shown performance can be improved by training on additional data , from either the test set ( Krause et al. , 2019 ) or large amounts of web text ( Shoeybi et al. , 2019 ) . We also experiment with a continuous cache model , a related but orthogonal technique from Grave et al . ( 2017c ) , in which the model saves and retrieves neighbors from earlier in the test document , 1Code is available at : https : //github.com/urvashik/knnlm rather than the training set . Gains from interpolating with the continuous cache are smaller than reported in the original setting that used LSTMs , perhaps because self-attentive language models can learn to perform such queries . Improvements from the continous cache are additive with the kNN-LM , pushing our state-of-the-art result to 15.79 , a gain of 2.86 over the base model . Finally , we repeat the experiment using text from a different domain , BOOKS , to control for the possibility that encyclopedic Wikipedia text is somehow uniquely good for caching . Table 2 shows an improvement in test set perplexity from 11.89 to 10.89 , suggesting that this is not the case .

Summary: In this paper, the authors proposed a simple but effective way to augmentation the language model through memorization. Specifically, after obtaining a language model on a dataset, the model further uses the dataset to build a lookup table and then a k-nearest neighbor is used to searching the closest tokens for a token during inference. Based on this, the output distribution of a target token during the inference time would be modified accordingly. Through a comprehensive experiments and ablation studies, the authors showed that the proposed strategy can improve the performance of language models significantly for both the in-domain and out-domain testing scenarios. This is very insightful considering recently a lot of language models are focusing on increasing the size of model and training data.