SlowGuess/NAACL_2022_OCR huggingface dataset:
Barlas Oğuz*, Kushal Lakhotia*, Anchit Gupta*, Patrick Lewis, Vladimir Karpukhin, Aleksandra Piktus, Xilun Chen, Sebastian Riedel, Wen-tau Yih, Sonal Gupta, Yashar Mehdad
{barlaso, kushall, anchit, plewis, vladk, piktus, xilun, sriedel, scottyih, sonalgupta, mehdad}@fb.com
Pre-training on larger datasets with ever increasing model size is now a proven recipe for increased performance across almost all NLP tasks. A notable exception is information retrieval, where additional pre-training has so far failed to produce convincing results. We show that, with the right pre-training setup, this barrier can be overcome. We demonstrate this by pre-training large bi-encoder models on 1) a recently released set of 65 million synthetically generated questions, and 2) 200 million post-comment pairs from a preexisting dataset of Reddit conversations. We evaluate on a set of information retrieval and dialogue retrieval benchmarks, showing substantial improvements over supervised baselines.
As a pre-training task, language modeling and its variants (causal (Radford et al., 2018), bi-directional (Peters et al., 2018; Baevski et al., 2019), masked (Devlin et al., 2018), seq2seq (Lewis et al., 2020; Raffel et al., 2019)) have proven to be extremely versatile and shown to transfer well to most, if not all NLP tasks. Nevertheless, in-domain fine tuning remains important, as there is still a gap between the pre-training task and the downstream tasks. Numerous approaches have been proposed to fill this gap, with an additional (intermediate) pre-training stage, mostly based on multi-task learning (Raffel et al., 2019; Aghajanyan et al., 2021). It's been generally accepted that the more similar the end task is to the pre-training task, the larger the gains (e.g., NLI tasks transfer better to other NLI tasks (Phang et al., 2018), QA tasks to QA tasks (Khashabi et al., 2020), inter alia).
From this perspective, information retrieval (IR), which is the task of identifying the most relevant document to a given query from a large corpus of
candidates, has a unique position. At the surface, IR looks similar to other NLP tasks in standard benchmarks, such as NLI or paraphrase detection. However, the need to accommodate large corpora imposes computational constraints, which leads to important practical differences. Most importantly, indexing needs to happen offline, therefore the candidate representations need to be calculated independently of the query representation. As a result, neural retrieval systems typically use a bi-encoder model (Figure 1), trained to minimize the similarity between the document representation and the query representation. This shallow interaction between document and query encoders makes neural IR models architecturally unique, compared to the block cross-attention transformers which are the universal choice for almost every other NLP task.
Researchers have therefore recognized the need to construct intermediate pre-training tasks that are better matched to retrieval. Lee et al. (2019) proposed the inverse cloze task (ICT), which treats sentences as pseudo-queries, and matches them to the paragraph they originate from. Chang et al. (2020) combined this with body first selection (BFS) (selecting the first paragraph given a sentence from the same document), and wiki link prediction. Guu et al. (2020) pre-trained a retrieval model jointly in an end-to-end system to minimize a language modelling objective.
In each of these cases, pre-training approaches were shown to improve over their respective baselines. However, subsequent work showed that a careful fine-tuning over a vanilla BERT model can outperform all of these approaches (Karpukhin et al., 2020). The findings for model scaling are also similar to those of data scaling. Published results show only modest improvements from larger models for retrieval, and retrieval models which top the most competitive document ranking leaderboards are still based on the relatively small BERT
Findings of the Association for Computational Linguistics: NAACL 2022, pages 1524 - 1534 July 10-15, 2022 ©2022 Association for Computational Linguistics
base architecture. $^{1}$  This is in sharp contrast to other NLP benchmarks, where data and model scaling has been extremely successful.
We hypothesise that previously proposed pretraining tasks might be still too distant from the target task, which limits useful transfer. We therefore investigate pre-training tasks for retrieval which are as closely matched to the the target task and domain as possible. To this end, we propose using two corpora for retrieval pre-training:
- A corpus of 65 million synthetically generated question-answer pairs from Wikipedia (PAQ, Lewis et al., 2021), which we target for open domain question answering and other passage retrieval tasks.
- A corpus of 220 million post-comment pairs from Reddit, which we use for dialogue retrieval tasks.
We conduct extensive evaluations on two popular information retrieval tasks, a benchmark composed of 8 knowledge-intensive retrieval tasks, and 3 dialogue retrieval benchmarks. We find that pretraining leads to strong improvements in all cases, and also demonstrate robust generalization. We compare different pre-training tasks, investigating the effect of domain and task similarity, and find both to be important. We also experiment with models of varying sizes, with and without pre-training, showing in some cases that retrieval can indeed benefit from larger models.
2 Dense retrieval
In this section we give an overview of dense retrieval models and how they are trained.
2.1 Bi-encoder architecture
A typical dense retrieval system consists of a query encoder  $E_{Q}$  and a passage encoder  $E_{P}$ , which output a fixed  $d$ -dimensional representation for each query and passage respectively. Passages are processed offline, and their representations are indexed using a fast vector similarity search library such as FAISS (Johnson et al., 2017). At runtime, an incoming query is encoded, and the top-  $k$  closest passages to its representation in vector distance are returned using the index. Dot-product similarity is most commonly used:
$$
\operatorname {s i m} (q, p) = E _ {Q} (q) ^ {\top} E _ {P} (p). \tag {1}
$$
Figure 1: Bi-encoder architecture for retrieval.
The resulting bi-encoder architecture is pictured in Figure 1. Crucially, this formulation allows passage representations to be calculated independently from the query encoder, making efficient retrieval possible.
Given a query, a relevant (positive) passage, and a list of non-relevant (negative) passages, the bi-encoder model is trained to minimize the negative log likelihood of picking the positive passage, where the probability assigned to each passage is proportional to  $e^{\mathrm{sim}(q,p)}$ . For efficiency reasons, positive passages are recycled as negative passages for queries they are not paired with in the batch, referred to as in-batch negatives. In addition, hard negatives have been found to be useful, which can either come from a standard retrieval system such as BM25, or an earlier iteration of the dense model (Xiong et al., 2020). We do training in two steps (which we refer to as iterative training). In the first step we use a single BM25 negative per query, following best practice from Karpukhin et al. (2020), and in the second step we use hard negatives obtained using the first round model. This procedure approximates the asynchronous model update, which was shown to be helpful in Xiong et al. (2020).
3.1 Pre-training tasks
In this section, we describe the datasets we used to pre-train our retrieval models.
3.1.1 PAQ
For open-domain question answering tasks, we employ the recently-released PAQ dataset (Lewis et al., 2021). This dataset consists of 65 mil
<sup>1</sup>https://microsoft.github.io/msmarco/
lion synthetic question-answer pairs, generated from Wikipedia passages. PAQ is generated by a pipeline of models trained on Natural Questions (NQ, Kwiatkowski et al., 2019) and TriviaQA (TQA, Joshi et al., 2017). PAQ's main distinguishing features relative to other QA-pair generation techniques are its large size, and the use of a novel global consistency filtering. This leads to higher quality, less ambiguous open-domain-style questions than can be achieved using standard consistency filtering with machine-comprehension models (Alberti et al., 2019).
PAQ has previously been employed as a semistructured knowledge base of facts extracted from Wikipedia, and used as data-augmentation for closed-book question answering models (Roberts et al., 2020). However, since QA-pairs in PAQ are generated from Wikipedia passages, we can repurpose PAQ as a source of training data for a passage retrieval task. Here, given a PAQ question, the task is to retrieve the Wikipedia passage that was used to generate said question from a pool of negatives. PAQ's size makes this a suitable large-scale pretraining task and represents a close proxy of the actual downstream open-domain QA retrieval task.
3.1.2 Reddit
For dialogue tasks, we use 200 million postcomment pairs mined from Reddit. This dataset was originally extracted and made available by pushshift.io and shown to be useful for dialogue and chit-chat applications previously (Humeau et al., 2019; Roller et al., 2020).
3.2 Evaluation tasks
In this section, we describe our evaluation setup.
3.2.1 Passage retrieval
We evaluate on a mix of standard information retrieval and open-domain question answering benchmarks, and a suite of knowledge-intensive retrieval tasks:
MSMARCO (Nguyen et al., 2016) is a suite of benchmarks created using real user queries to the Bing search engine, with human annotated search results. We evaluate on the passage retrieval task, which is widely reported on in the IR community.
Natural Questions (NQ, Kwiatkowski et al., 2019) is a popular open-domain QA dataset, with questions originating from Google users, and answers annotated from Wikipedia.
KILT (Petroni et al., 2020) is a benchmark consisting of a diverse set of 8 knowledge-intensive tasks, including fact-checking, entity linking, relation extraction, dialogue and question answering. All tasks are grounded in Wikipedia, and we report on the passage selection metrics.
3.2.2 Retrieval for dialogue
We also evaluate on a set of dialogue retrieval benchmarks, to see how far our conclusions generalize to a different domain.
ConvAI2 is based on the PersonaChat dataset (Zhang et al., 2018), and was presented for the NeurIPS ConvAI2 competition (Dinan et al., 2019). The task involves selecting the correct next utterance in a dialogue, out of 20 candidates, given the dialogue history as well as some context about the speakers persona.
Ubuntu v2 (Lowe et al., 2015; Kummerfeld et al., 2019) is a large corpus of 1 million conversations from Ubuntu chat logs, which document users receiving support from other users regarding Ubuntu-related issues.
DSTC7 (Gunasekara et al., 2019) is a challenge set consisting of  $100\mathrm{k}$  samples extracted from the Ubuntu dataset described above.
**Frameworks** We use the Pytorch Lightning (PL) framework (Falcon, 2019) for implementing our models. PL enables effortless scaling to hundreds of GPUs, with memory and speed optimizations such as half-precision training, and sharded gradients during distributed training. We add memory-mapped data loaders, which allow us to scale to datasets with hundreds of millions of query-passage pairs. We use pre-trained encoders provided by the Huggingface transformers library (Wolf et al., 2020).
In-batch negatives Following (Karpukhin et al., 2020) we implement in batch negatives by using the differentiable all gather primitive provided by PL. Unlike the original implementation in (Karpukhin et al., 2020) this lets us gather negatives across all nodes leading to higher training efficiency.
Validation Metrics Evaluating neural retrieval models requires embedding tens of millions of passages for indexing. This is a one-time, manageable cost for deployment systems, however for research
iteration and model selection purposes, it is prohibitively expensive. One option is to use proxymetrics such as validation cross-entropy loss, or in-batch accuracy to do model selection. Unfortunately such metrics often do not correlate well with end-to-end retrieval accuracy. As a middle-ground, we implement distributed in-memory validation using the all gather primitive. This allows us to use a fairly large proxy corpus of up to 300k passages, including up to 50 hard negative examples for each test query. We find that using mean reciprocal rank on this corpus as a model selection metric correlates well with full evaluation metrics.
Training details Pre-training on PAQ and Reddit are run for up to 10 epochs on 64 Nvidia V100 32GB GPUs, with the ADAM optimizer and triangular learning rate schedule. Learning rate and batch size vary for each model, and are presented in the appendix. We fine-tune for up to 40 epochs on the end task on 8 GPUs. For BERT and DeBERTa models, we use the [CLS] token directly as the representation, whereas for RoBERTa we add a linear projection of the same size and an additional layer normalization. BERT models use seperate encoders for query and passage, where RoBERTa and DeBERTa models use shared encoders.
Data preparation For PAQ pre-training, we mined negative examples using a publicly available DPR checkpoint. For Reddit, the engineering effort to setup an index of 200M documents was too large; therefore we pre-train it without negatives. For MSMARCO and KILT, we use standard preprocessing and splits, and for NaturalQuestions we follow (Karpukhin et al., 2020). For the dialogue tasks, we use the dataset-provided negative examples when available. We concatenate all dialogue context (including persona for ConvAI2) to form the query, and truncate from the beginning if it is longer than 256 tokens.
Our code, pre-trained checkpoints, and preprocessed data files are publicly available<sup>3</sup>.
In this section, we summarize our main results, before we dive into some analysis in the next section.
4.1 Passage retrieval results
Our main results for passage retrieval are presented in two tables, Table 1 for MSMARCO and NQ, and Table 2 for KILT. PAQ-based pre-training results in strong gains on almost all passage retrieval tasks. For NaturalQuestions, pre-training improves  $+3.2$  points over our non-pretrained baseline on top-20 accuracy, without using iterative training (Figure 2). Our setup with iterative training is most similar to (Xiong et al., 2020), on which pre-training improves by additional 2.1 points (81.9 vs. 84.0). We advance the best published results (Qu et al., 2021) by  $+1.7$  points on both top-20 and top-100 accuracy. We note that the main contribution of (Qu et al., 2021) is using a large cross-encoder model to pre-filter training data - an approach which is orthogonal to pre-training and could provide additional gains. On MSMARCO, we see similar gains, improving  $+3.8$  points over our best non-pretrained baseline.
On KILT, we advance passage retrieval SoTA on all tasks by 6.7 points of R-precision on average. This result shows that PAQ-based pre-training generalizes well across a wide variety of tasks.
4.2 Dialogue retrieval results
To further verify the matched-domain hypothesis, we conduct experiments in the dialogue retrieval domain, using Reddit chat threads as pre-training data. We see clear gains on all datasets over vanilla BERT baselines, affirming the usefulness of additional pre-training for retrieval. However, the gains are less pronounced for UbuntuV2, which has a much larger training dataset. Nevertheless, our best model (RoBERTa large) still outperforms the previous SoTA by a comfortable margin on two tasks. For DSTC7, the results also support our conclusions, however we were not able to reproduce previous baselines on this dataset, and our numbers are generally lower.
5 Pre-training retrieval models
In this section we cover our findings regarding how to best pre-train bi-encoder models for retrieval. We compare our pre-training approach with previous approaches, and emphasize the importance of picking the right pre-training task. We discuss the effects of data and model size for pre-training retrieval models.
$^{2}$ https://github.com/facebookresearch/dpr, with id checkpoint.retriever(single.nq.bert-base-encoder $^{3}$ https://github.com/facebookresearch/dpr-scale
Methods	Base model	MSMARCO MRR@10	Natural Questions
R@5	R@20	R@100
1	BM25 (anserini) (Yang et al., 2017)	-	18.7	-	59.1	73.7
2	DPR (single) (Karpukhin et al., 2020)	BERTbase	-	65.8	78.4	85.4
3	GAR (Mao et al., 2020)	-	-	-	74.4	85.3
4	ANCE (single) (Xiong et al., 2020)	RoBERTa base	33.0	-	81.9	87.5
5	RocketQA (Qu et al., 2021)	ERNIEbase	37.0	74.0	82.7	88.5
6	DPR(ours)	BERTbase	29.0	65.5	78.3	85.6
7	DPR(ours)	BERTlarge	28.8	69.14	80.19	86.73
8	DPR(ours)	RoBERTa base	29.5	67.00	79.03	85.42
9	DPR(ours)	RoBERTa large	30.2	69.67	81.27	87.01
10	DPR(ours)	DeBERTa xlarge-v2	-	72.66	82.38	87.56
11	DPR-PAQ	BERTbase	31.4	74.5	83.7	88.6
12	DPR-PAQ	BERTlarge	31.1	75.3	84.4	88.9
13	DPR-PAQ	RoBERTa base	32.3	74.15	84.01	89.2
14	DPR-PAQ	RoBERTa large	34.0	76.93	84.68	89.22
15	DPR-PAQ	DeBERTa xlarge-v2	-	73.38	83	88.61
Table 1: Passage retrieval results for MSMARCO development set and NaturalQuestions test set.
Methods	Base model	FEV	T-REx	zsRE	NQ	HoPo	TQA	WoW	Avg.
BM25	-	40.1	51.6	53.0	14.2	38.4	16.2	18.4	33.1
Multi-task DPR	BERTbase	52.1	61.4	54.1	40.1	41.0	34.2	24.6	43.9
DPR-PAQ	BERTbase	61.4	68.4	73.28	44.1	44.6	38.9	26.5	50.6
DPR-PAQ	BERTlarge	62.8	66.58	66.9	42.6	42.1	37.9	23.4	48.9
Table 2: Paragraph-level  $R$  -Precision on the KILT benchmark.
5.1 Picking the pre-training task
As pointed out earlier, previous attempts at pretraining dense retrieval models have largely been ineffective. In Table 4, we confirm this conclusion. We see that BFS and ICT do result in nontrivial zero-shot retrieval performance on the NQ dataset. However, after fine-tuning these gains disappear, and they do not outperform a vanilla BERT model. The performance of PAQ-pretrained retrieval is exceptionally strong even before finetuning. This is expected to an extent, since PAQ has been trained on NQ, and many NQ training questions might already appear verbatim in the PAQ generated questions. Nevertheless, pre-training with PAQ results in robust gains, which persist after fine-tuning. Note that both BFS and ICT were pre-trained on more data than PAQ (200 million pairs vs. 65 million). We conclude that PAQ pairs are higher quality, and better matched to the end task than previously proposed artificial pre-training tasks, resulting in better performance.
For the dialogue experiments, we compare
against (Humeau et al., 2019), who also pre-trains on the same Reddit corpus, but using a cross-encoder with masked-language-modeling and next-sentence-prediction objectives  $a la$  BERT (Devlin et al., 2019). This allows us to compare bi-encoder pre-training, with cross-encoder pre-training on the same dataset. Looking at Table 3, we see that bi-encoder pre-training (DPR-Reddit,  $\mathrm{BERT}_{\text {base }}$ ) performs significantly better than cross-encoder pre-training on the ConvAI2 dataset (rows 5&8). However, the same conclusion does not hold for the larger and more domain-mismatched Ubuntu corpus. (Our RoBERTa-large bi-encoder does improve over (Humeau et al., 2019), but we don't have a corresponding cross-encoder pre-trained baseline for this model.) We conclude that transfer is somewhat fragile for dense retrieval pre-training, and is sensitive to domain and task mismatch.
5.2 Effect of data size
In Figure 2 we investigate the effect of pre-training data size on retrieval performance. We randomly downsample the PAQ pre-training dataset, and plot
Methods	Base model	ConvAI2	DSTC7	Ubuntu v2
R@1	R@1	MRR	R@1	MRR
1	(Wolf et al., 2019)	BERTbase	82.1	-	-	-	-
2	(Chen and Wang, 2019)	BERTbase	-	64.5	73.5	-	-
3	(Dong and Huang, 2018)	BERTbase	-	-	-	75.9	84.8
4	(Humeau et al., 2019)	BERTbase	83.3	66.8	74.6	80.6	88.0
5	(Humeau et al., 2019) (Reddit)	BERTbase	86.9	70.9	78.1	83.6	90.1
6	DPR (ours)	BERTbase	82.4	53.1	62.6	80.6	87.9
7	DPR (ours)	RoBERTabase	84.6	58.4	68.2	84.2	90.4
8	DPR-Reddit	BERTbase	88.5	61.5	70.2	82.0	88.8
9	DPR-Reddit	BERTlarge	88.2	62.0	70.9	81.8	88.7
10	DPR-Reddit	RoBERTabase	88.4	66.5	75.1	85.1	90.9
11	DPR-Reddit	RoBERTalarge	90.7	68.2	76.4	86.3	91.7
Table 3: Dialogue retrieval results.
Pre-training data	w/o FT	w/ FT
None	-	78.4
BFS	37.0	75.7
ICT	25.5	77.0
PAQ	78.1	81.6
Table 4: Comparison of different pre-training data, with and without fine-tuning (FT). Metric is top-20 accuracy on NaturalQuestions test set. Baseline is vanilla BERT-base model.
top-20 accuracy on NQ after fine-tuning on the full NQ training set. We see that as little as 1 million pre-training examples can improve performance, with larger pre-training data resulting in more gains as expected. This suggests that expanding PAQ with even more questions could potentially be beneficial (though this could be contingent on the quality of additional generated questions).
It is interesting to note that additional MLM training is not generally helpful for retrieval on open-domain QA. RoBERTa (Liu et al., 2019) was trained on an order of magnitude more data for much longer compared to BERT, yet fine-tuning on RoBERTa results in little, if any improvement over BERT, in the absence of retrieval-specific pretraining (Table 1, rows 6&8). For dialogue retrieval, better MLM training does help, as was shown previously (Humeau et al., 2019).
5.3 Effect of model size
We experimented with pre-trained models of varying sizes, including BERT(base/large) (Devlin et al., 2019), RoBERTa(base/large) (Liu et al.,
Figure 2: Effect of different sizes of PAQ data for pretraining. Results show top-20 accuracy on NQ after fine-tuning. No iterative pre-training is used.
2019), and DeBERTa(xlarge-v2) (He et al., 2020). In terms of how better and larger pre-trained models interact with retrieval-specific pre-training, we get mixed results. For instance, for passage retrieval, DeBERTa-xlarge-v2 model does outperform the BERT-base baseline significantly in the fully supervised setting (Table 1, rows 6&10), yet this gain disappears after additional pre-training with PAQ (rows 11&15). The opposite is true when comparing RoBERTa vs. BERT, as we see RoBERTa performing better after intermediate pretraining, both for passage retrieval (rows 12&14) and for dialogue tasks (Table 3, rows 9&11). In contrast to the clear-cut conclusions for other NLP tasks, it is hard to conclude whether larger and better language models actually make better retrieval models.
DPR-PAQ
DPR	R@20✓	R@20×
R@20✓	2.5	3.2
R@20×	3.3	3.1
Table 5: Mean Levenshtein distance to most similar question in PAQ, for DPR-PAQ and a DPR baseline for NQ test questions, stratified by whether the model achieves Recall@20
5.4 Effects of PAQ on Retrieval
In section 4.1, we established that pretraining on PAQ is beneficial for passage retrieval for QA. However, it is worth considering where the source of this improvement lies. Lewis et al. (2021) note that QA-pairs in PAQ have substantial overlap with the test sets of NQ — indeed, this is intentional, given their aim of preempting a large number of probable questions for use as a cache for question answering models. In fact,  $\sim 9\%$  of the NQ test questions appear verbatim in PAQ.
It is worth investigating then, whether the gains we observe are due to simply memorizing the relevant passages for PAQ questions which overlap with test questions, or, whether they are due to learning a more robust, generalizable model behavior.
To investigate, we compare the predictions of the DPR-PAQ retriever with an otherwise equal baseline DPR model, without PAQ-pretraining. On the subset of the NQ test set that overlaps verbatim with PAQ questions, we find that DPR-PAQ achieves  $95.5\%$  R@20, whereas the baseline achieves  $94.8\%$  These are both remarkably high scores, indicating that these verbatim questions are very easy for models to solve, regardless of pretraining. Due to the very similar performance on this subset, the difference in overall performance cannot be attributed to simply memorising verbatim-overlapping questions.