# Learning to Search in Long Documents Using Document Structure

Mor Geva
Tel Aviv University
<EMAIL_ADDRESS>
&Jonathan Berant
Tel Aviv University
<EMAIL_ADDRESS>

###### Abstract

Reading comprehension models are based on recurrent neural networks that
sequentially process the document tokens. As interest turns to answering more
complex questions over longer documents, sequential reading of large portions
of text becomes a substantial bottleneck. Inspired by how humans use document
structure, we propose a novel framework for reading comprehension. We
represent documents as trees, and model an agent that learns to interleave
quick navigation through the document tree with more expensive answer
extraction. To encourage exploration of the document tree, we propose a new
algorithm, based on Deep Q-Network (DQN), which strategically samples tree
nodes at training time. Empirically we find our algorithm improves question
answering performance compared to DQN and a strong information-retrieval (IR)
baseline, and that ensembling our model with the IR baseline results in
further gains in performance.

##  1 Introduction

00footnotetext:  This work is licensed under a Creative Commons Attribution
4.0 International License. License details:
http://creativecommons.org/licenses/by/4.0/

Reading comprehension (RC), the task of reading documents and answering
questions about their content, has attracted immense attention recently. While
early work focused on simple questions and short paragraphs [Hermann et al.,
2015, Rajpurkar et al., 2016, Trischler et al., 2017, Onishi et al., 2016],
current work is shifting towards more complex questions that require reasoning
over long documents [Joshi et al., 2017, Hewlett et al., 2016, Welbl et al.,
2017, Kočisky et al., 2017].

Long documents pose a challenge for current RC models, as they are dominated
by recurrent neural networks (RNNs) [Chen et al., 2016, Kadlec et al., 2016,
Xiong et al., 2017]. RNNs process documents token-by-token, and thus using
them for long documents is prohibitive. A common solution is to retrieve part
of the document with an IR approach [Chen et al., 2017, Clark and Gardner,
2017] or a cheap model [Watanabe et al., 2017], and run an RNN over the
retrieved excerpts. However, as documents become longer and questions become
complex, two problems emerge: (a) retrieving all the necessary text with a
one-shot IR method when performing complex reasoning becomes harder, and thus
thousands of tokens are retrieved [Clark and Gardner, 2017]. (b) Running even
a cheap model over the document in its entirety becomes expensive [Choi et
al., 2017].

Humans, in lieu of a mental inverted index, use document structure to guide
their search for answers. E.g., the answer to _“What high school did Leonard
Cohen go to?”_ is likely to appear in _“Early life”_ , while the answer to
_“How hot is it in Melbourne in July?”_ is likely to appear in _“Climate”_. In
this work we investigate whether we can train a model to navigate through a
document using its structure and find the answer, while reading only a small
portion of the entire document.

Figure 1: An overview of our framework: an agent answers the question $q$ for
a document $d$ by starting at the title, performing navigation actions until
reaching the relevant paragraph, extracting the answer $a$, and then stopping.

We represent documents as trees and train an agent that navigates through the
document tree until returning a final answer. Figure 1 illustrates this
process. Our agent reads the question _“The vacation destinations of Pattaya
and Phuket are in which country?”_ and starts navigation at the title of the
document. After reading a paragraph, and skipping _“History”_ , it drills down
to _“Geography”_ until finally halting at a paragraph that specifies the
answer (_“Thailand”_). The agent observes at each step only a glimpse of the
local text to determine its next action, which can be movement to a tree node,
answering the question with a more expensive RC model, or terminating
navigation. Thus, the agent consumes only a small fraction of the entire
document.

Our training data is question-document-answer triplets, without gold
navigation paths, and thus we train our model with the Deep Q-Network (DQN)
algorithm. Because the dataset is biased towards answers appearing at the
beginning of the document, the algorithm tends to stop early and does not
explore the document well. To overcome this challenge, we propose DocQN: a
variant of DQN for tree navigation that improves exploration by sampling nodes
from multiple parts of the tree.

We evaluate the ability of our agent to navigate to paragraphs containing the
answer on a variant of TriviaQA [Joshi et al., 2017] and find that: (a) DocQN
navigates better than DQN in documents both quantitatively and qualitatively.
(b) While DocQN observes only 6% of the document tokens, it outperforms an IR
method in end-to-end QA performance. (c) An ensemble of DocQN and IR
substantially improves both navigation and end-to-end QA performance over the
ensemble components.

To summarize, in this paper we ask: can an agent use document structure and
learn to find answers for complex questions in long documents? We propose a
new model and training algorithm that overcomes an inherent bias in the data,
answering the aforementioned question in the affirmative. Our code and dataset
are available at https://github.com/mega002/DocQN.

##  2 Problem Overview

We work in the traditional RC setup, where we are given question-document-
answer triplets $\\{(q_{i},d_{i},a_{i})\\}_{i=1}^{N}$ as a training set, and
aim to learn a function that finds the answer for an unseen question-document
pair. Unlike prior work, we assume documents are trees, where every tree node
$u$ corresponds to a structural element and is labeled with text $l(u)$.
Specifically, the root is labeled with the document title, sections and
subsections are labeled by their title, and paragraphs and sentences are
labeled by the text they contain. In addition, we order all non-sentence tree
nodes by a pre-order traversal (which corresponds to the linear order of text
in the document), and denote the index of a node $u$ by $n(u)$. For sentence
nodes, $n(u)$ is the index of their parent (a paragraph).

Figure 1 shows an example tree, where for each node we show the relevant
structural element and index (sentence nodes are not shown in the figure).

With this document representation, answering questions can be viewed as a
Markov Decision Process (MDP), where in each state the agent is located in a
particular tree node, actions allow movement through the document tree,
answering the question with a RC model, or stopping, and a reward is based on
whether the agent locates a node that contains the answer.

Our agent interleaves actions that _navigate_ in the document, with an action
that runs an RC model in a certain document position and _extracts an answer_.
Thus, the agent can decide to continue navigation after extracting a certain
answer. This is strictly more expressive than existing approaches that combine
IR with RC, where some text is retrieved exactly once before applying an RC
model. As RC shifts to reasoning over complex questions, navigating and
reading multiple parts of the document will become necessary. We show in
Section 5 this approach indeed improves QA performance.

##  3 Data

To test our framework, we capitalize on the recently-released TriviaQA
dataset, which contains question-answer pairs, along with a small set of
documents that (in almost all cases) contain the answer. TriviaQA is suitable
for our purposes as it is a large scale dataset, where questions are
relatively complex and documents are fairly long. The dataset includes only
raw text, and thus for every evidence document, we built a tree representation
by extracting the html metadata from the corresponding Wikipedia page, and
constructing the document structure from it.

Because our goal is to investigate whether a model can learn to search through
a document, it is important that a non-negligible fraction of the questions
require navigation through the document. However, in Wikipedia each document
starts with a preface that summarizes the document, and thus often contains
the answer. Consequently, a model that ignores the question and document and
always stops in the first paragraph is likely to obtain good performance.
Figure 2 shows the distribution of the first answer occurrence (FAO) in a
document in TriviaQA over the node indices $n(u)$ (x-axis), where all tree
nodes, except sentences are considered (see Section 2). We find that in most
question-document pairs the FAO is in the first few paragraphs, and that in
60% of the cases it is in the preface section.

|  | TriviaQA | TriviaQA-NoP
---|---|---|---
Train | Questions | 61,888 | 57,220
Documents | 110,647 | 99,315
Development | Questions | 7,993 | 7,336
Documents | 14,229 | 12,706
Test | Questions | 7,701 | 6,507
Documents | 13,661 | 10,481

Table 1: Data statistics for TriviaQA vs. TriviaQA-NoP.

Average number of tokens | 5590.7
---|---
Average number of tree nodes | 332.2
Average number of high-level sections | 6.6

Table 2: Data statistics for the evidence documents of TriviaQA-NoP.

Figure 2: FAO node index distribution of TriviaQA and TriviaQA-NoP (median
values in red).

| TriviaQA | TriviaQA-NoP
---|---|---
Questions | 86.7% | 83.8%
Question-document pairs | 69.7% | 66.9%

Table 3: Portion of answerable samples in a random subset of the training set,
which contains 278 questions and 475 question-document pairs.

To alleviate this heavy bias, we derive a new dataset, termed TriviaQA-NoP,
where we remove the preface section from every document. After removing the
preface, 2,144 out of 77,582 (2.8%) questions and 6,124 out of 138,537 (4.4%)
question-document pairs are left without an answer and are removed. To further
reduce the number of cases where an answer can not be inferred from a
document, we drop question-document pairs where: (a) the answer appears only
in titles; (b) the answer is a single-character; (c) the FAO node index is
$>700$ (in most cases the answer is an item in a list). Finally, the dataset
includes 91.6% of the questions and 88.4% of the question-document pairs from
the original dataset. We provide full statistics on the dataset in Tables 1
and 2.

To verify that questions remain answerable in TriviaQA-NoP after removing the
preface, we perform manual analysis on a random sample of 278 questions and
475 documents from the training set (Table 3). We find that the portion of
answerable questions and question-document pairs remains high and is reduced
by less than 3% in comparison to TriviaQA . This demonstrates that indeed, for
most questions and documents, the context necessary for answering the question
also appears in the document body and not only in the preface.

Figure 2 shows the FAO node index distribution in TriviaQA-NoP. We observe,
compared to TriviaQA, that the first occurrence of an answer is much more
spread out across the document, and that the median increases from 3 to 14,
which will require more navigation from the agent. However, even in TriviaQA-
NoP answers tend to appear at the beginning of the document, because document
content is usually organized by importance. This bias results in an
exploration challenge for our training algorithm, which we will address in
Section 4.

##  4 Method

In this section, we describe a model for the agent and a training algorithm
based on DQN [Mnih et al., 2015]. Specifically, we introduce a tree-sampling
strategy, which addresses the exploration challenge stemming from the bias
towards answers early in the document.

###  4.1 Framework

We represent the MDP as a tuple $(\mathcal{S},\mathcal{A},R,T)$, where
$\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $R(s,a)$
is a reward function, and $T(s,a)$ is a deterministic transition function. Our
model implements an action-value function $Q(s,a)$, which takes a state
$s\in\mathcal{S}$ and returns a value for every action $a\in\mathcal{A}$. This
function defines a policy $\pi(s)=\arg\max_{a\in\mathcal{A}}Q(s,a)$. We now
describe the state space, actions and reward function.

| Example
---|---
$o$ | _“Phuket Province Name”_
$\phi_{n}^{1}$: height | 2
$\phi_{n}^{2}$: depth | 1
$\phi_{n}^{3}$: h_dist_start | 0
$\phi_{n}^{4}$: h_dist_end | 2
$\phi_{n}^{5}$: parent.h_dist_start | 0
$\phi_{n}^{6}$: parent.h_dist_end | 0
$\phi_{n}^{7}$: navigation step | 1
Table 4: An example observation $o$ and list of navigation features $\phi_{n}$
for node 1 in Figure 1. The features height and depth correspond to distance
from the farthest leaf and root, respectively (height is 2 since sentence
nodes are omitted from the figure). h_dist_start and h_dist_end measure the
horizontal distance from the first and last child of the node’s parent.
navigation step is a counter for the number of performed steps.

