On the Generalization Ability of Retrieval-Enhanced Transformers (2023)

Abstract: Recent work on the Retrieval-Enhanced Transformer (RETRO) model has shown impressive results: off-loading memory from trainable weights to a retrieval database can significantly improve language modeling and match the performance of non-retrieval models that are an order of magnitude larger in size. It has been suggested that at least some of this performance gain is due to non-trivial generalization based on both model weights and retrieval. In this paper, we try to better understand the relative contributions of these two components. We find that the performance gains from retrieval to a very large extent originate from overlapping tokens between the database and the test data, suggesting less of non-trivial generalization than previously assumed. More generally, our results point to the challenges of evaluating the generalization of retrieval-augmented language models such as RETRO, as even limited token overlap may significantly decrease test-time loss. We release our code and model at https://github.com/TobiasNorlund/retro

## introduction large-scale generative language models have shown promising results toward creating a general-purpose foundation for many natural language applications. while sheer scale-up has resulted in better language modeling performance, the immense costs are an inhibiting factor towards further improvements @xcite . recent work on retrieval-augmented language models, such as the retrieval-enhanced transformer (retro; @xcite , suggests that memory can be effectively off-loaded from the model parameters to an external database. in retro, the information retrieved from the database is used to augment the context from which the model predicts new tokens, reducing the need to memorize this information in the model parameters. this opens up for smaller language models with retained performance. specifically, @xcite report that, with a large enough retrieval * corresponding author, tobiasno@chalmers.se database, retro can achieve a performance comparable to gpt-3 @xcite and jurassic-1 @xcite on the pile @xcite , at only 4% of the parameters. similarly, retro achieves significantly lower bits-per-byte performance compared to a baseline of the same size without retrieval. @xcite conclude that retro has the capacity for non-trivial generalization based on both the model parameters and the retrieval database, even though they find that part of the performance gains can be attributed to lexical overlap between retrieval and test data. in this work, we want to better understand the nature and magnitude of this effect. our findings indicate that performance gains foot_0 originate almost exclusively from retro's ability to copy tokens verbatim from retrieved data, effectively exploiting any (small or large) overlap between training and test data. this suggests that the ability of retro to fuse retrieved and in-parameter information may be more limited than previously assumed. ## method to investigate gains from retrieval, we re-implement the retro model described by @xcite (with a few deviations; see below). we present the model here in brevity. ## related work equipping language models with a retrievable external memory has been extensively studied @xcite @xcite @xcite . explicitly leveraging the training data through retrieval to reduce perplexity is proposed in knn-lm @xcite . knn-lm matches the leftward context with the leftward context of all training data tokens, and explicitly interpolates between generating and copying the next token. a recent study analyzes knn-lm to better understand 4 we note a sudden increase in accumulated loss difference for 𝑛 > 64 which is expected considering the way in which we construct the buckets; see appendix c for more details. the causes of performance gains @xcite . similar to our findings in retro, lexical overlap has also been found to play a significant role in explaining retrieval performance gains in knn-lm as well @xcite . the idea of knn-lm is extended in spalm @xcite to instead learn a gating function that facilitates more dynamic interpolation. in both knn-lm and spalm, retrieval is incorporated at the top of the network. this might induce a bias towards surface-level rather than semantic augmentation. in contrast, retrieval in retro is incorporated in lower layers of the network, which opens up for more sophisticated integration of the retrieved information. our results suggest, however, that retrieval in retro also contributes at the surface rather than at the semantic level, similar to the previous works. ## conclusions and future work the capacity of language models for generalization is often measured intrinsically using perplexity, loss or bits-per-byte on held-out validation data. low perplexity language models perform well as few-shot learners on many downstream tasks due to their capacity to both memorize and non-trivially combine textual information from many sources @xcite @xcite . the hope is that we can externalize memory to reduce the footprints of language models without reducing generalization and downstream task performance. our results show that the low loss in retro almost exclusively originates from tokens overlapping between retrieval and validation data, rather than from more sophisticated generalization. to better understand this effect, it would be interesting to modify the retrieval component and deliver semantically similar but lexically different context during training. if the retrieved context is uninformative, the model will learn to ignore it, but if the context is too specific (e.g. literal overlap) the model will learn to copy. by better balancing between these two modes, models may become better at utilizing retrieved information at a deeper and more generalizable level.