Baleen: Robust Multi-Hop Reasoning at Scale via Condensed Retrieval

The FLIPR encoders and interaction mechanism are shown in Figure 2. Our query encoder reads $Q_{t-1}$ and outputs a vector representation of every token in the input. Each query embedding interacts with all passage embeddings via a MaxSim operator, permitting us to inherit the efficiency and scalability of ColBERT [16]. While vanilla ColBERT sums all the MaxSim outputs indiscriminately, FLIPR considers only the strongest-matching query embeddings for evaluating each passage: it sums only the top-$\hat{N}$ partial scores from $N$ scores, with $\hat{N}<N$.

We refer to this top-$k$ filter as a “focused” interaction. Intuitively, it allows the same query to match multiple relevant passages that are contextually unrelated, by aligning—during training and inference—a different subset of the query embeddings with different relevant passages. These important subsets of the query are not always clear in advance. For instance, $Q_{1}$ in Table 1 should find the page about the 1965 World Series. This requires ignoring many distractions (e.g., World Series 1970 and Baseball Hall of Fame) that seem important in isolation while focusing on other phrases, namely a prominent “MVP” and the World Series 1965, a distinction evident only after finding the document and attempting to match it against the query.

FLIPR applies this mechanism over the representation of the query and the facts separately, keeping the top-$\hat{N}$ and top-$\hat{L}$ from $N$ and $L$ embeddings, respectively. More formally, the FLIPR score $S_{q,d}$ of a passage $d$ given a query $q$ is the summation $S_{q,d}=S^{Q}_{q,d}+S^{F}_{q,d}$ over the scores corresponding to the original query and to the facts. Define $M^{Q}_{i}=\max_{j=1}^{m}E_{q_{i}}\cdot E_{d_{j}}^{T}$, that is, the maximum similarity of the $i^{th}$ query embedding against the passage embeddings. The partial score $S^{Q}_{q,d}$ is computed as

where $\textrm{top}_{\hat{N}}$ is an operator that keeps the largest $\hat{N}$ values in a set. We define $S^{F}_{q,d}$ similarly, using the fact matches $M^{F}_{i}$ instead of $M^{Q}_{i}$.

We confirm the gains of FLIPR against the state-of-the-art ColBERT retriever in §5.4. By decomposing retrieval into token-level maximum similarities and pre-computing document representations, FLIPR maintains the scalability of the late interaction paradigm of ColBERT, reducing the candidate generation stage of retrieval into highly-efficient approximate $k$-nearest neighbor searches. See Appendix §B.1 for the implementation details of FLIPR.

For every training example, our datasets supply unordered gold passages. However, dependencies often exist in retrieval as Table 1 illustrates: a system needs to retrieve and read the page of Red Flaherty to deduce that it needs information about 1965 World Series. Such dependencies can be complex for 3- and 4-hop examples, as in HoVer, especially as multiple passages may be appropriate to retrieve together in the same hop.

We propose a generic mechanism for identifying the best gold passage(s) to learn to retrieve next. The key insight is that among the gold passages, a weakly-supervised retriever with strong inductive biases would reliably “prefer” those passages it can more naturally retrieve, given the hops so far. Thus, given $Q_{t-1}$, we train the retriever with every remaining gold passage and, among those, label as positives (for $Q_{t-1}$) only those ranked highly by the model. We subsequently use these weakly-labeled positives to train another retriever for all the hops.

Latent hop ordering is summarized in Algorithm 1, which assumes we have already trained a single-hop (i.e., first-hop) retriever $R_{1}$ in the manner of relevance-guided supervision (see §2). To start, we use $R_{1}$ to retrieve the top-$k$ (e.g., $k=1000$) passages for each training question (Line 2). We then divide these passages into positives $P_{1}$ and negatives $N_{1}$ (Line 3). Positives are the highly-ranked gold passages (i.e., the intersection of gold passages and the top-$\hat{k}$ passages with, e.g., $\hat{k}=10$) and negatives are the non-gold passages. Subsequently, we expand first-hop queries with the oracle facts from $P_{1}$ to obtain the queries for the second hop (Line 4).

We train a second-hop retriever $R_{2}$ using the aforementioned queries and negatives (Line 6). As we do not have second-hop positives, as a form of weak supervision, we train the retriever with all gold passages (per query) besides those already in $P_{1}$. Once trained, we use $R_{2}$ to discover the second-hop positives $P_{2}$, to collect negatives $N_{2}$, and to expand the queries (Lines 7 and 8). We repeat this procedure for the third hop onward. Once this iterative procedure is complete, we have bootstrapped positives (and negatives) corresponding to every retrieval hop for each query. We combine these sets to train a single multi-hop retriever (Lines 10 and 11), which takes a query $Q_{t-1}$ and learns to discriminate the positives in $P_{t}$ from the negatives in $N_{t}$ for every $t\in[T]$. In §5.4, we indeed find that this generic latent procedure outperforms a typical hand-crafted heuristic that relies on trying to match titles against the text.

After each hop, the condenser proceeds in two stages. In the first stage, it reads $Q_{t-1}$ and each of the top-$K$ passages retrieved by FLIPR. The input passages are divided into their constituent sentences and every sentence is prepended with a special token. The output embedding corresponding to these sentence markers are scored via a linear layer, assigning a score to each sentence in every retrieved passage. We train the first-stage condenser with a cross-entropy loss over the sentences of two passages, a positive and a negative.