#### States

Given a tree node $u$, a state is a tuple $s^{u}=(q,o,z,\phi_{n},\phi_{z})$,
where $q$ is the question, $o$ is an observation, $z$ is an answer prediction,
and $\phi_{n}$, $\phi_{z}$ are navigation and answer prediction features. An
observation $o=(o_{1},\dots,o_{|o|})$ is a sequence of tokens produced by
recursively concatenating the first $k$ tokens of text in the label $l(u)$ to
the observation of $u$’s parent. An answer prediction
$z=(z_{1},\dots,z_{|z|})$ is the sequence of tokens that were extracted by an
RC model, if an RC model was already run on $l(u)$ (and a null token
otherwise). The answer prediction features $\phi_{z}=(z_{e},\,z_{l},\,z_{n})$
provide information on the distribution over answer spans provided by the RC
model, which reflects its confidence: $z_{e}$ is the entropy of the
distribution, $z_{l}$ is the logit value for $z$, and $z_{n}$ is the number of
tokens in $l(u)$. Navigation features $\phi_{n}$ provide information on the
relative location of $u$ in the document. An example observation and full list
of navigation features are given in Table 4.

Note that the state $s$ does not depend on the history of visited tree
nodes.111except for navigation step, which can be approximated by the shortest
path from the root to any node. While incorporating history could be
beneficial, a memory-less model enables us to explore tree-sampling
strategies, which is important for training (Section 4.2).

#### Actions

We define the following set of actions $\mathcal{A}$. Let $u$ be a node with
an ordered list of children $(v_{1},v_{2},v_{3})$, and $w$ be a child of
$v_{2}$. We define five movement actions (Figure 3), where Down moves from $u$
to its first child $v_{1}$, Right moves from $v_{2}$ to $v_{3}$, and Left
moves from $v_{2}$ to $v_{1}$. Because moving upwards reaches a node we
already visited, we define UpR, which moves from $w$ to $v_{3}$, and UpL,
which moves from $w$ to $v_{1}$. If an illegal action is chosen (e.g., Left
from $v_{1})$, then the agent stays in its current position. The action Answer
returns an answer (and a distribution over spans) by running a RC model on
$l(u)$, unless $u$ is a sentence, in which case it is run on the paragraph
containing $u$. After Answer, the agent can resume navigation. The action Stop
also returns the answer given the current node $u$, but also terminates
navigation.

Figure 3: Movement actions in our environment.

#### Reward

Our goal is to develop an agent that can navigate in the document, and thus we
define the reward based on whether the agent stops in a node that contains the
gold answer text (this is noisy, because the answer might be there
sporadically). While a simple reward would be an indicator for whether the
agent stopped at a correct node, such a reward would not capture the proximity
of the agent to the answer. Moreover, we would like to consider the overall
document length, rewarding successful navigations in long documents.
Therefore, we define the following reward:

$r(s,a)=\begin{cases}2&a=\textsc{Stop},|n(u)-n(u^{*})|=0\\\
1-\frac{|n(u)-n(u^{*}|)}{\max_{u}n(u)}&a=\textsc{Stop},|n(u)-n(u^{*})|>0\\\
-0.06&a=\textsc{Answer}\\\
-0.02&a\not\in\\{\textsc{Answer},\textsc{Stop}\\}\end{cases}$

where $u$ is the node where the agent is located and $u^{*}$ is the closest
tree node that contains the answer ($n(u)$ is the node index as defined
above). Thus, when stopping, the reward is proportional to the distance to the
closest answer location given the document length. An additional reward is
given if navigation is successful, and a penalty is given for any other
action, to encourage shorter trajectories. We further penalize the Answer
action to discourage frequent usage of the RC model.

###  4.2 DocQN: DQN with Tree Sampling

Training the navigation model is based on DQN [Mnih et al., 2015]. In DQN, at
every step an agent at state $s_{t}$ selects an action $a_{t}$ using
$\epsilon$-greedy policy, given the current action-value function
$Q_{\theta}(s_{t},a_{t})$ parameterized by $\theta$. The agent observes a
reward $r_{t}$ and a state $s_{t+1}=T(s_{t},a_{t})$ and adds a transition
$(s_{t},a_{t},r_{t},s_{t+1})$ to a replay memory buffer $\mathcal{D}$ that
holds a large number of recent transitions. The parameters are then optimized
so that the action-value function matches better the observed reward. This is
done by sampling a batch of random transitions from $\mathcal{D}$ and
minimizing the regression loss

$(r_{t}+\gamma\max_{a^{\prime}}Q_{\hat{\theta}}(s_{t+1},a^{\prime})-Q_{\theta}(s_{t},a_{t}))^{2},$
(1)

where $\hat{\theta}$ are the parameters of a _target network_ , which is a
periodic copy of $\theta$ that is not optimized, and $\gamma$ is a discount
factor.

We also add some of the recent enhancements to DQN [Hessel et al., 2017],
which have proved to be useful in our setup. Specifically we implement _Double
Q-Learning_ [van Hasselt et al., 2016], _Prioritized Experienced Replay_
[Schaul et al., 2015], and _Dueling Networks_ [Wang et al., 2016].

#### Reducing bias with DocQN

The DQN algorithm contains episodes, where in each episode the agent is placed
at an initial state $s_{0}$, from which it starts taking actions. In our
setup, this state corresponds to the root of the document tree. Because the
data has a bias towards answers appearing at the beginning of the document,
the agent learns that stopping early improves the reward and is stuck at a
local minimum, where it ceases exploration. Examining Figures 2 and 8, we
observe that DQN learns to stop very early in the document compared to the FAO
node index distribution of TriviaQA-NoP, amplifying the bias in the data.

To address this issue, we suggest DocQN, a variant of DQN aimed at increasing
exploration when navigating in a document structure. DocQN capitalized on two
properties. First, DQN is an off-policy algorithm that trains from transitions
$(s_{t},a_{t},r_{t},s_{t+1})$. Second, our model is memory-less, and thus we
can sample any node $u$ and compute the corresponding state $s^{u}$.
Therefore, we modify DQN, and instead of initializing every episode with
$s_{0}$ and performing a sequence of actions, we sample states from
distributions that explore the document better. By exploring transitions from
across the document, the model learns from more distant parts of the document.

Figure 4: Illustration of the different state distributions for a specific
documents tree. Nodes with darker color have higher sampling probability, and
nodes marked in red are paragraph nodes containing the answer.

Algorithm 1 DocQN

1:Let $f$ be a distribution over tree nodes

2:Initialize replay memory $\mathcal{D}$ and parameters $\theta,\hat{\theta}$

3:for episode $=1,M$ do

4: if random() $<\varepsilon_{s}$ then

5: for $i=1\dots K$ do

6: sample node $u\sim f$ and generate $s^{u}$

7: sample action $a$ from $s^{u}$ ($\epsilon$-greedily).

8: $r\leftarrow R(s^{u},a),s^{\prime}\leftarrow T(s^{u},a)$

9: store $(s^{u},a,r,s^{\prime})$ in $\mathcal{D}$

10: else

11: Initialize start state $s_{0}$

12: for $t=0\dots T-1$ do

13: sample action $a$ from $s_{t}$ ($\epsilon$-greedily).

14: $r\leftarrow R(s_{t},a),s_{t+1}\leftarrow T(s_{t},a)$

15: store $(s_{t},a,r,s_{t+1})$ in $\mathcal{D}$

16: $\textsc{UpdateParams}(\theta,\hat{\theta},\mathcal{D})$ (Equation 1)

Figure 5: High-level overview of the network architecture.

Algorithm 1 provides the details of DocQN. The algorithm initializes the
replay memory buffer, the model and target network parameters, and chooses a
distribution $f$ over tree nodes. At the beginning of every episode, with
probability $\epsilon_{s}$ (line 4),222$\epsilon_{s}$ is annealed from $1$ to
$0.5$ during training. we sample $K$ state transitions using $f$ and store
them in $\mathcal{D}$, and with probability $(1-\epsilon_{s})$ (line 10), we
start at the initial state $s_{0}$ and sample a trajectory as in DQN.
Parameters are updated as in DQN by sampling transitions from the replay
buffer and minimizing Equation 1. If $f$ samples nodes from various locations
in the document, the algorithm will explore better and will not getting stuck
at its beginning. We consider the following instantiations of $f$:

  1. 1.

$f_{U}$: _Uniform sampling_ over nodes, except we discourage sampling
sentences by uniformly sampling with probability $0.2$ a leaf (sentence), and
probability $0.8$ an inner node.

  2. 2.

$f_{B}$: _Backward sampling_ : We uniformly sample a paragraph node $p$ that
contains the gold answer, then uniformly a number of actions
$B\in\\{1,2,3\\}$, and perform $B$ random movement actions from $p$ to output
a node. This results in a node that is “close” to the answer $a$, and can be
viewed as similar to bi-directional search or backward search [Lao et al.,
2015].

Figure 4 shows the distribution over tree nodes for a specific document tree
where the answer appears in two paragraph nodes (ignoring sentence nodes). We
see that sequential sampling (as in DQN) puts most of the probability mass
close to the root, uniform sampling is uniform (across paragraphs, as sentence
nodes are omitted), and backward sampling is concentrated close to answer
nodes. This illustrates how different distributions result in different
exploration strategies.

#### Network Architecture

We briefly describe the neural architecture (Figure 5) and provide full
details in the supplementary material. As explained in Section 4.1, the input
to the network is the state $s$, which comprises the question tokens $q$,
observation tokens $o$, answer prediction tokens $z$ and features $\phi_{n}$,
$\phi_{z}$. The question, observation and answer prediction tokens are encoded
with pre-trained word embeddings and trained character embeddings, where
character embeddings are followed by a convolutional layer with max pooling,
yielding a single vector per token. Each token is then represented by
concatenating the word embedding with the character embedding.

Question tokens, observation tokens, and answer to are then fed into a BiLSTM
and LSTM [Hochreiter and Schmidhuber, 1997] respectively and the LSTM outputs
are compressed to a single vector through self attention [Cheng et al., 2016],
resulting in one vector for $q$ and one for $o$. Answer prediction tokens are
fed into an LSTM, where the last hidden state is concatenated to the features
$\phi_{z}$, thus creating a third vector for the answer prediction.

