Efficient Training of Retrieval Models Using Negative Cache
Year: 2021

Abstract:
Factorized models, such as two tower neural network models, are widely used for scoring (query, document) pairs in information retrieval tasks. These models are typically trained by optimizing the model parameters to score relevant “positive" pairs higher than the irrelevant “negative" ones. While a large set of negatives typically improves the model performance, limited computation and memory budgets place constraints on the number of negatives used during training. In this paper, we develop a novel negative sampling technique for accelerating training with softmax cross-entropy loss. By using cached (possibly stale) item embeddings, our technique enables training with a large pool of negatives with reduced memory and computation. We also develop a streaming variant of our algorithm geared towards very large datasets. Furthermore, we establish a theoretical basis for our approach by showing that updating a very small fraction of the cache at each iteration can still ensure fast convergence. Finally, we experimentally validate our approach and show that it is efficient and compares favorably with more complex, state-of-the-art approaches.

heading: 1 Introduction
text: 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.

heading: 2 Background

heading: 2.1 Dual Encoders
text: 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].

heading: 2.2 Training Dual Encoders with Cross-Entropy Loss
text: 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.

heading: 2.3 Related Work
text: 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 cannot 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].

heading: 3 Full Document Cache
text: The primary objective is to find the parameters θ∗ that minimize the softmax cross-entropy loss LCE(θ) (Equation 1). The standard algorithm to optimize the loss is stochastic gradient descent (or its variants). However, calculating stochastic gradients for this loss function is expensive in terms of both time and memory. Computation of the gradient involves: (1) a forward pass on every document in our dataset, (2) storing all intermediate activations (which are needed for the backward pass), and (3) calculating the gradients in the backward pass. Unfortunately, for large embedding encoder models such as those based on Transformers ([35]), this is prohibitively expensive. For instance, for a large BERT model on 1 million documents, the required memory would be hundreds of terabytes.
We propose an approach that can approximate the gradient of cross-entropy loss without needing to embed every document, thereby reducing the per-iteration computation and memory requirements for gradient computation. Our approach is based on cached embeddings [13, 38]. To understand our approach, we first start with gradient estimation with Gumbel-Max sampling. The following known fact shows that if we can sample from the softmax distribution then we can get an unbiased estimate of the stochastic gradient.
Fact 1. The gradient of the cross-entropy has the following form:
∇LCEi = −β∇siyi + m∑ j=1 pijβ∇sij ,
where LCEi = − log ( exp(βsiyi )∑m j=1 exp(βsij) ) and pij = exp(βsij)∑m j′=1 exp(βsij′ ) .
If J is a sample from pi = (pi1, pi2, . . . , pim), we have ∇LCEi = −β∇siyi + βE[∇siJ ]. Further, if N1, N2, . . . , Nm are i.i.d. standard Gumbel random variables, then the index arg maxj∈[m] βsij +Nj is a sample from the distribution pi = (pi1, pi2, . . . , pim).
We refer to this approach as Gumbel-Max sampling and GumbelMaxSample(βs) is used to represent this sampling procedure. A natural estimator of the gradient is∇L̂CEi = −β∇siyi + β∇siJ where J is the index in [m] obtained through Gumbel-Max sampling. Due to the feasibility of fast maximum inner product search (MIPS), Gumbel-Max sampling is an efficient way to sample from the softmax distribution; this, thereby, provides an efficient way to obtain an unbiased estimate of the gradient. In particular, the per-iteration memory requirement is reduced from O(mτ) to O(md) where τ and d are the model and embedding sizes respectively. This is due to the fact that the intermediate activations need to be stored only for documents yi and J , the negative obtained through Gumbel-Max sampling. When d τ (which is typically the case while using large transformer models), this leads to a significant reduction in memory requirements. While the memory requirements are reduced drastically, we still need to do a forward pass on all documents to compute the scores sij . Using Fact 1 , we would need to embed every document every iteration, which is still computationally intensive.
Algorithm 1 Cached Gumbel-Max Gradient Descent Input: Learning rate η, Document refresh fraction ρ ∈ (0, 1] Initialize parameters θ0. Initialize embeddings table E : Ej ← φD(zj ; θ0) for all j ∈ [m] for t ∈ 0, 1, . . . , T − 1 do
Sample qi, yi from the training set eqi ← φQ(qi; θt) ezyi ← φD(zyi ; θt) and update Eyi = ezyi Calculate scores s̃ij ← eqi · Ej for all j ∈ [m] J ← GumbelMaxSample(βs̃i) ezJ ← φD(zJ ; θt) siJ ← eqi · ezJ gt ← −β∇siyi + β∇siJ θt+1 ← θt − ηgt Select oldest ρm embeddings of E and update them to Ej ← φD(zj ; θt+1)
end for
This motivates an algorithm where we cache previously calculated embeddings to efficiently approximate the softmax distribution pi. We first note the structure of the sij = φQ(qi; θ) · φD(zj ; θ). The key computational challenge in Gumbel-Max sampling is computation of the embeddings φD(zj ; θ) for all j ∈ [m]. This is required to find the index J in Fact 1, which essentially renders it computationally intractable for large m. Our approach is to compute embeddings φD(zj ; θ) for only a few documents j ∈ [m] at each iteration and reuse the previously computed embeddings for the rest of the documents. In particular, let zj = φD(zj ; θ) be the current embedding of a document and z̃j = φD(zj ; θ̃) be its previous embedding. If θ and θ̃ are reasonably close, then zj and z̃j would be similar. In such a scenario, the scores sij and s̃ij , and, thereby the corresponding distributions pi = (pi1, pi2, . . . , pim) and p̃i = (p̃i1, p̃i2, . . . , p̃im) are also similar. As a consequence, the true gradient ∇LCEi = −β∇siyi + ∑m j=1 pijβ∇sij and the approximation to the gradient
−β∇siyi + ∑m j=1 p̃ijβ∇sij should be similar. Note that for the approximation we take the gradient of sij and not s̃ij ; we only replace the weight terms pij with p̃ij .
We described this algorithm in Algorithm 1. We present the algorithm with a batch size of 1—it can be easily extended to larger batches. We maintain an embedding table of all documents. At each iteration, we update only a small fraction ρ of the stale embeddings to ensure that pi and p̃i are similar. Using this approach, we can approximately sample from this distribution using Gumbel-Max sampling in an efficient manner.

