PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Nearest Neighbor Language Models

Junxian He, Graham Neubig, Taylor Berg-Kirkpatrick

arXiv:2109.04212v3cs.CL

TL;DR

Large datastores make non-parametric language models costly at inference, constraining practical deployment. The paper explores several efficiency strategies for kNN-LM and reports up to 6x faster inference with comparable performance.

  • Problem

    Large datastore retrieval adds substantial test-time overhead to non-parametric language models, limiting their practical deployment.

    Source: §1 · p. 1

  • Method

    The paper explores adaptive retrieval, datastore pruning, and dimension reduction to improve kNN-LM efficiency.

    Source: §1 · p. 2

  • Results

    Up to 6x speed-up in inference is reported with comparable performance across WikiText-103 and training-free domain adaptation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper’s empirical results and analysis may help researchers understand the speed-performance tradeoff in non-parametric language models.

    Source: §1 · p. 2

    Exact-distance recomputation is impractical at scale: storing full-precision vectors in memory is not scalable, and disk reads are too slow for practical use.

    Source: §2 · p. 3

Abstract

from arXiv · show

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.

1 Introduction

Large external datastores make effective non-parametric language models costly at test time, limiting their practical deployment. The paper explores adaptive retrieval, datastore pruning, and dimension reduction, reporting speedups with comparable perplexity.

1 Introduction

  • kNN-LM interpolates a parametric language model with a next-token distribution computed from the nearest context-token pairs in an external datastore.The datastore supports explicit memorization, which can be particularly effective for rarer patterns.Source: §1 · p. 1
  • Large kNN-LM datastores add test-time retrieval overhead; a 100M-entry datastore can cause over 10x slowdown versus parametric models.This overhead is described as a hurdle to practical deployment despite non-parametric language models’ effectiveness.Source: §1 · pp. 1–2
  • Up to 6x faster inference with comparable perplexity is reported for the proposed methods on WikiText-103 and training-free domain adaptation.The empirical analysis is intended to help researchers understand the speed-performance tradeoff in non-parametric language models.Source: §1 · p. 2
  • The paper explores a learned gate for skipping retrieval, datastore pruning, and vector dimension reduction to improve kNN-LM efficiency.The stated pruning approaches include clustering, importance-guided filtering, and greedy merging.Source: §1 · p. 2

2 k-Nearest Neighbors Language Model

kNN-LM combines a pretrained neural language model with a distribution derived from nearest datastore contexts. Its retrieval overhead grows with datastore size, while storing dense context vectors can also require substantial memory.

2 k-Nearest Neighbors Language Model

  • kNN-LM interpolates a pretrained autoregressive neural language model’s next-token distribution with a distribution computed from an external datastore.The model requires no additional training; its parameters remain unchanged and interpolation is applied at test time.Source: §2 · p. 2
  • A 100M-token datastore with 1024-dimensional context vectors at 16-bit precision could require 200GB of memory.The datastore can approach the number of training tokens because contexts are nearly unique, and its dense vectors are difficult to fit in memory.Source: §2 · pp. 2–3
  • At inference, kNN-LM encodes the current context, retrieves its k nearest datastore neighbors, and aggregates their tokens into a distribution for interpolation.The retrieved contexts’ following tokens provide explicit memorized examples that may help with patterns rarely seen by the pretrained model.Source: §2 · p. 3
  • Retrieval time grows with datastore size, whereas the neural language model’s per-token inference time remains constant across datasets.The paper identifies kNN search as the source of the additional inference overhead.Source: §2 · p. 3
  • The study treats its approximate-nearest-neighbor index as a black box and aims to improve efficiency independently of index choice.Exact-distance recomputation is avoided: keeping full-precision vectors in memory is not scalable, and reading them from disk on demand is too slow to be practical.Source: §2 · p. 3

3 The Efficiency of kNN-LM

The section evaluates kNN-LM in standard and domain-adaptation settings, specifying datasets, model configuration, and evaluation hardware. Its baseline comparison shows substantial perplexity gains alongside a large inference-speed penalty from retrieval.

