Efficient Training of Retrieval Models using Negative Cache

1 Introduction . Learning to represent objects as dense vectors , often called embeddings , has proved to be crucial in large scale information retrieval tasks from multiple domains including recommendation systems [ 39 ] , vision [ 18 ] and natural language processing [ 13 , 22 ] . A popular paradigm for such learning tasks involves training two separate neural networks ( often called two-towers or dual-encoders ) , each representing a query and a document . Given positive and negative ( query , document ) pairs , the learning task trains the two networks by minimizing a loss function , usually softmax cross-entropy , to encourage positive pairs to have higher similarity scores and negative pairs to have lower scores . While it is easy to sample positive pairs of examples through user feedback such as impressions or clicks , it is more challenging to sample good negative pairs from a pool of potentially millions or even billions of documents . A large number of negative pairs is often required to ensure high quality of the final model , which makes the training process expensive . A number of strategies have been proposed in the literature to address the problem of sampling good negative pairs from a large corpus . The most common approach is to use in-batch negatives , which treats random , non-positive pairs in a minibatch as negatives [ 15 , 22 ] . This approach is computationally efficient and works in a streaming setting , but the pool of negative examples is limited to the minibatch . Towards the later stages of the training , the in-batch negatives become less informative ( i.e. , have low gradients ) since they are sampled randomly without paying attention 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to which negatives are hard for a given query [ 22 ] . Another popular approach is to maintain an asynchronous retrieval index of the full dataset for negative sampling [ 13 , 38 ] . Negatives from the full dataset can be extracted based on approximate retrieval techniques such as ScaNN [ 11 ] , Faiss [ 20 ] or SPTAG [ 7 ] . However , it requires coordinating with a separate process for re-indexing and re-building the retrieval index , which is not only computationally expensive and hard to maintain but also suffers from the problem of stale index . In this work , we theoretically and experimentally analyze training dual encoders using a large cache of negative elements . The cached elements are stale , as they may be generated from a prior iteration ’ s parameters . Our contributions are the following . Main Contributions : • We propose an approach to train retrieval models with cross-entropy loss using a large negative cache . We utilize Gumbel-Max sampling on the cached embeddings to efficiently sample the negatives . • We analyze the convergence of our algorithm in terms of the refresh rate of the cache . We show that even for a small refresh rate we can obtain a first-order convergence rate comparable to that of getting exact gradients using the entire dataset . • We develop a streaming version of our approach . Our streaming algorithm allows us to scale to very large datasets and avoids needing to maintain an up-to-date index for nearest neighbor search ( or maximum inner product search ) . We analyze the bias induced by our method and how it affects the convergence . • We experimentally validate our method using the MS MARCO and TREC 2019 passage retrieval tasks . We show that our approach can be efficiently implemented and achieves statistical performance comparable to state-of-the-art benchmarks with a computationally simpler approach that requires only a fraction of the memory . 2 Background . 2.1 Dual Encoders . Given a query q , our goal is to match the query with the most relevant documents . The set of documents is represented by D = { z1 , z2 , . . . , zm } ∈ RdD , where m is the total number of documents . While we refer to them as documents in this paper , they can be any set of items that we wish to retrieve e.g . movies or products . The training data consists of n positive ( query , doc ) pairs T = { ( q1 , y1 ) , ( q2 , y2 ) , . . . ( qn , yn ) } , where qi is a feature vector in RdQ and the label yi is an integer in [ m ] , which indicates a relevant document for query qi . Our goal is to learn a dual encoder model [ 6 ] for embedding the queries and documents . Formally , a dual encoder consists of two components : 1 . Query encoder φQ : RdQ 7→ Rd , which maps a query to a d-dimensional embedding space . 2 . Document encoder φD : RdD 7→ Rd , which maps a document to the embedding space . The score for a particular query , document pair is simply computed as inner product of the corresponding query and document embeddings i.e. , s ( q , z ) = φQ ( q ) · φD ( z ) . Intuitively , this represents the “ similarity ” between the query and documents in the embedding space . The model is trained such that the relevant query , document pairs have high scores . Dual encoders are popular in large scale settings since they are highly efficient during training and inference as the high scoring items for a given query can be found using efficient nearest neighbors algorithms , often achieving more than 100x speed up over brute force search [ 10 ] . 2.2 Training Dual Encoders with Cross-Entropy Loss . In this paper , we focus on the popular cross-entropy loss for training dual encoders . Even though these loss functions have been found to be highly effective , the main challenge in large scale settings is computational . To observe the computational bottleneck , we look at the cross-entropy loss : LCEi ( θ ) = − log ( exp ( βsiyi ) ∑m j=1 exp ( βsjyj ) ) , ( 1 ) where β is the inverse temperature that scales the scores . Here , sij = sij ( θ ) = φQ ( qi ; θ ) · φD ( zyi ; θ ) and θ ∈ Rp . This computation involves computing the scores for all documents , which can be prohibitively expensive when the number of documents m is large . For simplicity , throughout this paper , we will assume that the embeddings have unit bounded norms , i.e. , ‖φQ ( · ; θ ) ‖2 ≤ 1 and ‖φD ( · ; θ ) ‖2 ≤ 1 for all θ ∈ Rp . Note that this can be assumed without a loss in generality since one can adjust the inverse temperature β accordingly to compensate for the scale . On large-scale datasets , it may not be possible to have the entire dataset in memory , especially with large embedding models where we need the activations of the intermediate layers to persist for back-propagation . An alternative approach is to first select a sample of documents S and use them to approximate the partition function . We thus have the following loss function : LSAMPLE−CEi ( θ ) = − log ( exp ( βsiyi ) ∑ j∈S exp ( βsjyj ) ) . ( 2 ) It has been observed that using a larger set of negative samples leads to better performance [ 16 ] . A simple approach is to uniformly sample this set from D ( popularly referred to as negative sampling ) . However , when m is large , negative sampling is typically inefficient since it is difficult to obtain high-scoring irrelevant documents , thereby , providing a poor approximation of the cross-entropy loss . In the following sections , we explore highly scalable solutions for optimizing cross-entropy loss . Random sampling is often implemented using in-batch negative sampling [ 15 , 22 , 16 ] . However , this approach is not scalable because huge amount of accelerator memory is required to achieve a bigger pool of in-batch negatives . For example , BERT [ 9 ] based transformers are typically used in NLP tasks but a single pair of ( query , document ) BERT-base consumes 600MB of accelerator memory during training . This further limits the effectiveness of random negative sampling . 2.3 Related Work . There have been many recent papers on using cached embeddings and nearest neighbors for improving negative mining during training . Several works consider caching embeddings for the entire dataset [ 13 , 38 ] . They utilize a separate task to recompute embeddings based on the current parameters and recreate an index of nearest neighbor search . This can be computationally challenging for very large datasets , the cost of inference and index refresh scales linearly with the number of documents . REALM [ 13 ] uses their embedding cache to optimize a softmax-cross entropy loss . They approximate the partition function by retrieving the k largest scoring elements . This can be a poor approximation if the distribution is not highly concentrated around the top elements . Additionally , it requires O ( k ) additional memory to perform backpropagation . ANCE [ 38 ] uses their cache to optimize a negative contrastive loss . MoCo [ 14 ] uses a streaming negative cache to increase the number of negative examples for feature embedding tasks . But it can not be extended to dual encoder training , as the document encoder does not contribute to gradient updates . Negative sampling is also widely used in classification problem with very large number of labels , where each document is considered a separate class . This setting is simpler than the dual encoder setting , since the model consists of one encoder and one large classification layer . The weights in the classification layer are only updated when their corresponding labels are included in a gradient step . In contrast , in the dual encoder setting , all document embeddings are changed during a gradient step . For training with a softmax cross-entropy loss in the classification setting , Bengio and Senécal [ 4 ] develop an importance sampling approach for estimating the gradient of the cross-entropy loss function in large output spaces . Many following works utilize sampling techniques for estimating the gradient [ 17 , 36 , 30 , 27 , 25 , 2 ] . Rawat et al . [ 30 ] analyzes the bias of different sampling techniques in gradient estimation . Mussmann et al . [ 28 ] use Gumbel-Max sampling to sample from the softmax distribution and accelerate gradient descent for certain problems . Zhang et al . [ 42 ] analyze the bias between gradients of the softmax loss and gradients generated from using hard negative mining . Other loss functions besides softmax cross-entropy include noise contrastive estimation , which was developed by [ 12 ] and utilized in [ 26 , 25 ] , ordered weighted losses [ 31 ] , and triplet losses [ 33 ] . Hard and semi-hard negative mining is a collection of techniques for selecting negatives that the model is scoring above or near positive elements . Works in this area include [ 34 , 37 ] . In particular , Dense Passage Retrieval [ 22 ] utilizes hard negatives generated by a BM25 ranker [ 32 ] . We note that hard negatives such as BM25 are often domain-specific and require finding a hard negative document for each query , which can be expensive with a large number of queries and documents . We note that there are several approaches that take a dense retrieval model and use its retrieval results to improve performance further . The LTRe method [ 41 ] and the STAR and ADORE method [ 40 ] uses the retrieval results of ANCE to further improve performance with a ranking loss function . The RocketQA method [ 29 ] starts by training a model with sample negatives ( in the cross-batch setting when training with multiple accelerators ) . This model is used to generate negatives for a cross-attention model , which is used to denoise the dataset to train an improved dense retriveal model . We compare to their Step1 model in the experimental section . Our setting consists of learning from observed click data , thus we only see positive pairs . In the related learning-to-rank setting , the goal is to learn a ranking model using a set of positive and negative relevance labels . Works in this setting include [ 19 , 1 , 5 , 21 ] .

Summary: This paper focuses on the problem of negative mining for training bi-encoders with cross-entropy loss in a retrieval setting.  Overall, I think this is a good paper. The authors provide a method that compares to ANCE on the MS MARCO dataset, but is simpler and much more efficient. The ideas are well argumented, and experiments seem to be consistent with the theoretical aspects. Additional experiments would have made the paper stronger, but results on MS MARCO might be sufficient for acceptance.