PaperlayerView exact revision on arXiv

Source-linked AI summary

End-to-End Training of Neural Retrievers for Open-Domain Question Answering

Devendra Singh Sachan, Mostofa Patwary, Mohammad Shoeybi, Neel Kant, Wei Ping, William L Hamilton, Bryan Catanzaro

arXiv:2101.00408v2cs.CLcs.AI

TL;DR

Neural retriever training for OpenQA lacks clear guidance on combining unsupervised and supervised methods. The paper studies unsupervised pre-training followed by supervised finetuning, then evaluates end-to-end reader-retriever training and scaling. It reports improved retrieval and answer extraction, including state-of-the-art results and consistent gains from larger models.

  • Problem

    The paper addresses uncertainty about how unsupervised and supervised methods can be used most effectively for neural retrievers.

  • Method

    The method pre-trains retrievers with ICT or masked salient spans, then finetunes them using supervised question-context pairs and explores two end-to-end reader-retriever designs.

  • Results

    The methods obtain state-of-the-art retrieval and answer-extraction results, including 84 top-20 retrieval accuracy on Natural Questions and gains over DPR, REALM, and RAG.

  • Takeaways & Limitations

    Pre-training strategy effectiveness varies with supervised-data availability, while larger end-to-end models consistently improve performance over smaller ones.

  • Takeaways & Limitations

    The system was not designed for robust quality assurance in sensitive settings, and training-data biases may remain without explicit controls or regularization.

Abstract

from arXiv · show

Recent work on training neural retrievers for open-domain question answering (OpenQA) has employed both supervised and unsupervised approaches. However, it remains unclear how unsupervised and supervised methods can be used most effectively for neural retrievers. In this work, we systematically study retriever pre-training. We first propose an approach of unsupervised pre-training with the Inverse Cloze Task and masked salient spans, followed by supervised finetuning using question-context pairs. This approach leads to absolute gains of 2+ points over the previous best result in the top-20 retrieval accuracy on Natural Questions and TriviaQA datasets. We also explore two approaches for end-to-end supervised training of the reader and retriever components in OpenQA models. In the first approach, the reader considers each retrieved document separately while in the second approach, the reader considers all the retrieved documents together. Our experiments demonstrate the effectiveness of these approaches as we obtain new state-of-the-art results. On the Natural Questions dataset, we obtain a top-20 retrieval accuracy of 84, an improvement of 5 points over the recent DPR model. In addition, we achieve good results on answer extraction, outperforming recent models like REALM and RAG by 3+ points. We further scale up end-to-end training to large models and show consistent gains in performance over smaller models.

1 Introduction

OpenQA retrieves relevant documents from a large knowledge source before a reader extracts an answer. This work studies how unsupervised pre-training, supervised finetuning, and end-to-end training improve neural retrievers and readers.

  • Motivation: The study addresses how unsupervised and supervised retriever training can be combined effectively when OpenQA evidence contains millions of documents.The authors identify this comparison and its conditions as a central unresolved question.
  • Contributions: Readers that process retrieved documents separately improve retrieval, whereas readers that process them jointly improve answer extraction.The paper compares these two end-to-end designs on both retrieval accuracy and answer extraction.
  • Contributions: End-to-end training achieves new state-of-the-art retrieval and answer-extraction results, including 84 top-20 accuracy on Natural Questions, 5 points above DPR.The reader and retriever are trained together in the proposed end-to-end approaches.
  • Contributions: Answer extraction gains exceed 3 points over REALM and RAG, and scaling retriever-reader models produces consistent gains over smaller configurations.These results extend the reported improvements beyond retrieval accuracy.

2 Neural Retriever

The neural retriever uses dual encoders to rank evidence documents by question-context relevance. Training combines supervised question-context examples with unsupervised ICT or masked-salient-span pre-training, followed by supervised finetuning.

  • Retriever architecture: The retriever ranks a large evidence collection and returns the top-ranked documents as question-specific context.Its question and context representations support ranking over millions of documents.
  • Retriever architecture: A dual-encoder architecture uses separate question and context encoders, modeled with BERT-style transformer networks.The encoders use the first sequence-token hidden state as their output representation.
  • Retriever architecture: The relevance score is the dot product between question and context representations, with a scaling factor used to improve optimization for large hidden sizes.The model is trained by maximizing relevance likelihood, while in-batch contexts approximate negatives for efficient normalization.
  • Training approaches: Supervised training uses annotated question-answer data, BM25-selected answer-containing contexts, and additional hard negatives when context annotations are unavailable.This setup follows the common distant-supervision approach for neural retrievers.
  • Training approaches: ICT treats a sentence as a query and the remaining paragraph sentences as context, while masked salient spans trains a reader to predict masked entities from retrieved documents.The proposed approach pre-trains with either objective and then finetunes the retriever on supervised data.

3 End-to-End Retriever and Reader Training

