Adversarial Retriever-Ranker for Dense Text Retrieval

Current dense text retrieval models face two typical challenges . First , they adopt a siamese dual-encoder architecture to encode queries and documents independently for fast indexing and searching , while neglecting the finer-grained termwise interactions . This results in a sub-optimal recall performance . Second , their model training highly relies on a negative sampling technique to build up the negative documents in their contrastive losses . To address these challenges , we present Adversarial Retriever-Ranker ( AR2 ) , which consists of a dual-encoder retriever plus a cross-encoder ranker . The two models are jointly optimized according to a minimax adversarial objective : the retriever learns to retrieve negative documents to cheat the ranker , while the ranker learns to rank a collection of candidates including both the ground-truth and the retrieved ones , as well as providing progressive direct feedback to the dual-encoder retriever . Through this adversarial game , the retriever gradually produces harder negative documents to train a better ranker , whereas the cross-encoder ranker provides progressive feedback to improve retriever . We evaluate AR2 on three benchmarks . Experimental results show that AR2 consistently and significantly outperforms existing dense retriever methods and achieves new state-of-the-art results on all of them . This includes the improvements on Natural Questions R @ 5 to 77.9 % ( +2.1 % ) , TriviaQA R @ 5 to 78.2 % ( +1.4 % ) , and MS-MARCO MRR @ 10 to 39.5 % ( +1.3 % ) . We will make our code , models , and data publicly available . 1 INTRODUCTION . Dense text retrieval ( Lee et al. , 2019 ; Karpukhin et al. , 2020 ) has achieved great successes in a wide variety of both research and industrial areas , such as search engines ( Brickley et al. , 2019 ; Shen et al. , 2014 ) , recommendation system ( Hu et al. , 2020 ) , open-domain question answering ( Chen et al. , 2017 ; Liu et al. , 2020 ) , etc . A typical dense retrieval model adopts a dual-encoder ( Huang et al. , 2013 ) architecture to encode queries and documents into low-dimensional embedding vectors , with the relevance between query and document being measured by the similarity between embeddings . In real-world dense text retrieval applications , it pre-computes all the embedding vectors of documents in the corpus , and leverages the approximate nearest neighbor ( ANN ) ( Johnson et al. , 2019 ) technique for efficiency . To train a dense retriever , contrastive loss with negative samples is widely applied in the literature ( Xiong et al. , 2021 ; Karpukhin et al. , 2020 ) . During training , the model utilizes a negative sampling method to obtain negative documents for a given querydocument pair , and then minimizes the contrastive loss which relies on both the positive document and the sampled negative ones ( Shen et al. , 2014 ; Chen et al. , 2017 ; Radford et al. , 2021 ) . Recent studies on contrastive learning ( Xiong et al. , 2021 ; Karpukhin et al. , 2020 ) show that the iterative “ hard-negative ” sampling technique can significantly improve the performance compared with “ random-negative ” sampling approach , as it can pick more representative negative samples to learn a more discriminative retriever . In the work ( Qu et al. , 2021 ) , it suggests leveraging crossencoder model to heuristically filter “ hard-negative ” samples to further improve performance and shows the importance of sampling technique in the contrastive learning . On the other hand , the model architecture of dual-encoders enables the encoding of queries and documents independently which is essential for document indexing and fast retrieval . However , this Under review as a conference paper at ICLR 2022 … … … … query document … … … … query document relevance ignores the modeling of finer-grained interactions between queries and documents which could be a sub-optimal solution in terms of retrieval accuracy . Motivated by these phenomena , we propose an Adversarial Retriever-Ranker ( AR2 ) framework . The intuitive idea of AR2 is inspired by the “ retriever-ranker ” architecture in the classical information retrieval systems . AR2 consists of two modules : a dual-encoder model served as the retrieval module in Figure 1a and a cross-encoder model served as the ranker module in Figure 1b . The crossencoder model takes the concatenation of a query and document as input text , and can generate more accurate relevance scores compared with the dual-encoder model , since it can fully explore the interactions between the query and document through a self-attention mechanism using a conventional transformer model ( Vaswani et al. , 2017 ) . Instead of training “ retriever-ranker ” modules independently in some IR systems ( Manning et al. , 2008 ; Mitra & Craswell , 2017 ) , AR2 constructs a unified minimax game for training the retriever and ranker models interactively , as shown in Figure 2 . In particular , AR2 adopts a minimax objective from the adversarial game ( Goodfellow et al. , 2014 ) where the retrieval model is optimized to produce relevant documents to fool the ranker model , whereas the ranker model learns to distinguish the ground-truth relevant document and retrieved ones by its opponent retrieval model . Within the adversarial “ retriever-ranker ” training framework , the retrieval model receives the smooth training signals from the ranker model which helps alleviate the harmful effects of “ false-negative ” issues . For example , a “ false-negative ” example which is rated as high-relevance by the ranker model , will also be granted with high probability by retrieval model in order to fool the ranker , meanwhile the ranker model with better generalization capability is more resistant to label noises compared to the retrieval model . In the empirical studies of AR2 , we further introduce a distillation regularization approach to help stabilize/improve the training of the retriever . Intuitively , the retriever would converge to sharp conditionalprobabilities over documents given a query within the adversarial training framework , i.e. , high retrieval probabilities for the top relevant documents and near-zero retrieval ones for the rest . However , it is not a desirable property as it might impede exploring different documents during training . Thus , we incorporate the distillation loss between the retriever and ranker models as a smooth term for further improvement . In experiments , we evaluate AR2 on three widely used benchmarks for dense text retrieval : Natural Questions , Trivia QA and MS-MARCO . Experimental results show that AR2 achieves state-of-theart performance on all these datasets . Meanwhile , we provide a comprehensive ablation study to demonstrate the advantage of different AR2 components . 2 PRELIMINARIES . Dense Text Retrieval : We mainly consider a contrastive-learning paradigm for dense text retrieval in this work , where the training set consists of a collection of text pairs . C = { ( q1 , d1 ) , ... , ( qn , dn ) } . In the scenario of open-domain question answering , a text pair ( q , d ) refers to a question and a corresponding document which contains the answer . A typical dense retrieval model adopts a dual encoder architecture , where questions and documents are represented as dense vectors separately and the relevance score sθ ( q , d ) between them is measured by the similarity between their embeddings : sθ ( q , d ) = 〈E ( q ; θ ) , E ( d ; θ ) ) 〉 ( 1 ) where E ( · ; θ ) denotes the encoder module parameterized with θ , and 〈·〉 is the similarity function , e.g. , inner-product , Euclidean distance . Based on the embeddings , existing solutions generally leverage on-the-shelf fast ANN-search ( Johnson et al. , 2019 ) for efficiency . A conventional contrastive-learning algorithm could be applied for training the dual encoders ( Shen et al. , 2014 ; Chen et al. , 2017 ; Liu et al. , 2020 ) . For example , given a training instance ( q , d ) , we select n negative irrelevant documents ( d−1 , ... , d − n ) ( denoted as D−q ) to optimize the loss function of the negative log likelihood of the positive document : Lθ ( q , d , D−q ) = −log eτsθ ( q , d ) eτsθ ( q , d ) + ∑n i=1 e τsθ ( q , d − i ) ( 2 ) where τ is a hyper-parameter to control the temperature . Previous works ( Shen et al. , 2014 ; Chen et al. , 2017 ; Liu et al. , 2020 ) present an effective strategy on negative document sampling , called “ InBatch Negatives ” where negative documents are randomly sampled from a collection of documents which are within the same mini-batch as question-document training pairs . Recently , some studies e.g. , ANCE ( Xiong et al. , 2021 ) and Condenser ( Gao & Callan , 2021b ) , have shown that selecting “ hard-negatives ” in the training can significantly improve the retrieval performance in open-domain question answering . For example , instead of sampling negative document randomly , “ hard-negatives ” are iteratively retrieved through previous checkpoints of the dual encoder model . However , a more recent work RocketQA ( Qu et al. , 2021 ) continues to point out that the retrieved “ hard-negatives ” could potential be “ false-negatives ” in some cases , which might limit the performance . Generative Adversarial Network : GANs have been widely studied for generating the realisticlooking images in computation vision ( Goodfellow et al. , 2014 ; Brock et al. , 2018 ) . In the past few years , the idea of GANs has been applied in information retrieval ( Wang et al. , 2017 ) . For example , IRGAN ( Wang et al. , 2017 ) , proposes a minimax retrieval framework which constructs two types of IR models : a generative retrieval model and a discriminative retrieval model . The two IR models are optimized through a minimax game : the generative retrieval model generates relevant documents that look like ground-truth relevant documents to fool the discriminative retrieval model , whereas the discriminative retrieval model learns to draw a clear distinction between the groundtruth relevant documents and the generated ones made by its opponent generative retrieval model . The minimax objective is formulated as : JG ∗ , D∗ = minθmaxφEd∼ptrue ( ·|q ) [ logDφ ( d , q ) ] + Ed−∼Gθ ( ·|q ) [ log ( 1−Dφ ( d− , q ) ) ] ( 3 ) where Gθ ( ·|q ) and Dφ ( d− , q ) denote the generative retrieval model and discriminative retrieval model in IRGAN , respectively . It is worth noting the original IRGAN model doesn ’ t work for dense text retrieval tasks as it doesn ’ t contain the dual-encoder model for document indexing or fast retrieval . 3 METHOD . In this section , we introduce the proposed adversarial retriever-ranker ( AR2 ) approach . It consists of two modules : the dual-encoder retriever module Gθ as in Figure 1a , and the cross-encoder ranker module Dφ as in Figure 1b . Gθ and Dφ computes the relevance score between question and document as follows : Gθ ( q , d ) = Eθ ( q ) TEθ ( d ) Dφ ( q , d ) = wφ TEφ ( [ q , d ] ) ( 4 ) where Eθ ( · ) and Eφ ( · ) are language model encoders which can be initialized with any pre-trained language model , wφ is the linear projector in Dφ , and [ q , d ] is the concatenation of question and document . In AR2 , the retriever and ranker modules are optimized jointly through a contrastive minimax objective : JG ∗ , D∗ = minθmaxφED−q ∼Gθ ( q , · ) [ logpφ ( d|q ; Dq ) ] ( 5 ) where D−q : { d−i } ni=1 is the set of n negative documents sampled by Gθ ( q , · ) given q , and pφ ( d|q ; Dq ) denotes the probability of selecting the ground-truth document d from the document set Dq ( Dq = { d } ∪ D−q ) by the ranker module Dφ ; pφ ( d|q ; Dq ) = eτDφ ( q , d ) ∑ d′∈Dq e τDφ ( q , d′ ) ( 6 ) According to the objective function ( Eqn . 5 ) , the dual-encoder retrieval model Gθ ( q , · ) would try to sample the high-relevant documents to fool the ranker model , whereas the ranker model Dφ ( q , · ) is optimized to draw distinctions between ground-truth passage and the ones sampled by Gθ ( q , · ) . We present the illustration of the AR2 framework in Figure 2 . In order to optimize the minimax objective function , we adopt a conventional iterative-learning mechanism to optimize the retriever and ranker modules coordinately . 3.1 TRAINING THE RANKER Dφ Given the fixed retriever Gθ , the ranker model Dφ is updated by maximizing the log likelihood of selecting ground-truth d from set Dq given a query q : φ∗ = argmaxφlogpφ ( d|q ; Dq ) ( 7 ) where Dq consists of ground-truth document d and negative document set D−q . D−q is sampled by Gθ according to Eqn . 5 . In experiments , we first retrieve top-100 negative documents , and then randomly sample n examples from them to obtain D−q . 3.2 TRAINING RETRIEVER Gθ With fixing the ranker Dφ , the model parameters θ∗ for the retriever Gθ is optimized by minimizing the expectation of log likelihood of function . In particular , by isolating θ from the minimax function ( Eqn . 5 ) , the objective for the retriever can be written as : θ∗ = argminθJ θ = ED−q ∼Gθ ( q , · ) [ logpφ ( d|q ; Dq ) ] ( 8 ) However , it is intractable to optimize θ directly through Eqn . 8 , as the computation of probability D−q ∼ Gθ ( q , · ) does not follow a close form . Thus , we seek to minimize an alternative upper-bound of the loss criteria : Jθ ≤ Ĵθ = Ed−∼pθ ( ·|q ; D−q ) [ logpφ ( d|q ; { d , d− } ) ] ( 9 ) The detailed deviation of Eqn . 9 is provided in the Appendix A.1 . Therefore , the gradient of parameter θ can be computed as the derivative of Ĵθ with respect to θ : ∇θĴθ = Ed−∼pθ ( ·|q ; D−q ) ∇θ log pθ ( d −|q ; D−q ) [ log pφ ( d|q ; { d , d− } ) ] ( 10 ) Here , the same approach is applied to obtain set D−q as in Eqn . 7 . Regularization : we further introduce a distillation regularization term in Gθ ’ s training , which encourages the retriever model to distill from the ranker model . JθR = H ( pφ ( ·|q ; D ) , pθ ( ·|q ; D ) ) ( 11 ) H ( · ) is the cross entropy function . pφ ( ·|q ; D ) and pθ ( ·|q ; D ) denote the conditional probabilities of document in the whole corpus D by the ranker and the retriever model , respectively . In practice , we also limit the sampling space over documents to a fixed set , i.e. , Dq = { d } ∪ D−q . Thus the regularization loss for the retriever model can be rewritten as : JθR = H ( pφ ( ·|q ; Dq ) , pθ ( ·|q ; Dq ) ) ( 12 ) Algorithm 1 Adversarial Retriever-Ranker ( AR2 ) Require : Retriever Gθ ; Ranker Dφ ; Document pool D ; Training dataset C. 1 : Initialize the retriever Gθ and the ranker Dφ with pre-trained language models . 2 : Train the warm-up retriever G0θ on training dataset C. 3 : Build ANN index on D 4 : Retrieve negative samples on D. 5 : Train the warm-up ranker D0θ 6 : while AR2 has not converged do 7 : for Retriever training step do 8 : Sample n documents { d−i } n from ANN index . 9 : Update parameters of the retriever Gθ . 10 : end for 11 : Refresh ANN Index . 12 : for Ranker training step do 13 : Sample n hard negatives { d−i } n from ANN index . 14 : Update parameters of the ranker Dφ . 15 : end for 16 : end while

Summary: This paper proposes an adversarial method to jointly train the dual encoder retriever and the cross-attention ranker. The key idea is to obtain harder negatives for the retriever by training a separate ranker with higher capacity and taking passages that the model is confused with. It also proposes a technique like distillation from the ranker to the retriever. Experiments on three datasets demonstrate improvements over a range of baselines.