Generalization through Memorization: Nearest Neighbor Language Models
Authors: Urvashi Khandelwal
Omer Levy
Dan Jurafsky
Luke Zettlemoyer
Mike Lewis
Venue:  (2020)
Keywords: language models
k-nearest neighbors

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 transformation 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.

Published as a conference paper at ICLR 2020
GENERALIZATION THROUGH MEMORIZATION: NEAREST NEIGHBOR LANGUAGE MODELS
Urvashi Khandelwal Omer Levy', Dan Jurafskyl, Luke Zettlemoyer! & Mike Lewis!
Stanford University
Facebook AI Research
{urvashik, wwwrafayyy
[omerlevy, 1sz, mikelewis
ABSTRACT
We introduce KNN-LMs, which extend pre-trained neural language model (LM) by linearly interpolating it with a k-nearest neighbors (kNN) model. The near- est neighbors computed according to distance in the pre-trained LM embed- ding space, and can be drawn from any text collection, including the original LM training data. Applying this augmentation neighbors drawn from the original training set, our KNN-LM achieves new state- of-the-art perplexity of 15.79 2.9 point improvement with no additional train- ing. 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 knowl- edge. Together, these results strongly suggest that learning similarity between se- quences of text is casier than predicting the next word, and that nearest neighbor search is effective approach for language modeling in the long tail.
strong WIKITEXT-103 LM, with
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 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, approach that extends pre-trained LM by linearly interpolating its next word distribution with a ie-nearest neighbors (kNN) model. The nearest neighbors are computed according distance in the preitrained embedding space and can be drawn from any text collec- tion, 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 ANN augmentation to a strong WIKITEXT- 1 03 LM using only the original dataset achieves new state- of-the-art perplexity of 15.79 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 model on 100-million tokens and using &NN search over a 3-billion token dataset can outperform training the same model on all 3-billion tokens, opening
Work done while the first author was interning Facchook AT Research.

---

Published as a conference paper at ICLR 2020
Training Contexte
Targets
Representations flog
Distances dig.
Nearest K
Normalization (K exp +3
Aggregation Rocil (6
Obeme senator for AWnois
4
Hawar INinols Hawaii
Hawai 0.7
Hawaii 0.8 ANhoos 0.2
Barack memned to Michelle
100
Milnols 0.2
Obama was born Wi
Hawaii
5
Häwsi 0.1
Obeme native of Hawaii
3
Classification PLArÍa)
Interpolation
Test Context
Target
Representation
.2*
7(x)
Hawai 0.2
Hewali 0.6
Mllnoos 0.2
Ilimois 0.2
birthplace
Figure I:
An illustration of KNN-LM. A datastore is constructed with entry for each training set token, and an encoding of its leftward context. For inference, a test context encoded, and the most similar training contexts are retrieved from the datastore, along with the corresponding targets. distribution over targets is computed based on the distance of the corresponding context from the test context. This distribution is then interpolated with the original model's 's output distribution.
new path for efficiently using large datasets in language models. Similarly, adding out-of-domain data to the datastore makes single LM useful across multiple domains, again without further train- ing. Qualitatively, find the model is particularly helpful for long-tail patterns, such factual knowledge, which might be easier to access via explicit memory.
2 NEAREST NEIGHBOR LANGUAGE MODELING
Language models (LMs) assign probabilities to sequences. Given context sequence of tokens a: (Wi.... 1t- autoregressive LMs estimate ((w]cc)), the distribution over the target token
The ANN-LM involves augmenting such pre-trained LM with a nearest neighbors retrieval mech- anism, without any additional training (the representations learned by the LM remain unchanged). This can be done with a single forward pass over text collection (potentially including the original LM training set), where the resulting vontext-targer pairs are stored in key-value datastore that is queried during inference, as illustrated in Figure
Datastore Let f(-) be the function that maps a context e to a fixed-length vector representation computed by the pre-trained LM. For instance, in a Transformer LM, [c) could map to an inter- mediate representation that is output by an arbitrary self-attention layer. Then, given the i-th training example (c, w.) D, define the key-value pair (k;, v.). where the key kę is the vector represen- tation of the context f(c;) ) and the value e is the target word wy. The datastore (AC, V) is thus the set of all key-value pairs constructed from all the training examples in D:
(CC, V) {(f(e:), w.))(c,, w:) E D}
Inference At test time, given the input context æ the model generates the output distribution over next words (w(x)) and the context representation f(x). The model queries the datastore with (xx) to retrieve its k-nearest neighbors A according to distance function d(.,)) (squared L22 distance in our experiments, making the similarity function an RBF kernel). it computes a distribution over neighbors based 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 probability):
PVNN(y
Σ (KWW)NN
1y=e,
f(x)))
Finally, we follow Grave et al. (201 and interpolate the nearest neighbor distribution PaNN with the model distribution PLM using a tuned parameter À to produce the final ANN-LM distribution:
p(yla) DANN(ylz2
())()
2

---

Published as a conference paper at ICLR 2020
Implementation The datastore contains 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 hased on the cluster centroids, while reducing memory usage by storing compressed versions of the vectors. We found in pre- liminary experiments that using L 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 similar ap- proach 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 à 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, improve domain adaptation. In our work, we only save training data, with the goal of explicitly memorizing training examples better generalize to similar cases test time.
EXPERIMENTAL SETUP
Data Experiments in this paper use the following English corpora:
WIKITEXT-I03 is a standard benchmark by Merity et al. (2017) for autoregressive language mod- eling with 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 al. 2013), 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- 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 al. 2019).
Model Architecture ANN-LM is compatible with any model that produces fixed size context representations. We use decoder-only Transformers (Vaswani et al., (017) for language modeling, which the current state of the art. Since the kNN-LM makes no changes to the underlying IMM, we take the exact architecture and optimization described by Baevski & Auli (2019) and use create a ANN-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-I0r 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, (017] 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 scored per test example, but up 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.
ANN-LM The keys used for ANN-LM are the 1024-dimensional representations fed to the feed- forward network in the final layer of the Transformer LM (after self-attention and layernorm; see Section for further explanation). We perform 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 minimum of 1536 tokens of prior context for WIKITEXT- 1 103 and minimum of 512
3

---

Published as a conference paper at ICLR 2020
Model
Perplexity (L)
Trainable Params
Dev
Test
Baevski Auli (0019)
17.96
18.65
247M
+Transformer-XI XL (Dai al.
2019
18.30
257M
+Phrase Induction (Luo et al.
2019
17.40
257M
Base LM Baevski & Auli 2019
17.96
18.65
247M
+kNN-LM
16.06
16.12
247M
+Continuous Cache Grave et al. 2017c)
17.67
18.27
247M
+ANN-LM + Continuous Cache
15.81
15.79
247M
Table 1: Performance WIKITEXT- 103. The &NN-LM substantially outperforms existing work. Gains are additive with the related but orthogonal continuous cache, allowing us to improve the base model by almost 3 perplexity points with no additional training. We report the median of three random seeds.
Model
Perplexity (L)
# Trainable Params
Dev
Test
Base LM (Baevski & Auli 2019)
14.75
11.89
247M
+ANN-LM
14.20
10.89
247M
Table Performance on BOOKs, showing that ANN-LM works well in multiple domains.
tokens for the rest of the corpora. 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 & 1024 neighbors, and the index looks up 32 cluster centroids while searching for the nearest neighbors. For WIKITEXT- 103 experiments, full precision keys, but for the other datasets we use the FAISS L distances (not squared) between quantized keys directly, for faster evaluation. We tune the interpolation parameter À on the validation set
compute squared L distances with
Computational Cost Although the ANN-LM requires no training given an existing LM,, does add some other computational overheads. Storing the keys and values requires a single forward pass over the training set, which amounts to fraction the cost of training for one epoch on the same examples. Once the keys are saved, for WIKITEXT-103 103 building the cache with 103M entries takes roughly two hours on single CPU. Finally, running on the validation set took approximately 25 minutes when retrieving 1024 keys. While the cost of building large cache grows linearly in the number entries, it 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 I shows that ANN-LM improves perplexity WIKITEXT- 103 from 18.65 (Baevski & Auli 2019) to new state-of-ther -fte--r-the-rtt 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 aug- menting the KNN-LM with these techniques. We compare with models trained only on the standard training set, but recent work has shown performance 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, related but orthogonal technique from Grave et al. (2017c), in which the model saves and retrieves neighbors from earlier in the test document,
'Code available http:p///gggihhh....orrreee/e///

---

Published as a conference paper at ICLR 2020
Training Data
Datastore
Perplexity ())
Dev
Test
WIKI-3B WIKI-10OM
16.11
15.17
20.99
19.59
WIKI-100M WIKI-3B
14.61
13.73
Table 3: Experimental results on WIKI-3B. The model trained on 100M tokens is augmented with a datastore that contains about 3B training examples, outperforming the vanilla LM trained on the entire WIKI-3B training set.
0.5
0.4
optinal
UNIK-1COM KNN
0.2
16
0.1
0.0
0.0
0.5
3.0
1.5
2.0
2.5
3.0
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Size dafastore (In billioss
Size datastore (In billonst
(a) Effect of datastore size on perplexities.
(b) Tuned values of X for different datastore sizes.
Figure 2: Varying the size of the datastore. (a) Increasing the datastore size monotonically improves performance, and has not saturated even at about 3B tokens. A &NN-LM trained on 100M tokens with datastore 1.6B tokens already outperforms the LM trained on all 3B tokens. (b) The optimal value of À increases with the size of the datastore.
rather than the training set. Gains from interpolating with the continuous cache are smaller than reported 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 İ-N--MM, pushing our state-of-the-art result 15.79, 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. T an improvement in test set perplexity from 11.89 to 10.89, suggesting that this is not the case.
4.2 MORE DATA WITHOUT TRAINING
Section has shown that retrieving neighbors from the training data can significantly improve language modeling performance. This raises the question: can retrieving nearest neighbors from data be a substitute for training on it? To test this, we train LM on WIKI-100M and use it to build a datastore from WIKI-3B, corpus 30 times larger than the training set. We then compare this -NN-LM to vanilla LM trained on the entire WIKI-3B corpus
Table shoows that, as expected, the model trained on 3B tokens dramatically outperforms the model trained on 100M tokens, improving perplexity from 19.59 to 15.17. However, adding nearest neigh- bors retrieval over those 3B examples to the model trained on 100M tokens improves perplexity from 19.59 to 13.73; i.e. retrieving nearest neighbors from the corpus outperforms training on it. This result suggests that rather than training language models on ever larger datasets, can use smaller datasets to learn representations and augment them with ANN-LM over a large corpus.
"The original LM (Baevski & Auli 2019 was trained for 286K steps corpus of similar size to WIKI- 100M. When scaling up WIKI-3B, tuned only the number of updates the validation set and found that training for 572K steps (double) produces slightly stronger baseline.
5

---

Published as a conference paper at ICLR 2020
Training Data
Datastore
Perplexity ())
Dev
Test
WIKI-3B
37.13
34.84
BOOKS
14.75
11.89
WIKI-3B
BOOKS
24.85
20.47
Table 4: Domain adaptation experiments, with results on Books. Adding an in-domain datastore
Wikipedia-traiieee trained model improves results by 23 points, approaching in-domain training.
Key Type
Dev ppl. (v)
No datastore
17.96
Feed Forward Network
Model output
17.07
Model output layer normalized
17.01
Layer Norm
FFN input after layer
16.06
FFN input before layer norm
17.06
MHSA input after layer norm
16.76
Multi Headed Self Attention +
MHSA input before layer norm
17.14
WIKITEXT- 100 validation results using dif- ferent states from the final layer of the LM as the rep- resentation function f-) for keys and queries. We trieve k= 1024 neighbors and tuned for each.
Table 5:
Layer Norm
Figure 3: Transformer LM layer.
To understand how the amount of data used for kNN retrieval affects performance, we use the WIKI- 100M model create datastores using different amounts of randomly sampled data from WIKI-3B. Figure [2 shows that using only 1.6B examples for the datastore already surpasses the performance of the model trained on all of WIKI-3B. in the datastore, suggesting that growing the datastore more could lead to further gains. Figure 2b shows the model relies more on the ANN component as the size the datastore increases.
addition, performance does not saturate at examples
4.3 DOMAIN ADAPTATION
We also experiment with domain adaptation by creating datastore the target domain training set. shows that an in-domain LM on BOOKS has relatively low perplexity (11.89), while a model trained WIKI-3B performs poorly the Books domain (34.84 perplexity). Adding ANN search over BOOKS to the WIKI-3B model reduces perplexity by 14 points (to 20.47), demonstrating that &NN-LM allows single model to be useful in multiple domains, by simply adding datastore per domain.
TUNING NEAREST NEIGHBOR SEARCH
While the ANN-LM is conceptually straightforward, and requires no additional training, a number of hyperparameters are introduced for nearest neighbor search. We experiment with different choices here.
Key Function For similarity search, we extract a representation of context using an intermediate state of the LM f()). Transformers compute a number different intermediate states, and we com- pare several choices depicted in Figure with results shown in Table While all the instantiations of f we tried are helpful, we achieved the largest improvement by using the input to the final layer's feedforward network. We also observe that normalized representations (i.e. taken immediately af- ter the layer norm) perform better. Repeating the experiment on the second-last transformer layer showed similar trends with slightly worse results (not shown), suggesting that the feedforward layer might focusing more on the prediction problem, while the onus of representing the input falls more the self- self-attention layer.

---

Published as a conference paper at ICLR 2020
Books Inddimairt
17.4
UM W 103
18-
W8-25 a Bopca Dutastors
-36
Doma i Adsppatppii Prepleslty
Inddmen Perteedyyy
17.2
17.9
2.5
16.6
184
26
14
neighter
A.rteractation paramwteer
Figure 4: Effect of the number of nearest neigh- bors returned per word on WIKITEXT-103 (val- idation set). Returning more entries from the datastore monotonically improves performance.
Figure 5: Effect of interpolation parameter À on in-domain (left y-axis) and out-of-domain (right y-axis) validation set performances. More weight on P&NN improves domain adaptation.
Number of Neighbors per Query Each query returns the top-k neighbors. Figure 4 shows that performance monotonically improves as more neighbors are returned, and suggests that even larger improvements may be possible with higher value of k. Nonetheless, even small number of neighbors (k = 8) is enough to achieve new state of the art.
Interpolation Parameter We use a parameter to interpolate between the base model distribution and the distribution from &NN search over the dataset. Figure 5 shows that 0.25 is optimal on WIKITEXT- 103. However, À = 0.65 works best for domain adaptation results (Figure 5
Precision of Similarity Function In FAISS, the nearest neighbor search computes T.2 distances against quantized keys. We found results were improved from 16.5 perplexity on WIKITEXT-I03 16.06 by computing squared 2 distances with full precision keys for Equation 2
ANALYSIS
Qualitative Analysis To understand why ANN-LM improves performance, we manually examine cases in which PeNx was significantly better than PLM 6 one such example, along with several others in Appendix A The example shows interesting case where the model matches the trigram impact on the in several retrieved neighbors, but puts almost all weight on the most relevant neighbor, thus adding more value than an n-gram LM.
In general, we find that examples where kNN-LM is most helpful typically contain rare patterns. Examples include factual knowledge, names, and near-duplicate sentences from the training set. In these cases, assigning train and test instances similar representations (via f(-)) appears to be an easier problem than implicitly memorizing the next word in model parameters.
Simple Neural Representation We observe that many long-tail phenomena manifest as rare n-grams (e.g. names). Is therefore possible to interpolate an n-gram model with a Transformer LM, an alternative to our &NN approach? Figure shows little improvement from using n-gram LMs 0.2 perplexity points (similarly to Bakhtin et al. (2018)). This result highlights the need to use the learned representation function f()) measure similarity between more varied contexts.
Implicit vs Explicit Memory If a neural representation function is crucial for ANN-LM, could implicitly memorizing the training dataset in the neural network parameters replace the explicit memory in the datastore? To test this, we train a Transformer LM with no dropout. Figure [8 shows that this model eventually reaches zero training loss, indicating that can make perfect predictions for all examples in the training set; the model has memorized the dataset. Naturally, the memorizing LM overfits, i.e. the training loss drops to 0 while the best validation perplexity is much higher at 28.59. For comparison, the vanilla Transformer LM (with dropout) has much higher training loss (shown in Figure 8). but also generalizes better with a validation perplexity of 17.96. This result shows that the Transformer has sufficient capacity to memorize the training set.

---

Published as a conference paper at ICLR 2020
Test Context (PENN = 0.998, PtM = 0.124)
Test Target
it was organised by New Zealand international player Joseph Warbrick, promoted by civil servant Thomas Eyton, and managed by James Scott, publican. The Natives were the first New Zealand team perform haka, and also the first wear all black. They played 107 rugby matches during the tour, as well as a small number of Victorian Rules football and associ- ation football matches in Australia. Having made a significant impact on the...
development
Training Set Target
Context Probability
Training Set Context
the captain and instigator of the 1888-89 Natives the first New Zealand team to tour the British Isles Warbrick had lasting impact on the...
development
0.998
promoted new first grade competition which started in 1900. immediately made a big impact on the...
Glebe
district
0,00012
centuries, few were large as other players managed. Howeven, others contend that his impact on the...
game
0.000034
Nearly every game the main series has either an anime or manga adap- lation, or both. The series has had a significant impact on the...
development
0.00000092
Figure 6:
Example where the &NN model has much higher confidence in the correct target than the LM. Although there are other training set examples with similar local n-gram matches, the nearest neighbour search is highly confident of specific and very relevant context.
18.00
17.75
Without Dropout
17.50
Trainin looe (baze
Poplotty
17.25
Whitext
37.00
16.50
16.25
16.00
10
150
(Site of gram)
Epoch
Figure 7: Interpolating the Transformer LM with n-gram LMs on WIKITEXT- 103 (validation set). Using &NN-LM gives much lower perplexity, suggesting that the representations are learning more than just matching local context.
Figure 8: Training curves for the Transformer LM with and without dropout. Turning off dropout allows the training loss go 0, in- dicating that the model has sufficient capacity to memorize the training data.
We consider whether the memorizing LM can be effective substitute for nearest neighbor search. Interpolating the memorizing LM with the original LM improves validation perplexity by just 0. compared to 1.9 from i-N--MM. This result suggests that although the Transformer is expressive enough to memorize all training examples, learning to do does not result in context representations that generalize. contrast, ANN-LM memorizes training data while improving generalization.
From these experiments, conjecture that ANN-LM improves performance because (1) the Trans- former LM is very good at learning a representation function for contexts with an implicit notion of similarity, and (2) while the Transformer has capacity to memorize all training examples, doing so causes its representation to generalize less effectively, but (3) the ANN-LM allows the model to memorize the training data while retaining effective similarity function.
8

---

Published as a conference paper at ICLR 2020
7 RELATED WORK
We discuss related uses of caches for language modeling in Section 2
Similar &NN models ours have been proposed for computer vision tasks (Paperoot McDaniel, 2018 Orhan 2018: Zhao & bustness to adversarial attacks. We hypothesize that our method may be particularly effective for language modeling, because plentiful unlabeled data allows datastores of billions of tokens, and language modeling often requires world knowledge to be learnt from few examples.
primarily motivated by improving interpretability and ro-
Nearest neighbor models have been applied to number of NLP problems in the past, such as part of speech tagging (Daelemans et al. 1996) and morphological analysis (Bosch et al. 2007). but the use of learned representations makes the similarity function much more effective in the case of neural models. More recently, Kaiser et al. (2017) have used a similarly differentiable memory that is learned and updated during training, and applied to one-shot learning tasks.
Several models have also improved language generation by using training examples directly at test time. Guu et al. (2018) propose model that samples training sentences random and edits them with a prquence-to-sequene
model, but does not use retrieval mechanism such as KNN Gu et al. (2018) introduce translation model that attends over retrieved training set examples, WWston al. 20188) improve a dialogue response generation model by refining similar instances from the training set. KNN-LM differs from these approaches by working at the level of individual tokens instead of whole training sentences, as well as not incorporating the retrieval mechanism into the training pipeline.
A general trend in machine learning, and in language modeling in particular, is that adding more data consistently improves performance (Devlin et al. 2019 Radford et al.. 2019] Yang al.j Liu et al.) 2019; Zellers et al. 2019] (Shoeybi et al. [219). Our work offers an alternative method for scaling language models, in which relatively small models learn context representations, and nearest neighbour search acts highly expressive classifier.
CONCLUSION AND FUTURE WORK
We have introduced ANN-LMs, which can significantly outperform standard language models by directly querying training examples at test time. The approach can be applied to any neural language model. The success of this method suggests that learning similarity functions between contexts may be an easier problem than predicting the next word from some given context. Future work should explore explicitly training similarity functions, and reducing the size of the datastore.
ACKNOWLEDGMENTS
The authors thank the anonymous reviewers well as Sida Wang, Kartikay Khandelwal, Kevin Clark and members of the FAIR Seattle team for helpful discussions and comments.
REFERENCES
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling, In ICLR, 2019.
Anton Bakhtin, Arthur Szlam, Mare' Aurelio Ranzato, and Edouard Grave. Lightweight adaptive mixture of neural and n- n-gram language models. arXiv preprint arXiv:1804.07705. 2018.
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Jauvin. A neural probabilistic language model. Journal of machine learning research,
137--1155, 2003.
Antal van den Bosch, Bertjan Busser, Sander Canisius, and Walter Daelemans. An efficient memory- based morphosyntactic tagger and parser for dutch. LOT Occasional Series, 7:191-206, 2007.
Walter Daelemans, Jakub Zavrel, Peter Berck, and Steven Gillis. Mbt: A memory-based part of speech tagger-generator. In WVLC, 1996.

---

Published as a conference paper at ICLR 2020
Zihang Dai, Zhilin Yang, Yiming Yang, William W Cohen, Jaime Carbonell, Quoe V Le, and Ruslan Salakhutdinov. Transformer-xi. Attentive language models beyond a fixed-length context. In ACL, 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
Edouard Grave, Moustapha M Cisse, and Armand Joulin. Unbounded cache model for online lan- guage modeling with open vocabulary. In NIPS, pp. 6042-6052, 2017a.
Edouard Grave, Armand Joulin, Moustapha Cissé, Hervé Jégou, et al. Efficient softmax approxima- tion for gpus. In Proceedings of the 34th International Conference Machine Learning-Vood Vooluee 70, pp. 1302-1310 1310. JMLR. org. 2017b.
Edouard Grave, Armand Joulin, and Nicolas Usunier. Improving neural language models with continuous cache. In ICLR, 2017c
Jiatao Gu, Yong Wang, Kyunghyun Cho, and Victor OK Li. Search engine guided neural machine translation. Pherot--ecood AAAI Conference on Artificial Intelligence, 2018.
Kelvin Guu, Tatsunori
Hashimoto, Yonatan Oren, and Percy Liang. Generating sentences by editing prototypes. Transactions of the Association for Computational Linguistics, 6:437-450, 2018.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion scale similarity search with gpus. arXiv preprint arXhv:1702.0873: 2017.
Lukasz Kaiser, Ofir Nachum, Aurko Roy, and Samy Bengio. Learning to remember rare events. In ICLR, 2017.
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of trans- former language models. arXiv preprint AXXx:1904.08376 2019.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint a9Xiv: 11692, 2019.
Hongyin Luo, Lan Jiang, Yonatan Belinkov, and James Glass. Improving neural language models by segmenting, attending, and predicting the future. ACL, 2019.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. ICLR, 2017.
Tomášs Mikolov, Martin Karafiát, Lukáš Burget, Jan Černocký, and Sanjeev Khudanpur. Recurrent neural network based language model. In Eleventh annual conference the international speech communication association, 2010.
A. Emin Orhan. A simple cache model for image recognition. In NeurIPS, 2018.
Nicolas Papernot and Patrick McDaniel. Deep k-nearest neighbors: Towards confident, interpretable and robust deep learning. arXiv preprint 1803. 047655, 2018.
Ofir Press and Lior Wolf. Using the output embedding to improve language models. In ICLR, 2017.
Alec Radford, Jeff Wu. Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. URL. http://ddmuuiikyyyyooooooorrrrrr 2019.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. arXiv preprint arXiv: 077999 2015.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron- Training multi-billion parameter language models using gpu model parallelism. arXiv preprint arXiv: /t00::/9900000000 08053, 2019.
10

---

Published as a conference paper at ICLR 2020
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp.
2017.
Jason Weston, Emily Dinan, and Alexander H Miller. Retrieve and refine: Improved sequence generation models for dialogue. arXiv preprint arXiv:1808.04776, 2018.
Zhilin Yang. Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoe V
Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv:1966.08337, 2019.
Rowan Zellers, Ari Holtzman, Hannah Rashkin, Ali Farhadi, Franziska Roesner, and Yejin Choi. Defending against neural fake news. In NeurlPS, 2019.
Jake Zhao and Kyunghyun Cho. Retrieval-augmente convolutional neural networks for improved robustness against adversarial examples. arXiv preprint arXiv:1802.09502, 2018.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. Proceedings of the IEEE international conference on computer vision, pp. 19-27, 2015.

---

Published as a conference paper at ICLR 2020
A APPENDIX
This section provides several examples where PeNN places higher probability mass on the true target, compared to PLM-
Test Context (PENN = 0.995, PLM = 0.025)
Test Target
For Australians and New Zealanders the Gallipoli campaign came to sym- bolise an important milestone in the emergence of both nations as indepen- dent actors the world stage and the development of a sense of national identity. Today, the date of the initial landings, 25 April, known as An- zac Day in Australia and New Zealand every year thousands of people gather at memorials in both nations, as well as Turkey, to...
honour
Training Set Target
Context Probability
Training Set Context
Despite this, for Australians and New Zealanders the Gallipoli campaign has come to symbolise an important milestone nations as independent actors on the world stage and the development of a sense of national identity. Today, the date of the initial landings, 25 April, public holiday known Anzuc Day in Australia and New Zealand and every year thousands of people gather at memorials both nations, and indeed in Turkey, t....
honour
0.995
emergence of both
On anniversary date of death, every year since 1997, thousands of people gather at his home Memphis to...
celebrate
0.0086
Twenty-five years after Marseille's 's death, fighter pilot veterans of World War II gathered to...
honour
0.0000041
Table 6: Another example where the ANN model places much higher probability mass the correct target, compared to the LM. The nearest neighbors search has retrieved a training set context that is extremely similar to the test context, while very rare and in the long-tail of patterns.
Test Context (pENN = 0.959, prm = 0.503)
Test Target
U2 do what they 're best at, slipping into epic rock mode, playing music made for the arena". , In two other local newspuper reviews, critics praised the song's inclusion in sequence of greatest hits. For the PopMart Tour of 1997-..
1998
Training Set Target
Context Probability
Training Set Context
Following their original intent, "Sunday Bloody Sunday" was not played during any the forty-seven shows the Lovetown Tour 1989. The song reappeared for brief period during the Zoo TV Tour, and late during the second half of PopMart Tour (1997-...
1998
0.936
They 6 times Champions and they won the Challenge Cup in 1938, and have experienced two previous stretches in the Super League, 1997
2002
0.0071
About $40 million (S61.4 4 million in 2018 dollars) was spent on the property acquisition. After weather-related construction delays due to Nino season of the winter of 1997-...
1998
0.001.5
This made the highest-rated season of The X-Files to air as well as the highest rated Fox program for the 1997-...
0.00000048
this example, the desired date pattern appears in many examples. Yet, the nearest neighbors search is able to identify the only training set context which is relevant to the test context and assigns it the highest probability mass.
Table 7:
12

---

Published as a conference paper at ICLR 2020
Test Context (PINN = 0.624, PIM = 0.167)
Test Target
Lord Strathcona awarded Guuthier scholarship in 1906 that allowed her return Europe and continue her vocal studies. She returned there and continued both to study and give performances. Her first operatic perfor- mance came in 1909 in Pavia, Italy as Micaela in Bizet's s...
Carmen
Training Set Target
Context Probability
Training Set Context
Despite poor relations with the orchestra, Mahler brought five new operas the theatre, including Bizet's... s...
Carmen
0.356
The fourth movement of An die Jugend (1909), for instance, uses two of Niccolo Paganini's Caprices for solo violin (numbers and 15), while the 1920 piece Piano Sonatina No. (Fantasia da camera super Carmen) based on themes from Georges Bizet s...
opera
0.0937
also hosted the Ballet of her Majesty's 's Theatre the mid-19n 19th19th century, before returning to hosting the London premieres of such operas as Bizet's s...
Carmen
0.0686
Table 8: In this case, the model is able memorize the fact that Georges Bizet wrote Carmen.
Test Context (PENN 0.031, PLM 0.007)
Test Target
Mycena maculata bears resemblance M. <unk>, but is only sociated with decaying hardwood logs and stumps, and is found in eastern North America, and sometimes on ouk on the West Coust. In age, it...
develops
Training Set Target
Context Probability
Training Set Context
Morchella tridentina (=Mocchella frustrata) is also rufescent and very sim- ilar M. rufobrannea. is found mountainous forests and maquis and forms a marked sinus at the attachment of the cap with the stem, which is pure white. At maturity, it...
develops
0.031
The winter bonnet (M. tintinnabulum) a northern European species that is much smaller (cap diameter up to 2.6 cm (1.0 in) across) and has a brown cap, and has ragged hairs at the base. It...
generally
0.029
The "bleeding" will distinguish Mycena atkinsoniana from most other Mycena species commonly encountered. The common and widely dis- tributed M. sanguinolenta is another "bleeder". but it smaller than M. atkinsonia, with a cap diameter ranging from to mm (0.1 to in). Additionally, it...
has
0.028
Mycenu flavoalba bears resemblance to some members of the genus Hemimycena, such as H. lactea and H. <unk>, It...
can
0.018
Table 9: This is an example where the PeNN distribution is relatively flat, as several words are plausible continuations. However, the nearest neighbors search assigns the highest probability to the correct target and a corresponding context that is particularly relevant. In contrast, the LM probability on the correct target is lower.
13


Reviews:
[{"review_id": "ryghQUjTYH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nSummary:\nThe authors extend a pretrained LM by interpolating its next word distribution with a KNN model. The authors show retrieving nearest neighbor from corpus achieve quite large perplexity decrease in several language modeling benchmarks.\n\nDecision:\nOverall, the idea seems simple but is quite effective. Even with some discussions on the related work with cache based LM and the work that use training examples explicitly, I feel it is a simple extension/usage of previous approaches. Hence I am borderline with my decision.\n\nSupporting argument:\n1. The proposed idea uses KNN to look up training examples for interpolating the prediction. As discussed by the authors, this approach is effective in factual knowledge, names, and near-duplicate sentences.\n2. There are several experiments and ablation study in showing the effectiveness of the approach.\n3. The related work that uses training examples explicitly is quite similar to the proposed approach, though the authors claim that one is at the level of individual tokens and the other is the whole training sentences.\n\nAdditional feedback:\n1. In reference, ‘Bert’ -> ‘BERT’\n2. Missing reference: Yogatama et al., Memory Architectures in Recurrent Neural Network Language Models, 2018, https://arxiv.org/abs/1410.3916, https://arxiv.org/abs/1803.02400", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "H1eELYG-cS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "[Overview]\n\nIn 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.\n\n[Pros]:\n\nOverall I think the paper is well-written and presents clearly. Detailed points below:\n\n1. the authors proposed a simple but effective method for increasing the generalization ability of language model through a memorization strategy. Specifically, the authors proposed to build a lookup table which memorizes the representation and output token pairs which are then used for the inference of language model. Different from conventional way, the proposed strategy does not introduce any more parameters in the model and also does not need any more training or fine-tuning on the target dataset.\n\n2. The authors showed that the proposed strategy can improve the performance of language generation model (i.e., transformer) without any extra training or data, as shown in Table 1. Also, using the continuous caches  with KNN-LM further improve the performance.\n\n3. Besides the main results shown in Table 1 and Table 2, the authors also showed using kNN-LM can probably outperforms the model which is directly trained on it. Also, it also supports domain adaptation from one language domain to another domain.\n\n4. Finally, the authors presented a number of ablation studies to investigate how the performance is affected by the method of building datastore, including the size of nearest neighbor, the interpolation parameter, etc. These results are also insightful and meaningful for the readers to understand the method.\n\n[Cons]:\n\nI think this paper is a solid paper. So I would have some suggestions below:\n\n1. The first concern about the method is the efficiency. At page 3, the authors mentioned that the proposed strategy will bring more time cost. It would be good if the authors can perform more systematical analysis on the time cost of building the datastore and inference for the proposed model. \n\n2. Second, the authors should not only evaluate the proposed method based on transformers. It would be good to test on various language models to verify the generalization ability across different models, including the old-fashioned one like RNN and CNN.\n\n3. Also, the authors should try to extend the proposed model to other language tasks, such as translation.\n\n[Summary]\n\nIn this paper, the authors introduced a simple but effective method to augment the pertained language model through memorizations. Though this is not absolutely new and relatively simple , the authors successfully demonstrate that it can be applied to improve the generation of language model much. The. thorough ablation studies help to understand the property of the proposed strategy. I think this paper overall is insightful and thoughtful. It would be good to see the authors add more analysis on the computational complexity and also evaluate on more type of language models.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "B1gMnYs8tS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work utilizes the kNN method on dense vectors to augment the LMs. The method is simple and straightforward, meanwhile, the performance seems great if only in terms of PPL.\nThree of my most concerns:\n1)\tIt seems that this approach heavily relies on the similarity of context distribution between the training and test set. Intuitively, higher performance will be achieved with more similar examples between training and test set. This question should be discussed more in this work. This similarity cannot always satisfied in practice, I thus quite doubt the proposed method can work for general case.\n2)\tThe evaluation is only done for PPL, I notice the LM was trained in a corpus scale as pre-trained BERT, though none of real downstream tasks were evaluated like BERT. Expect to see some MRC or NLI results with the proposed LM.\n3)\tFurthermore, though FAISS is very fast, it is hard to get great results with only a small datastore which makes the retrieving slow. So it seems not suitable for tasks such as generations but maybe open-domain QA can be the scene for this method. It would be great if there are some experiments on such tasks, and also combining with models such as BERT could be much better and convincing.\n\nQuestions\nWhat about other distance functions such as cosine distance? The author only said L2 is better but there is no analysis on it.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes an idea of using a pre-trained language model on a potentially smaller set of text, and interpolating it with a k-nearest neighbor model over a large datastore. The authors provide extensive evaluation and insightful results. Two reviewers vote for accepting the paper, and one reviewer is negative. After considering the points made by reviewers, the AC decided that the paper carries value for the community and should be accepted.
Decision: Accept (Poster)