Augmenting Zero-Shot Dense Retrievers with Plug-in Mixture-of-Memories (2023)

Abstract: In this paper we improve the zero-shot generalization ability of language models via Mixture-Of-Memory Augmentation (MoMA), a mechanism that retrieves augmentation documents from multiple information corpora (external memories), with the option to “plug in” unseen memory at inference time. We develop a joint learning mechanism that trains the augmentation component with latent labels derived from the end retrieval task, paired with hard negatives from the memory mixture. We instantiate the model in a zero-shot dense retrieval setting by augmenting strong T5-based retrievers with MoMA. With only T5-base, our model obtains strong zero-shot retrieval accuracy on the eighteen tasks included in the standard BEIR benchmark, outperforming some systems with larger model sizes. As a plug-in-play model, our model can efficiently generalize to any unseen corpus, meanwhile achieving comparable or even better performance than methods relying on target-specific pretraining. Our analysis further illustrates the necessity of augmenting with mixture-of-memory for robust generalization, the benefits of augmentation learning, and how MoMA utilizes the plug-in memory at inference time without changing its parameters. Our code can be found at https://github.com/gesy17/MoMA.

## introduction scaling up language models-with more parameters and pretraining data-improves model generalization ability on downstream applications @xcite @xcite , but with diminishing return: linear improvements on downstream metrics often require exponentially more parameters and computing cost @xcite . hence, scaling pretrained language models in this way is economically unsustainable @xcite @xcite . ## method in this section we first describe our mixture-of-memory augmentation. then we discuss how it is jointly learned with the end system and enables plug-in memory at inference time. ## experimental methodologies datasets. we choose the ms marco passage dataset @xcite as the source domain dataset, whereas the target domains are from the 18 datasets in beir @xcite ) benchmark, which include including biomedical, scientific and financial texts. more details can be found in appendix a.3. the evaluation metric ndcg@10 is the same with beir benchmark, which measures normalized discounted cumulative gain @xcite of top 10 prediction. the higher ndcg@10 value indicates better performance. augmenting corpora. during training, the mixture-of-memory is composed of source training corpus (marco), wikipedia and a medical knowledge graph. we use the wikipedia chunk prepossessed by @xcite without further processing foot_0 . the medical knowledge graph is extracted from the medical subject headings (mesh) foot_1 , an open-source database for indexing and cataloging of biomedical and health-related information. since it is hierarchical in structure, we linearize it by concatenating spans with text information. during testing, we directly replace marco with the corresponding document sets from beir. each task from beir is augmented independently. more dataset and preprocessing details can be found in appendix a.3. baselines and model choices. we compare our moma with standard sparse and dense retrieval models on beir. we also compare moma with advanced approaches that are specifically designed for zero-shot generalization. they involve techniques that are not directly comparable with this paper, including pretraining on extra data, in-domain continuous pretraining, and generating target pairs using another pretrained generative model. besides, some baselines use larger scale language model as their backbone. we list the details of baselines in appendix a.4. as a plug-in-and-play method, moma can be following cocondenser @xcite , we continuously trained the original t5 model on the marco document corpus using a sentencelevel contrastive loss, combined with the original masked language modeling loss. we then performed the same moma training on top of the continuously pretrained t5 checkpoint and denoted it as moma (cocondenser). the only difference between moma (t5-ance) and moma (cocondenser) is the initialized model start point. we compare their pretraining details with other models in ## conclusion in this paper we propose a new plug-in mixtureof-memory mechanism for the retrieval augmented language models to improve their zero-shot ability on the dense retrieval task. to learn the memory mixture we develop a new joint learning approach that trains the augmentation component using the positive signals from the end task, the language model's attention scores, and hard negatives retrieved from the mixture of augmentation corpora. this leads to our final model moma (t5-ance) and moma (cocondenser) that achieve strong zero-shot accuracy on 18 retrieval tasks included in beir. our analysis shows the importance of augmenting with diverse memory sources and indomain information for robust generalization. we hope our findings can inspire more future research in better augmenting language models, to provide other alternatives to achieve generalization ability beyond solely relying on model scale.