Efficient Nearest Neighbor Language Models (2021)

Abstract: Non-parametric neural language models (NLMs) learn predictive distributions of text utilizing an external datastore, which allows them to learn through explicitly memorizing the training datapoints. While effective, these models often require retrieval from a large datastore at test time, significantly increasing the inference overhead and thus limiting the deployment of non-parametric NLMs in practical applications. In this paper, we take the recently proposed k-nearest neighbors language model as an example, exploring methods to improve its efficiency along various dimensions. Experiments on the standard WikiText-103 benchmark and domain-adaptation datasets show that our methods are able to achieve up to a 6x speed-up in inference speed while retaining comparable performance. The empirical analysis we present may provide guidelines for future research seeking to develop or deploy more efficient non-parametric NLMs.

## introduction language models (lms) are one of the most fundamental technologies in nlp, with applications spanning text generation @xcite , representation learning @xcite @xcite , and few-shot learning @xcite . modern neural language models (nlms) based on recurrent @xcite or selfattentional @xcite neural networks are mostly parametric, where the predictions are solely dependent on the model parameters given the input data. ## k-nearest neighbors language model in this section, we overview knn-lm @xcite and its implementation details. formulation. knn-lm @xcite is an lm that estimates token distributions by interpolating a pre-trained autoregressive nlm's distribution with another distribution computed using an external datastore. specifically, given a sequence of context tokens c @xmath0, • • • , w t-1 ), the knn-lm's next token probability p(w t |c t ) is calculated through the interpolation of the probability estimated by a standard parametric nlm p nlm (w t |c t ) and a probability computed using an external datastore p kn n (w t |c t ) (detailed later): the size of such a datastore is almost equal to the number of training tokens because the context c t is (nearly) unique due to the large context window size in modern recurrent @xcite or self-attentional @xcite nlms. this suggests that the datastore can easily scale to hundreds of millions or even billions of records. also, each f (c t ) is a highdimensional dense vector, which makes the datastore difficult to fit in memory. for example, a datastore from a 100m-token training dataset, using 1024-dimension context vectors at 16-bit precision, could require 200gb of memory. 3 the nearest neighbor distribution p knn (w t |c t ). at inference time, the knn-lm (1) computes the context vector f (c) from the current sequence using the pretrained nlm, (2) uses f (c) as the query to retrieve k nearest neighbors @xmath1, v i )|@xmath2, • • • , k} from the datastore, and (3) aggregates the retrieved tokens to form the distribution p knn (w|c) to be used in eq. 1 as: sources of inference overhead. the extra inference overhead stems from the knn search process in p knn (w t |c t ) computation. we denote the inference time per token as @xmath3. while t nlm remains constant with different datasets, t knn unfortunately grows as the datastore scales. in practice, the knn search process is often performed only approximately (ann, @xcite ; @xcite ) to reduce computational cost. @xcite implemented ann search in knn-lm foot_3 using faiss @xcite , which combines inverted file systems @xcite and product vector quantization @xcite . this type of index reduces memory usage by only storing quantized vectors and accelerates knn search by pre-clustering the datastore vectors; interested readers can refer to @xcite for more details. for the purpose of this paper we study knn-lm using this indexing method as a black box, aiming to improve efficiency in an index-agnostic way. at the same time, we note that building fast and accurate indexing methods remains an active area of research @xcite , and selection or improvement of the index itself (possibly in concert with the methods proposed in this paper) is an interesting avenue for future work. distance recomputation. the distances to the nearest neighbors are required to compute p knn (w t |c t ) as shown in eq. 3. however, as described above, knn-lm's nearest neighbor search process performs search over quantized vectors, and as a result it can only return approximate distances. while it is possible to compute the accurate distances by reading the full-precision vectors from the datastore after retrieval, this presents challenges as well: (1) storing the entire datastore in memory is not scalable for large datastores, (2) reading the vectors from a large datastore on disk on-the-fly is too slow to be practical (< 1 token per second). 5 therefore, in this paper we use the approximate distances directly to compute p knn . this comes at the cost of a minor performance loss, as we will show in §3.3. similar approximations were adopted to apply knn-lm to machine translation tasks @xcite . ## the efficiency of knn-lm in this section, we first introduce the datasets and setup that we will use throughout the paper, and then compare the inference speed of knn-lm to parametric nlms. ## the remedies in this section we propose and explore several different methods that may potentially improve the efficiency of knn-lm along three axes: (1) adaptive retrieval, (2) datastore pruning, and (3) dimension reduction. we analyze the performance of each method on wikitext-103, trying to conclude the best practices that we will evaluate in §5. ## putting it all together based on the analysis results in §4, in this section we combine best practices in adaptive retrieval, datastore pruning, and dimension reduction to assess the performance. we select the retrieval pruning rate r, datastore pruning rate n, and the reduced dimensions d on the validation set, 12 so that they achieve the largest speed-up at the cost @xmath4.1 perplexity compared to vanilla knn-lm . we report the results on the test set. results: different from wikitext-103 where the datastore is contructed from the data that trains the pretrained nlm, in the law-mt domain adaptation setting the datastore represents the domain-specific knowledge that the pretrained nlm never sees during 11 the tool we use for pca, the faiss pca implementation, applies random rotation to the pca output vectors by default to re-balance variances of components of a vector @xcite , which may provide additional benefits over vanilla pca on product vector quantization inside the index. the best. while minor gains are from the automatically learned weights ("random mask, learned weights"), most of the superiority can be attained with the smart pruning strategy even with constant weights ("learned mask, constant weights").