You can’t pick your neighbors, or can you? When and How to Rely on Retrieval in the k NN - LM (2022)

Abstract: Retrieval-enhanced language models (LMs), which condition their predictions on text retrieved from large external datastores, have recently shown significant perplexity improvements compared to standard LMs. One such approach, the kNN-LM, interpolates any existing LM’s predictions with the output of a k-nearest neighbors model and requires no additional training. In this paper, we explore the importance of lexical and semantic matching in the context of items retrieved by kNN-LM. We find two trends: (1) the presence of large overlapping n-grams between the datastore and evaluation set plays an important factor in strong performance, even when the datastore is derived from the training data; and (2) the kNN-LM is most beneficial when retrieved items have high semantic similarity with the query. Based on our analysis, we define a new formulation of the kNN-LM that uses retrieval quality to assign the interpolation coefficient. We empirically measure the effectiveness of our approach on two English language modeling datasets, Wikitext-103 and PG-19. Our re-formulation of the kNN-LM is beneficial in both cases, and leads to nearly 4% improvement in perplexity on the Wikitext-103 test set.

## introduction recently, a new class of language models (lms) that are augmented with retrieval capabilities have led to substantial improvements over standard neural lms @xcite @xcite @xcite . furthermore, lms with retrieval warrant investigation as they provide benefits for many tasks @xcite . these approaches generally involve a backbone neural lm that interacts with a retrieval component of varying complexity to find relevant documents. in this work, we analyze and improve a specific and simple type of retrieval-enhanced language model, the knn-lm originally proposed by @xcite . the knn-lm is non-parametric -it works by retrieving instances from an external datastore at each decoding timestep, and it improves language model performance without requiring additional training. in essence, the knn-lm interpolates a base lm's predicted probability distribution of the next word with a distribution formed by retrieving vectors similar to the current hidden state. knn-lm includes two tunable hyperparameters: the number of items to retrieve (k) and an interpolation coefficient (λ). the method's effectiveness depends crucially on source and size of the retrieval datastore: it is most effective when using a very large datastore with orders of magnitude more tokens than seen in the training corpus, but @xcite also observe improvements with smaller datastores. modern neural models have massive capacity to memorize their training data @xcite . nonetheless, simply using an lm's training corpus as the source for the datastore works well for knn-lm, as test perplexity on the wikitext-103 dataset decreases substantially from 18.65 to 16.12. however, it remains unclear how and why the knn-lm achieves these improvements. which types of tokens and contexts does it improve most on? as an effort to answer this question and motivate new more effective methods to enhance lms with retrieval we analyze the knn-lm's behavior with respect to parts of speech, semantic similarity between context and retrievals, and lexical overlap. among others, our analysis reveals the knn-lm is helpful beyond factual knowledge (i.e. proper nouns), and improves perplexity across many word types, so it would be difficult to extend knn-lm using syntactic information alone. on the other hand, we find the performance of the knn-lm highly correlates with lexical similarity between the context and retrieved items, although this is somewhat domain specific and does not fully explain its strong performance. semantic similarity is nearly as accurate a predictor of knn-lm performance as lexical similarity, making it a strong candidate to extend the knn-lm. based on our analysis, we devise a simple scheme to extend the knn-lm following the intuition that when retrieval quality is high (measured by semantic similarity), then the model should rely more heavily on the knn-based prediction. since retrieval in the knn-lm is latent, we use semantic similarity as a proxy to measure retrieval relevance. concretely, our method is an adaptive version of knn-lm that assigns the interpolation coefficient according to retrieval quality (see figure 1 ). while it introduces new hyperparameters, we show that the additional hyperparameter tuning comes at negligible cost. importantly, our empirical results demonstrate that our newly introduced re-formulation of knn-lm is beneficial for both encylopedic text and book data, and leads to an improvement of nearly 4% perplexity over the the vanilla knn-lm, measured on the english language modeling wikitext-103 test set. broadly, we hope our insights and methods helps to facilitate future development of retrieval-augmented lms. ## language modeling with knn-lm the knn-lm improves over a base language model by explicitly memorizing the lm's training data. it stores exact sentences from the training data in its datastore that can be accessed during language model inference to produce a k-nearest neighbor next word distribution that is interpolated with the base model's prediction. interpolation is preferred for similar reasons as approximate matrix factorization in collaborative filtering -the universe of text patterns is sparse and lossless compression of the training data alone is not sufficient to model new patterns. in this section, we explain the specifics of the knn-lm's inner workings in order to guide our analysis. ## a new formulation for knn-lm in the previous section, we analysed when knn-lm is most helpful. we use this information to design a new formulation of knn-lm that can exploit this behavior. the original knn-lm uses the same interpolation coefficient (λ) for every example, which may not be desirable. as our analysis reveals, we can predict when the knn-lm is most beneficial, which naturally leads us to a new formulation with an adaptive λ: ## experiments and results to measure the importance of retrieval quality in the knn-lm, we evaluate our approach ( §4) on two english language modeling datasets. the first is the wikitext-103 corpus @xcite used by @xcite . the second is pg-19 @xcite , which we include because it consists of books and is thematically distinct from the encyclopedic documents in wikitext-103. ## discussion in previous sections we use observations of knn-lm to motivate our new approach that adapts the interpolation coefficient to retrieval quality. here we analyze results with our new method to see how they compare with baselines and deepen our understand of retrieval-enhanced language modeling. 6.1 can we adapt to lexical similarity? the original knn-lm has similar performance when its results are stratified by either semantic or lexical similarity ( §3.1), but in our new formulation we adaptive the coefficient only according to semantic similarity. what if we use lexical similarity instead? we explore this possible alternative and report the results for wikitext-103 in in general, we find that both semantic and lexical similarity foot_6 yield similar results when used to bucket queries. for the best setting, when @xmath0, the learned vectors work better, reflecting recent findings that dense vectors outperform sparse representations for various retrieval-related tasks @xcite . hence, throughout this paper we adapt the coefficient using semantic similarity and @xmath1. interestingly, for lower values of k the bagof-words representation has an edge over semantic similarity. perhaps this suggests lexical similarity is more precise, and if retrieving many items is costly, then adapting the coefficient according to lexical similarity might be particularly helpful. ## related work we extend the knn-lm by adapting the interpolation coefficient to retrieval quality (measured by semantic similarity). adaptret @xcite models the interpolation coefficient as a function of the query. this is convenient, since one can skip retrieval if the coefficient is below a threshold, although requires training a separate adaptor network. crucially, their coefficient predictions are based solely on query features, and does not take into account whether retrieval is successful. our approach incorporates the quality of retrieval, and improves language modeling results. it is simple and effective, and only needs lightweight hyperparameter tuning without any additional training. retomaton @xcite provides an alternative means to bypass retrieval. they build a graph over the datastore, and at each time step they either retrieve like the original knn-lm or re-use the previously retrieved neighbors to traverse the graph. this is more efficient than adaptret, providing better results at lower cost. both adaptret and retomaton are designed with efficiency in mind. they rely on approximate distance using product quantization and perform about as well as the exact distance version of the knn-lm. we improve upon knn-lm by about 4% perplexity. there are many recent works that use retrieval components for language tasks besides language modeling, such as question answering @xcite @xcite , dialogue generation @xcite , conversational search @xcite , semantic parsing @xcite , data augmentation @xcite , and machine translation @xcite @xcite . there are alternatives to knn-lm that incorporate document structure @xcite , but their experimental setup is not comparable with ours. in our baselines we only consider models matching the original knn-lm backbone, although alternative architectures show promise for retrieval-enhanced language modeling @xcite @xcite . scaling the datastore @xcite or model size @xcite have shown to effectively improve language modeling. alternatively, text generation may be improved through more advanced ranking @xcite or decoding @xcite algorithms. researchers have explored fundamental extensions to knn that are agnostic to language data. @xcite spatially partition the datastore, adapting the value of k for each region. keeping k fixed, @xcite instead adapt the shape of the neighborhood based on local information. ## conclusion in this paper, we have proposed a novel and effective re-formulation of the knn-lm. our approach adapts the interpolation coefficient to the quality of retrieved documents measured by semantic similarity. we motivate our approach through extensive analysis, which also provides insights on the types of tokens and contexts knn-lm is most helpful for. importantly, we empirically demonstrate the effectiveness of our approach through experiments on two domains, wikitext-103 (encyclopedic text) and pg-19 (book data), and outperform the original knn-lm by 4% test perplexity on the wikitext-103 language modeling corpus.