Great Memory, Shallow Reasoning: Limits of k NN - LM s (2025)

Abstract: K-nearest neighbor language models (kNN-LMs), which integrate retrieval with next-word prediction, have demonstrated strong performance in language modeling as well as some downstream NLP benchmarks. These results have led researchers to argue that models trained on poor quality or outdated data could perform well by employing a kNN extension that has access to a higher-quality datastore. In this work, we ask whether this improved ability to recall information really translates into downstream abilities. We extensively evaluate kNN-LMs on a diverse set of tasks, ranging from sentiment classification and commonsense reasoning to multi-hop reasoning. Results show that kNN-LMs excel at memory-intensive tasks, where utilizing the patterns in the input is sufficient for determining the output, but struggle with reasoning tasks that require integrating multiple pieces of information to derive new knowledge. We further demonstrate through oracle experiments and qualitative analysis that even with perfect retrieval, kNN-LMs still fail to determine the correct answers, placing an upper bound on their reasoning performance.

## introduction a foundational property of pretrained language modeling @xcite has been that improvements to the perplexity of the model lead to improvements on downstream tasks. this property is central to the scaling of large language models (llms) where researchers focus nearly exclusively on perplexity as a proxy metric for improved general purpose abilities @xcite . in recent years, this research has centered primarily on high-quality text data at greater quantities as the limiting component for producing better language models @xcite . this increasing need for training data has led to significant challenges. on one hand, including as much high-quality data as possible results in improved downstream performance. on the other hand, this data is often protected by licenses or copyright, which means training on such data brings legal issues. it would be ideal to circumvent this issue entirely with alternative approaches. if a model could be trained on lower-quality data but adapted to perform well on real tasks, it might provide a technical workaround. non-parametric language models (nplms), such as knn-lms, have emerged as a promising approach in this space @xcite . knn-lms extend neural lms by linearly interpolating with simple k-nearest neighbor lms. this approach can improve language modeling with its memory over a massive collection of texts, usually referred to as a datastore. @xcite and @xcite validate that knn-lms achieve better performance on downstream tasks compared to standard lms. the silo model of @xcite applies this approach further by training a lm exclusively on license-permissive data and using a non-parametric datastore to improve the models during inference. in this work, we study the limits of how knn-lms can be used to improve llms. specifically, we are interested in whether the improvements in perplexity seen with knn-lms are equivalent to other improvements in lm ability. this question relates to debates about whether memory is separable from other language abilities and how they interact in nlp benchmarks. we summarize our contributions as follows. first, we evaluate knn-lms on 20 nlp tasks, with experimental results revealing that lower perplexity does not necessarily lead to better reasoning in non-parametric settings. to investigate the performance degradation, we conduct extensive analyses in appendix f, which shows that knn-lms are not sensitive to semantic information and can be distracted by irrelevant tokens. figure 1 illustrates such limitations using a multi-hop reasoning exthe knn approach finds several document, both irrelevant and relevant, that may help. however, two issues occur: first, an irrelevant document increases the probability of a random wrong answer; second, even though a relevant document has been found, it may not upweight the actual answer (ouse). we study how these issues may impact task performance as compared to perplexity. ## experimental setup we use llama-2-7b @xcite , llama-3-8b (ai@meta, 2024), and mistral-7b @xcite as our inference models. for each inference model, we build the corresponding datastores. the keys are the 4096-dimensional hidden representations before the final mlp which predicts the token distribution at each generation step, produced by executing forward passes over the datastore corpora. for efficient similarity search, we create a faiss index @xcite and search for nearest-neighbor tokens using euclidean distance. due to the scale of the datastores, we perform approximate search instead of exact search. we base our implementation on @xcite . hyperparameters include λ, k, and σ. λ determines the weight of the datastore, and we consider λ ∈ {0.1, 0.2, 0.3}. we retrieve @xmath0, 2048} neighbors and smooth the knn distribution with a temperature σ ∈ {1, 3, 5, 10}. for each inference model, we use math and wiki datastores for language modeling on the corresponding evaluation datasets: wikitext and math textbooks. each datastore represents a specific domain, and we evaluate the performance of knn-lms on a domain by measuring the perplexity of each evaluation dataset. we conduct a grid search to find the hyperparameters that yield the lowest ppl for each datastore. the optimal hyperparameters for each datastore are later applied across all downstream tasks in our experiments. we provide eight demonstrations for gsm8k and three demonstrations for bbh. for the other datasets, we perform zero-shot inference. details of the experiments are in appendix c. ## knn-lms help in-domain perplexity to explore how different sources of external knowledge impact downstream task performance, we experiment with two datastores. first, we follow the choice made by @xcite , where they identify heterogeneous data sources broadly relevant to common downstream nlp tasks. in particular, they mix wikitext103 @xcite , with other sources including the english portion of amazon review @xcite , cc-news @xcite and imdb @xcite . we call this datastore wiki. we begin by validating past results of knn-lms on language modeling. we present results in ## knn-lms can help memory-intensive tasks we begin by looking at a set of memory-intensive tasks, which we believe can be solved by pattern matching at scale without complex reasoning. we incorporate three types of tasks: sentiment classification, which aims to predict whether the sentiment of a text is positive or negative; textual entailment, which assesses the relationship between two sentences, determining if it constitutes entailment, contradiction, or neutrality; and topic classification, which involves identifying the main topic of a text. we describe dataset details in appendix d. for classification and multiple-choice questionanswering (qa) tasks, we utilize domain conditional pointwise mutual information (dcpmi) @xcite to predict answers. we then calculate accuracy metrics to compare performance across different models. we measure the performance using f1 scores at the token level for text generation. additionally, whenever feasible, we employ fuzzy verbalizers @xcite to maximize the performance of knn-lms. ## knn-lms hurt reasoning performance for reasoning tasks, we consider three types: knowledge-intensive reasoning, which focuses on utilizing world knowledge for making (potential) multi-hop inferences; commonsense reasoning, which involves leveraging commonsense knowledge to understand social and physical interactions; and mathematical reasoning, which includes arithmetic, logical, and discrete reasoning abilities. we describe dataset details in appendix d. we present the results for knowledge-intensive tasks in ## conclusions we investigate whether the improved perplexity observed in knn-lms models can be translated into enhanced reasoning capabilities. our findings indicate that while knn-lms improve perplexity and can achieve better performance on memory-intensive tasks, they struggle with reasoning-intensive tasks, showing a disconnect between lm ability and task ability.