Efficient Nearest Neighbor Language Models

By Junxian He, Graham Neubig, Taylor Berg-Kirkpatrick. 2021

124 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Natural Language Processing, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Natural Language Processing

In Efficiency & Compression

Citations in the corpus

Language Generation Models Can Cause Harm: So What Can We Do About It? An Actionable Survey (2022, background): “The first set of works treat the task of rewriting as a sequence labeling task, where each token in the output sequence is either substituted, deleted, or kept the same (Pryzant et al., 2020; He et al., 2021b).”
The Efficiency Misnomer (2021, background): “…Speed is one the most informative indicator for comparing the efﬁciency of different models (So et al., 2021; Dosovitskiy et al., 2020; Arnab et al., 2021; Tay et al., 2021b; Kim et al., 2021; He et al., 2021a; Narayanan et al., 2021; Lagunas et al., 2021; Liu et al., 2021; Tay et al., 2020c;b).”
Retrieval-based Language Models and Applications (2023, methodology): “We also plan to cover techniques for refining data stores and improving similarity search (He et al., 2021; Alon et al., 2022).”
When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories (2022, background): “Prior work investigates the effectiveness of deciding when to use non-parametric memories at the token level in k NN LM (He et al., 2021; Drozdov et al., 2022).”
SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore (2023, background): “…work can improve k NN-LM, e.g., by training the model to output a nonparametric distribution (Zhong et al., 2022; Lan et al., 2023; Min et al., 2023), by having a vocabulary-specific λ (Huang et al., 2023b), or by modeling λ as a function of the input x (He et al., 2021; Drozdov et al., 2022).”
Retrieval-Augmented Generation for AI-Generated Content: A Survey (2024, unknown): “Efficient-KNNLM [38] combines the generation probability of KNN-LM [37] and NPM [162] with a hyperparameter λ to determine the proportion of generation and retrieval.”
Training Language Models with Memory Augmentation (2022, background): “Potential solutions to this issue include (1) constructing the memory using a coarser granularity (e.g., text blocks) (Borgeaud et al., 2021); (2) compressing the external memory set and reducing the dimension of memory representations (He et al., 2021).”
In-Context Retrieval-Augmented Language Models (2023, methodology): “The first LM dataset is WikiText-103 (Merity et al., 2016), which has been extensively used to evaluate RALMs (Khandelwal et al., 2020; He et al., 2021; Borgeaud et al., 2022; Alon et al., 2022; Zhong et al., 2022).”
A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models (2024, unknown): “In this part, we will briefly review studies with these two types of generators: parameter-accessible (white-box) and parameter-inaccessible (black-box).”