We concatenate these three vectors, and pass them through a one layer feed-
forward network that then branches to two networks according to the Dueling
DQN architecture [Wang et al., 2016]. In each branch we also concatenate the
navigation features $\phi_{n}$. One branch predicts the value of the state
$V_{\theta}(s)\in\mathbb{R}$, and the other branch predicts the advantage of
every action $A_{\theta}(s,a)\in\mathbb{R}$ for every possible action $a$. The
output of the network is
$Q_{\theta}(s,a)=V_{\theta}(s)+(A_{\theta}(s,a)-\frac{1}{\mathcal{A}}\sum_{a^{\prime}}A_{\theta}(s,a^{\prime}))$
as in ?).

##  5 Experimental Evaluation

Our experimental evaluation aims to answer the following questions: (a) Can
document structure be used to learn to navigate to an answer in a document?
(b) How does DocQN compare to DQN? (c) How does DocQN compare to IR methods
that observe the entire document?

###  5.1 Experimental Setup

We evaluate on TriviaQA-NoP. Because our focus is on the navigation ability of
the agent, we train a single RC model and fix it in all experiments.
Specifically, we download RaSoR [Lee et al., 2017, Salant and Berant,
2018],333https://github.com/shimisalant/RaSoR and exactly follow the procedure
described by the authors of TriviaQA for training a RC model, i.e., we train
RaSoR on the first $400$ tokens of each document in TriviaQA-NoP. As a sanity
check for our RC model, we also train and evaluate RaSor on the original
TriviaQA dataset. Indeed, Rasor obtains $48.6\%$ EM and $53.4\%$ F1, which is
substantially higher than the baseline reported by ?).

#### Evaluation

We use two metrics: First, we measure _navigation accuracy_ , i.e., for a
question-document pair, whether a method returns text containing a gold answer
(if the agent stops at a sentence node we evaluate the encompassing
paragraph). Because questions in TriviaQA-NoP often have more than one
evidence document, we also measure _aggregated navigation accuracy_ , where we
give credit if the agent navigated correctly in any of the documents. This
gives performance assuming an oracle that always chooses the best document for
the question. Because the test set in TriviaQA is hidden, we evaluated
navigation accuracy on the development set only.

In addition, we measure end-to-end QA performance with the official Exact
Match (EM) and F1 metrics. The EM metric measures the percentage of
predictions that match exactly any of the answer aliases, and the F1 metric
measures the average overlap between the prediction and answer. To aggregate
evidence from multiple documents we follow ?) and define the score of an
answer to be the sum of probabilities for that answer across all documents.

#### Models

We compare the following models:

  * •

DocQN: This is our main model, where we use the sampling distribution
$f_{U+B}=\frac{1}{2}f_{U}+\frac{1}{2}f_{B}$.

  * •

DQN: DQN algorithm without state sampling.

  * •

{DocQN$\mid$DQN}-Coupled: A less expressive version of the model, where the
actions ANSWER and STOP are coupled, i.e., the agent can use the RC model
exactly once and then stops. This is similar to a setup where retrieval is
performed once and not interleaved with navigation.

  * •

RandomWalk An agent that selects an action at each step uniformly at random.

  * •

RandomPara An agent that randomly selects a non-sentence tree node.

  * •

Doc-Tf-Idf: An IR baseline, where a paragraph is selected based on its tf-idf
score. We implement the tf-idf scheme of DocQA [Clark and Gardner, 2017],
which is a high-performing system on TriviaQA. Here, idf is computed for each
paragraph in the context of the current document, and the paragraph with
highest cosine similarity to the question is chosen. Note that Doc-Tf-Idf
processes the entire document (up to tens of thousands of tokens), while DocQN
processes only a small fraction.

  * •

Tf-Idf: Vanilla tf-idf, where the idf score is computed from all documents in
TriviaQA-NoP.

  * •

Ensemble: We ensemble DocQN$f_{U+B}$ with Doc-Tf-Idf in two ways: (1) for
finding the final answer, we aggregate scores as described above, that is, we
sum the probabilities from both models over all documents and choose the
answer with highest score. (2) For navigation, we simply tune on the
development set a threshold $l$, where we take the prediction of
DocQN$f_{U+B}$ if it stopped at a node with index $\leq l$ and use Doc-Tf-Idf
otherwise.

  * •

ReadTop: Following ?), a model that runs the RC model on the first $800$
tokens of the document. Note that running the RC model on the text retrieved
by DocQN involves consuming far fewer tokens, namely, RaSoR consumes only
$160$ tokens on average.

We report the value of all hyper-parameters in the supplementary material (all
fixed without tuning).

###  5.2 Results

Figure 6: Navigation example of DocQN.

Figure 7: Model performance (top) and portion of samples in the development
set (bottom) as functions of the node index of the FAO.

Figure 8: Distribution of node index at navigation stopping point (median
value in red) for DQN and DocQN.

Tables 6 and 6 show the results of our experiments. Focusing on navigation
accuracy, we see that randomly walking or choosing a paragraph yields low
performance. Vanilla Tf-Idf performs considerably better than the random
baselines, but is outperformed by all other models. Comparing DQN to DocQN, we
see that DocQN outperforms DQN. Allowing DQN and DocQN to run the RC model
during navigation improves their performance, but the accuracy gain is
substantially higher for DocQN (2.3% accuracy and 3.1% aggregated accuracy)
than for DQN (0.4% accuracy and 0.7% aggregated accuracy). Doc-Tf-Idf, which
has access to the entire document outperforms DocQN, which consumes on average
6% of the entire document. Nonetheless, the two models obtain the same
aggregated accuracy. This good performance of Doc-Tf-Idf shows that in
TriviaQA-NoP answer paragraphs share a lot of lexical material with the
question. Importantly, an ensemble of Doc-Tf-Idf with DocQN substantially
improves the overall performance, reaching accuracy of 35% and aggregated
accuracy of 48%. This is in sharp contrast to an ensemble with DQN, where for
any value of $l$, Doc-Tf-Idf performs better on its own.

| Dev. | Dev. Agg.
---|---|---
RandomWalk | 1.8 | 3.1
RandomPara | 13.4 | 20.9
DQN-Coupled | 26.8 | 38.8
DQN | 27.2 | 39.5
DocQN-Coupled | 28.1 | 40.5
DocQN | 30.4 | 43.6
Tf-Idf | 25.3 | 35.4
Doc-Tf-Idf | 32.4 | 43.5
Ensemble (threshold, $l=5$) | 35.0 | 48.0
Table 5: Navigation accuracy for all models on the development set of
TriviaQA-NoP.

| Development | Test
---|---|---
| EM | F1 | EM | F1
DQN | 21.7 | 26.5 | 19.1 | 24.1
DocQN | 23.6 | 27.9 | 21.0 | 25.5
Doc-Tf-Idf | 21.4 | 27.1 | 18.2 | 23.5
Ensemble (threshold, $l=5$) | 26.8 | 32.0 | 24.2 | 29.4
Ensemble (answer) | 28.4 | 33.4 | 25.4 | 30.5
ReadTop | 32.5 | 36.7 | 28.1 | 32.4
Table 6: End-to-end QA performance of all models on the development and test
sets of TriviaQA-NoP. For the ensemble model, we choose the best threshold of
$l=5$.

Examining end-to-end performance, DocQN again outperforms DQN, but DocQN is
now better than Doc-Tf-Idf. This suggests that when Doc-Tf-Idf selects a
paragraph with the answer, it is often difficult to extract with the RC model.
Again, the ensembles leads to a dramatic increase in performance, showing that
Doc-Tf-Idf and DocQN are complementary. However, ReadTop, which consumes the
first $800$ tokens of each document compared to $\sim 160$ for DocQN,
substantially outperforms all navigation models. This shows that the bias for
answers at the beginning of documents is still strong in TriviaQA-NoP.

To further elucidate the differences between navigation models, Figure 7 shows
navigation accuracy of different models and the proportion of samples for
different node indices of the FAO. We see that DQN outperforms DocQN when the
answer is at the top of the document, but DocQN dominates DQN when the answer
is further down, showing that DocQN learns to find answers deeper in the
document. Doc-Tf-Idf has a more balanced navigation accuracy across the
document, which explains why an ensemble of Doc-Tf-Idf with DocQN works, as
they are complementary to one another.

#### Analysis

Figure 6 shows a navigation example, which includes the navigation step, node
index, observation $o$, and the action $a$ taken. For the observation we also
highlight the attention distribution from the self-attention component. In
this figure the question is about culture, and we see the agent going into
multiple sections, reading them, running the RC model and continuing forward,
until finally stopping at a paragraph that contains the answer. We provide
many more examples in the supplementary material.

Table 7 highlights some differences between DQN and DocQN. DocQN has longer
trajectories, and stops at sentence nodes more frequently, suggesting it reads
the document at a finer granularity. Additionally, it leverages the RC model
to collect more information and confidence during navigation, by choosing the
action ANSWER more frequently. Both models consume less than 7% of the entire
document. Figure 8 illustrates the navigation stopping point, and shows that
DocQN navigates deeper into the document.

|  DQN |  DocQN
---|---|---
Path length |  avg. 7.7, range [1,36] |  avg. 15.2, range [3,100]
Answer predictions |  avg. 2.8 |  avg. 4.1
Tokens consumed |  3.4% |  6.2%
Stopping node |  0.5% title, 2.2% headline, 89.0% paragraph, 8.3% sentence |  0% title, 0.4% headline, 75.1% paragraph, 24.5% sentence
Table 7: Comparing navigation properties of DQN and DocQN on the development
set. The average and range of navigation path length in steps, relative amount
of consumed tokens, and distribution of stopping node types.

##  6 Related Work

Handling the challenges of reasoning over multiple long documents is gaining
fast momentum recently [Shen et al., 2017]. As mentioned, some approaches use
IR for reducing the amount of processed text [Chen et al., 2017, Clark and
Gardner, 2017], while others use cheap or parallelizable models to handle long
documents [Hewlett et al., 2017, Swayamdipta et al., 2018, Wang et al.,
2018a]. Searching for answers while using a trained RC model as a black-box
was also implemented recently in ?), for open-domain questions and multiple
short evidence texts from the Web. Another thrust has focused on skimming text
in a sequential manner [Yu et al., 2017], or designing recurrent architectures
that can consume text quickly [Bradbury et al., 2017, Seo et al., 2018, Campos
et al., 2018, Yu et al., 2018]. However, to the best of our knowledge no work
has previously applied these methods to long documents such as Wikipedia
pages.

In this work we use TriviaQA-NoP for evaluation of our navigation based
approach and comparison to an IR baseline. While there are various aspects to
consider in such evaluation setup, our choice of data was derived mainly by
the requirements for long and structured context. Recently, several new
datasets such as WikiHop and NarrativeQA were published. These datasets try to
focus on the tendency of RC models to match local context patterns, and are
designed for multi-step reasoning. [Welbl et al., 2017, Wadhwa et al., 2018,
Kočisky et al., 2017].

Our work is also related to several papers which model an agent that navigates
in an environment to find objects in an image [Ba et al., 2015], relations in
a knowledge-base [Das et al., 2018], or documents on the web [Nogueira and
Cho, 2016].