3.1 Datasets

  • Disk random I/O remains an aspect that may benefit from further engineering.The authors identify it as future work.Source: §3.1 · p. 3
  • The standard setting uses WikiText-103’s 103M training tokens for a 103M-record datastore, while Law-MT supplies 19M domain-training tokens for adaptation.WikiText-103 uses a pretrained model trained on its training split; Law-MT uses English text and a model trained on WMT News Crawl.Source: §3.1 · pp. 3–4

3.2 Setup

  • Evaluation uses 1,024 neighbors, with 1,024-dimensional WikiText-103 and 1,536-dimensional Law-MT context vectors; the interpolation weight is tuned on validation data.The experiments use the original kNN-LM hyperparameters and index; the authors observed no speed effect from changing k between 8 and 1,024.Source: §3.2 · p. 4
  • Speed tests use 32 AMD EPYC 7282 CPU cores for index search and one NVIDIA RTX 3090 GPU for neural-network computation.The 32-core retrieval setup is also used by the FAISS repository for benchmarking large-scale retrieval.Source: §3.2 · p. 4

3.3 Baseline Speed

  • kNN-LM gains 2 perplexity points on WikiText-103 and over 90 on Law-MT, but is 10x–30x slower than parametric models.The baseline table reports evaluation performance and speed; exact-distance retrieval from disk takes over one second per token.Source: Table 1 · p. 4, §3.3 · p. 4

4 The Remedies

The section explores three kNN-LM efficiency remedies: adaptive retrieval, datastore pruning, and vector-dimension reduction.

4.1 Adaptive Retrieval

  • The retrieval adaptor learns context-dependent interpolation weights and removes retrievals by setting the lowest-weight cases to zero.It is trained to maximize kNN-LM log likelihood with an L1 regularizer that encourages sparse datastore weights.Source: §4.1 · p. 5
  • At 50% retrieval removal, adaptive retrieval (AR) matches original kNN-LM perplexity (16.67 vs. 16.65) at nearly 2x speed (530 vs. 277 tokens/s).AR’s gains come from both its pruning mask and optimized interpolation weight λ.Source: Figure 3 · p. 5, §4.1 · p. 6

4.2 Datastore Pruning

  • Pruning methods include random selection, target-aware k-means, rank-based filtering, and local greedy merging of nearby records sharing a target token.Target-aware k-means clusters context vectors separately by target word; greedy merging avoids propagating merges through neighbors-of-neighbors.Source: §4.2 · pp. 6–7
  • Greedy merging has a relatively flat perplexity–speed curve, while random pruning surprisingly outperforms k-means and rank-based pruning.At 20% of the original datastore, k-means can underperform the vanilla NLM baseline, which the authors suggest may reflect centroid distance errors.Source: §4.2 · p. 7, Figure 4 · p. 7

4.3 Dimension Reduction

  • At 256 and 512 dimensions, PCA gives comparable or better performance than original 1024-dimensional vectors with 3x–4x speed-up.Dimensions below 256 incur large perplexity costs; below 128, PCA loses its advantage over the NLM baseline.Source: Figure 5 · p. 7, §4.3 · p. 8

5 Putting it All Together

Combining adaptive retrieval, datastore pruning, and dimension reduction produced faster evaluation on WikiText-103 and Law-MT while retaining comparable or better perplexity. The techniques were selected using validation-set speed and perplexity trade-offs.

