Efficient Nearest Neighbor Language Models

TL;DR

This work addresses the practical inefficiency of non-parametric, datastore-backed language models by proposing a three-pronged efficiency framework: adaptive retrieval to skip unnecessary datastore queries, datastore pruning to aggressively reduce storage, and dimension reduction to speed up vector comparisons. Through comprehensive experiments on WikiText-103 and Law-MT, it demonstrates that careful combination of these techniques can yield up to 6x speed-ups while maintaining perplexity on par with standard kNN-LM baselines. The findings provide concrete guidance for deploying non-parametric LMs in real-world settings and offer a set of scalable techniques for future work in efficient retrieval-based language modeling. Overall, the paper clarifies the speed–performance trade-offs in non-parametric LM deployments and shows practical paths to bridge the gap with parametric models.

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 (Khandelwal et al., 2020) 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.

Figures (5)

• Figure 1: Perplexity (ppl) and evaluation speed for different models. Circled points represent the neural language model (NLM) and $k$-nearest neighbors language model ($k$NN-LM ) baselines respectively, while others are the methods that we propose and explore in this paper.
• Figure 2: Perplexity and speed results of adaptive retrieval on WikiText-103 test set. We annotate the coordinates of some points and the third number in the annotation is the fraction of retrieval operations that are removed.
• Figure 3: Perplexity and speed results of datastore pruning methods on WikiText-103 validation set. We annotate the coordinates of some points and the third number in the annotation is the compression rate (fraction of records remained).
• Figure 4: PCA dimension reduction results on WikiText-103 validation set. We annotate the coordinates of the points and the third number in the annotation is the PCA dimensions.
• Figure 5: Perplexity and speed results of adaptive retrieval on WikiText-103 test set. This figure includes different variants of adaptive retrieval for ablation analysis.