##  7 Conclusions

We investigate whether document structure can be leveraged to train an agent
that finds answers to questions in long documents while reading only a small
fraction of it. We show that an agent that reads 6% of the document can
improve QA performance compared to an IR method that utilizes the entire
document, and that ensembling the two substantially improves performance. We
also present DocQN, an algorithm that promotes better exploration of the
document, and show it outperforms DQN qualitatively and quantitatively.

Our approach represents a conceptual departure from previous methods for
reading long documents, as it interleaves searching for an answer in the
document with extracting the answer from a particular paragraph, which we show
improves both navigation and QA performance. We expect that as RC models
tackle longer documents that require reasoning and reading text that is spread
in multiple parts of the document, models that can efficiently navigate and
collect evidence will become more and more crucial. Our agent provides a first
step in this important research direction.

## Acknowledgments

We thank Eunsol Choi from Washington University for helpful discussions and
comments on the paper. This research was supported by the Yandex Initiative in
Machine Learning and by The Israel Science Foundation grant 942/16. This work
was completed in partial fulfillment for the Ph.D degree of the first author.

## References

  * [Ba et al., 2015]  J. Ba, V. Mnih, and K. Kavukcuoglu.  2015\.  Multiple object recognition with visual attention.  In International Conference on Learning Representations (ICLR).
  * [Bradbury et al., 2017]  J. Bradbury, S. Merity, C. Xiong, and R. Socher.  2017\.  Quasi-recurrent neural networks.  In International Conference on Learning Representations (ICLR).
  * [Campos et al., 2018]  V. Campos, B. Jou, X. Giro i Nieto, J. Torres, and S. Chang.  2018\.  Skip RNN: Learning to skip state updates in recurrent neural networks.  In International Conference on Learning Representations (ICLR).
  * [Chen et al., 2016]  D. Chen, J. Bolton, and C. D. Manning.  2016\.  A thorough examination of the CNN / Daily Mail reading comprehension task.  In Association for Computational Linguistics (ACL).
  * [Chen et al., 2017]  D. Chen, A. Fisch, J. Weston, and A. Bordes.  2017\.  Reading Wikipedia to answer open-domain questions.  In Association for Computational Linguistics (ACL).
  * [Cheng et al., 2016]  J. Cheng, L. Dong, and M. Lapata.  2016\.  Long short-term memory-networks for machine reading.  In Empirical Methods in Natural Language Processing (EMNLP).
  * [Choi et al., 2017]  E. Choi, D. Hewlett, A. Lacoste, I. Polosukhin, J. Uszkoreit, and J. Berant.  2017\.  Coarse-to-fine question answering for long documents.  In Association for Computational Linguistics (ACL).
  * [Clark and Gardner, 2017]  C. Clark and M. Gardner.  2017\.  Simple and effective multi-paragraph reading comprehension.  arXiv preprint arXiv:1710.10723.
  * [Das et al., 2018]  R. Das, S. Dhuliawala, M. Zaheer, L. Vilnis, I. Durugkar, A. Krishnamurthy, A. Smola, and A. McCallum.  2018\.  Go for a walk and arrive at the answer: Reasoning over paths in knowledge bases using reinforcement learning.  In International Conference on Learning Representations (ICLR).
  * [Hermann et al., 2015]  K. M. Hermann, T. Kočiský, E. Grefenstette, L. Espeholt, W. Kay, M. Suleyman, and P. Blunsom.  2015\.  Teaching machines to read and comprehend.  In Advances in Neural Information Processing Systems (NIPS).
  * [Hessel et al., 2017]  M. Hessel, J. Modayil, H. V. Hasselt, T. Schaul, G. Ostrovski, W. Dabney, D. Horgan, B. Piot, M. Azar, and D. Silver.  2017\.  Rainbow: Combining improvements in deep reinforcement learning.  arXiv preprint arXiv:1710.02298.
  * [Hewlett et al., 2016]  D. Hewlett, A. Lacoste, L. Jones, I. Polosukhin, A. Fandrianto, J. Han, M. Kelcey, and D. Berthelot.  2016\.  Wikireading: A novel large-scale language understanding task over Wikipedia.  In Association for Computational Linguistics (ACL).
  * [Hewlett et al., 2017]  D. Hewlett, L. Jones, A. Lacoste, et al.  2017\.  Accurate supervised and semi-supervised machine reading for long documents.  In Empirical Methods in Natural Language Processing (EMNLP), pages 2011–2020.
  * [Hochreiter and Schmidhuber, 1997]  S. Hochreiter and J. Schmidhuber.  1997\.  Long short-term memory.  Neural Computation, 9(8):1735–1780.
  * [Joshi et al., 2017]  M. Joshi, E. Choi, D. Weld, and L. Zettlemoyer.  2017\.  TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension.  In Association for Computational Linguistics (ACL).
  * [Kadlec et al., 2016]  R. Kadlec, M. Schmid, O. Bajgar, and J. Kleindienst.  2016\.  Text understanding with the attention sum reader network.  In Association for Computational Linguistics (ACL).
  * [Kočisky et al., 2017]  T. Kočisky, J. Schwarz, P. Blunsom, C. Dyer, K. M. Hermann, G. Melis, and E. Grefenstette.  2017\.  The NarrativeQA reading comprehension challenge.  arXiv preprint arXiv:1712.07040.
  * [Lao et al., 2015]  N. Lao, E. Minkov, and W. Cohen.  2015\.  Learning relational features with backward random walks.  In Association for Computational Linguistics (ACL).
  * [Lee et al., 2017]  K. Lee, S. Salant, T. Kwiatkowski, A. Parikh, D. Das, and J. Berant.  2017\.  Learning recurrent span representations for extractive question answering.  arXiv.
  * [Mnih et al., 2015]  V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al.  2015\.  Human-level control through deep reinforcement learning.  Nature, 518(7540):529–533.
  * [Nogueira and Cho, 2016]  R. Nogueira and K. Cho.  2016\.  End-to-end goal-driven web navigation.  In Advances in Neural Information Processing Systems (NIPS).
  * [Onishi et al., 2016]  T. Onishi, H. Wang, M. Bansal, K. Gimpel, and D. McAllester.  2016\.  Who did what: A large-scale person-centered cloze dataset.  In Empirical Methods in Natural Language Processing (EMNLP).
  * [Rajpurkar et al., 2016]  P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang.  2016\.  SQuAD: 100,000+ questions for machine comprehension of text.  In Empirical Methods in Natural Language Processing (EMNLP).
  * [Salant and Berant, 2018]  S. Salant and J. Berant.  2018\.  Contextualized word representations for reading comprehension.  In North American Association for Computational Linguistics (NAACL).
  * [Schaul et al., 2015]  T. Schaul, J. Quan, I. Antonoglou, and D. Silver.  2015\.  Prioritized experience replay.  In International Conference on Learning Representations (ICLR).
  * [Seo et al., 2018]  M. Seo, S. Min, A. Farhadi, and H. Hajishirzi.  2018\.  Neural speed reading via skim-RNN.  In International Conference on Learning Representations (ICLR).
  * [Shen et al., 2017]  Y. Shen, P. Huang, J. Gao, and W. Chen.  2017\.  Reasonet: Learning to stop reading in machine comprehension.  In International Conference on Knowledge Discovery and Data Mining (KDD).
  * [Swayamdipta et al., 2018]  S. Swayamdipta, A. P. Parikh, and T. Kwiatkowski.  2018\.  Multi-mention learning for reading comprehension with neural cascades.  In International Conference on Learning Representations (ICLR).
  * [Trischler et al., 2017]  A. Trischler, T. Wang, X. Yuan, J. Harris, A. Sordoni, P. Bachman, and K. Suleman.  2017\.  NewsQA: A machine comprehension dataset.  In Workshop on Representation Learning for NLP.
  * [van Hasselt et al., 2016]  H. van Hasselt, A. Guez, and D. Silver.  2016\.  Deep reinforcement learning with double Q-learning.  In Association for the Advancement of Artificial Intelligence (AAAI), volume 16, pages 2094–2100.
  * [Wadhwa et al., 2018]  S. Wadhwa, V. Embar, M. Grabmair, and E. Nyberg.  2018\.  Towards inference-oriented reading comprehension: Parallelqa.  In Workshop on New Forms of Generalization in Deep Learning and Natural Language Processing at NAACL.
  * [Wang et al., 2016]  Z. Wang, T. Schaul, M. Hessel, H. V. Hasselt, M. Lanctot, and N. D. Freitas.  2016\.  Dueling network architectures for deep reinforcement learning.  In International Conference on Machine Learning (ICML).
  * [Wang et al., 2018a]  S. Wang, M. Yu, X. Guo, Z. Wang, T. Klinger, W. Zhang, S. Chang, G. Tesauro, B. Zhou, and J. Jiang.  2018a.  R3: Reinforced ranker-reader for open-domain question answering.  In Association for the Advancement of Artificial Intelligence (AAAI).
  * [Wang et al., 2018b]  S. Wang, M. Yu, J. Jiang, W. Zhang, X. Guo, S. Chang, Z. Wang, T. Klinger, G. Tesauro, and M. Campbell.  2018b.  Evidence aggregation for answer re-ranking in open-domain question answering.  In International Conference on Learning Representations.
  * [Watanabe et al., 2017]  Y. Watanabe, B. Dhingra, and R. Salakhutdinov.  2017\.  Question answering from unstructured text by retrieval and comprehension.  arXiv preprint arXiv:1703.08885.
  * [Welbl et al., 2017]  J. Welbl, P. Stenetorp, and S. Riedel.  2017\.  Constructing datasets for multi-hop reading comprehension across documents.  arXiv preprint arXiv:1710.06481.
  * [Xiong et al., 2017]  C. Xiong, V. Zhong, and R. Socher.  2017\.  Dynamic coattention networks for question answering.  In International Conference on Learning Representations (ICLR).
  * [Yu et al., 2017]  A. W. Yu, H. Lee, and Q. V. Le.  2017\.  Learning to skim text.  In Association for Computational Linguistics (ACL).
  * [Yu et al., 2018]  K. Yu, Y. Liu, A. G. Schwing, and J. Peng.  2018\.  Fast and accurate text classification: Skimming, rereading and early stopping.  In International Conference on Learning Representations (ICLR).

##  Appendix A Supplementary Material

###  A.1 Network Architecture Details

Here, we elaborate on the network architecture, which was briefly described in
the paper. Given an input state $s=(q,o,z,\phi_{n},\phi_{z})$, we denote by
$q=(q_{1},...,q_{n})$, $o=(o_{1},...,o_{m})$ and $z=(z_{1},...,z_{r})$ the
question tokens, observation tokens and answer prediction tokens,
respectively.

#### Word-level embedding

For every $i=1,...,n$, every $j=1,...,m$ and every $k=1,...,r$, we create an
embedding of the input tokens $q_{i},o_{j},z_{k}$ in two steps. First, we
embed every token with a pre-trained GloVe word embedding matrix $W_{w}$:

