PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalization through Memorization: Nearest Neighbor Language Models

Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, Mike Lewis

arXiv:1911.00172v2cs.CL

TL;DR

Neural language modeling must both represent contexts and predict their next words, but the paper asks whether representation learning is easier than prediction. It augments a pre-trained LM with nearest-neighbor retrieval from a configurable datastore, achieving stronger WIKITEXT-103 performance without additional training and supporting larger-data use and domain adaptation.

  • Problem

    The paper investigates whether learning prefix representations is easier than predicting the next word, a distinction important for improving language-model performance.

  • Method

    kNN-LM linearly interpolates a pre-trained LM’s next-word distribution with a distance-based nearest-neighbor distribution retrieved from a key-value datastore.

  • Results

    15.79 WIKITEXT-103 perplexity is reported as a new state of the art, a 2.86-point improvement over the base model with no additional training.

  • Takeaways & Limitations

    The results suggest that learning similarity between contexts may be easier than next-word prediction, while explicit retrieval helps access long-tail patterns and larger text collections.

  • Takeaways & Limitations

    The approach adds datastore construction and retrieval costs that grow linearly with the number of entries, although retrieval is parallelizable and requires no GPU-based training.

Abstract

from arXiv · show

We introduce $k$NN-LMs, which extend a pre-trained neural language model (LM) by linearly interpolating it with a $k$-nearest neighbors ($k$NN) model. The nearest neighbors are computed according to distance in the pre-trained LM embedding space, and can be drawn from any text collection, including the original LM training data. Applying this augmentation to a strong Wikitext-103 LM, with neighbors drawn from the original training set, our $k$NN-LM achieves a new state-of-the-art perplexity of 15.79 - a 2.9 point improvement with no additional training. We also show that this approach has implications for efficiently scaling up to larger training sets and allows for effective domain adaptation, by simply varying the nearest neighbor datastore, again without further training. Qualitatively, the model is particularly helpful in predicting rare patterns, such as factual knowledge. Together, these results strongly suggest that learning similarity between sequences of text is easier than predicting the next word, and that nearest neighbor search is an effective approach for language modeling in the long tail.

1 INTRODUCTION

The paper proposes augmenting neural language models with nearest-neighbor retrieval because prefix representations may be easier to learn than next-word predictions. On WIKITEXT-103, this approach improves perplexity without additional training and supports larger-data scaling and domain adaptation.

  • Neural LMs both encode sentence prefixes into fixed-sized representations and predict the next word from those representations.
  • kNN-LM linearly interpolates a pre-trained LM’s next-word distribution with a nearest-neighbor model computed in its embedding space.Neighbors can be drawn from any text collection, including the original training data.
  • 15.79 perplexity on WIKITEXT-103 sets a new state of the art and improves the base model by 2.86 points with no additional training.
  • Using nearest-neighbor retrieval over a 3-billion-token dataset can outperform training the same model on all 3 billion tokens.A 100-million-token model with retrieval over the larger corpus is described as outperforming the model trained on 3 billion tokens.
  • Varying the nearest-neighbor datastore enables domain adaptation without further training, while helping with long-tail patterns such as factual knowledge.

2 NEAREST NEIGHBOR LANGUAGE MODELING

kNN-LM augments a frozen language model with a key-value datastore of context representations and target words. At inference, it retrieves nearby contexts, forms a distance-weighted target distribution, and interpolates that distribution with the model’s output.

  • The method stores context-target pairs from a text collection in a key-value datastore using one forward pass and does not change the pre-trained LM representations.The collection may include the original LM training set.
  • Each datastore key is the pre-trained LM’s fixed-length representation of a context, and each value is its target word.
  • At inference, the model encodes the test context, retrieves k nearest keys by distance, and aggregates distance-weighted probability mass over their target words.The experiments use squared L2 distance and a softmax over negative distances.
  • The final distribution interpolates the nearest-neighbor distribution with the original LM distribution using a tuned λ parameter.
  • FAISS accelerates retrieval in the potentially billion-entry datastore by clustering keys and storing compressed vectors.