Section 3 presents two end-to-end supervised strategies that train the retriever and reader jointly: one processes retrieved documents separately, while the other aggregates them together.

  • Overview: The two approaches jointly train retriever and reader parameters from question-answer data, aiming to improve both components through answer prediction.The retriever first selects top-k context documents, which the reader then uses to predict the answer.
  • Approach 1: Individual Top-k: Individual Top-k computes reader likelihoods separately for each retrieved document and combines them using retrieval probabilities normalized over the retrieved set.Its objective is the negative marginal log-likelihood.
  • Approach 1: Individual Top-k: Compared with RAG, the method updates both retriever encoders and uses T5 rather than BART as the reader.The paper attributes substantial gains over RAG to these differences.
  • Approach 2: Joint Top-k: The likelihood formulation conditions answer generation on the question, all retrieved documents, and their retrieval scores.The approach follows a reader likelihood over the retrieved set and uses a final loss defined by the paper’s objective.
  • Approach 2: Joint Top-k: Joint Top-k concatenates each retrieved document with the question, then lets the decoder jointly attend to all document representations for information aggregation.Retriever similarity scores are incorporated to bias encoder-decoder attention through a trainable parameter λ.

4 Experimental Setup

The experiments use Natural Questions and TriviaQA with a December 2018 English Wikipedia evidence collection, evaluating base and large model configurations.

  • Implementation: Training details and hyperparameters are provided in Appendix A for reproducibility.
  • Datasets: Experiments use Natural Questions and TriviaQA, two widely used open-domain question-answering datasets.Natural Questions contains Google search questions with Wikipedia answer annotations, while TriviaQA contains web-sourced trivia questions and answers.
  • Evidence: The evidence collection contains 21,015,324 English Wikipedia documents, each 100 words long.The corpus is the December 2018 preprocessed dump released by Karpukhin et al. (2020).
  • Models: The study compares base and large model configurations with substantially different depths, hidden sizes, attention heads, and parameter counts.BERT-base and T5-base contain 110M and 220M parameters, whereas BERT-large and T5-large contain 330M and 770M.

5 Results: Retriever Training

The study compares supervised, unsupervised, and end-to-end strategies for training neural retrievers, finding gains from pre-training, score scaling, longer training, hard negatives, and selective encoder updates.

  • Supervised Retriever Training: Relevance-score scaling, training for 80 epochs, and adding one hard negative improve supervised retriever accuracy over the original DPR settings.The reported top-5 and top-20 gains from scaling and longer training are 1.5–2 points before the additional hard-negative improvement.
  • Effect of Retriever Initialization: More than 8 points of zero-shot retrieval improvement are obtained from masked salient spans over ICT initialization, while BERT initialization performs poorly.Both unsupervised strategies bootstrap the retriever before supervised training.
  • Effect of Retriever Initialization: 2–3 points of absolute retrieval improvement over strong supervised training are obtained by combining ICT and masked salient spans pre-training with supervised finetuning.The gains are consistent across Natural Questions and TriviaQA, indicating that supervised finetuning does not erase the pre-training benefits.
  • Effect of Amount of Training Data: Masked salient spans pre-training is more effective with limited supervision, while its gains converge with ICT beyond 40% of Natural Questions training data.The comparison varies supervised data from 1% through 50%, with the low-resource advantage strongest for masked salient spans.
  • Effect of End-to-End Training: Individual Top-k improves retrieval when encoders are updated, whereas Joint Top-k yields little higher-k retrieval improvement and is mainly useful for answer extraction.For Joint Top-k, query updates improve top-1 but not substantially higher top-k scores; the context encoder was not updated after initial experiments.

6 Results: Answer Extraction

End-to-end training improves answer extraction across model sizes and datasets, with Individual Top-k and Joint Top-k showing complementary strengths in retrieval and document aggregation.

  • Individual Top-k Approach: The model outperforms RAG by 3.5+ points on Natural Questions and 2.8 points on TriviaQA in the large configuration.The reported improvements are attributed to a stronger initial retriever, stronger reader, and updating both query and context encoders.
  • Joint Top-k Approach: Joint Top-k outperforms FiD by 1 point on Natural Questions in the base configuration, highlighting the value of end-to-end training.FiD uses DPR and T5 but finetunes only the reader, whereas the proposed approach trains the retriever and reader together.
  • Joint Top-k Approach: Joint Top-k answer-extraction EM improves as more documents are retrieved, unlike Individual Top-k, whose performance can decrease with larger top-k values.Joint Top-k better aggregates information across retrieved documents.
  • Joint Top-k Approach: Similarity-enriched attention adds 0.8–1 EM points for top-k values of 5, 10, and 25, with a smaller gain at top-k=50.The benefit of this enrichment diminishes as more documents are retrieved.
  • Comparison of Approaches: The two end-to-end objectives have complementary effects: Individual Top-k mainly improves retrieval, while Joint Top-k mainly improves answer extraction.This division follows the reported retrieval and answer-extraction analyses.

7 Related Work

