On Retrieval Augmentation and the Limitations of Language Model Training
This work addresses the limitations of language model training for researchers and practitioners by improving generalization and reducing storage overhead, though it is incremental in nature.
The authors tackled the problem of understanding why retrieval augmentation improves language model performance and created a new dataset to test generalization when training data includes irrelevant information. They showed that kNN retrieval consistently boosts performance for GPT-2 and Mistral 7B, and proposed a method that reduces storage costs by over 25x.
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 performance in this setting. Finally, to make $k$NN retrieval more accessible, we propose using a multi-layer perceptron model that maps datastore keys to values as a drop-in replacement for traditional retrieval. This reduces storage costs by over 25x.