3 EXPERIMENTAL SETUP

The experiments use several English corpora and decoder-only Transformer LMs, evaluate perplexity on held-out data, and build kNN datastores from training text without additional training.

  • Data: Experiments cover WIKITEXT-103, BOOKS, WIKI-3B, and the WIKI-100M subset, with complete books or articles held out for validation and testing where specified.WIKITEXT-103 contains 103M training tokens; BOOKS contains 0.7B tokens; WIKI-3B contains about 2.87B tokens; WIKI-100M is a random 100M-token subset.
  • Data: Except for WIKITEXT-103, corpora use BERT's 29K-token byte-pair encoding vocabulary.
  • Model Architecture: The model uses decoder-only Transformers and retains the underlying pretrained LM architecture and optimization.The kNN-LM is compatible with any model producing fixed-size context representations.
  • Evaluation: Evaluation minimizes training-corpus negative log-likelihood and reports perplexity on held-out data, scoring 512 tokens per example with extra prior context.Extra prior context is up to 2560 tokens for WIKITEXT-103 and up to 512 tokens for the other corpora.
  • kNN-LM Setup: A datastore is constructed by one forward pass over training text, storing pretrained-LM context representations as keys and target tokens as values.The keys used in the kNN-LM are 1024-dimensional final-layer representations, and no additional training is performed.
  • Implementation: Experiments retrieve 1024 neighbors with FAISS, using clustered and compressed keys to make high-dimensional nearest-neighbor search efficient.The index learns 4096 cluster centroids from 1M sampled keys and searches 32 centroids during inference.
  • Computational Cost: Building the WIKITEXT-103 cache takes roughly two hours on one CPU, while validation retrieval takes approximately 25 minutes and scales linearly with datastore size.The overhead is parallelizable and does not require GPU-based training.

4 EXPERIMENTS

Experiments show that kNN retrieval can improve language modeling using the training data, scale performance beyond a 100M-token model, and support domain adaptation without additional training.

  • 4.1 USING THE TRAINING DATA AS THE DATASTORE: 15.79 perplexity is achieved on WIKITEXT-103 after combining kNN-LM with a continuous cache, improving 2.86 points over the base model.Using the training data as the datastore alone improves perplexity from 18.65 to 16.12.
  • 4.2 MORE DATA WITHOUT TRAINING: 13.73 perplexity results when a 100M-token model retrieves from WIKI-3B, outperforming the 15.17 perplexity of a model trained on all 3B tokens.The comparison suggests retrieval can exploit a larger corpus without training the LM on that corpus.
  • 4.2 MORE DATA WITHOUT TRAINING: 1.6B datastore tokens already outperform a model trained on all 3B tokens, while performance remains unsaturated at about 3B datastore tokens.The model increasingly relies on the kNN component as datastore size grows.
  • 4.2 MORE DATA WITHOUT TRAINING: Adding a BOOKS datastore to a WIKI-3B model reduces perplexity from 34.84 to 20.47, enabling one model to serve multiple domains.The in-domain BOOKS model has perplexity 11.89, so datastore adaptation approaches but does not match in-domain training.

5 TUNING NEAREST NEIGHBOR SEARCH

The tuning experiments examine representation choices, neighbor counts, interpolation weights, and distance precision, finding that each design choice affects kNN-LM performance and adaptation.

  • Key Function: The input to the final layer’s feedforward network gives the largest improvement among tested representation functions, with normalized representations performing better.The representation function f(c) supplies keys and queries for similarity search.
  • Number of Neighbors per Query: More retrieved neighbors monotonically improve performance, and k = 8 is already sufficient to achieve a new state of the art.The results suggest that still larger values of k may yield further improvements.
  • Interpolation Parameter: λ = 0.25 is optimal on WIKITEXT-103, whereas λ = 0.65 works best for domain adaptation.λ interpolates between the base model distribution and the distribution obtained from datastore search.
  • Precision of Similarity Function: 16.06 perplexity is obtained with squared L2 distances and full-precision keys, improving over 16.5 perplexity from quantized-key L2 retrieval.The comparison is reported for WIKITEXT-103.