Related work developed neural dual-encoder retrievers through unsupervised and supervised training, but prior studies did not systematically compare these strategies in the most challenging OpenQA setting.

  • Neural Retriever Training: Earlier work used word-frequency dense representations, BERT-based dual encoders, ICT pre-training, or supervised hard negatives to train retrievers with fewer annotated pairs.These approaches include Yih et al. (2011), Lee et al. (2019), and Karpukhin et al. (2020).
  • Pre-training for Retrieval: Prior paragraph-level pre-training studies evaluated evidence collections up to 1M documents, whereas this work studies OpenQA with 21M documents.The larger evidence collection makes the retrieval setting more challenging.
  • This Work: This work combines ICT and masked salient spans pre-training with supervised finetuning and then updates the retriever end to end using question-answer pairs.The combined strategy distinguishes the work from prior pre-training and supervised-retrieval studies.

8 Conclusion

The paper studies pre-training and end-to-end training strategies for neural OpenQA retrievers, reporting state-of-the-art retrieval and answer-extraction results.

  • Contributions: The paper systematically investigates ICT and masked salient spans pre-training for supervised dual-encoder retriever training.It also evaluates two end-to-end reader–retriever approaches that process retrieved documents individually or jointly.
  • Overall Results: The proposed methods achieve state-of-the-art results on both retrieval and answer extraction.The conclusion summarizes gains from both individual-document and joint-document reader designs.

Broader Impact and Ethics Statement

The paper discusses potential applications and deployment boundaries for its OpenQA systems, emphasizing that dataset-specific training limits off-the-shelf use and that sensitive applications require quality assurance.

  • Broader Impact and Ethics Statement: Models trained on Natural Questions and TriviaQA are not deployable off-the-shelf because they target those datasets’ specific training distributions.Real-world applications should retrain systems on custom data matching practical queries.
  • Broader Impact and Ethics Statement: Sensitive uses such as legal or health-care applications require extensive quality-assurance testing because the prototype was not designed to meet those criteria.
  • Broader Impact and Ethics Statement: The authors warn that training data may introduce social biases, which users should evaluate and address through training-data and procedural modifications.
  • Training Infrastructure: The implementation uses large-scale mixed-precision and distributed training with BERT and T5 models on NVIDIA A100 GPUs.The supplied passages describe Megatron-LM and PyTorch-based distributed training on the Selene cluster.
  • Retriever Pre-training: Masked salient-span training uses masked sentences as queries, retrieves evidence documents, and trains the reader to predict the masked words.The method refreshes evidence embeddings asynchronously during training.

A.3 End-to-End Supervised Training

The appendix describes inference and training details for end-to-end OpenQA models, including separate-document and joint-document reader configurations and exact retrieval over Wikipedia-scale evidence.

  • End-to-End Training: The Individual Top-k and Joint Top-k approaches use separate training configurations, with evidence-index refreshes and different GPU allocations.Individual Top-k training used 240 GPUs, whereas Joint Top-k used 64 GPUs.
  • Retrieval: The system retrieves top-k documents from approximately 21M evidence documents using exact matrix-multiplication search rather than LSH or another approximate index.The authors report that exact search was not a training bottleneck on GPUs.
  • Dataset-specific Details: For Natural Questions and TriviaQA, training uniformly samples a target answer, and answer extraction omits the Wikipedia article title from each retrieved document’s reader input.
  • Model Configuration: Table 8 reports the hyperparameters used to pre-train the BERT and T5 models.
  • Inference: During inference, the reader greedily generates an answer for each retrieved document and selects the answer with the highest likelihood score.Generated answers are scored using Eq. 5 before final selection.
  • Qualitative Examples: Table 9 presents top-1 retrieved documents for Natural Questions test examples, highlighting documents that contain the answer.

B.1 For all reported experimental results

The reproducibility statement identifies the code, data, infrastructure, evaluation procedures, and hyperparameter-selection practices used for the experiments, while noting limitations in validation reporting and repeated runs.

  • Code Availability: The authors provide source code based on Megatron-LM, but the submitted codebase lacks an exhaustive README file.
  • Computing Infrastructure: Experiments ran on NVIDIA Selene infrastructure with nodes containing 256 CPUs, 2.2TB memory, and eight 80GB NVIDIA A100 GPUs.
  • Compute Cost: The code was not carefully optimized for runtime or optimal hardware use, despite reporting average runtime and compute in Appendix A.
  • Validation Reporting: Validation-set performance is not reported in the main paper, although model selection used validation performance under a stated experimentation protocol.
  • Experimental Repetition: Because the experiments were compute-intensive and took tens of hours to several days, the authors did not run multiple trials and used seed 1234 for all runs.
  • Data and Splits: The experiments use standard NQ, TriviaQA, and Wikipedia data, with English-language datasets and standard train, development, and test splits.Dataset statistics and preprocessing details are provided in the paper’s referenced sections and tables.
Loading 2101.00408v2…