$\displaystyle e_{q_{i}}^{w}=W_{w}\cdot q_{i}\quad;\quad
e_{o_{j}}^{w}=W_{w}\cdot o_{j}\quad;\quad e_{z_{k}}^{w}=W_{w}\cdot z_{k}$

Next, we apply character-level embeddings with a learned embedding matrix
$W_{c}$. The embedded characters are then summarized with a max-pooling
convolutional neural network:

$\displaystyle e_{q_{i}}^{c}=\mathrm{CNN}(\\{W_{c}\cdot
q_{ix}\\}_{x=1}^{|q_{i}|})$ $\displaystyle
e_{o_{j}}^{c}=\mathrm{CNN}(\\{W_{c}\cdot o_{jx}\\}_{x=1}^{|o_{j}|})$
$\displaystyle e_{z_{k}}^{c}=\mathrm{CNN}(\\{W_{c}\cdot
o_{kx}\\}_{x=1}^{|z_{k}|})$

Concatenation of the two components yields the final word-level embeddings:

$\displaystyle e_{q_{i}}=[e_{q_{i}}^{w}\,;\,e_{q_{i}}^{c}]\quad;\quad
e_{o_{j}}=[e_{o_{j}}^{w}\,;\,e_{o_{j}}^{c}]\quad;\quad
e_{z_{k}}=[e_{z_{k}}^{w}\,;\,e_{z_{k}}^{c}]$

#### Question sequence encoding

The question is encoded with a BiLSTM, where for every timestamp $i$, the
forward and backward outputs are concatenated:

$\displaystyle\\{u_{q_{1}},...,u_{q_{n}}\\}=\mathrm{BiLSTM}(e_{q_{1}},...,e_{q_{n}})$
$\displaystyle
u_{q_{i}}:=[\overrightarrow{\mathrm{LSTM}}(\overrightarrow{h}_{q_{i-1}},e_{q_{i}})\;;\;\overleftarrow{\mathrm{LSTM}}(\overleftarrow{h}_{q_{i+1}},e_{q_{i}})]$

The outputs are then summarized with a self-attention:

$h_{q}=\sum_{i=1}^{n}{\alpha_{i}u_{q_{i}}}$

where the coefficients $\alpha_{1},...,\alpha_{n}$ are obtained by feeding the
outputs to a two-layer feed-forward network, and normalizing the output logits
with softmax:

$\displaystyle
a_{i}=\mathrm{FFNN}(u_{q_{i}})\quad;\quad\alpha_{i}=\frac{e^{a_{i}}}{\sum_{k=1}^{n}{e^{a_{k}}}}$

#### Observation sequence encoding

The observation sequence encoding $h_{o}$ is obtained in an analogous manner
to $h_{q}$, except that we use a $\mathrm{LSTM}$ rather than a
$\mathrm{BiLSTM}$.

#### Answer prediction encoding

The answer prediction encoding $h_{z}$ is obtained by running the answer
prediction tokens through a LSTM, and concatenating the last hidden state with
the input feature vector $\phi_{z}$:

$\displaystyle\\{u_{z_{1}},...,u_{z_{r}}\\}=\mathrm{LSTM}(e_{z_{1}},...,e_{z_{r}})$
$\displaystyle
u_{z_{k}}:=\overrightarrow{\mathrm{LSTM}}(\overrightarrow{h}_{z_{k-1}},e_{z_{k}})$
$\displaystyle h_{z}=[u_{z_{r}}\;;\;\phi_{z}]$

#### State representation

The final state representation is formed by concatenating the encoded question
$h_{q}$ with the encoded observation $h_{o}$. Concretely:

$h_{s}=[h_{q}\;;\;h_{o}\;;\;h_{z}]$

The input node features $\phi_{o}$ are concatenated to an upper layer, as we
describe next , to increase their weights in the final predictions. We have
found that incorporation of these features in this way accelerates the
navigation learning process.

#### Q-values prediction

We implement a Dueling DQN architecture, where the final Q-values are composed
of a state value $V(s)$ prediction and advantage predictions $A(s,a)$ for
every possible action $a$. Denoting by FFNN a single-layer feed-forward neural
network, predictions are derived as follows:

$\displaystyle v_{0}=\mathrm{FFNN}(h_{s})$ $\displaystyle
v_{1}^{V}=\mathrm{FFNN}(v_{0})\;;\;v_{1}^{A}=\mathrm{FFNN}(v_{0})$
$\displaystyle
v_{2}^{V}=\mathrm{FFNN}([v_{1}^{V}\;;\;\phi_{n}])\;;\;v_{2}^{A}=\mathrm{FFNN}([v_{1}^{A}\;;\;\phi_{n}])$

Where $v_{2}^{V}\in\mathbb{R}$, $v_{2}^{A}\in\mathbb{R}^{|\mathcal{A}|}$, and
$\phi$ are the navigation features. The final Q-values prediction is obtained
by averaging:

$v^{Q}=v_{2}^{V}+\Big{(}v_{2}^{A}-\frac{\sum_{i=1}^{|\mathcal{A}|}v_{2i}^{A}}{|\mathcal{A}|}\Big{)}$

###  A.2 Hyper Parameters

Table 8 summarizes the hyper-parameters used for building and training the
DocQN models.

| Parameter | Value
---|---|---
Network | Maximum node observation length | 20 tokens
Maximum observation length | 120 tokens
| Word embedding dimension | 300
| Character embedding dimension | 20
| Convolution filter size | 5
| BiLSTM and LSTM hidden dimension | 300
| First feed-forward layer dimension | 512
| Second feed-forward layer dimension | 256
| Dropout rate | 0.2
Training | RMSprop learning rate | 0.0001
Batch size | 64
Target network period | 10K steps
| Initial memory size | 50K transitions
| Maximal memory size | 300K transitions
| Action sampling $\epsilon$ | $1.0\rightarrow 0.1$
| Discount factor $\gamma$ | $0.996$
| Prioritization usage $\alpha$ | 0.6
| Prioritization importance sampling $\beta$ | $0.4\rightarrow 1.0$
Sampling | State sampling $\epsilon_{s}$ | $1.0\rightarrow 0.5$
Annealing steps for $\epsilon_{s}$ |  $1.2$M steps
| Sampling repetitions $K$ | 5
| Maximum navigation length $T$ (train) | 30 steps
| Maximum navigation length $T$ (test) | 100 steps
| Interpolation coefficient for $f_{U+B}$ | 0.5
Table 8: DocQN hyper-parameters

###  A.3 Navigation Examples

Figures 9,10,11,12,13 show sample navigations, performed by the DocQN model.
In each example, $q,a,n$ denote the question, answer aliases and answer node
numbers. The observation tokens are highlighted according to the self-
attention weights, given by the model. By choosing the action ANS, the agent
executes the RC model to obtain an answer prediction, which is part of the
observation in the next step.

Figure 9: Navigation example 1. Although the document is more general than the
question subject and the answer appears multiple times across the document,
the agent finds the correct context for answering the question. Figure 10:
Navigation example 2. The agent explores several sections before reaching the
most relevant one.

Figure 11: Navigation example 3. The agent quickly navigates to the answer,
without running the RC model. Figure 12: Navigation example 4. The agent
leverages the RC model and decides to stop after observing enough information
to answer.

Figure 13: Navigation example 5. After reading through the ”Television
episodes” section, the agent goes back to the ”Production” section to find the
answer.


$V_{\text{rag}}(\rho)$ with a linear decay factor $m(\rho)$. As we require
$V_{\text{rag}}(\rho)<1$ to ensure high retrieval quality, large distribution
shift radius $\rho$ must be compensated by small $V_{\text{rag}}$. This is
consistent with the existing observations that low-variance models can
generalize better under distribution shifts (Lam, 2016; Gotoh et al., 2018;
Namkoong & Duchi, 2017). (R2) Compared to Proposition 3, Proposition 4 removes
the dependency on varying label portions $r^{(c)}_{\text{cal}}$ during
distribution shifts, as long as the size of external knowledge base is
sufficiently large $N_{\text{ext}}$ to offset the worst-case long-tail
distributions, a condition often met in practice with large knowledge bases.

Proof sketch. We first formulate the expectation of similarity difference
between positive pairs and negative pairs
$\mathbb{E}[s_{\theta_{r}}(x,x^{+})-s_{\theta_{r}}(x,x^{-})]$ as a function of
the contrastive loss of the retrieval model $L_{\tau}$. Then, we apply
Chebyshev’s inequality to upper bound the failure probability
${\mathbb{P}}[s_{\theta}(x,x^{+})<s_{\theta}(x,x^{-})]$ as a function of the
variance and expectation of the similarity difference. Considering the
distribution shifts, the variance and expectation bound can be derived via
Grammian bound as (Weber et al., 2022). We then plug in the failure rate
corrected by distribution shifts and follow the proof structure of Proposition
3.

###### Proof of Theorem 2.

Recall that the calibration samples $(Z_{1},Z_{2},...,Z_{N_{\text{cal}}})$ are
drawn from the distribution ${\mathcal{D}}$ and the test sample
$Z_{\text{test}}=(X_{\text{test}},Y_{\text{test}})$ is sampled from the
distribution ${\mathcal{Q}}$. The distribution ${\mathcal{Q}}$ and
${\mathcal{D}}$ have bounded Hellinger distance
$H({\mathcal{Q}},{\mathcal{D}})\leq\rho$ and the risk is upper bounded by $1$.
Let $R_{\mathcal{Q}}$, $R_{\mathcal{D}}$ be the population risk on
distribution ${\mathcal{Q}}$, ${\mathcal{D}}$ and $V_{\mathcal{D}}$ be the
variance of risk on distribution ${\mathcal{D}}$. We can apply Equation 21 and
get the following:

$R_{\mathcal{Q}}\leq
R_{\mathcal{D}}+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}\sqrt{V_{\mathcal{D}}}+\rho^{2}(2-\rho^{2})\left(1-R_{\mathcal{D}}-\dfrac{V_{\mathcal{D}}}{1-R_{\mathcal{D}}}\right),$
(93)

with any given distance bound $\rho>0$ that satisfies:

$\rho^{2}\leq
1-\left(1+\dfrac{(1-R_{\mathcal{D}})^{2}}{V_{\mathcal{D}}}\right)^{-1/2}.$
(94)

Since the variance is non-negative, Equation 93 further implies that:

$R_{\mathcal{Q}}\leq
R_{\mathcal{D}}+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}\sqrt{V_{\mathcal{D}}}+\rho^{2}(2-\rho^{2})\left(1-R_{\mathcal{D}}\right).$
(95)

Then we will consider the finite-sample error of the calibration set. From
Hoeffding’s inequality (Hoeffding, 1994), with probability $1-\delta/4$, we
have:

