Regularized Training of Nearest Neighbor Language Models (2022)

Abstract: Including memory banks in a natural language processing architecture increases model capacity by equipping it with additional data at inference time. In this paper, we build upon k NN-LM (CITATION), which uses a pre-trained language model together with an exhaustive k NN search through the training data (memory bank) to achieve state-of-the-art results. We investigate whether we can improve the k NN-LM performance by instead training a LM with the knowledge that we will be using a k NN post-hoc. We achieved significant improvement using our method on language modeling tasks on WIKI-2 and WIKI-103. The main phenomenon that we encounter is that adding a simple L2 regularization on the activations (not weights) of the model, a transformer, improves the post-hoc k NN classification performance. We explore some possible reasons for this improvement. In particular, we find that the added L2 regularization seems to improve the performance for high-frequency words without deteriorating the performance for low frequency ones.

## introduction the problem of language modeling (lm) usually consists of two main challenges. firstly, mapping the context, i.e. the sentence prefixes, to a vector representation, and secondly using this representation to predict the subsequent word. in @xcite , the authors claim that the first problem is much easier to solve. hence, given a pre-trained lm, they post-hoc modify the representation using a k-nearest neighbor scheme (knn) and achieve significant improvements on challenging datasets, such as wiki-103. given that knn improves the overall language modeling of a pre-trained network, we examine training strategies that can make the underlying network's representations more amenable to the * work done as an intern at apple knn step. our results show improvements over applying knn to a generic lm network. we first explore a simple learning scheme for the language model, where during training we intentionally push representations that predict the same word closer together in the l2 sense, using a momentum contrastive (moco) @xcite style implementation. we go on to note that this moco style learning can be replaced by simply adding l2 regularization to the activation of the layer used for knn, eliminating implementation complexity. lastly, we present some initial experiments toward understanding why this l2 regularization brings improved performance. ## background our work builds upon knn-lm @xcite . in essence, knn-lm tackles the problem of how to improve a trained lm's representations, and how to adapt lms to capture non-frequent sentences that are usually forgotten by the model during training. knn-lms achieve significantly higher performance through a simple interpolation between the original lm predictions and the knn predictions. at inference time, given a new context sentence, knn-lm works as follows: this simple post-hoc implementation allows @xcite to improve upon the sota in lm by a significant margin. one thing to note about knn-lm is that they do not need to retrain the lm and hence the whole algorithm can be run on cpu only. furthermore, knn-lm uses faiss, which is an efficient library that allows them to quickly find knns. one detail to note in @xcite ) which was crucial for this work was that the authors tried both the inner product and the l2 for their distance metric in knn. they concluded that l2 worked significantly better. this observation implies the fact that the default training recipe of lms implicitly prefers one distance over the other. given that we know that a post-hoc knn adaptation significantly improves the performance, it is natural to ask whether we could train a lm with this in mind. in the next section, we describe how to adapt the training of the lm for this purpose. ## proposed method in our initial attempt, we experimented with the idea of explicitly minimizing the l2 distance between context vectors that predict the same target word. this strategy directly mirrors the use of context vectors at the knn step, and we hoped that training the representations in a way similar to testing will further improve the effectiveness of knn lm. however, a naïve implementation of it is infeasible due to having to store all the representation in memory. we then resorted to a moco- @xcite style training scheme. specifically, for each target word w, we construct a queue q of fixed length l, which stores the recent l context representations for w. during training, we optimize a regularized objective as follows: empirically, we found that equation 1 provides a practical solution and yields improved representations for the knn lm, as shown in fig 1 . in particular, we see from the figure that there is an optimal value for ω for which the added regularization seems to improve the knn lm model perplexity significantly i.e. from 76 to 70 at ω = 2 (orange line). the interesting part to note in this case, is the fact that the standard lm (without post-hoc knn) does not vary much up to ω = 5, leading us to conclude that the added regularization has changed the representation in a way that knn can more effectively exploit the neighbors. however, the use of the queue and momentum target network still adds overhead to a large-scale model training as we are required to access the queue for each batch. hence we tried to decrease q and l, which interestingly did not decrease the performance at all and therefore, to promote efficiency, we tested an even simpler formulation, where we replace q with all zero vectors. this eliminates the need to explicitly construct and update the queue, while instead encouraging the model to learn conservative representations w.r.t. the l2 norms of its context representations. the corresponding loss is as follows: to our surprise, equation 2 yields similar performance to equation 1 in practice see ## experiments we tested our method on the wiki-2 and wiki-103 datasets, which are widely used benchmarks for language modeling. we are interested in demonstrating two empirical results: improved performance using our approach over that of knn-lm, and exploring a possible mechanism for this improved performance. ## conclusion in conclusion, we propose a useful training mechanism that is inspired by the fact that the post-hoc application of knn seems to significantly improve the performance of standard lms. we have found that training a lm with l2 regularization at the final layer, i.e. layer which is used for the posthoc knn search, improves validation performance. we have also found initial indications that the l2 regularization mostly improves performance for the most frequent, lower-loss words. in addition, we have found further evidence for the hypothesis proposed (dauphin and cubuk, 2021) which states that l2 regularization helps generalization in vision tasks. this paper found similar results when working with transformer models in nlp tasks. there are, however, some shortcomings in our work. firstly, we have only given a preliminary explanation for why the added l2 regularization significantly improves upon standard knn lm, but we believe that we have given sufficient evidence that our proposed method promotes clustering of the representations which subsequently improves the knn. secondly, even though we have found great and promising improvement using our findings on wiki-2, further work with more compute should be done on wiki-103. we however leave this for future work due to computational constraints. lastly, we believe that training models with post-hoc knn in mind is a promising area and hence future work will consider more diverse datasets from the nlp literature. these findings motivate exploring various regularizations in different transformer architectures and lm tasks.