Efficient Nearest Neighbor Language Models

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 5703–5714 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 5703 Efficient Nearest Neighbor Language Models Junxian He†, Graham Neubig†, Taylor Berg-Kirkpatrick‡ †Language Technologies Institute, Carnegie Mellon University ‡Department of Computer Science and Engineering, University of California San Diego {junxianh, gneubig}@cs.cmu.edu, tberg@eng.ucsd.edu 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., 2019) 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 de- velop or deploy more efficient non-parametric NLMs.1 1 Introduction Language models (LMs) are one of the most fundamental technologies in NLP, with applica- tions spanning text generation (Bahdanau et al., 2015; Rush et al., 2015), representation learn- ing (Peters et al., 2018; Devlin et al., 2019; Yang et al., 2019), and few-shot learning (Radford et al., 2019; Brown et al., 2020). Modern neural lan- guage models (NLMs) based on recurrent (Mikolov et al., 2010; Sundermeyer et al., 2012) or self- attentional (Vaswani et al., 2017; Al-Rfou et al., 2019) neural networks are mostly parametric, where the predictions are solely dependent on the model parameters given the input data. In contrast, recent non-parametric LMs (Guu et al., 2018; Khandelwal et al., 2019; He et al., 1Code is available at https://github.com/jxhe/efficient- knnlm. 300 500 1000 2000 3000 tokens/s 16 17 18 19 ppl x x x NLM Vanilla kNNLM Datastore Compression Adaptive Retrieval Dimension Reduction Combined Figure 1: Perplexity (ppl) and evaluation speed for different models. Circled points represent the neural language model (NLM) and k-nearest neighbors language model (kNN-LM ) baselines respectively, while others are the methods that we propose and explore in this paper. 2020) model text distributions by referencing both the parameters of the underlying model and exam- ples from an external datastore. Non-parametric LMs are appealing since they allow for effective language modeling – particularly for rarer patterns – through explicit memorization via a datastore, which mitigates the burden on model parameters to learn to encode all information from a large dataset. One effective and representative example is the k-nearest neighbors LM (kNN-LM, Khandelwal et al. (2019)). The kNN-LM computes the probabil- ity of the next token by interpolating a parametric LM with a distribution calculated from the k near- est context-token pairs in the datastore, as demon- strated in Figure 2. This model is particularly no- table for its large improvements in performance – it outperforms …

Content truncated to 500 words