$R_{\mathcal{D}}\leq\hat{R}+\sqrt{\dfrac{\ln(4/\delta)}{2N_{\text{cal}}}}.$
(96)

From sample variance bound in (Maurer & Pontil, 2009), with probability
$1-\delta/4$, we have:

$\sqrt{V_{\mathcal{D}}}\leq\sqrt{\hat{V}}+\sqrt{\dfrac{2\ln(4/\delta)}{N_{\text{cal}}-1}}.$
(97)

Note that the RHS of Equation 95 monotonically increases in $R_{\mathcal{D}}$
and $V_{\mathcal{D}}$, combining Equations 96, 97 and 95 and applying the
union bound, with probability $1-\delta/2$, we have:

$\displaystyle
R_{\mathcal{Q}}\leq\hat{R}+\rho^{2}(2-\rho^{2})\left(1-\hat{R}\right)+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}\sqrt{\hat{V}}+(1-\rho^{2})(\dfrac{1-\rho^{2}}{\sqrt{2N_{\text{cal}}}}+\dfrac{2\sqrt{2}\rho\sqrt{2-\rho^{2}}}{\sqrt{N_{\text{cal}}-1}})\sqrt{\ln(4/\delta)}.$
(98)

By two-sided Hoeffding’s inequality, with probability $1-\delta/4$, we have:

$\hat{R}_{\rho}\leq
R_{\mathcal{Q}}+\sqrt{\dfrac{\ln{(8/\delta)}}{2N_{\text{cal}}}}.$ (99)

Combining Equations 98 and 99 and applying the union bound, with probability
$1-3\delta/4$, we have:

$\hat{R}_{\rho}\leq\hat{R}+\rho^{2}(2-\rho^{2})\left(1-\hat{R}\right)+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}\sqrt{\hat{V}}+(1-\rho^{2})(\dfrac{1-\rho^{2}}{\sqrt{2N_{\text{cal}}}}+\dfrac{2\sqrt{2}\rho\sqrt{2-\rho^{2}}}{\sqrt{N_{\text{cal}}-1}})\sqrt{\ln(4/\delta)}+\sqrt{\dfrac{\ln{(8/\delta)}}{2N_{\text{cal}}}}.$
(100)

Recall that the controlled conformal risk $\hat{\alpha}$ can be formulated as
a function of the empirical risk $\hat{R}$ as Equation 14. When
$(X_{\text{test}},Y_{\text{test}})$ is sampled from ${\mathcal{D}}$, the
guarantee of conformal risk is as follows:

${\mathbb{P}}\left[R(T_{\hat{{\bm{\lambda}}}};X_{\text{test}},Y_{\text{test}})\leq\hat{\alpha}:=\min\left\\{h^{-1}\left(\dfrac{8/\delta}{N_{\text{cal}}};\hat{R}\right),\Phi^{-1}_{\text{bin}}\left(\dfrac{\delta}{8e};N_{\text{cal}},\hat{R}\right)\right\\}\right]\geq
1-\delta/4,$ (101)

where $h^{-1}(\cdot;\cdot)$ is the partial inverse function such that
$h^{-1}(h(a,b);a)=b$ with $h_{1}(a,b)=a\log(a/b)+(1-a)\log((1-a)/(1-b))$. Note
that $\hat{\alpha}$ monotonically increases in $\hat{R}$. By Equation 100,
with probability $1-3\delta/4$, the following holds about the conformal risk
on distribution ${\mathcal{Q}}$ (denoted by $\hat{\alpha}_{\rho}$), which is
within bounded Hellinger distance $\rho$ to the original distribution
${\mathcal{D}}$:

$\hat{\alpha}_{\rho}\leq\min\left\\{h^{-1}\left(\dfrac{8/\delta}{N_{\text{cal}}};\overline{\hat{R}_{\rho}}\right),\Phi^{-1}_{\text{bin}}\left(\dfrac{\delta}{8e};N_{\text{cal}},\overline{\hat{R}_{\rho}}\right)\right\\},$
(102)

where $\overline{\hat{R}_{\rho}}$ is formulated as:

$\overline{\hat{R}_{\rho}}=\hat{R}+\rho^{2}(2-\rho^{2})\left(1-\hat{R}\right)+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}\sqrt{\hat{V}}+(1-\rho^{2})(\dfrac{1-\rho^{2}}{\sqrt{2N_{\text{cal}}}}+\dfrac{2\sqrt{2}\rho\sqrt{2-\rho^{2}}}{\sqrt{N_{\text{cal}}-1}})\sqrt{\ln(4/\delta)}+\sqrt{\dfrac{\ln{(8/\delta)}}{2N_{\text{cal}}}}.$
(103)

Combining Equations 102 and 101 and applying the union bound, we can finally
conclude that:

${\mathbb{P}}\left[R(T_{\hat{{\bm{\lambda}}}};X_{\text{test}},Y_{\text{test}})\leq\hat{\alpha}_{\rho}:=\min\left\\{h^{-1}\left(\dfrac{8/\delta}{N_{\text{cal}}};\overline{\hat{R}_{\rho}}\right),\Phi^{-1}_{\text{bin}}\left(\dfrac{\delta}{8e};N_{\text{cal}},\overline{\hat{R}_{\rho}}\right)\right\\}\right]\geq
1-\delta,$ (104)

where $\overline{\hat{R}_{\rho}}$ is formulated as:

$\overline{\hat{R}_{\rho}}=\hat{R}+\rho^{2}(2-\rho^{2})\left(1-\hat{R}\right)+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}\sqrt{\hat{V}}+(1-\rho^{2})(\dfrac{1-\rho^{2}}{\sqrt{2N_{\text{cal}}}}+\dfrac{2\sqrt{2}\rho\sqrt{2-\rho^{2}}}{\sqrt{N_{\text{cal}}-1}})\sqrt{\ln(4/\delta)}+\sqrt{\dfrac{\ln{(8/\delta)}}{2N_{\text{cal}}}}.$
(105)

∎

###  I.2 Proof and detailed remark of Theorem 3

######  Remarks.

Our result rigorously characterizes the effect of distribution shift on the
reduced risk guarantee of RAG. (R1) Compared to Theorem 1, only the
uncertainty of retrieval model $p_{r}(\rho)$ is affected by the distribution
shift $\rho$. This affect is reflected on the the retrieval quality
$V_{\text{rag}}(\rho)$. In particular, a large distance radius $\rho$ will
downgrade the retrieval quality $V_{\text{rag}}(\rho)$ and thus lead to a
higher uncertainty $p_{r}(\rho)$. However, the influence of $\rho$ on
$p_{r}(\rho)$ can be reduced by $N_{\text{rag}}$ inverse proportionally and by
$N_{\text{ext}}$ exponentially, demonstrating the robustness of RAG with more
retrieval knowledge. (R2) Since $V_{\text{rag}}(\rho)$ is proportional to
model variance $V_{\text{rag}}$, a low-variance retrieval model demonstrates
better robustness against distribution drifts, aligning with existing
empirical observations (Lam, 2016; Gotoh et al., 2018; Namkoong & Duchi,
2017), which evaluate the generalization ability of low-variance retrieval
models under distribution shifts. Different from Proposition 3, the quality of
retrieval models under distribution shifts is decreased from $V_{\text{rag}}$
to $V_{\text{rag}}(\rho)$ with a linear decay factor $m(\rho)$. As we require
$V_{\text{rag}}(\rho)<1$ to ensure high retrieval quality, large distribution
shift radius $\rho$ must be compensated by small $V_{\text{rag}}$. This is
consistent with the existing observations that low-variance models can
generalize better under distribution shifts (Lam, 2016; Gotoh et al., 2018;
Namkoong & Duchi, 2017). (R3) Compared to Theorem 1, Theorem 3 has no
dependence on varying label portions $r^{(c)}_{\text{cal}}$ during
distribution shifts, as long as the size of external knowledge base
$N_{\text{ext}}$ is moderately large, ($N_{\text{ext}}>2\sqrt{2\ln
10}/\min_{c}r^{(c)}_{\text{ext}}$) to offset the worst-case long-tail
distributions, a condition often met in practice with large knowledge bases.

Proof sketch. We apply Proposition 4 to provide the lower bound of the
retrieved positive examples under distribution shifts. We plug in the term and
analyze the functionality (logit difference statistics) of the self-attention
transformer as proof of Theorem 1. Connecting the logit difference statistics
to the empirical risks and the distribution-shifted conformal risk bound in
Theorem 2 finally concludes the proof.

###### Proof of Theorem 3.

From Proposition 4, we prove that:

$\mathbb{E}\left[N_{\text{pos}}\right]\geq\underline{N}_{\text{pos}}:=\dfrac{9}{10}N_{\text{rag}}\left(1-1.5N_{\text{ext}}V_{\text{rag}}(\rho)^{0.25\left(\min_{c}r^{(c)}_{\text{ext}}N_{\text{ext}}\right)}\right).$
(106)

Since $N_{\text{pos}}$ is a binomial random variable with $N_{\text{rag}}$
trials, we have the upper bound of the variance
$\mathbb{V}[N_{\text{pos}}]\leq\dfrac{N_{\text{rag}}}{4}$. Applying
Chebyshev’s inequality to the random variable $N_{\text{pos}}$, the following
holds $\forall n_{\text{pos}}<\underline{N}_{\text{pos}}$:

${\mathbb{P}}\left[N_{\text{pos}}\geq n_{\text{pos}}\right]\geq
1-\dfrac{\mathbb{V}[N_{\text{rag}}]}{(\underline{N}_{\text{pos}}-n_{\text{pos}})^{2}}\geq
1-\dfrac{N_{\text{rag}}}{4(\underline{N}_{\text{pos}}-n_{\text{pos}})^{2}},$
(107)

which implicates that we can do the analysis with $N_{\text{pos}}\geq
n_{\text{pos}}$ with probability
$1-\dfrac{\mathbb{V}[N_{\text{rag}}]}{(\underline{N}_{\text{pos}}-n_{\text{pos}})^{2}}$.

According to the proof of Theorem 1, by Equation 60, we have the following:

$\mathbb{E}\left[R-R_{\text{rag}}\right]\geq\Phi_{M}\left(\left((n_{\text{pos}}+1)d^{+}-1\right)n_{\text{pos}}\left(\int_{-1}^{1}\Phi_{M}(v)dv-1\right)\right)-\Phi_{M}(0).$
(108)

.

Let $n_{\text{rag}}=N_{\text{rag}}/2$. Combining Equation 108 and Equation
107, we get that if $\underline{N}_{\text{pos}}>{N_{\text{rag}}}/{2}>1/d^{+}$,
with probability
$1-\dfrac{N_{\text{rag}}}{4(\underline{N}_{\text{pos}}-{N_{\text{rag}}}/{2})^{2}}$,
we have:

$\mathbb{E}\left[R-R_{\text{rag}}\right]\geq\Phi_{M}\left(\dfrac{d^{+}\left(\int_{-1}^{1}\Phi_{M}(v)dv-1\right)N_{\text{rag}}}{2}\right)-\Phi_{M}(0).$
(109)