5 Putting it All Together

  • The retrieval-pruning rate, datastore-pruning rate, and reduced dimension were selected on validation data to maximize speed-up while keeping perplexity within 0.1 of vanilla kNN-LM.The retrieval rate was selected on its own held-out validation data before combination with the other settings.Source: §5 · p. 8
  • On WikiText-103, adaptive retrieval removed 50% of retrieval for nearly 2x speed-up, greedy merging pruned 40% of the datastore at a 0.2-perplexity-point cost, and PCA was 3.6x faster with a minor perplexity improvement.These are individual-strategy results reported alongside the combined approach.Source: §5 · p. 8
  • WikiText-103 combination achieved comparable perplexity (16.67 v.s. 16.65) and a 6.6x speed-up (1835 v.s. 277); Law-MT evaluation was 5.4x faster with superior perplexity.Table 3 reports test-set perplexity and speed for both datasets and the individual strategies.Source: Table 3 · p. 8, §5 · p. 8
  • On Law-MT, adaptive retrieval removed 10% of retrieval at a 0.1-perplexity-point cost, greedy merging pruned 40% of the datastore at a 0.7-point cost, and PCA gained 1 perplexity point while achieving 3.3x speed-up.The authors suggest the smaller 19M-entry Law-MT datastore may have less redundancy than WikiText-103's 103M-entry datastore.Source: §5 · p. 8
  • The Law-MT oracle comparison found p_kNN(w|c) ≥ p_NLM(w|c) 76% of the time, versus 39% on WikiText-103.The authors relate the larger datastore-retrieval gains in Law-MT to the difficulty of adaptively removing retrieval operations.Source: §5 · p. 8

6 Implications and Future Work

The paper explores efficiency improvements for kNN-LM and reports up to 6x speed-up with comparable performance. It identifies datastore-side retrieval cues and optimized indexing as directions for further work.

6 Implications and Future Work

  • The explored kNN-LM efficiency methods achieved up to 6x speed-up while retaining comparable performance.The paper suggests combining its approaches with more optimized indexing methods could further reduce the gap between retrieval-based and parametric NLMs.Source: §6 · p. 9
  • Future work could use datastore-side features to better determine when retrieval is needed.The paper presents this as a direction for further exploration.Source: §6 · p. 9
  • The paper identifies combining its approaches with more optimized indexing methods as a possible way to further reduce the gap between retrieval-based and parametric NLMs.Source: §6 · p. 9

A.1 General Setup Details

The interpolation weight λ was tuned separately on each dataset's validation split, with a 0.05 search interval.

A.1 General Setup Details

  • The interpolation weight λ was tuned on the validation split of each dataset separately.Source: §A.1 · p. 11
  • The λ search used an interval of 0.05.Source: §A.1 · p. 11
  • The selected λ values were 0.25 for WikiText-103 and 0.9 for Law-MT.Source: §A.1 · p. 11

A.2 Adaptive Retreival

The adaptive retrieval configuration used a feature-based MLP to produce interpolation weights, with its hyperparameters validated on WikiText-103.

A.2 Adaptive Retreival

  • The retrieval adaptor was a four-hidden-layer MLP with 128 hidden units per layer and dropout probability 0.2, except at the output layer.Its output layer mapped the hidden representation to two dimensions and used log softmax to produce log λ and log(1 − λ).Source: §A.2 · p. 11
  • Before concatenation, scalar input features other than f(c) were mapped to m-dimensional vectors, where m = dim(f(c))/n.The paper used this transformation to balance the context-vector feature against the other features.Source: §A.2 · p. 11
  • The L1 regularizer coefficient was fixed at 0.05 for both datasets, and the adaptor was trained with Adam at learning rate 0.0005.The authors report that an LSTM alternative was unstable and failed to converge; the selected checkpoint had the best validation perplexity at 50% pruning.Source: §A.2 · p. 11

B Ablation Analysis

The ablations find that perplexity is robust to removing individual retrieval-adaptor input features, and that learned pruning masks account for most of the gains over random pruning. The experiments drop the log-frequency feature to save memory while retaining comparable perplexity.

B Ablation Analysis

  • At 50% retrieval pruning, perplexity is relatively robust to removing a single input feature, and feature choice has only a marginal effect on evaluation speed.All features together produce the best perplexity.Source: B Ablation Analysis · p. 11, Table 4 · p. 11
  • Dropping the log-frequency feature saves memory while achieving comparable perplexity to using all features.Source: B Ablation Analysis · p. 11
  • Learned mask, learned weights performs best in Figure 6, but most of the advantage comes from learned masks, with only minor gains from learned weights.The figure compares perplexity and speed across adaptive-retrieval variants for ablation analysis.Source: B Ablation Analysis · pp. 11–12, Figure 6 · p. 11
Loading 2109.04212v3…