SlowGuess/NAACL_2022_OCR huggingface dataset:
$\mathbf{Re}^2 \mathbf{G}$ : Retrieve, Rerank, Generate
Michael Glass $^{1}$ , Gaetano Rossiello $^{1}$ , Md Faisal Mahbub Chowdhury $^{1}$ , Ankita Rajaram Naik $^{1,2}$ , Pengshan Cai $^{1,2}$ , Alfio Gliozzo $^{1}$
<sup>1</sup> IBM Research AI, Yorktown Heights, NY, USA
$^{2}$  University of Massachusetts Amherst, MA, USA
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, nonparametric 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  $\mathrm{Re}^2\mathrm{G}$ , 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<sup>1</sup>.
GPT-3 [Brown et al., 2020] and T5 [Raffel et al., 2020] are arguably the most powerful members in a family of deep learning NLP models called transformers. Such models store surprising amount of world knowledge. They have been shown to produce good performance on a range of demanding tasks, especially in generating human like texts. However, such large transformers' capability is tied to the increasingly larger parameter spaces on which they are trained.
Recently, there has been work towards transformers that make use of non-parametric knowledge. REALM (Retrieval Augmented Language Model) [Guu et al., 2020] and RAG (Retrieval Augmented Generation) [Lewis et al., 2020b] both use an indexed corpus of passages to support conditional generation. By using the corpus as a source of knowledge these models can extend the information available to the model by tens or even hundreds of gigabytes with a sub-linear scaling in computation cost.
These recent advancements, in turn, have been inspired by BART (Bidirectional and Auto-Regressive Transformer) [Lewis et al., 2020a] that combines a Bidirectional Encoder (e.g. BERT [Devlin et al., 2019]) with an Autoregressive decoder (e.g. GPT [Brown et al., 2020]) into one sequence-to-sequence model.
We build on this line of research, pioneered by REALM and RAG, and propose a new approach that we call  $\mathrm{Re}^2\mathrm{G}$  (Retrieve, Rerank, Generate), which combines both neural initial retrieval and reranking into a BART-based sequence-to-sequence generation.
There are two particular aspects on which our approach is different from the previous works. Firstly, our reranking approach permits merging retrieval results from sources with incomparable scores, e.g. enabling an ensemble of BM25 and neural initial retrieval. Secondly, 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.
The KILT benchmark [Petroni et al., 2021] has been recently introduced to evaluate the capabilities of pre-trained language models to address NLP tasks that require access to external knowledge. We evaluate on four diverse tasks from KILT: slot filling, question answering, fact checking and dialog. Figure 1 shows examples of these tasks.  $\mathrm{Re}^{2}\mathrm{G}$
https://github.com/one-qa
Human Language Technologies, pages 2701 - 2715
makes significant gains on all four tasks, reaching the top of the KILT leaderboards and establishing a new state-of-the-art.
- We introduce  $\mathrm{Re}^2\mathrm{G}$ , demonstrating the effectiveness of reranking for generative language models that incorporate retrieval.
- We further extend  $\mathrm{Re}^2\mathrm{G}$  by ensembling initial retrieval methods, combining neural and traditional keyword-based approaches.
-  $\mathrm{Re}^{2}\mathrm{G}$  improves the current state-of-the-art of  $9\%$ ,  $31\%$ ,  $34\%$ ,  $22\%$  and  $10\%$  relative gains on the headline KILT metrics for T-REx (slot filling), Natural Questions (question answering), TriviaQA (question answering), FEVER (fact checking), and Wizard of Wikipedia (dialog), respectively.
- We publicly release our code as open source to support continued development.
The KILT benchmark and public leaderboard<sup>2</sup> combines eleven datasets across five tasks. The main advantage of the KILT distribution of these datasets is that the provenance information from each dataset is realigned to reference the same snapshot of Wikipedia. A unified evaluation script and set of metrics is also provided. In this work, we focus on four tasks, such as Slot Filling [Levy et al., 2017, Elsahar et al., 2018], Question Answering [Kwiatkowski et al., 2019, Joshi et al., 2017], Fact Checking [Thorne et al., 2018a,c], and Dialog [Dinan et al., 2019] (see Figure 1).
A set of baseline methods have been proposed for KILT. GENRE [Cao et al., 2021] is trained on BLINK [Wu et al., 2020] and all KILT tasks jointly using a sequence-to-sequence language model to generate the title of the Wikipedia page where the answer can be found. This method is a strong baseline to evaluate the retrieval performance, but it does not address the downstream tasks. On the other hand, generative models, such as BART [Lewis et al., 2020a] and T5 [Raffel et al., 2020], show interesting performance when finetuned on the downstream tasks relying only on the implicit knowledge stored in the weights of the
neural networks, without the use of any explicit retrieval component.
RAG [Lewis et al., 2020b], an end-to-end retrieval-based generative model, is the best performing baseline in KILT and it incorporates DPR [Karpukhin et al., 2020] to first retrieve relevant passages for the query, then it uses a model initialized from BART [Lewis et al., 2020a] to perform a sequence-to-sequence generation from each evidence passage concatenated with the query in order to generate the answer. Figure 2 shows the architecture of RAG.
Multi-task DPR [Maillard et al., 2021] exploits multi-task learning by training both DPR passage and query encoder on all KILT tasks. DensePhrases [Lee et al., 2021] addresses the knowledge intensive tasks with a short answer, such as slot filling. It indexes the phrases in the corpus that can be potential answers. The extracted phrases are represented by their start and end token vectors from the final layer of a transformer initialized from SpanBERT [Joshi et al., 2020].
Knowledge Graph Induction (KGI) [Glass et al., 2021] combines DPR and RAG models, both trained with task and dataset specific training. KGI employs a two phase training procedure: first training the DPR model, i.e. both the query and context encoder, using the KILT provenance ground truth. Then, KGI trains the sequence-to-sequence generation and further trains the query encoder using only the target output as the objective. This results in large improvements in retrieval performance and, as a consequence, in the downstream tasks.
KILT-WEB 2 [Piktus et al., 2021] addresses the KILT tasks by broadening the knowledge source used. Rather than rely only on KILT's Wikipedia snapshot, KILT-WEB 2 creates SPHERE as a knowledge source. SPHERE is built from CCNet [Wenzek et al., 2020] and over twenty times the size of the Wikipedia corpus. It can use either BM25 or DPR retrieval (though not both combined) followed by a 'reader' component, but not trained end-to-end. The reader component is the Fusion-in-Decoder [Izacard and Grave, 2021] model, where retrieved documents are encoded independently, then their encoded representations are concatenated for the decoder.
SEAL [Bevilacqua et al., 2022] introduces a novel generative approach to retrieval. Rather than generating the unique document identifier like GENRE, SEAL can generate any ngrams present
2https://eval.ai/web/challenges/challenge-page/689/leaderboard
T-REx
Input:
Dracula [SEP] narrative location
Output: Transylvania
Provenance: 7923-2
Natural Questions
Input: when did bram stoker's dracula come out
Output: 1897
Provenance: 7923-1
FEVER
Input: Dracula is a novel by a Scottish author.
Output: REFUTES
Provenance: 7923-1
Dracula (7923)
Dracula is an 1897 Gothic horror novel by Irish author Bram Stoker. It introduced the character of Count Dracula, and established many conventions of subsequent vampire fantasy.
The novel tells the story of Dracula's attempt to move from Transylvania to England so that he may find new blood and spread the undead curse, and of the battle between Dracula and a small group of men and a woman led by Professor Abraham Van Helsing.
Wizard of Wikipedia
- I really like vampires!!
- Vampires are intense and based on European folklore. Do you have any favorite vampires?
- I think dracula is the best one!!!
Output: He's one of the best! He's based on the character from the 1897 horror book of the same name.
Provenance: 7923-1
Figure 1: KILT tasks of slot filling, question answering, fact checking and dialog
Figure 2: RAG Architecture
Figure 3:  $\mathrm{Re}^2\mathrm{G}$  Architecture
in the corpus, which are then mapped to passages. The neural retrieval generator is based on BART and constrained to generate ngrams that appear in the corpus with an FM-Index [Ferragina and Manzini, 2000]. Like KILT-WEB 2, SEAL uses Fusion-in-Decoder as the component responsible for generating the output conditioned on the retrieved passages.
Multi-stage or cascade approaches to retrieval have received ample attention in Information Retrieval (IR) research. The multi-stage approach begins with the initial retrieval phase, where an initial set of documents or passages form the pool of candidates to be considered for ranking. Then one or more phases of increasingly computationally demanding rerankers are applied. Early approaches in learning to rank [Liu, 2009] used features and linear classifiers. Pre-trained language models, especially BERT [Devlin et al., 2019], have shown state-of
the-art performance when applied to the task of relevance ranking. Transformers may be applied as classifiers to each query and passage pair independently [Nogueira and Cho, 2019] or as generators to produce labels for passages in a sequence-to-sequence model [Nogueira et al., 2020].
The approach of RAG, Multi-DPR, and KGI is to train a neural IR (Information Retrieval) component and further train it end-to-end through its impact in generating the correct output. Figure 2 illustrates the end-to-end RAG system.
It has been previously established that results from initial retrieval can be greatly improved through the use of a reranker [Liu, 2009, Wang et al., 2011]. Therefore we hypothesized that natural language generation systems incorporating retrieval can benefit from reranking.
Figure 4: Interaction Model Reranker
Figure 5: Representation Model for Initial Retrieval
In addition to improving the ranking of passages returned from DPR, a reranker can be used after merging the results of multiple retrieval methods with incomparable scores. For example, the scores returned by BM25 [Robertson and Zaragoza, 2009] are not comparable to the inner products from DPR. Using the scores from a reranker, we can find the top-k documents from the union of DPR and BM25 results. Figure 3 illustrates our extension of RAG with a reranker. We call our system  $\mathrm{Re}^2\mathrm{G}$  (Retrieve, Rerank, Generate).
3.1 Reranker
The reranker we use is based on the sequence-pair classification of Nogueira and Cho [2019]. This model is shown in Figure 4. The query and passage are input together to a BERT [Devlin et al., 2019] transformer. Cross attention is applied over the tokens of both sequences jointly. This is called an interaction model.
This model contrasts with the representation model used for initial retrieval. Figure 5 shows the bi-encoder representation model for DPR. The representation vectors for the query and passage are produced independently. This allows for efficient retrieval by pre-computing vectors for all passages in the corpus and indexing them with an ANN (Approximate Nearest Neighbors) index. By using an interaction model to rerank the top-N passages from the representation model, we can get the advantages of both model types: accuracy and scalability.
We initialize the reranker from the BERT model trained on MS MARCO [Nguyen et al., 2016] by NBoost [Thienes and Pertschuk, 2019] and avail
able through Hugging Face3.
As Figure 1 illustrates, KILT tasks are provided with two types of ground truth: the target output sequence and the provenance information indicating the passage or passages in the corpus that support the output.
Our training is carried out in four phases: DPR training, generation training, reranking training, and full end-to-end training. The initial DPR and reranking phases make use of the provenance ground truth. The generation and full end-to-end training make use of only the target output.
- The original KILT instances are a tuple:  $\langle q, t, \mathbf{Prov} \rangle$  where  $q$  is the input or prompt,  $t$  is the target output, and  $\mathbf{Prov}$  is the set of provenance passages that support the target output.
- DPR training is a tuple:  $\langle q, p^{+}, p^{-} \rangle$  where  $p^{+} \in \mathbf{Prov}$  and  $p^{-}$  where  $p^{-} \in \mathrm{BM25}(q) \land p^{-} \notin \mathbf{Prov}$
- Reranking training begins with the application of DPR and BM25, producing tuples:  $\langle q, \mathbf{P}, \mathbf{Prov} \rangle$  where  $\mathbf{P} = \mathrm{BM25}(q) \cup \mathrm{DPR}(q)$
- Generation and end-to-end training instances are pairs of query and target:  $\langle q, t \rangle$
The first two phases, DPR and generation, are identical to KGI, specifically  $\mathrm{KGI}_0$ . We use the codes from Glass et al. [2021] $^4$ .
DPR Stage 1 training is the same training used by Karpukhin et al. [2020]. The triplets of query, positive passage and "hard negative" passages from BM25 are put into batches of 128 instances. The positives and hard negatives from other instances form the "batch negatives" for each instance. The DPR bi-encoder model gives each query a probability distribution over the positive, hard negative, and batch negatives. The loss is the negative log-likelihood for the positive. After DPR Stage 1 training the passages from the corpus are indexed with a Hierarchical Navigable Small World (HNSW) [Malkov and Yashunin, 2018] using FAISS [Johnson et al., 2017].
3https://huggingface.co/nboost/ pt-bert-base-uncased-msmarco 4https://github.com/IBM/ kgi-slot-filling
Generation training extends the training of the query encoder and trains the BARTLarge sequence-to-sequence model on the target sequence output. This training is the same as that described by Lewis et al. [2020b].
3.3 Reranking Training
The next phase, training the reranking in isolation, begins with gathering the initial retrieval results from DPR and BM25 on the training set. These results are merged and used as training data for the reranker.
In some datasets there are multiple positive passages. Therefore, we use the negative of the summed log-likelihood for the positive passages as the loss function. The logits given by the reranker are  $\mathbf{z}_{\mathbf{r}}$  and the indices for the correct passages (from the ground truth provenance) are Prov.
$$
l o s s = - \sum_ {i \in \mathbf {P r o v}} l o g (\operatorname {s o f t m a x} (\mathbf {z _ {r}}) _ {i})
$$
Training end-to-end poses a special challenge. In RAG, the gradient propagates to the query encoder because the inner product between the query vector and the passage vector is used to weight the influence of each sequence, a process RAG calls marginalization. The inputs to the BART model are sequences  $(s_j = p_j$  [SEP]  $q$ ) that comprise a query  $q$  plus retrieved passage  $p_j$ . The probability for each sequence is determined from the softmax over the retrieval (or reranker) scores for the passage. The probability for each target token  $t_i$  given the sequence  $s_j$  is a softmax over BART's token prediction logits. The loss therefore is a negative log-likelihood summed over all target tokens and sequences, weighted by each sequence's probability.
Consider that in  $\mathrm{Re}^2\mathrm{G}$  the score from the reranker, not the initial retrieval, is used to weight the impact of each sequence in generation. This allows the reranker to be trained through the ground truth on target output, but it means the gradient for the query encoder will be zero since the marginalization no longer depends on the inner product from the query and passage representation vectors.
$$
P (s _ {j}) = \operatorname {s o f t m a x} (\mathbf {z} _ {\mathbf {r}}) _ {j}
$$
$$
P \left(t _ {i} \mid s _ {j}\right) = \operatorname {s o f t m a x} \left(\operatorname {B A R T} \left(s _ {j}\right) _ {i}\right) _ {t _ {i}}
$$
$$
l o s s = - \sum_ {i, j} l o g \left(P \left(t _ {i} \mid s _ {j}\right) \cdot P \left(s _ {j}\right)\right)
$$
We consider three possible resolutions to this issue.
- Combine the DPR and reranker scores
- Freeze the query encoder
Online Knowledge Distillation
The first candidate solution is tempting but fatally flawed. By adding the log softmax from DPR and the reranker we can ensure that both systems are trained through impact in generation. However, if the DPR score is added to the reranker score, then the DPR score is being trained to provide a complementary signal to the reranker. Therefore, when DPR is used to gather the candidate passages, it does not give the highest scores to the passages that are most likely to be relevant, but instead gives the highest scores to the passages the reranker is most likely to underrate. We find that this theoretical concern is also a practical concern, as DPR performance (and overall system performance) declines greatly when trained in this way.
The simplest solution is to freeze the parameters of the query encoder, training only the reranker and generation components. We find this is indeed the best solution for one of our datasets, Wizard of Wikipedia. Note that DPR has already been trained in two phases, first from the provenance ground truth and then again in generation training in the RAG model.
The third solution is our novel application of knowledge distillation [Hinton et al., 2015]. We use the reranker as a teacher model to provide labels to the DPR student model. We distill the knowledge across architectures: from an interaction model to a representation model. Further, this knowledge distillation occurs online, while the reranker is being trained. The loss for the initial retrieval is therefore the KL-divergence between the probability distribution it gives over the retrieved passages and the reranker's probability distribution over the same passages. A temperature hyperparameter  $T$  smoothes these distributions to prevent excessive loss and stabilize training.
T-REx	(Slot Filling)
R-Prec	Recall@5	Accuracy	F1	KILT-AC	KILT-F1
Re2G (ours)	80.70	89.00	87.68	89.93	75.84	77.05
KGI1 [Glass et al., 2021]	74.36	83.14	84.36	87.24	69.14	70.58
KILT-WEB 2 [Piktus et al., 2021]	75.64	87.57	81.34	84.46	64.64	66.64
SEAL [Bevilacqua et al., 2022]	67.80	81.52	83.72	86.53	60.08	61.72
KGI0 [Glass et al., 2021]	59.70	70.38	77.90	81.31	55.54	56.79
Natural Questions	(Question Answering)
R-Prec	Recall@5	Accuracy	F1	KILT-AC	KILT-F1
Re2G (ours)	70.78	76.63	51.73	60.97	43.56	49.80
SEAL [Bevilacqua et al., 2022]	63.16	68.19	53.74	62.24	38.78	44.40
KGI0 [Glass et al., 2021]	63.71	70.17	45.22	53.38	36.36	41.83
KILT-WEB 2 [Piktus et al., 2021]	59.83	71.17	51.59	60.83	35.32	40.73
RAG [Petroni et al., 2021]	59.49	67.06	44.39	52.35	32.69	37.91
TriviaQA	(Question Answering)
R-Prec	Recall@5	Accuracy	F1	KILT-AC	KILT-F1
Re2G (ours)	72.68	74.23	76.27	81.40	57.91	61.78
SEAL [Bevilacqua et al., 2022]	68.36	76.36	70.86	77.29	50.56	54.99
KILT-WEB 2 [Piktus et al., 2021]	58.85	71.55	72.73	79.54	45.55	49.57
KGI0 [Glass et al., 2021]	60.49	63.54	60.99	66.55	42.85	46.08
MultiDPR [Maillard et al., 2021]	61.49	68.33	59.60	66.53	42.36	46.19
FEVER	(Fact Checking)
R-Prec	Recall@5	Accuracy	KILT-AC
Re2G (ours)	88.92	92.52	89.55	78.53
SEAL [Bevilacqua et al., 2022]	81.45	89.56	89.54	71.28
KILT-WEB 2 [Piktus et al., 2021]	74.77	87.89	88.99	65.68
KGI0 [Glass et al., 2021]	75.60	84.95	85.58	64.41
MultiDPR [Maillard et al., 2021]	74.48	87.52	86.32	63.94
Wizard of Wikipedia	(Dialog)
R-Prec	Recall@5	Rouge-L	F1	KILT-RL	KILT-F1
Hindsight [Paranjape et al., 2021]	56.08	74.27	17.06	19.19	11.92	13.39
Re2G (ours)	60.10	79.98	16.76	18.90	11.39	12.98
SEAL [Bevilacqua et al., 2022]	57.55	78.96	16.65	18.34	10.45	11.63
KGI0 [Glass et al., 2021]	55.37	78.45	16.36	18.57	10.36	11.79
RAG [Petroni et al., 2021]	57.75	74.61	11.57	13.11	7.59	8.75
KILT-WEB 2 [Piktus et al., 2021]	41.54	68.25	13.94	15.66	6.55	7.57
Table 1: KILT leaderboard top systems
$$
l o s s = D _ {K L} \left( \right.\operatorname {s o f t m a x} \left(\frac {\mathbf {Z _ {s}}}{T}\right)\left\| \right. \operatorname {s o f t m a x} \left(\frac {\mathbf {Z _ {t}}}{T}\right)\left. \right) \cdot T ^ {2}
$$
The knowledge distillation has the usual advantage of providing signal not only of positive and negative instances, but degrees of negativity. In addition, since we retrieve  $n = 12$  passages from DPR but only use the top- $k$  ( $k = 5$ ) for generation, the knowledge distillation loss is providing a (soft) label for more passages.
At inference time the query is encoded using the DPR query encoder and the top-12 passages from the HNSW index are returned. The query is also passed to BM25 search, specifically Anserini $^{5}$ , gathering the top-12 BM25 results. Both sets of passages are passed to the reranker and scored. The top-5 passages are then joined with the query and passed to BART $_{\text{LARGE}}$  to generate the output. The five output sequences are weighted according to the softmax over the reranker scores to produce the
<sup>5</sup>https://github.com/castorini/anserini
We test our model on five datasets, over four distinct tasks in the KILT benchmark: slot filling, question answering, fact checking and dialog. Figure 1 shows an example of these four tasks.
The slot filling dataset, T-REx [Elsahar et al., 2018], provides as input a head entity and relation, and expects as output the entity or term that fills the slot, also called the tail entity. The T-REx dataset contains 2.3M instances. We use only 370k training instances by downsampling the relations that occur more than 5000 times. This reduces the training time required while keeping state-of-the-art performance. The development and test sets each have 5k instances.
The question answering datasets are "open" versions of Natural Questions [Kwiatkowski et al., 2019] and TriviaQA [Joshi et al., 2017]. Unlike the original versions, the relevant Wikipedia page must be found by a retrieval step. The training sets for Natural Questions and TriviaQA contain  $87\mathrm{k}$  and  $62\mathrm{k}$  questions, with another  $3\mathrm{k}$  and  $5\mathrm{k}$  for the development and  $1.4\mathrm{k}$  and  $6.5\mathrm{k}$  for test.
The fact checking dataset in KILT is FEVER (Fact Extraction and VERIFICATION). It is a combination of the two FEVER versions [Thorne et al., 2018b, 2019] omitting the NOTENOUGHINFO class. There are approximately 10k instances in the development and test sets, and 100k for training. FEVER is a classification task, but we cast it as a generation task by training the model to generate either the token "SUPPORTS" or "REFUTES".