Adaptation Approaches for Nearest Neighbor Language Models (2023)

Abstract: Semi-parametric Nearest Neighbor Language Models (kNN-LMs) have produced impressive gains over purely parametric LMs, by leveraging large-scale neighborhood retrieval over external memory datastores. However, there has been little investigation into adapting such models for new domains. This work attempts to fill that gap and suggests the following approaches for adapting kNN-LMs — 1) adapting the underlying LM (using Adapters), 2) expanding neighborhood retrieval over an additional adaptation datastore, and 3) adapting the weights (scores) of retrieved neighbors using a learned Rescorer module. We study each adaptation strategy separately, as well as the combined performance improvement through ablation experiments and an extensive set of evaluations run over seven adaptation domains. Our combined adaptation approach consistently outperforms purely parametric adaptation and zero-shot (kNN-LM) baselines that construct datastores from the adaptation data. On average, we see perplexity improvements of 17.1% and 16% for these respective baselines, across domains.

## introduction natural language processing (nlp) has observed large performance improvements with recent advancements in neural language models (lms). these models have enabled learning rich, semantic text representations @xcite that have facilitated a wide range of downstream language tasks @xcite . for the task of next-word prediction, parametric lms utilize the rich contextual text representations as input to a classifier (output layer), which produces a distribution over the possible next words. in contrast to parametric lms, k-nearest neighbor lms (knn-lms) are semi-parametric models * work done during an internship at aws ai labs. figure 1 : an illustration of the proposed knn-lm adaptation approach. the current context is used as a query (q) for nearest-neighbor retrieval. the context is passed through the lm to obtain the query vector representation ⃗ c q , which is then used to retrieve nearest neighbors from a large pretraining datastore and a smaller adaptation datastore (displayed in pink and yellow, respectively). the function f (•) represents merging of datastores (merge), followed by rescoring (rank) of the retrieved neighbors to obtain p knn . the probability distribution over the candidate next words is computed by the mixture of probabilities p knn and p lm , where p lm denotes probabilities obtained from domainadapted lm. we posit that large, general-purpose datastores, referred to as the pretraining datastore, contain a significant amount of relevant information which can still be applied to specialized domains. this information can be leveraged through nearestneighbor retrieval and should prove especially useful in situations where there is an insufficient amount of domain-specific data to generate an effective standalone datastore. unlike parametric neural architectures which can employ gradient-based finetuning for domain adaptation, it is less obvious how to adapt knn-lms primarily because of the non-parametric nature of datastores. one simple approach would be to reconstruct the datastore, using domainadapted lm representations. however, this comes at the cost of incurring a large memory footprint for each adaptation domain. in this work, we instead choose to focus on adaptation strategies that are parameter and memory efficient. given the complementary nature of the parametric and nonparametric components in a knn-lm, we pursue adaptation strategies separately for each component and analyze their impact on the knn-lm system's adaptation performance. in line with previous works, we focus our experiments solely on the core language modeling task of next-word prediction @xcite . results on seven adaptation domains ranging from science and books, to conversational text, demonstrate that our componentlevel strategies consistently improve over respective parametric and semi-parametric baselines, and produce even better results when combined together. specifically, we find that adaptation of the parametric component increases recall of groundtruth labels found in the retrieved neighbors. we also confirm that the large-scale pretraining datastore contains relevant information for adaptation domains, via its performance edge over models that exclude it. finally, we observe that expanding the nearest neighbor search to include elements from the adaptation datastore contributes to the best overall performing strategy. figure 1 demonstrates the overall approach using wikipedia and us news as example pretraining and adaptation domains, respectively. ## knn-lms for a context c t defined by the sequence of words (w 1 , . . . , w t-1 ), the causal language modeling task aims to model a probability distribution over the next word foot_0 w t . let p lm (w t |c t ) and p knn (w t |c t ) be the probability masses computed by the lm and knn components, respectively. details on how p knn (w t |c t ) is computed and combined with p lm (w t |c t ) to produce the final knn-lm predictions, are outlined in the following sections. datastore creation: given a source domain training set x s , let c @xmath0, . . . , w t-1 ) be a sequence in x s . the datastore is defined as a set of d s tuples {(⃗ c i , w i )} ds @xmath1, where the key ⃗ c @xmath2, produced by the lm and value w i denotes the next word label in the sequence. thus, the knn probability of the next word can be obtained via: unifying knn and lm: the probability distribution of the knn-lm system can be obtained by interpolating the component probabilities since each probability distribution lies on a simplex spanning the token vocabulary, performing a convex combination of the two maintains a valid probability distribution. ## related work our proposed work investigates the intersection of techniques used for parametric language model adaptation with semi-parametric systems (knn-lms). therefore we discuss the related works in each of these areas and contrast our respective contributions. ## conclusion we proposed a multi-pronged strategy for adapting knn-lm systems. through our studies, we demonstrated that a general-purpose pretraining datastore contains relevant information, which can be utilized for downstream domains. we showed that parametric and non-parametric adaptation methods complement each other and that using the complete semi-parametric adaptation strategy outperforms adapting just one of the knn-lm components. our methods could further be extended by noting that the recall of retrieved neighborhoods is often imperfect. thus, a gate could be learned to predict whether knn retrieval should be triggered. while our study focused on the language modeling task, our approach could be applied towards other nlp tasks such as text generation and translation.