Source-linked AI summary
Re2G: Retrieve, Rerank, Generate
Michael Glass, Gaetano Rossiello, Md Faisal Mahbub Chowdhury, Ankita Rajaram Naik, Pengshan Cai, Alfio Gliozzo
TL;DR
Knowledge-intensive generation needs access to large external knowledge without relying solely on increasingly large parameter spaces. Re2G combines neural retrieval, reranking, and BART generation, using target-output distillation and mixed BM25-neural retrieval; it reports gains across four KILT tasks.
Problem
Large transformers store substantial knowledge, but expanding their parameter spaces motivates using non-parametric memory for knowledge-intensive tasks.
Method
Re2G combines neural initial retrieval and reranking with BART-based generation, merges BM25 and neural results, and trains components using target-sequence knowledge distillation.
Results
Re2G improves headline KILT metrics across five datasets, with relative gains of 9% to 34% over previous state of the art.
Takeaways & Limitations
Reranking improves retrieval and end-to-end performance while enabling multiple retrieval sources, and better passages support the generation component.
Takeaways & Limitations
The paper identifies further domain-adaptation experiments for question answering and dialog as useful for understanding real-world applications.
Abstract
from arXiv · showhide
As demonstrated by GPT-3 and T5, transformers grow in capability as parameter spaces become larger and larger. However, for tasks that require a large amount of knowledge, non-parametric memory allows models to grow dramatically with a sub-linear increase in computational cost and GPU memory requirements. Recent models such as RAG and REALM have introduced retrieval into conditional generation. These models incorporate neural initial retrieval from a corpus of passages. We build on this line of research, proposing Re2G, which combines both neural initial retrieval and reranking into a BART-based sequence-to-sequence generation. Our reranking approach also permits merging retrieval results from sources with incomparable scores, enabling an ensemble of BM25 and neural initial retrieval. To train our system end-to-end, we introduce a novel variation of knowledge distillation to train the initial retrieval, reranker, and generation using only ground truth on the target sequence output. We find large gains in four diverse tasks: zero-shot slot filling, question answering, fact-checking, and dialog, with relative gains of 9% to 34% over the previous state-of-the-art on the KILT leaderboard. We make our code available as open source at https://github.com/IBM/kgi-slot-filling/tree/re2g.
1 Introduction
Re2G extends retrieval-augmented generation with reranking and mixed-source retrieval in a BART-based sequence-to-sequence model. Across four diverse KILT tasks, it reaches the top of the leaderboards with substantial relative gains.
- Re2G combines neural initial retrieval and reranking within BART-based sequence-to-sequence generation.
- Its reranker merges retrieval results with incomparable scores, enabling ensembles of BM25 and neural retrieval.
- A novel knowledge-distillation variation trains initial retrieval, reranking, and generation end-to-end using only target-sequence ground truth.
- Re2G evaluates slot filling, question answering, fact checking, and dialog, reaching the top of the KILT leaderboards.
- 9%, 31%, 34%, 22%, and 10% relative gains improve headline KILT metrics for T-REx, Natural Questions, TriviaQA, FEVER, and Wizard of Wikipedia, respectively.
2 Related Work
Related work spans retrieval, reranking, and generative approaches for knowledge-intensive KILT tasks. Re2G builds on systems such as RAG, DPR, KGI, and BM25 while addressing multi-source retrieval through reranking.
- KILT unifies eleven datasets across five tasks with shared Wikipedia provenance, evaluation metrics, and a public leaderboard; this work focuses on four tasks.
- GENRE evaluates retrieval by generating Wikipedia page titles, whereas BART and T5 can perform downstream tasks using only implicit parametric knowledge.
- RAG retrieves passages with DPR and generates answers using a BART-initialized sequence-to-sequence model, making it a strong KILT baseline.
- KGI trains DPR and sequence generation in two phases, using provenance ground truth for retrieval and target output for later training.
3 Methodology
Re2G combines neural and traditional initial retrieval with reranking before BART generation, and trains these components through staged and end-to-end procedures. Its end-to-end design addresses the zero-gradient problem caused when reranker scores, rather than initial retrieval scores, weight generation.
- Retrieval and reranking: Re2G merges DPR and BM25 results, then uses a reranker to select passages for BART generation despite incomparable retrieval scores.The reranker scores the union of DPR and BM25 candidates before the top passages are passed to generation.
- Retrieval and reranking: The interaction-model reranker jointly processes each query and passage, while the DPR representation model independently encodes them for scalable indexed retrieval.The two model types combine reranking accuracy with efficient initial retrieval over precomputed passage vectors.
- Training procedure: Training proceeds through DPR, generation, reranking, and full end-to-end phases, using provenance supervision initially and only target outputs in the latter phases.Generation training uses a BARTLARGE sequence-to-sequence model on target output.
- Training procedure: End-to-end generation weights retrieved query-passage sequences by retrieval or reranker probabilities and token probabilities from BART across target tokens.The loss is a weighted negative log-likelihood over target tokens and retrieved sequences.
- End-to-end training: Using reranker scores for generation makes the query-encoder gradient zero, so Re2G considers score combination, encoder freezing, and online knowledge distillation.The score-combination solution is rejected because it causes DPR to seek passages complementary to the reranker rather than most relevant to the query.
- End-to-end training: Online knowledge distillation trains DPR as a student from the reranker teacher across interaction and representation architectures using a KL-divergence loss.A temperature parameter smooths the distributions and stabilizes training.
4 Experiments
Re2G is evaluated on four KILT tasks using retrieval, reranking, and generation, with strong leaderboard gains and mixed effects across retrieval-training stages. Ablations show that online knowledge distillation and BM25 ensembling generally improve performance, but not uniformly.
- Re2G achieves relative gains of 9%, 31%, 34%, 22%, and 10% on the headline KILT metrics for T-REx, Natural Questions, TriviaQA, FEVER, and Wizard of Wikipedia, respectively.
- Re2G leads the headline KILT metrics on every evaluated dataset except Wizard of Wikipedia, where it ranks second.
- Retrieval: The retrieval-training stages improve performance across all datasets initially, while online distillation further benefits two of the three datasets examined in the later step.
- Retrieval: Reranker training improves over DPR across all five datasets on both retrieval metrics, while end-to-end reranker gains are mixed.
- Ablations: Both online knowledge distillation and BM25 ensembling improve performance in four of five datasets, with distillation failing on Wizard of Wikipedia and ensembling failing on Natural Questions.
5 Analysis
The analysis attributes much of Re2G’s output improvement to better retrieval, but shows that retrieval changes do not explain all gains. Error analysis identifies incomplete provenance and retrieval failures as important sources of error.
- 67.73%, 61.08%, and 66.86% of output gains for T-REx, Natural Questions, and FEVER, respectively, were accompanied by improved ranking of the correct passage.
- 36.86% and 27.74% of output improvements for TriviaQA and Wizard of Wikipedia, respectively, were accompanied by improved ranking for the correct passage.
- At least a third of the output-quality gain is not attributable to improved recall, even when retrieved passages are equivalent at test time.
- Error analysis: 33/50 sampled zero-score T-REx errors were caused by incomplete ground truth, often involving ambiguous head entities or multiple relation fillers.
- Error analysis: Among genuine errors, 9/17 resulted from failing to retrieve the necessary passage, while 6/17 selected an entity related in a different way.
6 Conclusions
Re2G advances performance across five KILT datasets by combining reranking, retrieval-source ensembling, and online knowledge distillation. The authors identify domain adaptation as an area for further study and release the code for research use.
- Re2G holds the top position on four of five KILT datasets and substantially improves retrieval and end-to-end performance across slot filling, question answering, fact checking, and dialog.
- The reranker enables combining BM25 with DPR results, while online knowledge distillation improves DPR in four of five datasets.
- Further domain-adaptation experiments on question answering and dialog could clarify Re2G’s application to real-world use cases.
- The authors release Re2G as open-source software under the Apache 2.0 license to enable further research.
A Hyperparameters
The experiments use hyperparameters largely inherited from prior work rather than extensively tuned for Re2G components. Results come from a single run with a fixed random seed.
- DPR Stage 1, generation, and reranking hyperparameters were not tuned; settings similar to the original training procedures were used instead.
- Knowledge distillation used generation settings, with reported temperature 10.0 and learning-rate scaling 1.0.
- Online distillation uses a separate query-encoder optimizer during generation training with the same hyperparameter settings.
- All results are from a single run using random seed 42 for Python, NumPy, and PyTorch.
C Model Details
Re2G uses three BERTBASE transformers for query encoding, passage encoding, and reranking, alongside a BARTLARGE generator. Training and indexing require substantial compute and memory resources.
- Model architecture: Re2G uses three BERTBASE transformers, each with 110M parameters, plus a 400M-parameter BARTLARGE generator.The total model size is 730M parameters.
- Training and infrastructure: DPR training on one NVIDIA V100 takes approximately 24 hours for T-REx and less than 12 hours for FEVER and WoW.
- Training and infrastructure: Generation training on two NVIDIA P100 GPUs takes two days for 370k T-REx instances and half a day for FEVER and WoW.
- Training and infrastructure: The FAISS index for the KILT knowledge source requires a large-memory machine; the authors use 128GB of memory.
D Generation Analysis
Manual analysis finds that WoW generation quality is difficult to assess because benchmark targets can be inconsistent, while Re2G outputs are sometimes coherent despite low automatic scores. Retrieval failures also explain some poor generations.
- Manual evaluation: 5/20 ground-truth WoW target texts were judged inconsistent, indicating limitations in annotation quality.
- Manual evaluation: Re2G and KGI0 received similar GOOD/OK/INCONSISTENT judgments: Re2G 8/2/10 and KGI0 9/2/9.
- Low-score cases: Among 20 bottom-quintile Re2G F1 cases, 8/20 were inconsistent, and supporting ground-truth passages were not retrieved in 4 of those 8 cases.
- Low-score cases: One Re2G response was consistent with the conversation despite low F1 and Rouge-L scores, stating that red has a dominant wavelength of approximately 625-740 nanometres.
D.1 Generation Quality
The generation-quality examples compare system outputs with retrieved evidence and ground truth, showing coherent system responses alongside cases where reference texts are factually or semantically inconsistent.
- Evaluation setup: Table 7 contains randomly selected WoW instances evaluated by a human evaluator, with comments documenting the judgments.
- Example evaluations: In the first example, both Re2G and KGI0 retrieved a passage that may have supported outputs judged coherent by evaluators.
- Example evaluations: The ground-truth response in that example was marked factually inconsistent based on a retrieved Niagara Falls passage.
- Example evaluations: In the second example, all three texts were marked inconsistent because system outputs nearly repeated earlier conversation and the ground truth was semantically incoherent.