SlowGuess/NAACL_2022_OCR huggingface dataset:
Patrick Lewis*, Barlas Oğuz*, Wenhan Xiong, Fabio Petroni, Wen-tau Yih, Sebastian Riedel
{plewis,barlaso,xwhan,fabiopetroni,scottyih,sriedel}@fb.com
We propose DrBoost, a dense retrieval ensemble inspired by boosting. DrBoost is trained in stages: each component model is learned sequentially and specialized by focusing only on retrieval mistakes made by the current ensemble. The final representation is the concatenation of the output vectors of all the component models, making it a drop-in replacement for standard dense retrievers at test time. DrBoost enjoys several advantages compared to standard dense retrieval models. It produces representations which are 4x more compact, while delivering comparable retrieval results. It also performs surprisingly well under approximate search with coarse quantization, reducing latency and bandwidth needs by another 4x. In practice, this can make the difference between serving indices from disk versus from memory, paving the way for much cheaper deployments.
Identifying a small number of relevant documents from a large corpus to a given query, information retrieval is not only an important task in-and-of itself, but also plays a vital role in supporting a variety of knowledge-intensive NLP tasks (Lewis et al., 2020; Petroni et al., 2021), such as open-domain Question Answering (ODQA, Voorhees and Tice, 2000; Chen et al., 2017) and Fact Checking (Thorne et al., 2018). While traditional retrieval methods, such as TF-IDF and BM25 (Robertson, 2008), are built on sparse representations of queries and documents, dense retrieval approaches have shown superior performance recently on a range of retrieval and ranking tasks (Guu et al., 2020; Karpukhin et al., 2020; Reimers and Gurevych, 2019; Hofstätter et al., 2021b). Dense retrieval
involves embedding queries and documents as low-dimensional, continuous vectors, such that query and document embeddings are similar when the document is relevant to the query. The embedding function leverages the representational power of pretrained language models and is further finetuned using any available training query-document pairs. Document representations are computed offline in an index allowing dense retrieval to scale to millions of documents, with query embeddings being computed on the fly.
When deploying dense retrievers in real-world settings, however, there are two practical concerns: the size of the index and the retrieval time latency. The index size is largely determined by the number of documents in the collection, as well as the embedding dimension. Whilst we cannot generally control the former, reducing the embedding size is an attractive way to reduce index size. On lowering latency, Approximate Nearest-Neighbor (ANN) or Maximum Inner Product Search (MIPS) techniques are required in practice. This implies that it is far more important for retrieval models to perform well under approximate search rather than in the exact search setting. Developing a dense retrieval model that produces more compact embeddings and are more amenable to approximate search is thus the focus of this research.
In this paper, we propose DrBoost, an ensemble method for learning a dense retriever, inspired by boosting (Schapire, 1990; Freund and Schapire, 1997). DrBoost attempts to incrementally build compact representations at training time. It consists of multiple component dense retrieval models ("weak learners" in boosting terminology), where each component is a BERT-based bi-encoder, producing vector embeddings of the query and document. These component embeddings are in lower dimensions (e.g., 32 vs. 768) compared to those of regular BERT encoders. The final relevance function is implemented as a linear combination
1Our code is available at https://github.com/facebookresearch/drboost.
Human Language Technologies, pages 3102 - 3117
of inner products of embeddings produced by each weak learner. This can be efficiently calculated by concatenating vectors from each component and then performing a single MIPS search, which makes DrBoost a drop-in replacement for standard dense retrievers at test time. Component models are trained and added to the ensemble sequentially. Each model is trained as a reranker over negative examples sampled by the current ensemble and can be seen as specializing on retrieval mistakes made previously. For example, early components focus on high-level topical information, whereas later components can capture finer-grained tail phenomena. Through this mechanism, individual components are disentangled and redundancy minimized, leading to more compact representations.
There are a couple of noticeable differences in training DrBoost when compared to existing dense retrieval models. Although iterative training using negatives sampled by models learned in the previous rounds has been proposed (Xiong et al., 2020; Qu et al., 2021; Oğuz et al., 2021; Sachan et al., 2021, inter alia.), existing methods keep only the final model. In contrast, each weak learner in DrBoost is preserved and added to the ensemble. The construction of the embedding also differs. DrBoost can be viewed as a method of slowly "growing" overall dense vector representations, lending some structure to otherwise de-localized representations, while existing retrieval models encode queries and documents in one step.
More importantly, DrBoost enjoys several advantages in real-world settings. Because each weak learner in DrBoost produces very low-dimensional embeddings to avoid overfitting (32-dim in our experiments), many components can be added whilst the index stays small. Our experiments demonstrate that DrBoost produces very compact embeddings overall, achieving accuracy on par with a comparable non-boosting baseline with 4-5x smaller vectors, and strongly outperforming a dimensionally-matched variant. Probing DrBoost's embeddings using a novel technique, we also show that the embeddings can be used to recover more topical information from Wikipedia than a dimensionally-matched baseline.
Empirically, DrBoost performs superbly when using approximate fast MIPS. With a  $K$ -mean inverted file index (IVF), the simple and widely used approach, especially in hierarchical indices and Web-scale settings (Jégou et al., 2011; Johnson
et al., 2019; Matsui et al., 2018), DrBoost greatly outperforms the baseline DPR model (Karpukhin et al., 2020) by 3-10 points. Alternatively, it can reduce bandwidth and latency requirements by  $4 - 64\mathrm{x}$  while retaining accuracy. In principle, this allows for the approximate index to be served on disk rather than in expensive and limited RAM (which is typically  $25\mathrm{x}$  faster), making it feasible to deploy dense retrieval systems more cheaply and at much larger scale. We also show that DrBoost's index is amenable to compression, and can be compressed to  $800\mathrm{MB}$ ,  $2.5\mathrm{x}$  smaller than a recent state-of-the-art efficient retriever, whilst being more accurate (Yamada et al., 2021).
2 Dense Retrieval
We give here the background of dense retrieval and boosting, as well as our proposed method. More extensive related work can be found in §5.5.
Dense Retrieval involves learning a scalable relevance function  $h(q, c)$  which takes high values for passages  $c$  that are relevant for question  $q$ , and low otherwise. In the popular dense bi-encoder framework,  $h(q, c)$  is implemented as the dot product between  $\mathbf{q}$  and  $\mathbf{c}$ , dense vector representations of passages and questions respectively, produced by a pair of neural network encoders,  $E_{Q}$  and  $E_{C}$ ,
$$
h (q, c) = E _ {Q} (q) ^ {\top} E _ {C} (c) = \mathbf {q} ^ {\top} \mathbf {c} \tag {1}
$$
where  $\mathbf{q} = E_{Q}(q)$  and  $\mathbf{c} = E_{C}(c)$ . At inference time, retrieval from a large corpus  $\mathcal{C} = \{c_1,\dots ,c_{|\mathcal{C}|}\}$  is accomplished by solving the following MIPS problem:  $c^{*} = \arg \max_{c\in \mathcal{C}}\mathbf{q}^{\top}\mathbf{c}$ .
In standard settings, we assume access to a set of  $m$  gold question-passage pairs  $\mathcal{D} = \{(q_i,c_i^+)\}_{i = 1}^m$ . It is most common to learn models by training to score gold pairs higher than sampled negatives. Negatives can be obtained in a variety of ways, such as by sampling at random from corpus  $\mathcal{C}$ , or by using some kind of importance sampling function on retrieval results (see §2.1). When augmented by  $n$  negatives per gold passagedocument pair, we have training data of the form  $\widetilde{\mathcal{D}} = \{(q_i,c_i^+,c_{i,1}^-,\dots c_{i,n}^-)\}_{i = 1}^m$ , which we use to train a model, e.g., using a ranking or margin objective, or in our case, by optimizing negative log-likelihood (NLL) of positive pairs
$$
\mathcal {L} _ {\theta} = - \log \frac {e ^ {h (q _ {i} , c _ {i} ^ {+})}}{e ^ {h (q _ {i} , c _ {i} ^ {+})} + \sum_ {j = 1} ^ {n} e ^ {h (q _ {i} , c _ {i , j} ^ {-})}}
$$
2.1 Iterated Negatives for Dense Retrieval
The choice of negatives is an important factor for what behaviour dense retrievers will learn. Simply using randomly-sampled negatives has been shown to perform poorly, because they are too easy for the model to discriminate. Thus it is common to mix in some hard negatives along with random negatives, which are designed to be more challenging to distinguish from gold passages (Karpukhin et al., 2020). Hard negatives are typically collected by retrieving passages using an untrained retriever, such as BM25, and filtering out any unintentional golds. This ensures the hard negatives are at least topically-relevant.
Recently, it has become common practice to run a number of rounds of dense retrieval training to bootstrap hard negatives (Xiong et al., 2020; Qu et al., 2021; Oğuz et al., 2021; Sachan et al., 2021, inter alia.). Here, we first train a dense retriever following the method we describe above, and then use this retriever to produce a new set of hard negatives. This retriever is discarded, and a new one is trained from scratch, using the new, "harder" negatives. This process can then be repeated until performance ceases to improve. This approach, which we refer to dense retrieval with iteratively-sampled negatives is described in Algorithm 1.
2.2 Boosting
Boosting is a loose family of training algorithms for machine learning problems, based on the principle of gradually ensembling "weak learners" into a
strong learner. Boosting can be described by the following high-level formalism (Schapire, 2007). For a task with a training set  $\{(x_1,y_1),\dots ,(x_m,y_m)\}$  where  $(x_{i},y_{i})\in X\times Y$  we want to learn a function  $h:X\to Y$  , such that  $h(x_{i}) = \hat{y}_{i}\approx y_{i}$  . This is achieved using an iterative procedure over  $R$  steps:
- For round  $r$ , we construct an importance distribution  $D_r$  over the training data, based on where error  $\epsilon$  of our current model  $h$  is high
- Learn a "weak learner"  $h_r$  to minimize error  $\epsilon_r = \sum_i D_r(i) \mathcal{L}(h_r(x_i), y_i)$  for some loss function  $\mathcal{L}$  measuring the discrepancy between predictions and real values.
- Combine  $h$  and  $h_r$  to form a new, stronger overall model, e.g., by a linear combination  $h_{\mathrm{new}} = \alpha h_r + \beta h$ . The iteration can now be repeated.
The initial importance distribution  $D_0$  is usually assumed to be a uniform distribution, and  $h_0$  models a constant function. Note that how each additional model added to  $h$  is specifically designed to solve instances that  $h$  currently struggles with.
2.3 Boosted Dense Retrieval: DrBoost
We note similarities between the boosting formulation, and the dense retrieval with iteratively-sampled negatives. We can adapt a boosting-inspired approach to dense retrieval with minimal changes, as shown in Algorithm 1. Algorithmically, the only difference (lines 10-13) is that in the case of iterative negatives, the model  $h$  after  $r$  rounds is replaced by the new model  $h_r$ , whereas in the boosting case, we combine  $h_r$  and  $h$ .
In this paper, we view the boosted "weak learner" models  $h_r$  as rerankers over the retrieval distribution from the current model  $h$ . That is, when training dense boosted Retrievers, we only train using hard negatives, and do not use any random or in-batch negatives. Using the construction of negatives as a mechanism to define the importance distribution, each new model is directly trained to solve the retrieval mistakes that the current ensemble makes. Each model  $h_r$  is implemented as a bi-encoder, as in Eq. (1). We combine models as linear combinations:
$$
\text {C o m b i n e M o d e l s} (h, h _ {r}) = h _ {\text {n e w}} = \alpha h _ {r} + \beta h
$$
The coefficients could be learnt from development data, or, simply by setting all coefficients to 1,
which we find to be empirically effective. The overall model after  $R$  rounds can be written as:
$$
\begin{array}{l} h (q, c) = \alpha h _ {R} (q, c) + \beta \left(h _ {R - 1} (q, c) + \gamma (\dots)\right) \\ = \sum_ {r = 1} ^ {R} \alpha_ {r} ^ {\prime} h _ {r} (q, c) = \sum_ {r = 1} ^ {R} \alpha_ {r} ^ {\prime} \mathbf {q} _ {r} ^ {\top} \mathbf {c} _ {r} \\ = [ \mathbf {q} _ {R}, \dots , \mathbf {q} _ {1}, \mathbf {q} _ {0} ] ^ {\top} [ \alpha_ {R} ^ {\prime} \mathbf {c} _ {R}, \dots , \alpha_ {1} ^ {\prime} \mathbf {c} _ {1}, \alpha_ {0} ^ {\prime} \mathbf {c} _ {0} ] \\ = \bar {\mathbf {q}} ^ {\top} \bar {\mathbf {c}} \\ \end{array}
$$
where  $[\ldots]$  indicates vector concatenation. Thus  $h$  can be computed as a single inner product, making it a drop-in replacement for standard MIPS dense retrievers at test time.
One downside of the boosting approach is that we must maintain  $R$  encoders for both passages and questions. Since passages are embedded offline, this does not create additional computational burden on the passage side at test time. However, on the query side, for a question  $q$ , boosted dense retrieval requires  $R$  forward passes to compute the full representation, one for each subvector  $\mathbf{q}_r$ . While this step is fully parallelizable, it is still undesirable. We can remedy this for low-latency, low-resource settings by distilling the question encoders of  $h$  into a single encoder, which can produce the overall question representations  $\bar{\mathbf{q}}$  directly. Here, given the training dataset  $\mathcal{D}_{\mathrm{train}}$  of gold question-passage pairs, and a model  $h$  we want to distill, we first compute overall representations  $\bar{\mathbf{q}}$  and  $\bar{\mathbf{c}}$  for all pairs using  $h$  as distillation targets, then train a new question encoder  $E_{Q}^{\mathrm{dist}}$  with parameters  $\phi$ , by minimizing the objective:
$$
\mathcal {L} _ {\phi} = \sum_ {(q, c ^ {+}) \in \mathcal {D} _ {\text {t r a i n}}} \| E _ {Q} ^ {\text {d i s t}} (q) - \bar {\mathbf {q}} \| ^ {2} + \| E _ {Q} ^ {\text {d i s t}} (q) - \bar {\mathbf {c}} \| ^ {2}
$$
Natural Questions (NQ) We evaluate retrieval for downstream ODQA using the widely-used NQopen retrieval task (Kwiatkowski et al., 2019). This requires retrieving Wikipedia passages which contain answers to questions mined from Google search logs. We use the preprocessed and gold pairs prepared by Karpukhin et al. (2020), and report recall-at-K  $(\mathbf{R}@\mathbf{K})$  for  $\mathrm{K}\in \{20,100\}$
MSMARCO We evaluate in a Web-text setting using the widely-used passage retrieval task from MSMARCO (Bajaj et al., 2016). Queries consist of user search queries from Bing, with human-annotated gold relevant documents. We use the
preprocessed corpus, training and dev data (gold pairs and data splits) from Oğuz et al. (2021). We follow the common practice of reporting the Mean-Reciprocal-Rank-at-10 (MRR@10) metric for the public development set.
In this section, we will describe the experiments we perform, and the motivations behind them.
Exact Retrieval We are interested in understanding whether the boosting approach results in superior performance for exhaustive (exact) retrieval. Here, no quantization or approximations are made to MIPS, which results in large indices, and slow retrieval, but represents the upper bound of accuracy. This setting is the most commonly-reported in the literature.
Approximate MIPS: IVF Exact Retrieval does not evaluate how a model performs in practically-relevant settings. As a result, we also evaluate in two approximate MIPS settings. First, we consider approximate MIPS with an Inverted File Index (IVF, Sivic and Zisserman, 2003). IVF works by first clustering the document embeddings offline using  $K$ -means (Lloyd, 1982) resulting  $K$  cluster centroids. At test time, for a given query vector, rather than compute an inner product for each document in the index, we instead compute inner products to the  $K$  centroids. We then visit the n-probes highest scoring clusters, and compute inner products for only the documents in these clusters. This technique increases the speed of search significantly, at the expense of some accuracy. Increasing  $K$ , the number of centroids, increases speed, at the expense of accuracy, as does decreasing the value of n-probes. A model is preferable if retrieval accuracy remains high with very fast search, i.e., low n-probes and high  $K^2$ . In our experiments we fit  $K = 65536$  clusters and sweep over a range of values of n-probes from  $2^{0}$  to  $2^{15}$ . Other methods such as HNSW (Malkov and Yashunin, 2020) are also available for fast search, but are generally more complex and can increase index sizes significantly. IVF is a particularly popular approach due it its simplicity, and as a first coarse quantizer in hierarchical indexing (Johnson et al., 2019), since it is straightforward to apply
$^2$ Up to the point in  $K$  where the first stage search becomes the bottleneck. This happens when  $K$  is in the order of  $\sqrt{|\mathcal{C}|}$ , which is how we pick  $K = 65536$ . We also include results with  $K \in \{4092, 16384\}$  in the Appendix.
sharding to the clusters, and further search indices can be built for each cluster.
Approximate MIPS: PQ Whilst IVF will increase search speeds, it does not reduce the size of the index, which may be important for scalability, latency and memory bandwidth considerations. To investigate whether embeddings are amenable to compression, we experiment with applying Product Quantization (PQ, Jégou et al., 2011). PQ is a lossy quantization method that works by 1) splitting vectors into subvectors 2) clustering each subvector space and 3) representing vectors as a collection cluster assignment codes. We apply PQ using 4-dimensional sub-vectors and 256 clusters per subspace, leading to a compression factor of 16x over uncompressed float32.
All MIPS retrieval is implemented using FAISS (Johnson et al., 2019).
Generalization Tests In addition to in-domain evaluation, we also perform two generalization tests. These will determine whether the boosting approach is superior to iteratively-sampling negatives in out-of-distribution settings. We evaluate MSMARCO-trained models for zero-shot generalization using selected BEIR (Thakur et al., 2021) datasets that have binary relevance labels. Namely, we test on the SciFact, FiQA, Quora and ArguAna subsets. This will test how well models generalize to new textual domains and different query surface forms. We also evaluate NQ-trained models on EntityQuestions (Sciavolino et al., 2021), a dataset of simple entity-centric questions which has been recently shown to challenge dense retrievers. This dataset uses the same Wikipedia index as NQ, and tests primarily for robustness and generalization to new entities at test time.
We compare a model trained with iteratively-sampled negatives to an analogous model trained with boosting, which we call DrBoost. There are many dense retrieval training algorithms available which would be suitable for training with iteratively-sampled negatives and boosting with DrBoost. Broadly-speaking, any dense retriever could be used if utilizes negative sampling, and could be trained in Step 9 of Algorithm 1. We choose Dense Passage Retriever (DPR, Karpukhin et al., 2020) with iteratively-sampled negatives due to its comparative simplicity and popularity.
3.3.1 Iteratively-sampled negatives baseline: DPR
DPR follows the dense retrieval paradigm outlined in §2. It is trained with a combination of in-batch negatives, where gold passages for one question are treated as negatives for other questions in the batch (which efficiently simulates random negatives), and with hard negatives, sampled initially from BM25, and then from the previous round, as in Algorithm 1. We broadly follow the DPR training set-up of Oğuz et al. (2021). We train BERT-base DPR models using the standard 768 dimensions, as well as models which match the final dimension size of DrBoost. We use parametersharing for the bi-encoders, and layer-norm after linear projection. Models are trained to minimize the negative log-likelihood of positives, and the number of training rounds is decided using development data, as in Algorithm 1, using an initial  $h_0$  retriever BM25.
3.3.2 DrBoost Implementation
For our DrBoost version of DPR, we keep as many experimental settings the same as possible. There are two exceptions, which are required for adapting dense retrieval to boosting. The first is that each component "weak learner" model has a low embedding dimension. This is to avoid overfitting, and to make sure the final index size is manageable. We report using models of 32 dims (cf. the standard 768 dims), but note that training with dimension as low as 8 is stable. The second is that, as motivated in §2.3, we train each weak learner using only hard negatives, and no in-batch negatives. In effect, this choice of negatives means that each model is essentially trained as a reranker. $^{3}$  DrBoost models are fit following Algorithm 1, and we stop adding models when the development set performance stops improving. The initial retriever  $h_0$  for DrBoost is a constant function, and thus the initial negatives for DrBoost are sampled at random from the corpus, unlike DPR, which uses initial hard negatives collected from BM25.
DrBoost  $\alpha$  Coefficients DrBoost combines weak learners as a linear combination. We experiment with learning the  $\alpha$  coefficients using development data, however this does not significantly improve results over simply setting them all to 1. There-
<sup>3</sup>Note: We sample negatives from the model's retrieval distribution rather than taking the top-K retrieved negatives. We find this improves results for early rounds
fore, for the sake of simplicity and efficiency, we report DrBoost numbers with all  $\alpha = 1.0$ . Empirically, we find the magnitudes of embeddings for DrBoost's component models to be similar. In other words, one component does not dominate over others.
DrBoost Distillation We experiment with distilling DrBoost ensembles into a single model for latency-sensitive applications using the L2 loss at the end of §2.3. We distill a single BERT-base query encoder, and perform early stopping and model selection using development L2 loss.
4.1 Exact Retrieval
Exact Retrieval results for MSMARCO and Natural Questions are shown in Table 1 in the "Exact Search" Column. We find that our DrBoost version of DPR reaches peak accuracy after 5 or 6 rounds when using 32-dim weak learners (see §A.1), leading to overall test-time index of 160/192-dim. In terms of Exact Search, DrBoost outperforms the iteratively-sampled negatives DPR baseline on MSMARCO by  $2.2\%$ , and trails it by only  $0.3\%$  on NQ R@100, despite having a total dimension 4-5x smaller. It also strongly outperforms a dimensionally-matched DPR, by  $3\%$  on MSMARCO, and  $1\%$  on NQ in R@100, demonstrating DrBoost's ability to learn high-quality, compact embeddings. We also quote recent state-of-the-art results, which generally achieve stronger exact search results (AR2, Zhang et al., 2021). Our emphasis, however, is on comparing iteratively-sampled negatives to boosting, and we note that state-of-the-art approaches generally use larger models and more complex training strategies than the "inner loop" BERT-base DPR we report here. Such strategies could also be incorporated into DrBoost if higher accuracy was desired, as DrBoost is largely-agnostic to the training algorithm used.