On Retrieval Augmentation and the Limitations of Language Model Training (2024)

Abstract: Augmenting a language model (LM) with k -nearest neighbors ( k NN) retrieval on its training data alone can decrease its perplexity, though the underlying reasons for this remain elusive. In this work, we rule out one previously posited possibility — the “softmax bottleneck.” We then create a new dataset to evaluate LM generalization ability in the setting where training data contains additional information that is not causally relevant. This task is challenging even for GPT-3.5 Turbo. We show that, for both GPT-2 and Mistral 7B, k NN retrieval augmentation consistently improves per formance in this setting. Finally, to make k NN retrieval more accessible, we propose using amulti-layer perceptron model that maps datastore keys to values as a drop-in replacement for traditional retrieval. This reduces storage costsby over 25x.

## introduction recent efforts to improve the performance of language models (lms) have focused on scaling up model @xcite and training data size @xcite . the resulting models have reached near-human or even super-human performance on some tasks @xcite , though with steep accompanying energy and compute resource costs @xcite @xcite . another approach for improving lm performance has been retrieval augmentation. @xcite proposed to build a datastore using lm training data. the datastore associates the next token of prefixes in the training data with the representations of the prefixes extracted from an intermediate layer of an lm. they found that when predicting the next token for a given prefix, using k-nearest neighbor (knn) retrieval, which retrieves the next token based on the intermediate representation of a given prefix, reduced language models' perplexity. because the datastore is drawn entirely from the lm's training data, the success of knn augmentation suggests the standard lm training setup does not yield models that best utilize their parametric capacity or training data. studying why lms augmented with knn retrieval (knn-lms) outperform vanilla lms may shed light on ways to improve the standard lm training setup. in this work, we base our study on the analyses of knn-lms by @xcite . among the aspects they explore are the limitations of model architecture and memorization. @xcite suggest the knn component may be able to map intermediate representation of context to distributions in a more flexible way, while the last layer of lms has a softmax bottleneck @xcite that restricts lms from generating certain distributions. this discrepancy of expressiveness may thus cause the performance gap. they also show that replacing the knn component with an overfitted lm performs worse than knn-lm, suggesting that knn augmentation does not perform better solely because it memorizes the training data better. in this work, we start with inspecting the bottlenecks in the model as suggested by @xcite . we propose an experiment that shows that the softmax bottleneck is not the cause of the performance gap between knn and vanilla lm. our experimental results show that the last linear layers of lms can generate distributions that approximate the distribution from a knn-lm well. therefore, we conclude that the bottleneck issues in the last layers, including the softmax bottleneck issue, are not the cause of the performance gap. we then investigate the performance gap from the perspective of generalization. this explains why an overfitted lm is less effective than a knn retrieval component @xcite . we identify a scenario which we refer to as over-specification. that is, when a statement about certain knowledge (e.g., relational knowledge @xcite or commonsense @xcite @xcite ) contains redundant information. we create a synthetic dataset macondo and use it to show that over-specification in training data prevents lms from learning the knowledge in a robust way, i.e., lms cannot generalize to test data which is not over-specified. even gpt-3.5 turbo, fails, indicating it is a fundamental limitation of lm training. it may be crucial when the size of training data is limited, because in this scenario, it is likely that there are only few statements about certain knowledge and all of them are overspecified. decounfounding the effect of having redundant information also requires more training examples. this may explain why we need to scale up the training data size. because the better generalization ability may be what makes the knn component helpful, we explore alternatives to a knn component by looking for components that also generalize well. it turns out that we can close the generalization gap on moncodo by training another neural model that maps the intermediate representation to the target token. we also show that on the wikitext dataset, this approach reduces the perplexity by 1.45 while requiring less than 4% storage space of knn augmentation. we suggest it is a promising future direction for improving lms. ## generalization from over-specification as the last layers do not have a bottleneck issue that explains the performance gap, we turn to study the efficacy of knn augmentation from the perspective of generalization. in this section, we identify a limitation of lm training that may cause the performance gap: the failure to generalize from over-specified descriptions. ## an alternative to knn-augmentation motivated by the results in §4.3, we explore whether using a datastore is necessary to improve perplexity. the success of knn-augmentation in §4.3 shows that it is possible to generalize better by utilizing the intermediate representation with a knn module. we wonder whether we can use a classification model instead of a knn module. because deep models have been known for their generalization ability @xcite , we explore using an mlp model to replace knn retrieval. we use the key-value pairs in the datastore for knn retrieval to train an mlp model to map the keys to the values (details in §d). results in figure 2 : gpt-3.5-turbo fine-tuned with macondo-conv using openai api. the results are the average of 5 runs with 5 datasets generated with 5 random seeds. note that the presented loss involves special tokens, e.g., endof-string tokens, so the theoretical perfect likelihood is greater than log 0.5. the gray line is the test loss we achieve when we use the test data to train the model. ## related work lms that solely rely on parametric knowledge learned during training time are known to hallucinate @xcite @xcite @xcite , suffer to learn long-tail knowledge @xcite , and fail to adapt to new knowledge over time @xcite @xcite . to overcome these limitations, recent works @xcite @xcite @xcite @xcite include an external datastore with the parametric model, resulting in a retrieval-augmentated model paradigm. meanwhile, @xcite and @xcite analyzes the effect of knn-lm on generation tasks, while @xcite focuses on using knn-lm on few-and zero-shot classification tasks. the traditional lm training setup has been shown to yield models that fail to generalize to test data with reversed relations @xcite , respective readings @xcite , and longer tasks @xcite . these models can also struggle with linguistic generalization between unseen but related contexts @xcite and learn shortcuts that harm generalization @xcite . @xcite have also argued that such models will necessarily be limited due to the ungrounded nature of their training data. ## conclusion we study the performance gap between vanilla and knn-augmented lms. we develop an experiment that allows us to directly inspect the bottleneck is-232 sue and exclude the possibility that it causes the performance gap ( §3). we further identify the overspecified scenario where vanilla lms fail to generalize while knn-lms generalize better ( §4). we also show with gpt-3.5-turbo that this failure of generalization can not be solved by scaling up the model size, suggesting that this is a fundamental limitation of lm training. finally, we show the potential of augmenting lms with an mlp model, indicating a promising future direction ( §5).