Joint Passage Ranking for Diverse Multi-Answer Retrieval

In a typical single-answer retrieval problem, a model is given a natural language question $q$ and retrieves $k$ passages $\{p_{1}...p_{k}\}$ from a large text corpus $\mathcal{C}$ (Voorhees et al., 1999; Ramos et al., 2003; Robertson and Zaragoza, 2009; Chen et al., 2017; Lee et al., 2019; Karpukhin et al., 2020; Luan et al., 2020). The goal is to retrieve at least one passage that contains the answer to $q$. During training, question-answer pairs $(q,a)$ are given to the model.

We now formally define the task of multi-answer retrieval. A model is given a natural language question $q$ and needs to find $k$ passages $\{p_{1}...p_{k}\}$ from $\mathcal{C}$ that contain all distinct answers to $q$. Unlike in single-answer retrieval, question-and-answer-set pairs $(q,\{a_{1}...a_{n}\})$ are given during training.

JPR makes use of the encoder-decoder architecture, where the encoder processes candidate passages and the decoder autoregressively generates a sequence of $k$ passage references (indexes). Intuitively, dependencies between passages can be modeled by the autoregressive architecture.

We extend the architecture from Izacard and Grave (2021); we reuse the encoder but modify the decoder. Each candidate passage $p_{i}$ is concatenated with the question $q$ and the number $i$ (namely index). It is fed into the encoder to be transformed to $\mathbf{p}_{i}\in\mathbb{R}^{L\times h}$, where $L$ is the length of the input text and $h$ is a hidden size. Next, $\mathbf{p}_{1}...\mathbf{p}_{|\mathcal{B}|}$ are concatenated to form $\bar{\mathbf{p}}\in\mathbb{R}^{L|\mathcal{B}|\times h}$, and then fed into the decoder. The decoder is trained to autoregressively output a sequence of indexes $i_{1}...i_{k}$, representing a sequence of passages $p_{1}...p_{k}$. As the generation at step $t$ ($1\leq t\leq k$) is dependent on the generation at step $1\dots t-1$, it can naturally capture dependencies between selected passages. As each index occupies one token, the length of the decoded sequence is $k$.

A standard way of training the encoder-decoder is teacher forcing which assumes a single correct sequence. However, in our task, a set of answers can be retrieved through many possible sequences of passages, and it is unknown which sequence is the best. To this end, we dynamically form the supervision data which pushes the model to assign high probability to a dynamic oracle—any positive passage covering a correct answer that is not in the prefix, i.e., previously selected passages.

We first precompute a set of positive passages $\mathcal{\tilde{O}}$ and a prefix $\tilde{p}_{1}...\tilde{p}_{k}$. A set of positive passages $\mathcal{\tilde{O}}$ includes up to $k$ passages with maximal coverage of the distinct answers.³³3 $|\mathcal{\tilde{O}}|<k$ if fewer than $k$ passages are sufficient to cover all distinct answers; $|\mathcal{\tilde{O}}|=k$ otherwise. A prefix $\tilde{p}_{1}...\tilde{p}_{k}$ is a simulated prediction of the model, consisting of $\mathcal{\tilde{O}}$ and $k-|\mathcal{\tilde{O}}|$ sampled negatives. At each step $t$ ($1\leq t\leq k$) , given a set of positive passages $\mathcal{\tilde{O}}$ and a prefix $\tilde{p}_{1}...\tilde{p}_{t}$ (denoted as $\mathcal{P}_{t}$), JPR is trained to assign high probabilities to the dynamic oracle $\mathcal{\tilde{O}}-\mathcal{P}_{t}$. The objective is defined as follows:

A typical autoregressive decoder makes the top 1 prediction at each step to decode a sequence of $k$ (SeqDecode in Algorithm 1),⁴⁴4We explored beam search decoding but it gives results that are the same as or marginally different from SeqDecode. which, based on our training scheme, asks the decoder to find a new answer at every step. However, when $k$ is larger than the number of correct answers, it would be counter-productive to ask for $k$ passages that each covers a distinct answer. Instead, we want the flexibility of decoding fewer timesteps and take multiple predictions from each timestep.

