Why do Nearest Neighbor Language Models Work?

TL;DR

This paper investigates why retrieval-augmented k-nearest neighbor language models outperform standard neural LMs, identifying key factors like input representation, approximate search, and softmax temperature, and proposes improvements to standard models.

Contribution

The study provides a detailed analysis of the reasons behind kNN-LMs' superior performance and introduces methods to enhance standard LMs based on these insights.

Findings

kNN-LMs outperform standard LMs due to input representation differences

Approximate kNN search impacts model performance

Softmax temperature plays a crucial role in kNN distribution effectiveness

Abstract

Language models (LMs) compute the probability of a text by sequentially computing a representation of an already-seen context and using this representation to predict the next word. Currently, most LMs calculate these representations through a neural network consuming the immediate previous context. However recently, retrieval-augmented LMs have shown to improve over standard neural LMs, by accessing information retrieved from a large datastore, in addition to their standard, parametric, next-word prediction. In this paper, we set out to understand why retrieval-augmented language models, and specifically why k-nearest neighbor language models (kNN-LMs) perform better than standard parametric LMs, even when the k-nearest neighbor component retrieves examples from the same training set that the LM was originally trained on. To this end, we perform a careful analysis of the various dimensions over which kNN-LM diverges from standard LMs, and investigate these dimensions one by one. Empirically, we identify three main reasons why kNN-LM performs better than standard LMs: using a different input representation for predicting the next tokens, approximate kNN search, and the importance of softmax temperature for the kNN distribution. Further, we incorporate these insights into the model architecture or the training procedure of the standard parametric LM, improving its results without the need for an explicit retrieval component. The code is available at https://github.com/frankxu2004/knnlm-why.