Across the $K$ passages, the top-$k$ facts are identified and concatenated into the second-stage condenser model. As above, we include special tokens between the sentences, and the second-stage condenser assigns a score to every fact while attending jointly over all of the top-$k$ facts, allowing a direct comparison within the encoder. Every fact whose score is positive is selected to proceed, and the other facts are discarded. We train the second-hop condenser over a set of 7–9 facts, some positive and others (sampled) negative, using a linear combination of cross-entropy loss for each positive fact (against all negatives) and binary cross-entropy loss for each individual fact.

In practice, this condensed architecture offers novel tradeoffs against a re-ranking pipeline. On one hand, condensed retrieval scales better to more hops, as it represents the facts from $K$ long passages using only a few sentences. This may provide stronger interpretability, as the inputs to the retriever (e.g., the example in Table 1) become more focused. On the other hand, re-ranking requires only passage-level labels, which can be cheaper to collect, and provides the retriever with broader context, which can help resolve ambiguous references. Empirically, we find the two approaches to be complementary. In §5.4, we show that condensing is competitive with re-ranking (despite much shorter context; Appendix §C) and that a single retriever that combines both pipelines together substantially outperforms a standard rerank-only pipeline.

After all hops are conducted, a reader is used to extract answers or verify the claim. We use a standard Transformer encoder, in particular ELECTRA-large [5], which we feed the final output $Q_{t}$ of our multi-hop retrieval and condensing pipeline. We train for question answering similar to Khattab et al. [17] or claim verification similar to Jiang et al. [13].

The input to the model is a claim, which is often one or a few sentences long. The model is to attempt to verify this claim, outputting “Supported” if the corpus contains enough evidence to verify the claim, or otherwise “Unsupported” if the corpus contradicts—or otherwise fails to support—this claim. Alongside this binary label, the model must extract facts (i.e., sentences) that justify its prediction. In particular, HoVer contains two-, three-, and four-hop claims that require facts from 2–4 different Wikipedia pages, and the model must find and extract one or more supporting facts from each of these pages for every claim. Models are not given the number of hops per claim.

Table 3 reports the retrieval and passage extraction quality of the systems. We report three metrics, each sub-divided by the number of hops of each claim. The first is Retrieval@100, which is the percentage of claims for which the system retrieves all of the relevant passages within the top-100 results. To compare with Jiang et al. [13], we report the Retrieval@100 results only for supported claims. All other metrics are reported for all claims. The second metric is Passage EM, which is the percentage of claims for which the system can provide the exact set of relevant passages. The third metric is Passage F1, which uses the standard F1 measure to offer a relaxed version of EM.

At the top of the table, the TF-IDF baseline retrieves top-100 results in one round of retrieval. At the bottom of the table, we include the results reported by Jiang et al. [13] for their BERT ranking model when supplied with oracle retrieval, and also their reported human performance. We show Baleen’s performance after four hops, where Baleen’s FLIPR retrieves 25 passages in each hop. Note that we exclude passages retrieved in hop $t$ from the retrieval in further hops $t^{\prime}>t$. As the table shows, Baleen achieves strong results, outperforming the official baseline model by 47.6 points in Retrieval@100, 51.1 points in Passage EM, and 29.0 points in Passage F1. In fact, Baleen’s performance also consistently exceeds “Oracle + BERT” at passage extraction.

Next, we evaluate the performance of fact/sentence extraction and thereby the end-to-end claim verification of Baleen. The results are in Table 4, where Sentence EM and Sentence F1 are defined similar to the corresponding metrics in Table 3 but focus on sentence-level extraction. Moreover, Verification Accuracy is the percentage of claims that are labelled correctly as “Supported” or “Unsupported”. In Table 4, we see that Baleen’s sentence-level results mirror its document-level results, with large performance gains across the board against the baseline model. Baleen also outperforms the oracle-retrieval BERT model of Jiang et al. [13], emphasizing the strong retrieval and condensing performance of Baleen. In the leaderboard test-set evaluation (scores expected to be posted soon), Baleen increases the overall “HoVer Score” metric from 15.3 for the baseline to 57.5.

Lastly, Table 5 reports various settings and ablations of Baleen, corresponding to its contributions. Here, models A and F are simply the four-hop Baleen and ColBERT-Hop in the previous experiments, respectively. Models B and C test the effect of different architectural decisions on Baleen. In particular, model B replaces our condensed retrieval architecture with a simpler re-ranking architecture (Appendix §B.2), which reflects a common design choice for multi-hop systems. In this case, FLIPR is given the top-ranked passage as context in each hop, helping expose the value of extracting facts in condensed retrieval, which exhibits much shorter contexts as we show in Appendix §C. Second, model C trains a FLIPR retriever that uses both a condenser pipeline and a reranking pipeline independently, and then retains the overall top-100 unique passages retrieved (Appendix §B.2). The results suggest that condensing offers complimentary gains, and the hybrid model attains the highest scores.

Moreover, models D and E test modifications to the retriever, while using the condenser of model A. Specifically, model D replaces our FLIPR retriever with a simpler ColBERT retriever, while retaining all other components of our architecture. This allows us to contrast our proposed focused late interaction with the “vanilla” late interaction of ColBERT. Further, model E replaces our weak-supervision strategy for inferring effective hop orders with a hand-crafted rule (Appendix §B.3) that deliberately exploits a construction bias in HotPotQA and HoVer, namely that passage titles can offer a strong signal for hop ordering.