Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval

1 INTRODUCTION . Many language systems rely on text retrieval as their first step to find relevant information . For example , search ranking ( Nogueira & Cho , 2019 ) , open domain question answering ( OpenQA ) ( Chen et al. , 2017 ) , and fact verification ( Thorne et al. , 2018 ) all first retrieve relevant documents for their later stage reranking , machine reading , and reasoning models . All these later-stage models enjoy the advancements of deep learning techniques ( Rajpurkar et al. , 2016 ; Wang et al. , 2019 ) , while , the first stage retrieval still mainly relies on matching discrete bag-of-words , e.g. , BM25 , which has become the pain point of many systems ( Nogueira & Cho , 2019 ; Luan et al. , 2020 ; Zhao et al. , 2020 ) . Dense Retrieval ( DR ) aims to overcome the sparse retrieval bottleneck by matching in a continuous representation space learned via neural networks ( Lee et al. , 2019 ; Karpukhin et al. , 2020 ; Luan et al. , 2020 ) . It has many desired properties : fully learnable representation , easy integration with pretraining , and efficiency support from approximate nearest neighbor ( ANN ) search ( Johnson et al. , 2017 ) . These grant dense retrieval an intriguing potential to fundamentally overcome some intrinsic limitations of sparse retrieval , for example , vocabulary mismatch ( Croft et al. , 2009 ) . One challenge in dense retrieval is to construct proper negative instances when learning the representation space ( Karpukhin et al. , 2020 ) . Unlike in reranking ( Liu , 2009 ) where the training and testing negatives are both irrelevant documents from previous retrieval stages , in first stage retrieval , DR models need to distinguish all irrelevant ones in a corpus with millions or billions of documents . As illustrated in Fig . 1 , these negatives are quite different from those retrieved by sparse models . Recent research explored various ways to construct negative training instances for dense retrieval ( Karpukhin et al. , 2020 ) , e.g. , using contrastive learning ( Oord et al. , 2018 ; He et al. , 2020 ; Chen et al. , 2020a ) to select hard negatives in current or recent mini-batches . However , as observed in recent research ( Karpukhin et al. , 2020 ) , the in-batch local negatives , though effective in learning word or visual representations , are not significantly better than spare-retrieved negatives in representation learning for dense retrieval . In addition , the accuracy of dense retrieval models often underperform BM25 , especially on documents ( Gao et al. , 2020b ; Luan et al. , 2020 ) . ∗Lee and Chenyan contributed equally . In this paper , we first theoretically analyze the convergence of dense retrieval training with negative sampling . Using the variance reduction framework ( Alain et al. , 2015 ; Katharopoulos & Fleuret , 2018 ) , we show that , under conditions commonly met in dense retrieval , local in-batch negatives lead to diminishing gradient norms , resulted in high stochastic gradient variances and slow training convergence — the local negative sampling is the bottleneck of dense retrieval ’ s effectiveness . Based on our analysis , we propose Approximate nearest neighbor Negative Contrastive Estimation ( ANCE ) , a new contrastive representation learning mechanism for dense retrieval . Instead of random or in-batch local negatives , ANCE constructs global negatives using the beingoptimized DR model to retrieve from the entire corpus . This fundamentally aligns the distribution of negative samples in training and of irrelevant documents to separate in testing . From the variance reduction perspective , these ANCE negatives lift the upper bound of per instance gradient norm , reduce the variance of the stochastic gradient estimation , and lead to faster learning convergence . We implement ANCE using an asynchronously updated ANN index of the corpus representation . Similar to Guu et al . ( 2020 ) , we maintain an Inferencer that parallelly computes the document encodings with a recent checkpoint from the being optimized DR model , and refresh the ANN index used for negative sampling once it finishes , to keep up with the model training . Our experiments demonstrate the advantage of ANCE in three text retrieval scenarios : standard web search ( Craswell et al. , 2020 ) , OpenQA ( Rajpurkar et al. , 2016 ; Kwiatkowski et al. , 2019 ) , and in a commercial search engine ’ s retrieval system . We also empirically validate our theory that the gradient norms on ANCE sampled negatives are much bigger than local negatives , thus improving the convergence of dense retrieval models.1 2 PRELIMINARIES . In this section , we discuss the preliminaries of dense retrieval and its representation learning . Task Definition : Given a query q and a corpus C , the first stage retrieval is to find a set of documents relevant to the query D+ = { d1 , ... , di , ... , dn } from C ( |D+| |C| ) , which then serve as input to later more complex models ( Croft et al. , 2009 ) . Instead of using sparse term matches and inverted index , Dense Retrieval calculates the retrieval score f ( ) using similarities in a learned embedding space ( Lee et al. , 2019 ; Luan et al. , 2020 ; Karpukhin et al. , 2020 ) : f ( q , d ) = sim ( g ( q ; θ ) , g ( d ; θ ) ) , ( 1 ) where g ( ) is the representation model that encodes the query or document to dense embeddings . The encoder parameter θ provides the main capacity . The similarity function ( sim ( ) ) is often simply cosine or dot product to leverage efficient ANN retrieval ( Johnson et al. , 2017 ; Guo et al. , 2020 ) . BERT-Siamese Model : A standard instantiation of Eqn . 1 is to use the BERT-Siamese/twotower/dual-encoder model ( Lee et al. , 2019 ; Karpukhin et al. , 2020 ; Luan et al. , 2020 ) : f ( q , d ) = BERT ( q ) · BERT ( d ) = MLP ( ~ [ CLS ] q ) ·MLP ( ~ [ CLS ] d ) . ( 2 ) It encodes the query and document separately with BERT as the encoder g ( ) , using their last layer ’ s [ CLS ] token representation , and applied dot product ( · ) on them . This enables offline precomputing of the document encodings and efficient first-stage retrieval . In comparison , the BERT reranker ( Nogueira et al. , 2019 ) applies BERT on the concatenation of each to-rerank query-document pair : BERT ( q ◦ d ) , which has explicit access to term level interactions between query-document with transformer attentions , but is often infeasible in first stage retrieval as enumerating all documents in the corpus for each query is too costly . 1Our code and trained models are available at http : //aka.ms/ance . Learning with Negative Sampling : The effectiveness of DR resides in learning a good representation space that maps query and relevant documents together , while separating irrelevant ones . The learning of this representation often follows standard learning to rank ( Liu , 2009 ) : Given a query q , a set of its relevant document D+q and irrelevant ones D − q , find the best θ ∗ that : θ∗ = argminθ ∑ q ∑ d+∈D+q ∑ d−∈D−q l ( f ( q , d+ ) , f ( q , d− ) ) . ( 3 ) The loss l ( ) can be binary cross entropy ( BCE ) , hinge loss , or negative log likelihood ( NLL ) . A unique challenge in dense retrieval , targeting first stage retrieval , is that the irrelevant documents to separate are from the entire corpus ( D−q = C \ D+q ) . This often leads to millions of negative instances , which have to be sampled in training : θ∗ = argminθ ∑ q ∑ d+∈D+ ∑ d−∈D̂− l ( f ( q , d+ ) , f ( q , d− ) ) . ( 4 ) Here we start to omit the subscript q in Dq . All D+ and D− are query dependent . A natural choice is to sample negatives D̂− from top documents retrieved by BM25 . However , they may bias the DR model to merely mimic sparse retrieval ( Luan et al. , 2020 ) . Another way is to sample negatives in local mini-batches , e.g. , as in contrastive learning ( Oord et al. , 2018 ) , however , these local negatives do not significantly outperform BM25 negatives ( Karpukhin et al. , 2020 ; Luan et al. , 2020 ) . 3 ANALYSES ON THE CONVERGENCE OF DENSE RETRIEVAL TRAINING . In this section , we theoretically analyze the convergence of dense retrieval training . We first show the connections between learning convergence and gradient norms ( Sec . 3.1 ) , then we discuss how non-informative negatives in dense retrieval yield less optimal convergence ( Sec . 3.2 ) . 3.1 ORACLE NEGATIVE SAMPLING ACCORDING TO PER-INSTANCE GRADIENT-NORM . Let l ( d+ , d− ) = l ( f ( q , d+ ) , f ( q , d− ) be the loss function on the training triple ( q , d+ , d− ) , PD− the negative sampling distribution for the given ( q , d+ ) , and pd− the sampling probability of negative instance d− , a stochastic gradient decent ( SGD ) step with importance sampling ( Alain et al. , 2015 ) is : θt+1 = θt − η 1 Npd− ∇θt l ( d+ , d− ) , ( 5 ) with θt the parameter at t-th step , θt+1 the one after , and N the total number of negatives . The scaling factor 1Npd− ensures Eqn . 5 is an unbiased estimator of the non-stochastic gradient on the full data . Then we can characterize the converge rate of this SGD step as the movement to the optimal θ∗ . Following derivations in variance reduction ( Katharopoulos & Fleuret , 2018 ; Johnson & Guestrin , 2018 ) , let gd− = 1Npd−∇θt l ( d + , d− ) the weighted gradient , the convergence rate is : E∆t = ||θt − θ∗||2 − EPD− ( ||θt+1 − θ ∗||2 ) ( 6 ) = ||θt||2 − 2θTt θ∗ − EPD− ( ||θt − ηgd− || 2 ) + 2θ∗TEPD− ( θt − ηgd− ) ( 7 ) = −η2EPD− ( ||gd− || 2 ) + 2ηθTt EPD− ( gd− ) − 2ηθ ∗TEPD− ( gd− ) ( 8 ) = 2ηEPD− ( gd− ) T ( θt − θ∗ ) − η2EPD− ( ||gd− || 2 ) ( 9 ) = 2ηEPD− ( gd− ) T ( θt − θ∗ ) − η2EPD− ( gd− ) TEPD− ( gd− ) − η 2Tr ( VPD− ( gd− ) ) . ( 10 ) This shows we can obtain better convergence rate by sampling from a distribution PD− that minimizes the variance of the stochastic gradient estimator EPD− ( ||gd− || 2 ) , or Tr ( VPD− ( gd− ) ) as the estimator is unbiased . The variance reflects how good the stochastic gradient from negative sampling represents the full gradient on all negatives—the latter is ideal but infeasible . Intuitively , we prefer the stochastic estimator to be stable and have smaller variances . A well known result in importance sampling ( Alain et al. , 2015 ; Johnson & Guestrin , 2018 ) is that there exists an optimal distribution that : p∗d− = argminpd− Tr ( VPD− ( gd− ) ) ∝ ||∇θt l ( d + , d− ) ||2 . ( 11 ) To prove this , one can apply Jensen ’ s inequality on the gradient variance and verify that Eqn . 11 achieves the minimum . The detailed derivations can be find in Johnson & Guestrin ( 2018 ) . Eqn . 11 shows that the convergence rate can be improved by sampling negatives proportional to their per-instance gradient norms ( though too expensive to calculate ) . Intuitively , an negative instance with larger gradient norm is more likely to reduce the non-stochastic training loss , thus should be sampled more frequently than those with diminishing gradients . The correlation of larger gradient norm and better training convergence is also observed in BERT fine-tuning ( Mosbach et al. , 2021 ) .

Summary: This paper studies the problem of dense text retrieval, which represents texts as dense vectors for approximate nearest neighbors (ANN) search. Dense text retrieval has two phases. The first phase learns a representation model to project semantically similar texts to vectors of large similarity scores (e.g. inner products or cosine similarity scores). The second phase adopts an ANN search algorithm to index these vectors and process queries. The paper claims key contributions at the first phase. Specifically, (1) The paper introduces a better negative sampling method to sample good dissimilar text pairs for training.  (2) The new method enables faster converge of model learning. (3) The new method leads to 100x faster efficiency than a BERT-based baseline, while achieving almost the same accuracy as the baseline.