Let $R(Z)$ be the risk of $Z$ sampled from the distribution ${\mathcal{Q}}$.
Define the empirical risk $\hat{R}$ and $\hat{R}_{\text{rag}}$ as the
following:

$\displaystyle\hat{R}=\dfrac{1}{N_{\text{cal}}}\sum_{i=1}^{N_{\text{cal}}}R(Z_{i}),\quad\hat{R}_{\text{rag}}=\dfrac{1}{N_{\text{cal}}}\sum_{i=1}^{N_{\text{cal}}}R_{\text{rag}}(Z_{i})$
(110)

Note that since the risk $R(\cdot)$ is bounded in $[0,1]$, the variance
estimator $\hat{V}=\dfrac{1}{N_{\text{cal}}(N_{\text{cal}}-1)}\sum_{1\leq
i<j\leq N_{\text{cal}}}(R(Z_{i})-R(Z_{j}))^{2}$ is bounded in $[0,1]$.
Leveraging the fact and according to Theorem 2, the statistical guarantee of
conformal risk $\hat{\alpha}_{\rho}$ and $\hat{\alpha}_{\rho}^{\text{rag}}$
with confidence $1-\delta$ can be formulated as:

$\displaystyle\hat{\alpha}_{\rho}:=\min\left\\{h^{-1}\left(\dfrac{8/\delta}{N_{\text{cal}}};\overline{\hat{R}_{\rho}}\right),\Phi^{-1}_{\text{bin}}\left(\dfrac{\delta}{8e};N_{\text{cal}},\overline{\hat{R}_{\rho}}\right)\right\\}$
(111) where
$\displaystyle\overline{\hat{R}_{\rho}}={\hat{R}+\rho^{2}(2-\rho^{2})(1-\hat{R})}+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}+C(\rho,N_{\text{cal}}),$
$\displaystyle\hat{\alpha}_{\rho}^{\text{rag}}:=\min\left\\{h^{-1}\left(\dfrac{8/\delta}{N_{\text{cal}}};\overline{\hat{R}_{\rho}^{\text{rag}}}\right),\Phi^{-1}_{\text{bin}}\left(\dfrac{\delta}{8e};N_{\text{cal}},\overline{\hat{R}_{\rho}^{\text{rag}}}\right)\right\\}$
(112) where
$\displaystyle\overline{\hat{R}_{\rho}^{\text{rag}}}={\hat{R}_{\text{rag}}+\rho^{2}(2-\rho^{2})(1-\hat{R}_{\text{rag}})}+2\rho(1-\rho^{2})\sqrt{2-\rho^{2}}+C(\rho,N_{\text{cal}}),$

where
$C(\rho,N_{\text{cal}})={(1-\rho^{2})\left(\dfrac{1-\rho^{2}}{\sqrt{2N_{\text{cal}}}}+\dfrac{2\sqrt{2}\rho\sqrt{2-\rho^{2}}}{\sqrt{N_{\text{cal}}-1}}\right)\sqrt{\ln(4/\delta)}+\sqrt{\dfrac{\ln{(8/\delta)}}{2N_{\text{cal}}}}}$.

Noting that $\hat{\alpha}_{\rho}$ is monotonically increasing in
$\overline{\hat{R}_{\rho}}$ and $\overline{\hat{R}_{\rho}}$ is monotonically
increasing in $\hat{R}$, $\hat{\alpha}_{\rho}$ is monotonically increasing in
$\hat{R}$. Then, the following holds by Hoeffding’s inequality:

${\mathbb{P}}\left[\hat{\alpha}_{\rho}^{\text{rag}}<\hat{\alpha}_{\rho}\right]\geq{\mathbb{P}}\left[\hat{R}_{\text{rag}}<\hat{R}\right]\geq
1-\exp\left\\{-2N_{\text{cal}}\mathbb{E}\left[R-R_{\text{rag}}\right]^{2}\right\\}.$
(113)

Combining Equation 109 and Equation 113 and using the union bound, under the
condition that $\underline{N}_{\text{pos}}>{N_{\text{rag}}}/{2}>1/d^{+}$, we
have:

${\mathbb{P}}\left[\hat{\alpha}_{\text{rag}}<\hat{\alpha}\right]\geq
1-\exp\left\\{-2N_{\text{cal}}\left[\Phi_{M}\left(\dfrac{d^{+}\left(\int_{-1}^{1}\Phi_{M}(v)dv-1\right)N_{\text{rag}}}{2}\right)-\Phi_{M}(0)\right]^{2}\right\\}-\dfrac{N_{\text{rag}}}{4(\underline{N}_{\text{pos}}-{N_{\text{rag}}}/{2})^{2}}.$
(114)

Combining Equations 106 and 114, under the condition that
$N_{\text{ext}}>\dfrac{2\sqrt{2\ln 10}}{r_{\text{ext}}^{m}}$,
$N_{\text{ext}}V_{\text{rag}}(\rho)^{0.25r_{\text{ext}}^{m}N_{\text{ext}}}<\dfrac{8}{17}$,
and $N_{\text{rag}}>\dfrac{2}{d^{+}}$, we can finally conclude that:

$\displaystyle{\mathbb{P}}\left[\hat{\alpha}_{\text{rag}}<\hat{\alpha}\right]\geq$
$\displaystyle
1-\exp\left\\{-2N_{\text{cal}}\left[\Phi_{M}\left(\dfrac{d^{+}\left(\int_{-1}^{1}\Phi_{M}(v)dv-1\right)N_{\text{rag}}}{2}\right)-\Phi_{M}(0)\right]^{2}\right\\}$
(115)
$\displaystyle-\dfrac{100}{N_{\text{rag}}}\left(8-17N_{\text{ext}}V_{\text{rag}}(\rho)^{0.25\left(\min_{c}r^{(c)}_{\text{ext}}N_{\text{ext}}\right)}\right)^{-2}.$

∎

##  Appendix J Additional evaluation results

Algorithm 3 Test distribution sampling protocol.

1: Input: original test set ${\mathcal{D}}$, test sample pool
${\mathcal{D}}_{\text{pool}}$, Risk function
$R(\cdot,\cdot):{\mathcal{X}}\times{\mathcal{Y}}\mapsto{\mathbb{R}}$

2: Output: empirical risk on the sampled test set ${\mathcal{Q}}$
$\hat{R}_{\mathcal{Q}}$

3: Randomly sample ${\mathcal{Q}}$ from ${\mathcal{D}}_{\text{pool}}$ with
equalized set size as ${\mathcal{D}}$: $|{\mathcal{Q}}|=|{\mathcal{D}}|$

4: Evaluate empirical risks for all samples in ${\mathcal{Q}}$:
$\hat{{R}}_{\mathcal{Q}}=\dfrac{1}{|{\mathcal{Q}}|}\sum_{(x,y)\in{\mathcal{Q}}}R(x,y)$

5: Return $\hat{{R}}_{\mathcal{Q}}$

Algorithm 4 Shifted distribution sampling protocol.

1: Input: original test set ${\mathcal{D}}$, test sample pool
${\mathcal{D}}_{\text{pool}}$, Risk function
$R(\cdot,\cdot):{\mathcal{X}}\times{\mathcal{Y}}\mapsto{\mathbb{R}}$

2: Output: empirical risk on the shifted test set ${\mathcal{Q}}$
$\hat{R}_{\mathcal{Q}}$, Hellinger distance between ${\mathcal{D}}$ and
${\mathcal{Q}}$ $H_{{\mathcal{P}}{\mathcal{Q}}}$

3: Randomly sample ${\mathcal{Q}}$ from ${\mathcal{D}}_{\text{pool}}$ with
equalized set size as ${\mathcal{D}}$: $|{\mathcal{Q}}|=|{\mathcal{D}}|$

4: Randomly sample the sample weight vector of ${\mathcal{Q}}$
${\bm{w}}\in\Delta^{|{\mathcal{D}}|}$

5: Compute Hellinger distance as
$H_{{\mathcal{P}}{\mathcal{Q}}}=\sqrt{1-\sum_{i=1}^{|{\mathcal{D}}|}\sqrt{{{\bm{w}}_{i}}/{|{\mathcal{D}}|}}}$

6: Evaluate risks for all samples in ${\mathcal{Q}}$ with risk function
$R(\cdot,\cdot)$:
$\hat{\bm{R}}_{\mathcal{Q}}\in{\mathbb{R}}^{|{\mathcal{D}}|}$

7: Compute the empirical risk on ${\mathcal{Q}}$ with weight vector
${\bm{w}}$: $\hat{R}_{\mathcal{Q}}={\bm{w}}^{T}\hat{\bm{R}}_{\mathcal{Q}}$

8: Return $\hat{R}_{\mathcal{Q}}$, $H_{{\mathcal{P}}{\mathcal{Q}}}$

AESLC

Hellinger Distance $\rho$

Evaluation Risk

CommonGen

Hellinger Distance $\rho$

DART

Hellinger Distance $\rho$

E2E

Hellinger Distance $\rho$

Figure 8: Conformal generation risk $\hat{\alpha}_{\text{rag}}$ and empirical
risks with BM25 under distribution shifts with
$N_{\text{rag}}=15,\lambda_{g}=1,\lambda_{s}=1.0$.

AESLC

Hellinger Distance $\rho$

Evaluation Risk

CommonGen

Hellinger Distance $\rho$

DART

Hellinger Distance $\rho$

E2E

Hellinger Distance $\rho$

Figure 9: Conformal generation risk $\hat{\alpha}_{\text{rag}}$ and empirical
risks with BAAI/bge under distribution shifts with
$N_{\text{rag}}=15,\lambda_{g}=1,\lambda_{s}=1.0$.

AESLC

Hellinger Distance $\rho$

Conformal Risk $\alpha_{\text{rag}}$

CommonGen

Hellinger Distance $\rho$

DART

Hellinger Distance $\rho$

E2E

Hellinger Distance $\rho$

Figure 10: Conformal generation risk $\alpha_{\text{rag}}$ vs. Hellinger
distance $\rho$ for different retrieval models under distribution shifts with
$N_{\text{rag}}=15,\lambda_{g}=1,\lambda_{s}=1.0$.

DART

E2E

Figure 11: Conformal generation risk $\hat{\alpha}_{\text{rag}}$ and simulated
empirical risks with different $\lambda_{g}$ and $N_{\text{rag}}$ for
OpenAI/ada.

AESLC

Hellinger Distance

Evaluation Risk

CommonGen

Hellinger Distance

DART

Hellinger Distance

E2E

Hellinger Distance

Figure 12: Conformal generation risk $\hat{\alpha}_{\text{rag}}$ and empirical
risks with $N_{\text{rag}}=5,\lambda_{g}=20,\lambda_{s}=1.0$ under
distribution shifts for Biencoder-SFT.

AESLC

