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 to learn a retriever (to find relevant documents in a dataset / dual encoder) and a ranker (to re-rank retrieved documents / cross-encoder) using an adversarial framework (the retriever trying to fool the ranker). This training allows to reach state-of-the-art performance for both the retriever and the ranker. The paper leverages all the recent techniques of neural retrieval models (pre-training, teacher training) and paper mostly focus on the training of the 1st stage (retriever), but also presents result on the full pipeline too (the 2nd stage is more standard, so it is normal the authors did not focus on this part).  Overall, this is a good paper that is inspired by many recent techniques on neural retrievers, bringing back the IRGAN adversarial training but with more success. The experiments are well designed and the analysis interesting.