In this context, we introduce a new decoding algorithm TreeDecode, which decodes a tree from an autoregressive model. TreeDecode iteratively chooses between the depth-wise and the width-wise expansion of the tree—going forward to the next step and taking the next best passage within the same step, respectively—until it reaches $k$ passages (Figure 3). Intuitively, if the model believes that there are many distinct answers covered by different passages, it will choose to take the next step, being closer to SeqDecode. On the other hand, if the model believes that there are very few distinct answers, it will choose to take more predictions from the same step, resulting in behavior closer to independent scoring.

The formal algorithm is as follows. We represent a tree $\mathcal{S}$ as a list of ordered lists $[s_{1}...s_{T}]$ where $s_{1}$ is an empty list and $s_{i}$ is one element appended to any of $s_{1}...s_{i-1}$. The corresponding set $\mathcal{O}$ is $\cup_{s\in\mathcal{S}}\mathrm{Set}(s)$. We define a score of a tree $\mathcal{S}$ as

We form $\mathcal{S}$ and $\mathcal{O}$ through an iterative process by (1) starting from $\mathcal{O}=\emptyset$ and $\mathcal{S}=[\texttt{null}]$, and (2) updating $\mathcal{O}$ and $\mathcal{S}$ by finding the best addition of an element that maximizes the gain in $f(\mathcal{S})$, until $|\mathcal{O}|=k$, as described in Algorithm 1.

We train and evaluate on three datasets that provide a set of distinct answers for each question. Statistics of each dataset are provided in Table 2.

JPR can obtain candidate passages $\mathcal{B}$ from any first-stage retrieval model. In this paper, we use DPR⁺, our own improved version of DPR (Karpukhin et al., 2020) combined with REALM (Guu et al., 2020). DPR and REALM are dual encoders with a supervised objective and an unsupervised, language modeling objective, respectively. We initialize the dual encoder with REALM and train on supervised datasets using the objective from DPR. More details are provided in Appendix A.

We compare JPR with three baselines, all of which are published models or enhanced versions of them. All baselines independently score each passage.

We use the English Wikipedia from 12/20/2018 as the retrieval corpus $\mathcal{C}$, where each article is split into passages with up to 288 wordpieces, All rerankers are based on T5 (Raffel et al., 2020), a pretrained encoder-decoder model; T5-base is used unless otherwise specified. We use $|\mathcal{B}|=100,k=\{5,10\}$. Models are first trained on NQ (Kwiatkowski et al., 2019) and then finetuned on multi-answer datasets, which we find helpful since all multi-answer datasets are relatively small. During dynamic oracle training, $k-|\mathcal{\tilde{O}}|$ negatives are sampled from $\mathcal{B}-\mathcal{\tilde{O}}$ based on $s(p_{i})+\gamma g_{i}$, where $s(p_{i})$ is a prior logit value from IndepPR, $g_{i}\sim\mathrm{Gumbel}(0,1)$ and $\gamma$ is a hyperparameter. In TreeDecode, to control the trade-off between the depth and the width of the tree, we use a length penalty function $l(y)=\left(\frac{5+y}{5+1}\right)^{\beta},$ where $\beta$ is a hyperparameter, following Wu et al. (2016). More details are in Appendix B.2.

Table 3 reports MRecall on all questions and on questions with more than one answer.

This section discusses experiments on downstream question answering: given a question and a set of passages from retrieval, the model outputs all valid answers to the question. We aim to answer two research questions: (1) whether the improvements in passage retrieval are transferred to improvements in downstream question answering, and (2) whether using a smaller number of passages through reranking is better than using the largest possible number of passages given fixed hardware memory.

We use an answer generation model based on Izacard and Grave (2021) which we train to generate a sequence of answers, separated by a [SEP] token, given a set of retrieved passages. Our main model uses JPR to obtain passages fed into the answer generation model. The baselines obtain passages from either DPR⁺ only or IndepPR, described in Section 4.3.

We compare different models that fit on the same hardware by varying the sizes of T5 (base, large, 3B) and use the maximum number of passages ($k$).⁵⁵5The memory requirement is $O(k\times\text{T5 size})$. This results in three settings: $\{k=140,$ base$\}$, $\{k=40,$ large$\}$ and $\{k=10,$ 3B$\}$.