# Generations $\lambda_{g}$

Evaluation Risk

CommonGen

# Generations $\lambda_{g}$

DART

# Generations $\lambda_{g}$

E2E

# Generations $\lambda_{g}$

Figure 13: Conformal generation risk $\hat{\alpha}_{\text{rag}}$ and simulated
empirical risks for different generation set size $\lambda_{g}$ and fixed
$N_{\text{rag}}=5,\lambda_{s}=1.0$ for Biencoder-SFT.

AESLC

diversity threshold $\lambda_{s}$

Evaluation Risk

CommonGen

diversity threshold $\lambda_{s}$

DART

diversity threshold $\lambda_{s}$

E2E

diversity threshold $\lambda_{s}$

Figure 14: Conformal generation risk $\hat{\alpha}_{\text{rag}}$ and simulated
empirical risks for different diversity threshold $\lambda_{s}$ with fixed
$N_{\text{rag}}=5,\lambda_{g}=20$ for Biencoder-SFT.

###  J.1 Experiment setup

#### Datasets

We evaluate C-RAG on four NLP datasets with retrieval augmented generation
utilizing an external knowledge base. (1) AESLC: The Annotated Enron Subject
Line Corpus (AESLC) dataset (Zhang & Tetreault, 2019) contains a collection of
email messages from employees of the Enron Corporation. It contains two
primary features: the “email_body“, which is the text of the email body, and
the “subject_line“, which is the text of the email subject. The task is to
generate the email subject given the email body. (2) The Generative
Commonsense Reasoning (CommonGen) dataset (Lin et al., 2019) is a collection
of commonsense descriptions, amounting to 79k descriptions over 35k unique
concept sets, constructed via crowdsourcing and existing caption corpora. The
task is to generate a sentence that uses given concepts in a coherent and
commonsense way. (3) The Data Record to Text (DART) dataset (Nan et al., 2020)
is a large-scale dataset designed to facilitate the generation of text from
structured data records. It comprises 82,191 examples spanning various
domains, with each example being a set of Resource Description Framework (RDF)
triples derived from data records in tables and tree ontologies of schemas.
These are annotated with sentence descriptions that encapsulate all the facts
presented in the RDF triplet. (4) The End-to-End Generation (E2E) dataset
(Novikova et al., 2017) contains about 50k comments in the restaurant domain.
The task is to generate text from meaning representations (MRs), which are
structured inputs that describe various aspects of a restaurant. These MRs
consist of slots and values that a model needs to convert into natural
language descriptions that are coherent and fluent.

External knowledge base: Following (Wang et al., 2023b; Wei et al., 2021;
Cheng et al., 2023), we construct the external knowledge base as the union of
a total of 30 publicly available datasets from 9 distinct categories with over
6 million documents.

#### Metrics

For generation tasks, we leverage ROUGE-L to quantify the quality of
generations. ROUGE-L measures the longest common subsequence between a
candidate generation and reference texts and is typically adopted for
generation quality evaluations across the literature (Lin, 2004; Gatt &
Krahmer, 2018). A low ROUGE-L implies poor quality generations and accordingly
a high generation risk. Therefore, we adopt the risk function as
$1-\text{ROUGE-L}$ to bound the risk in $[0,1]$. Note that C-RAG is agnostic
to selection of risk functions, and thus, practitioners can specify any
function that suits their specific use cases.

#### Retrieval models

We consider four types of retrieval models. (1) BM25 (Robertson et al., 2009)
is a sparse encoding metric used to rank the candidate documents given a query
in information retrieval tasks. BM25 scores are linearly weighted combinations
of token-level scores and can be analytically formulated as a function of term
frequency, inverse document frequency, and length of documents. (2) BAAI
general embedding (BAAI/bge) (Zhang et al., 2023a) trains the SOTA embedding
model in the MTEB benchmark (Muennighoff et al., 2022) and supports the
diverse retrieval augmentation needs of LLMs by a reward formulation based on
LLMs’ feedback, the stabilization of knowledge distillation, multi-task fine-
tuning with explicit instructions, and homogeneous in-batch negative sampling.
(3) OpenAI ada-text-embedding-02 (OpenAI/ada) is a close source text embedding
models designed to convert text into high-dimensional vectors, which capture
the semantic meaning of the input text and can be used for a variety of tasks,
such as text similarity, classification, and clustering. (4) Biencoder-
supervised fine-tuning (Biencoder-SFT) (Wang et al., 2023b) is a bi-encoder
based dense retriever trained with contrastive loss and hard negative sampling
strategies. It iteratively trains the retrieval model with hard negative
samples identified by computing similarity scores by the current retrieval
model.

#### Implementation details

We leverage Llama-7b for inference without specification. We perform conformal
calibration on validation sets for different NLP datasets and fix the
confidence levels $1-\delta=0.9$ across all evaluations. Lastly, we
concatenate the retrieved in-context examples and the query sample with line
break tokens.

###  J.2 Conformal risk bounds evaluation

Soundness and tightness of generation risk guarantees in C-RAG. To achieve the
generation risk guarantee in Equation 1, C-RAG computes the upper confidence
bound of generation risk by Proposition 1, which takes the empirical risk,
calibration size and confidence level as input. We evaluate the conformal
risks of RAG models $\alpha_{\text{rag}}$ with variations of the number of
retrieved examples $N_{\text{rag}}$ by calibration statistics on the
validation set. To validate the soundness and tightness of the risk bounds, we
randomly sample multiple test sets from a pool of test samples and compute the
empirical risks on the sampled test sets. The sampling protocol is detailed in
Algorithm 3 in Section J.2. We provide the results for OpenAI/ada, BM25,
BAAI/bge, Biencoder-SFT in Figure 3. The results show that (1) the certified
conformal risks $\alpha_{\text{rag}}$ (black up-pointing triangles) are larger
than the empirical risks of sampled test sets (grey points) and some empirical
risks approach the risk bounds, demonstrating the soundness and tightness of
risk bounds in C-RAG, and (2) the conformal risks decrease much as the number
of retrieved examples $N_{\text{rag}}$ increases, which shows the
effectiveness of retrieved in-context examples in RAG model and aligns with
our theoretical analysis in Theorem 1.

Comparisons of C-RAG risk bounds for SOTA retrieval models. We also compare
the conformal risk bounds of C-RAG for different retrieval models, including
sparse encoding method BM25, and SOTA dense encoding models BAAI/bge,
OpenAI/ada, Biencoder-SFT. The results in Figure 4 show that (1) RAG benefits
in achieving a lower conformal risks of generations for different retrieval
models, and (2) Biencoder-SFT is the most performant generally since the model
is trained in the same domain as the test sets, while OpenAI/ada trained with
an open corpus also demonstrates impressive effectiveness and even outperforms
Biencoder-SFT on CommonGen dataset.

###  J.3 Conformal risk bounds evaluations under distribution shifts

Soundness and tightness of distribution-shifted conformal risk bounds in
C-RAG. The test user input text can be out of the calibration distribution in
practice, which needs correction of the conformal risk bounds considering the
distribution drift. We provide the first distribution-shfted conformal risk
bounds for general bounded risk functions in Theorem 2. We evaluate the bounds
in Equation 6 as a function of the distribution shifting distance measured by
Hellinger distance $\rho$. To validate the bounds, we also construct various
test sets with covariate shift induced by sample weight shifting.
Specifically, different weights can be assigned to test samples in the shifted
test sets and the Hellinger distance can also be explicitly computed by the
original sample weights and shifted sample weights. We provide the detailed
procedure in Algorithm 4 in Section J.3. We provide the results of conformal
risk bounds and simulated empirical risks with $N_{\text{rag}}=15$ in Figures
8, 9 and LABEL:fig:bound_and_simulation_llmr_dist_shft with BM25, BAAI/bge,
and Biencoder-SFT. The results demonstrate that (1) the distribution-shifted
conformal risk bounds in Theorem 2 is sound and tight for different retrieval
models, and (2) the conformal risk bounds increases linearly with the
Hellinger distance and are non-trivial for a Hellinger distance up to $0.2$.

Comparisons of C-RAG distribution-shifted risk bounds for SOTA retrieval
models. We compare the distribution-drift conformal risk bounds for different
retrieval models in Figure 10. The results show that (1) the bounds for
different retrieval models increases linearly with a same slope as the
Hellinger distance $\rho$ increases, and (2) Biencoder-SFT and OpenAI/ada
demonstrate lower conformal risk bounds for different distances due to a lower
initial risk without distribution shifts.

###  J.4 Conformal risk bounds with multi-dimensional RAG configurations

We already theoretically and empirically demonstrate the effectiveness of
retrieved in-context examples quantified by $N_{\text{rag}}$. To further
improve the conformal risk bounds, we can addtionally consider more RAG
parameters such as the number of generations $\lambda_{g}$ in the generation
set and a similarity threshold $\lambda_{s}$ to control the diversity of
generations as Algorithm 1. We follow the RAG generation protocol in Algorithm
1 and define the risk function as the minimal risks among all candidate
generations. We similarly construct random test sets and provide the results
on DART and E2E in Figure 7. The results show that (1) the conformal risk
bounds for multi-dimensional RAG configurations are still sound and tight, and
(2) a larger $N_{\text{rag}}$ can reduce the conformal risks more sensitively
compared to the number of generations $\lambda_{g}$, demonstrating the
effectiveness of more retrieved in-context examples. We also fix the number of
retrieved examples $N_{\text{rag}}=5$ and the diversity threshold
$\lambda_{s}=1.0$, and evaluate the certified conformal risk and empirical
risk of randomly sampled test set in Figure 13, which demonstrates that
although not effective as the number of retrieved examples, a larger
generation set size also benefits a low generation risk. Further more, we fix
$N_{\text{rag}}=5,\lambda_{g}=20,\lambda_{s}=1.0$ and evaluate the risks for
different distribution drift distance in Figure 12.

###  J.5 Valid configurations identification with specified risk levels

In conformal analysis (2) in Section 3.3, given a desired risk level $\alpha$,
C-RAG can certify a set of valid configurations $\hat{\Lambda}_{\alpha}$ such
that any configuration in the set results in a conformal risk below $\alpha$.
We use Bonferroni correction for family-wise error rate control and evaluate
empirical risks on randomly sampled test sets for the identified valid
configurations. We provide the results on DART and ECE datasets in Figure 11.
The results demonstrate that (1) the certification is empirically sound since
the empirical risks of valid configurations $\hat{\Lambda}_{\alpha}$ are
always below the desired level $\alpha$, and (2) a large number of retrieved
examples $N_{\text{rag}}$ and a large generation set size $\lambda_{g}$ are
effective in achieving a low generation risk since the valid configuration set
includes the region with a large $N_{\text{rag}}$ and $\lambda_{g}$. We also
individually demonstrate the effectiveness of generation set size
$\lambda_{g}$ and diversity threshold $\lambda_{s}$ in Figures 13 and 14.