Zero-Shot Dense Retrieval with Momentum Adversarial Domain Invariant Representations (2022)

Abstract: Dense retrieval (DR) methods conduct text retrieval by first encoding texts in the embedding space and then matching them by nearest neighbor search. This requires strong locality properties from the representation space, e.g., close allocations of each small group of relevant texts, which are hard to generalize to domains without sufficient training data. In this paper, we aim to improve the generalization ability of DR models from source training domains with rich supervision signals to target domains without any relevance label, in the zero-shot setting. To achieve that, we propose Momentum adversarial Domain Invariant Representation learning (MoDIR), which introduces a momentum method to train a domain classifier that distinguishes source versus target domains, and then adversarially updates the DR encoder to learn domain invariant representations. Our experiments show that MoDIR robustly outperforms its baselines on 10+ ranking datasets collected in the BEIR benchmark in the zero-shot setup, with more than 10% relative gains on datasets with enough sensitivity for DR models’ evaluation. Source code is available at https://github.com/ji-xin/modir .

## introduction rather than matching texts in the bag-of-words space, dense retrieval (dr) methods first encode texts into a dense embedding space @xcite @xcite and then conduct text retrieval using efficient nearest neighbor search @xcite @xcite . with pre-trained language models and dedicated fine-tuning techniques, the learned representation space has significantly advanced the first stage retrieval accuracy of many language systems, including web search @xcite ., * work partly done during ji's internship at microsoft. source positive source negative target positive target negative source query target query 2021), grounded generation @xcite , open domain question answering @xcite , etc. purely using the learned embedding space for retrieval has raised concerns on the generalization ability, especially in scenarios without dedicated supervision signals. many have observed diminishing advantages of dr models in various datasets if they are not fine-tuned with task-specific labels, i.e., in the zero-shot setup @xcite . however, in many scenarios outside commercial web search, zero-shot is the norm. obtaining training labels is difficult, expensive, and sometimes infeasible, especially in special domains (e.g., medical) where annotation requires strong expertise or is even prohibited because of privacy constraints. the lack of zero-shot ability hinders the democratization of advancements in dense retrieval from data-rich domains to everywhere else. many equally, if not more important, real-world search scenarios still rely on unsupervised exact match methods that have been around for decades, e.g., bm25 @xcite . within the search pipeline, the generalization of first stage dr models is notably worse than subsequent reranking models @xcite . reranking models, similar to many classification models, only require a decision boundary between relevant and irrelevant query-document pairs (q-d pairs) in the representation space. in comparison, dr needs good local alignments across the entire space to support nearest neighbor matching, which is much harder to learn. in figure 1 , we use t-sne (van der @xcite to illustrate this difference. we show learned representations of a bert-based reranker @xcite and a bertbased dense retriever @xcite , in zeroshot transfer from web @xcite to medical domain @xcite . the representation space learned for reranking yields two manifolds with a clear decision boundary; data points in the target domain naturally cluster with their corresponding classes (relevant or irrelevant) from the source domain, leading to good generalization. in comparison, the representation space learned for dr is more scattered. target domain data points are grouped separately from those of the source domain; it is much harder for the learned nearest neighbor locality to generalize from source to the isolated target domain region. in this paper, we present momentum adversarial domain invariant representations learning (modir), to improve the accuracy of zero-shot dense retrieval (zerodr). we first introduce an auxiliary domain classifier that is trained to discriminate source embeddings from target ones. then the dr encoder is not only updated to encode queries and relevant documents together in the source domain, but also trained adversarially to confuse the domain classifier and to push for a more domain invariant embedding space. to ensure stable and efficient adversarial learning, we propose a momentum method that trains the domain classifier with a momentum queue of embeddings saved from previous iterations. our experiments evaluate the generalization ability of dense retrieval with modir using 15 retrieval tasks from the beir benchmark @xcite . on these retrieval tasks from various domains including biomedical, finance, scientific, etc., modir improves the zero-shot accuracy of two standard models, dpr @xcite and ance @xcite . on tasks where evaluation labels have sufficient coverage for dr @xcite ), modir's improvements are robust and significant, despite not using any target domain training labels. we also verify the necessity of the proposed momentum approach, without which the domain classifier fails to capture the domain gaps, and the adversarial training does not learn domain invariant representations, resulting in little improvement in zerodr. we conduct further analyses to reveal interesting properties of modir and its learned embedding space. during the adversarial training process, the target domain embeddings are gradually pushed towards the source domain and eventually absorbed as a subgroup of the source. in the learned representation space, our manual examinations find various cases where a target domain query is located close to source queries with similar information needs. this indicates that zerodr's generalization ability comes from the combination of information overlaps of source/target domains, and modir's ability to identify the right correspondence between them. ## related work in this section, we recap related work in dense retrieval and adversarial domain adaptation. dense retrieval different from sparse first stage retrieval models, dense retrieval with transformerbased models @xcite such as bert @xcite conducts retrieval in the dense embedding space @xcite @xcite @xcite . compared with its sparse counterparts, dr improves retrieval efficiency and also provides comparable or even superior effectiveness for in-domain datasets. one important research question for dr is how to obtain meaningful negative training instances. dpr @xcite uses bm25 to find stronger negatives in addition to in-batch random negatives. rocketqa @xcite uses crossbatch negatives and also filters them with a strong reranking model. ance @xcite uses an asynchronously updated negative index built from the being-trained dr model to retrieve global hard negatives. recently, challenges of zerodr have attracted much attention @xcite @xcite . one way to improve zerodr is query generation @xcite , which first trains a doc2query model in the source domain and then applies the nlg model on target domain documents to generate queries. the target domain documents and generated queries form weak supervision labels for dr models. our method differs from them and focuses on directly improving the generalization ability of the learned representation space. adversarial domain adaptation unsupervised domain adaptation (uda) has been studied extensively for computer vision applications. for example, maximum mean discrepancy @xcite @xcite ) measures domain difference with a pre-defined metric and explicitly minimizes the difference. following the advent of gan @xcite , adversarial training for uda is proposed: an auxiliary domain classifier learns to discriminate source and target domains, while the main classifier model is adversarially trained to confuse the domain classifier @xcite @xcite @xcite @xcite . the adversarial method does not require pre-defining the domain difference metric, allowing more flexible domain adaptation. modir builds upon the success of uda methods and introduces a new momentum learning technique that is necessary to learn domain invariant representations in the zerodr setting. ## training domain invariant representations for dense retrieval in this work, we aim to improve generalization in zerodr under the unsupervised domain adaptation setting (uda) @xcite . given a source domain with sufficient training signals, the goal is to transfer the dr model to a target domain, with access to its queries and documents, but without any relevance label. this is the common case when applying dr in real-world scenarios: in target domains (e.g., medical), example queries and documents are available but annotating relevance is expensive and may require domain expertise; on the other hand, in the source domain (e.g., web search), training signals are available at large scale @xcite . our method, modir, improves zerodr in the uda setup by encouraging the dr models to learn a domain invariant representation space that facilitates the generalization from source to target. in this section, we describe (1) how to train a vanilla dense retrieval model, (2) how to train a momentum domain classifier to distinguish the two domains, and (3) how to adversarially train the dr model for domain invariant representations. ## experiments this section describes experimental setups and evaluates the effectiveness of modir. furthermore, we dive deep into the importance of momentum training and properties of domain invariant embedding space, which provides new insights for zerodr. ## conclusion and future work in this paper, we present modir, a new representation learning method that improves the zero-shot generalization ability of dense retrieval models. we first show that dense retrieval models differ from classification models in that they emphasize locality properties in the representation space. then we present a momentum-based adversarial training method that robustly pushes text encoders to provide a more domain invariant representation space for dense retrieval. our experiments demonstrate that, compared with ance, a recent sota dr model, modir's improvements are robust overall and significant on datasets where zerodr's evaluation is more accurate. we conduct a series of studies to show the effects of our momentum method in learning domain invariant representations. without momentum, the adversarial learning is unstable. the inherent variance of the dr embedding space hinders the convergence of the domain classifier. with momentum training, the model fuses the target domain data into the source domain representation space and discovers related information from the source domain, thus improving generalization of zerodr. we view modir as an initial step of zero-shot dense retrieval, an area that democratizes the rapid advancements in search technologies to many realworld scenarios. our approach inherits the success of domain adaptation techniques and upgrades them by addressing the unique challenges of zerodr. understanding the dynamics of dense retrieval is an im-portant future direction for not only representation learning research but also real-world applications.