heading: 3.1 Theoretical Results
text: In this section, we establish error guarantees on the error of our Cached Gumbel-Max gradient approximation. Before delving into the technical details, we state the following key assumptions on the query and document encoder. Assumption 1. The following conditions hold for the query φQ and document encoder φD:
A1 The query and document encoder functions are both L-Lipschitz in the parameters θ. In particular, we have ‖φ(qi; θ) − φ(qi; θ′)‖2 ≤ L‖θ − θ′‖2 for all i ∈ [n] and ‖φ(zi; θ) − φ(zi; θ ′)‖2 ≤ L‖θ − θ′‖2 for all i ∈ [m].
A2 The query and document embeddings are bounded i.e., we have eq, ez = φQ(q; θ), φD(z; θ) satisfy ‖eq‖2, ‖ez‖2 ≤ 1.
A3 The score functions have bounded gradients i.e., we have ‖∇sij‖2 ≤M for all i ∈ [n] and j ∈ [m].
All of these assumptions are fairly mild and are common in optimization literature. As noted earlier, the second assumption does not lead to much loss of generality as a larger bound on the norm can be absorbed into the inverse temperature parameter β. In the following result, we first show that the error can be bounded by the `∞ error between the true scores and the scores with the cached embeddings.
We then show that if the cached embeddings are generated from Algorithm 1 with learning rate η and refresh rate ρ, we can bound the gradient error at each iteration. Theorem 2. Let θt and index i be the parameters and training point selected at tth iteration of Algorithm 1, respectively. Let J represent the Gumbel-Max index selected at that iteration. Then, under Assumption 1, we have the following gradient approximation with the cached embeddings
∇L̃CEi(θt) = EJ [gt] = −β∇siyi + m∑ j=1 p̃ijβ∇sij ,
where p̃ij = exp(βs̃ij)∑m
j′=1 exp(βs̃ij′ ) and s̃ij = eqi · Ej . Furthermore, we have the following bound on the
gradient approximation of ∇LCEi(θt):
‖∇LCEi(θt)−∇L̃CEi(θt)‖2 ≤ 2β2M‖s̃i − si‖∞. When parameter updates are generated by Algorithm 1 with step size η and update rate ρ, we have
‖∇LCEi(θt)−∇L̃CEi(θt)‖2 ≤ 4ηβ3LM2 ( 1 ρ − 1 ) .
We see that the error can be controlled by either increasing the refresh rate or decreasing the gradient norm. This provides a bound on the bias in the gradient approximation. Using the above result, we show the following first-order convergence guarantees. For proving convergence, we need the following additional assumption. Assumption 2. We assume loss function ∇LCEi is S-smooth i.e., we have ‖∇LCEi(θ) − ∇LCEi(θ′)‖2 ≤ S‖θ − θ′‖2 holds for all θ, θ′ ∈ Rp and i ∈ [n].
Under the above assumption, we have the following convergence result in general non-convex settings. Theorem 3. Suppose we run Algorithm 1 for T iterations with stepsize η = √ LCE(θ0)−LCE(θ∗)√
2TSM .
Then under Assumption 1 and 2, we have the following:
1
T T∑ t=0 E[‖∇LCE(θt)‖22] ≤ 4M √ S(LCE(θ0)− LCE(θ∗)) T + 4β6L2M2(LCE(θ0)− LCE(θ∗)) ST ( 1 ρ − 1 )2 .
We have that −2β + logm ≤ LCE(θ) ≤ 2β + logm. Thus the term LCE(θ0)− LCE(θ∗) ≤ 4β. We observe that the bias introduced due to stale embeddings is a lower order term in the bound of Theorem 3. In particular, one can use ρ = 1
1+T 1/4 without affecting the convergence rate of the
standard SGD algorithm. For a large T (which is typical in machine learning settings), this can have a significant impact on the computational complexity since a very small fraction of the documents need to be updated at each iteration.

heading: 3.2 Computational Discussion
text: We store our cache using accelerator memory. This prevents the need to have a separate task that constantly reindexes the embeddings as they change throughout training. This is feasible for moderate size datasets. For instance, 1 million training points with an embedding dimension of 512 and feature vector dimension of 1024 will use about 6 GB memory, which can fit on a single accelerator. To calculate the perturbed nearest neighbor for Gumbel-Max sampling, we simply brute force calculate the largest dot product. Accelerators such as GPUs and TPUs can do this very efficiently—in our experiments we see that the steps/second increases only a small amount as we increase the cache size. This is because the cost of computing the embeddings is much larger than the cost of doing the search due to the complexity of large transformer models.
If faster sampling is needed, it can be accelerated by using fast nearest neighbor search. Mussman et al. [28] consider an approach that applies Gumbel perturbations to the k highest scoring elements plus a small number of random elements to perform Gumbel-Max sampling in sublinear time.
However, having the cache on an accelerator has limitations due to the accelerator memory limits. In the next section, we modify Algorithm 1 to be a streaming algorithm.

heading: 3.3 Conditional Sampling Negatives
text: Note that Gumbel-Max sampling has a chance to sample the positive element. If this happens, then our gradient approximation is zero. Since we are training our model to make the score of the positive element large, this can happen often. We develop an approach to force the sampled element to be a negative element while maintaining a similar expected gradient.
Assume that the positive element is z1 and the negative elements are z2, z3, . . . , zm. We have that
∇LCEi = −β∇si1 + m∑ j=1 pijβ∇sij = −β(1− pi1)∇si1 + (1− pi1) m∑ j=2 pij 1− pi1 β∇sij .
Note that ( pi21−pi1 , pi3 1−pi1 , · · · , pim 1−pi1 ) is the conditional distribution where we condition on not sampling the first element. Thus we can sample from the conditional distribution as long as we properly scale the gradient by 1 − pi1. We note that the value of pi1 depends on the negative embeddings. Since we have all the negatives on our accelerator, we can also calculate the partition function without much additional compute, as it has the same complexity as of our nearest neighbor search.

heading: 4 Streaming Cache
text: For very large datasets, the Gumbel-Max sampling step can become difficult. If the embedding vectors are stored in accelerator memory, then only a finite number of vectors can be stored on a fixed computational budget. If the vectors are stored in CPU memory with a CPU-based retrieval system, then the nearest neighbor index needs to be constantly recreated, which is computationally expensive for large datasets. Additionally, we see in Theorem 2 that as the fraction of elements refreshed each iteration ρ decreases, then the learning rate needs to decrease as well to maintain a given bias in gradient estimation. Alternatively, more embeddings will need to be refreshed each iteration as the size of the dataset grows.
To address this issue, we develop a streaming variant of our algorithm. The key difference with respect to the full dataset setting is that instead of storing all the document embeddings in memory, we only store a sample multiset S of size αm. After every iteration we remove the oldest ραm elements in S and replace them with ραm new elements sampled uniform i.i.d. from the dataset. The fraction α can be tuned to fit a given computational budget.
Our gradient estimator approximates gradients from a cache cross-entropy loss. We define the cache cross-entropy loss to be the following:
LCacheCEi = − log exp(βsyi)
exp(βsiyi) + 1 α ∑ j∈S,j 6=yi exp(βsij)
Following [4], we scale the weight of the negative elements by 1α , otherwise the partition function would underestimate the true partition function. Thus is equivalent to shifting the scores of the negative elements by 1β log 1 α .
The use of the cache cross-entropy loss LCacheCE instead of the true cross entropy loss induces a bias in gradient estimation. We are able to show that the bias scales inversely with the cache size. The proof of Lemma 4 is a modification of the proof of Theorem 1 in [30]. Lemma 4. Assume Assumptions A2 and A3 hold, i.e., that the norm of the encoder embeddings is bounded by 1 and the norm of the gradients of the scores is bounded by M .
We have that
‖∇LCEi − E [∇LCacheCEi] ‖2 ≤ exp(O(β))M
αm ,
where the expectation is taken over the randomness of the elements in the cache.
In Algorithm 2 we present our Algorithm for training with the streaming cache. We use GumbelMaxSample(βs,S) to denote the Gumbel-Max sampling described in Fact 1 with scores restricted to documents in set S. There are two sources of bias in our gradient estimate: first due to the staleness of the cache and the second due to using a sampled set of negatives rather than the entire dataset. We handle the latter in
Algorithm 2 Streaming Cached Gumbel-Max Gradient Descent Input: Learning rate η, Cache fraction α ∈ (0, 1], Document refresh fraction ρ ∈ (0, 1] Initializes parameters θ0. Random sample S ⊆ [m]αm. Initialize embeddings table E : Ej ← φD(zj ; θ0) for all j ∈ S for t ∈ 0, 1, . . . , T − 1 do
Sample qi, yi from the training set eqi ← φQ(qi; θt) ezyi ← φD(zyi ; θt) and update Eyi = ezyi Calculate scores s̃ij ← eqi · Ej + log(1/α)/β for all j ∈ S and s̃iyi = eqi · ezyi Ŝ ← S with all instances of yi removed J ← GumbelMaxSample(βs̃i, Ŝ) ezJ ← φD(zJ ; θt) siJ ← eqi · ezJ pyi ← probability of yi under softmax(s̃i) gt ← β(1− pyi)(−∇siyi +∇siJ) θt+1 ← θt − ηgt Select oldest ραm embeddings set O of E and remove them from E Sample S ′ ⊆ [m], |S ′| = ραm Update Ej ← φD(zj ; θt+1) for j ∈ S ′ Update set S ← S −O ∪ S ′
end for
Lemma 4. For the former, note that our error guarantees in Theorem 2 still apply to the streaming cache, as we can bound the error between true gradient using this set of negatives with our gradient approximation. This allows us to bound the bias due to the staleness of the cache.
Since we were able to control the bias of our gradient estimator even with the streaming cache, we can establish a first order convergence theorem for training with Algorithm 2.
Theorem 5. Assume that Assumptions 1 and 2 hold and we use a learning rate η =√ LCE(θ0)−LCE(θ∗)√
2TSM . Running Algorithm 2 with a cache of size αm with a refresh rate ρ creates
updates θ1, θ2, . . . , θT such that
1
T T∑ t=0 E[‖∇LCE(θt)‖22] ≤ O ( 1√ T + 1 α2m2 + 1 T ( 1 ρ − 1 )2) ,
where we omit terms that depend on S, L, M , and β.
If the cache size satisfies αm ≥ T 1/4, then we do not asymptotically affect the rate of convergence. We see that using a streaming cache smaller than the entire dataset adds only a small amount of bias in convergence and reduces the computational burden of storing embeddings and calculating nearest neighbors significantly.

heading: 5 Experiments
text: We analyse the performance of our approach on the MS MARCO passage retrieval task [3] and the TREC 2019 passage retrieval task [8]. Both tasks utilize the same dataset consisting of 500,000 (query, passage) pairs and 8.8 million passages in the database. MS MARCO measures the mean reciprocal rank at 10(MRR@10) with a holdout set of query, passage pairs. TREC 2019 measures the normalized discounted cumulative gains at 10 (NDCG@10) compared to a set of human labeled relevance scores.

heading: 5.1 MSMARCO / TREC Experimental Set Up
text: Following related work [22, 38], we start with a pretrained BERT-base model [9] 1. Both our query and document encoder are initialized with this model and share parameters during training. We add a fully connected layer to the end of the encoder that projects to 512 dimensions. We do no additional pre-training. We use a global batch size of 8, the Adam optimizer [23], and we train for 250, 000 steps. We normalize the output embeddings and have a trainable parameter β scale the scores. We use a learning rate of 1× 10−5 for all experiments except when we train with a cache of 2 million elements, where we use a learning rate of 5× 10−5. This is due to the increased staleness of the large cache. Our experiments use 8 V2 Cloud TPUs. Each replica on the TPU has 8GB memory, for a total of 64GB memory.

heading: 5.2 MSMARCO / TREC Computational Performance
text: We first examine the computational performance of our method. Since we use brute force to calculate the negative with the highest perturbed score, we are interested in seeing how this affects the runtime performance of our algorithm. Randomly sampled passages are used as negatives for sample cross-entropy and the updates for the negative cache.
Due to the increased staleness of large caches, we do 512 updates when using a cache with 2 million elements. For all other cache sizes we use 256 elements.
We see in Table 1 that the cache size can scale to hundreds of thousands of elements with minimal effect on the speed of training. The cost of calculating embeddings and gradients with large transformer models is significantly more expensive than the cost of the nearest neighbor search, even if we use a brute force approach.
Further, we see that using a cache with 2 million elements requires less memory than using just 256 sample negatives, even when 512 elements are updated every iteration. This is because the memory required for the backward step is about 10x larger than the memory needed for the forward step.
Additionally, we see that there is a limit to the cache size. If we were to use all 8.8M negatives, we would not be able to fit them on 8 V2 Cloud TPUs and would have to move to a separate nearest neighbor system. Our streaming approach allows us to adapt the cache size to the available memory.

heading: 5.3 MSMARCO / TREC Experimental Results
text: We now analyze the statistical improvement of using a large negative cache. In Table 2 we present the output MRR@10 on the development set of MS MARCO passage retrieval task and NDCG@10 on the TREC 2019 passage retrieval task for training with stochastic negative mining, our negative cache approach, as well as two external benchmarks, ANCE [38] and DPR [22].
We see that our streaming cache method achieves similar statistical performance to ANCE without needing to store and search all elements in the dataset. Additionally we see that it outperforms all
1Model obtained from https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/4.
sample cross-entropy approaches with a smaller memory footprint. Further, we see that our approach outperform the BM25 negatives utilized by DPR.

heading: 5.4 Natural Question Experimental Results
text: We now experimentally compare our method on the Natural Question dataset [24]. We follow the methodology used by [22]. The goal of the experiment is to learn a model that matches questions with passages that contain the answer to the question. The dataset consists of about 60,000 query, positive passage pairs and there are about 20 million passages. In evaluation, we consider a retrieved passage a positive if it contains an exact match to one of the labeled answers to the given question.
In Table 3 we compare training with a negative cache with 2 million negatives to several state-of-theart dual encoder training methods. We see that training with a negative cache compares favorably with DPR without the use of a domain-specific hard negative mining techniques such as BM25. The use of such a techniques requires an inference step to find the hard negative among all passages, which can be expensive for large datasets. Additionally we note that ANCE is initialized with the DPR model and we use only 10% of the negatives in our cache.

heading: 6 Conclusion
text: In this paper we developed an algorithm for efficiently estimating gradients for dual encoder models with a softmax cross-entropy loss function. We see that we can efficiently obtain gradient estimations with low bias with an appropriately chosen learning rate and a sufficiently large negative cache. We analyze first order convergence of training with a streaming negative cache and establish nearoptimal convergence bounds compared to SGD for reasonable parameter choices. In dense retrieval experiments we see that our approach is efficient and can compare favorably to state-of-the-art dual encoder training methods.

Review summaries:
This paper is about finding good negative samples for contrastive training. The authors take a caching approach, and show how to make it work in a streaming fashion. Lots of theoretical results are part of the contribution.

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.  

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.