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 addresses the technique of negative sampling in IR tasks. The author proposes to utilize Gumbel-Max sampling on the cached embeddings to efficiently sample the negatives. For a very large dataset, the paper develops a streaming variant of the algorithm. The author also conducts some theoretical analysis to demonstrate its convergence and bias. Moreover, the experimental results show its efficiency.