Simple and Efficient ways to Improve REALM

Open Domain QA is typically modeled as a Machine Reading Comprehension (MRC) model $M$ which takes question $Q$ to predict an answer $A$ using a large corpus $D$ of $m$ text documents/passages $D=\{D_{1},D_{2},...D_{m}\}$. REALM Guu et al. (2020); Lee et al. (2019) specifically uses dense retrieval by learning document representations $d=\{h_{d_{1}},h_{d_{2}},...h_{d_{m}}\}$ and question representation $h_{q}$ and using approximate MIPS to retrieve a large set of top $c$ relevant passages $P=\{p_{1},....,p_{c}\}$ with their dense representations $\{h_{p_{1}},....,h_{p_{c}}\}$. A smaller set of top $k$ passages based on the retrieval score, $S_{\mathit{retr}}(p_{i},q)=h_{q}^{\top}h_{p_{i}}$, are passed to the transformer based reader module (e.g BERT) to process the tokens in the question and passages to find the answer as a span in the passage. When finetuning, the retriever is trained using distant supervision considering any passage that contains the target answer as positive. The reader is supervised by predicting the target answer span for each passage. The retriever and reader are jointly trained using the distant passage supervision and strong span level short answer supervision.

For our experiments, we follow the same setup of finetuning REALM on QA with the same optimization and evaluation setup. REALM retrieves $c=5000$ documents for distant passage supervision and $k=5$ documents are processed by the reader for finding the target answer. Similar to REALM, we use pre-trained checkpoints based on masked language modeling on the CC-News corpus. While REALM uses TPUs for distributed training for pretraining on MLM objective, the system is finetuned on a single machine with a 12GB GPU with batch size 1.

For our study, we use three open-domain QA datasets following Guu et al. (2020):
Natural Questions (NQ) Kwiatkowski et al. (2019) contains real user queries from Google Search. We consider questions with short answers ($<$=5 tokens) and the long answers for passage supervision.
WebQuestions (WQ) Berant et al. (2013) is a collection of questions collected from Google Suggest API, with Freebase entity answers whose string forms are the target short answers.
CuratedTREC (CT) Baudiš and Šedivỳ (2015) contains curated questions from TREC QA track with real user questions and answer as regular expression matching all acceptable answers.

Table 1 reports the results from our experiments and compares them to published results from REALM Guu et al. (2020). We find that our experiments produce similar results on NQ and WQ with slightly lower results on CT on the test set. We believe that this could be due to varying checkpoints due to early stopping. For fair evaluation, we use results from our experiments as a comparison for the remainder of the study.

We additionally present a breakdown of REALM’s retriever and reader performances on the development set across the three datasets. While REALM retrieves $c=5000$ documents for distantly supervising the retriever, only the top $k=5$ documents are processed by the reader for finding the right answer. Comparing the recall of answers in the retrieved documents at different subsets of documents we observe very high ($>90\%$) recall@5000 for all three datasets but the recall@5 effectively drops to $\sim$70%, showing that the document that contains the answer is not necessarily present in the top-5 highlighting limitations in the retriever.

Comparing the Exact Match accuracy of the system with the upper bound (the system is right if the passage contains the answer) shows wide gaps in the performance of the reader. While $\sim$63% of the questions from NQ have the answer in the top retrieved documents, REALM is only able to get the exact span of the answer for $\sim$36% of them showing the limitations of the reader in identifying the exact answer span in the document.

Table 5 shows that while REALM++ significantly improves over the baseline REALM model by $4.4$% on the NQ dataset, it additionally outperforms prior methods of similar size (models based on BERT_base) with no modifications to the model design. When transferred to WQ and CT which do not have human annotation for passage supervision, REALM++ shows an improvement of $3.8$% on WQ and $4.3$% on CT over base REALM. This shows that while sourcing human annotations for evidence passages can be expensive, their benefit can be transferred well to other datasets increasing their use beyond a single dataset. REALM++ produces state-of-art results on extractive Open Domain QA among models of similar size in all three datasets using a single end-to-end model.

From Table 5 we observe that our REALM++ model which uses BERT_base ($\sim$ 110M parameters) performs comparable to large models based on BERT_large and BART_large Lewis et al. (2020a) ($\sim$ 340M parameters) with 3x lesser parameters. While ReConsider_large Iyer et al. (2020) has higher accuracy numbers, their approach of using an additional model for answer span focused reranking is orthogonal to our model and can be directly applied to our outputs.

By using 8 TPUv3 cores and increased batch size for training our REALM++ model, we can process 4x more examples/sec as compared to REALM and reduce training time from 2 days to 12hrs. Further, as we train for the same number of epochs, similar improvements in accuracy could potentially be obtained using distributed training on multiple GPUs or gradient accumulation on a single GPU. Exact top-K increases the memory utilization by 5GB since we store the entire document index in memory. Our entire model fits within 12GB memory which is the equivalent of an Nvidia Titan X. This demonstrates that our REALM++ model is efficient and can improve training time by leveraging distributed training.

Open Domain QA Voorhees et al. (1999) aims to answers questions directly based on a large corpus of unstructured documents. Chen et al. (2017) use heuristic TF-IDF based retrievers for identifying relevant documents from a corpus and and RNN based reader for identifying the right answer within the document. Early systems used sparse retrieval like BM25 or TF-IDF for retrieving passages/documents Chen et al. (2017); Yang et al. (2019); Nie et al. (2019); Min et al. (2019b). Min et al. (2019a); Asai et al. (2020) explore augmenting sparse retrieval with structured knowledge via connecting graphs to improve passage retrieval. Seo et al. (2018); Das et al. (2019); Lee et al. (2019); Seo et al. (2019) explore dense embedding based methods for retrieving relevant passage using semantic match between question and passage. DPR Karpukhin et al. (2020) revisits pipelined QA systems where independent retrieval and reader systems are trained. The retrieval is a dense retrieval system trained using distant supervision from BM25. Iyer et al. (2020) extend DPR by performing answer focused reranking. We focus on REALM Guu et al. (2020), which is an end-to-end dense-retrieval system, for the purpose of our study and reranking is orthogonal and can be applied to our system as well. Lewis et al. (2020b); Roberts et al. (2020) use an alternative approach to generate the answer instead of extracting it from the document. While this reduces the constraints on the model, generated answers are not grounded to the document creating potential for hallucinations. In contrast, REALM is constrained in ensuring the answer is always a span in the document reducing any hallucinatory effects.

Passage retrieval has been an important component for open-domain QA. Das et al. (2019); Lin et al. (2018); Wang et al. (2018) use distant supervision for training the retriever by considering a paragraph that contains the ground truth answer as a positive example. Karpukhin et al. (2020) use a different form of distant supervision considering top retrieved BM25 passages as positive examples. Datasets like WikiQA Yang et al. (2015) and MSMarco Nguyen et al. (2016) provide gold supervision of candidate sentences or passages relevant to the question for supervision. Our system leverages similar candidate passage annotations from Natural Questions dataset for supervising the dense passage retrieval. While distant supervision has been shown to be effective for dense retrieval, our work experimentally shows the benefit of gold candidate passage supervision for training dense retrieval systems.