Long-Tail Crisis in Nearest Neighbor Language Models

2025-03-28

AI Summary

This paper identifies a fundamental failure of kNN-LM in predicting low-frequency (long-tail) target words: although it effectively retrieves long-tail contextual neighbors, it fails to improve probability estimates for such tokens—and instead predominantly boosts predictions for high-frequency words, even introducing negative bias for low-frequency ones. This challenges the prevailing assumption that explicit memory inherently benefits long-tail modeling. Method: We conduct a frequency-aware analysis integrating retrieval accuracy evaluation, token-frequency statistics over the datastore, product quantization error modeling, and perplexity decomposition experiments. Contribution/Results: Our analysis reveals a significant absence of predictive gain for low-frequency target words, exposing a structural limitation of kNN-LM in out-of-distribution long-tail modeling. These findings provide critical diagnostic insights for advancing memory-augmented language models, particularly in addressing distributional skew and improving tail coverage.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

The $k$-nearest-neighbor language model ($k$NN-LM), one of the retrieval-augmented language models, improves the perplexity for given text by directly accessing a large datastore built from any text data during inference. A widely held hypothesis for the success of $k$NN-LM is that its explicit memory, i.e., the datastore, enhances predictions for long-tail phenomena. However, prior works have primarily shown its ability to retrieve long-tail contexts, leaving the model's performance remain underexplored in estimating the probabilities of long-tail target tokens during inference. In this paper, we investigate the behavior of $k$NN-LM on low-frequency tokens, examining prediction probability, retrieval accuracy, token distribution in the datastore, and approximation error of the product quantization. Our experimental results reveal that $k$NN-LM does not improve prediction performance for low-frequency tokens but mainly benefits high-frequency tokens regardless of long-tail contexts in the datastore.

Problem

• Investigates kNN-LM's performance on low-frequency tokens
• Examines prediction accuracy for long-tail target tokens
• Assesses datastore impact on high vs low-frequency token predictions

Innovation

• Investigates kNN-LM on low-frequency tokens
• Examines prediction probability and retrieval accuracy
• Analyzes datastore token distribution and quantization error