6 ANALYSIS

The analysis finds that kNN-LM is especially useful for rare, long-tail patterns because learned contextual similarity identifies relevant examples better than local n-grams or implicit parameter memorization alone.

  • Qualitative Analysis: Retrieved neighbors can assign most probability to a particularly relevant context even when several examples share local n-gram matches.This provides more value than an n-gram language model in the examined cases.
  • Qualitative Analysis: kNN-LM is most helpful on rare patterns, including factual knowledge, names, and near-duplicate training sentences.These cases appear easier to handle by assigning similar representations than by implicitly memorizing next words in model parameters.
  • Simple vs Neural Representation: kNN-LM improves on n-gram interpolation because learned representations measure similarity across more varied contexts.Interpolating an n-gram model provides only a 0.2-point perplexity improvement.
  • Implicit vs Explicit Memory: The Transformer can memorize the training set, but parameter memorization does not replace nearest-neighbor retrieval.Removing dropout drives training loss to zero, while kNN-LM improves validation perplexity by 1.9 points versus only 0.1 for the memorizing LM.
  • Implicit vs Explicit Memory: The authors conjecture that kNN-LM retains an effective context-similarity function while storing training examples explicitly.Their explanation separates representation learning from memorization in the model parameters.

7 RELATED WORK

Related work includes nearest-neighbor methods across NLP and vision, retrieval-based language generation, and data scaling; kNN-LM combines learned representations with explicit datastore retrieval.

  • Nearest-Neighbor Methods: Nearest-neighbor methods have been used for tasks such as part-of-speech tagging, morphological analysis, and computer vision.The paper distinguishes its approach by using learned representations as the similarity function for neural models.
  • Retrieval at Test Time: Prior language-generation systems use training examples at test time through sampling, editing, attention, or instance refinement rather than kNN retrieval.These examples span generation, translation, and dialogue response modeling.
  • Scaling Language Models: kNN-LM offers an alternative scaling strategy in which relatively small models learn context representations and nearest-neighbor search acts as an expressive classifier.This complements the broader trend that adding data improves language-model performance.

8 CONCLUSION AND FUTURE WORK

The conclusion presents kNN-LM as a broadly applicable way to improve neural language models by querying training examples at test time, while motivating future work on similarity learning and datastore efficiency.

  • Conclusion: kNN-LMs can significantly outperform standard language models by directly querying training examples at test time.The approach is applicable to any neural language model.
  • Conclusion: The results suggest that learning similarity functions between contexts may be easier than predicting the next word.This is presented as the paper’s interpretation of kNN-LM’s success.
  • Future Work: Future work should explicitly train similarity functions and reduce the size of the datastore.These directions target the representation and memory components of the approach.

A APPENDIX

The appendix illustrates how kNN-LM improves target probabilities by retrieving highly relevant examples, including rare contexts, date patterns, factual knowledge, and cases with multiple plausible continuations.

  • Examples: kNN-LM can retrieve an extremely similar training context for a rare long-tail pattern and assign greater probability to the correct target.The corresponding example contrasts the nearest-neighbor search with the language model.
  • Examples: For repeated date patterns, nearest-neighbor search can identify the single training context relevant to the test context.That context receives the highest probability mass.
  • Examples: The model can memorize the factual relation that Georges Bizet wrote Carmen.This example represents factual knowledge stored in the retrieved training data.
  • Examples: When several continuations are plausible, kNN-LM can still assign the highest probability to the correct target through a particularly relevant retrieved context.The language-model probability on that target is lower in the example.
Loading 1911.00172v2…