SlowGuess/ACL_2020_OCR huggingface dataset:
Hongyu Gong $^{1*}$  Yelong Shen $^{2\dagger}$  Dian Yu $^{3}$  Jianshu Chen $^{3}$  Dong Yu $^{3}$
<sup>1</sup>University of Illinois at Urbana-Champaign, IL, USA
$^{2}$ Microsoft Dynamics 365 AI, Redmond, WA, USA
hgong6@illinois.edu
yeshe@microsoft.com
{yudian, jianshuchen, dyu} @tencent.com
In this paper, we study machine reading comprehension (MRC) on long texts, where a model takes as inputs a lengthy document and a question and then extracts a text span from the document as an answer. State-of-the-art models tend to use a pretrained transformer model (e.g., BERT) to encode the joint contextual information of document and question. However, these transformer-based models can only take a fixed-length (e.g., 512) text as its input. To deal with even longer text inputs, previous approaches usually chunk them into equally-spaced segments and predict answers based on each segment independently without considering the information from other segments. As a result, they may form segments that fail to cover the correct answer span or retain insufficient contexts around it, which significantly degrades the performance. Moreover, they are less capable of answering questions that need cross-segment information.
We propose to let a model learn to chunk in a more flexible way via reinforcement learning: a model can decide the next segment that it wants to process in either direction. We also employ recurrent mechanisms to enable information to flow across segments. Experiments on three MRC datasets - CoQA, QuAC, and TriviaQA - demonstrate the effectiveness of our proposed recurrent chunking mechanisms: we can obtain segments that are more likely to contain complete answers and at the same time provide sufficient contexts around the ground truth answers for better predictions.
Teaching machines to read, process, and comprehend natural language is a coveted goal of machine reading comprehension (MRC) problems
(Hermann et al., 2015; Hill et al., 2016; Rajpurkar et al., 2016; Trischler et al., 2017; Zhang et al., 2018; Kocisky et al., 2018). Many existing MRC datasets have a similar task definition: given a document and a question, the goal is to extract a span from the document (in most cases) or instead generate an abstractive answer to answer the question.
There is a growing trend of building MRC readers (Hu et al., 2019; Xu et al., 2019; Yang et al., 2019a; Keskar et al., 2019) based on pre-trained language models (Baker et al., 2019; Yang et al., 2019b), such as GPT (Radford et al., 2018) and BERT (Devlin et al., 2019). These models typically consist of a stack of transformer layers that only allow fixed-length (e.g., 512) inputs. However, it is often the case that input sequences exceed the length constraint, e.g., documents in the TriviaQA dataset (Joshi et al., 2017) contain 2,622 tokens on average. Some conversational MRC datasets such as CoQA (Reddy et al., 2018) and QuAC (Choi et al., 2018) often go beyond the length limit as we may need to incorporate previous questions as well as relatively long documents into the input to answer the current question.
To deal with long text inputs, a commonly used approach firstly chunks the input text into equally-spaced segments, secondly predicts the answer for each individual segment, and finally ensembles the answers from multiple segments (Devlin et al., 2019). However, there are two major limitations of this approach: first, a predetermined large stride size for chunking may result in incomplete answers, and we observe that models are more likely to fail when the answer is near the boundaries of a segment, compared to the cases when an answer is in the center of a segment surrounded by richer context (Figure 1); second, we empirically observe that chunking with a smaller stride size contributes little to (sometimes even hurts) the model performance. A possible explanation is that predicting answer
* The work was performed during an internship at Tencent AI Lab, Bellevue, WA, USA.
† The work was performed when Yelong Shen was at Tencent AI Lab, Bellevue, WA, USA.
6751
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6751-6761 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
for each segment independently may cause incomparable answer scores across segments. A similar phenomenon is also observed in open-domain question answering tasks (Clark and Gardner, 2017).
Considering the limitations mentioned above, we propose recurrent chunking mechanisms (RCM) on top of the transformer-based models for MRC tasks. There are two main characteristics of RCM. First, it could let the machine reader learn how to choose the stride size intelligently when reading a lengthy document via reinforcement learning, so it helps prevent extracting incomplete answers from a segment and retain sufficient contexts around the answer. Second, we apply recurrent mechanisms to allow the information to flow across segments. As a result, the model can have access to the global contextual information beyond the current segment.
Figure 1: The influence of the distance between the center of the answer span and the center of the segment. The test performance (in F1 score) is evaluated on CoQA using a BERT-Large reader. The best performance is achieved when the chunk center coincides with the answer span center. Within the distance of  $\pm 80$  (in tokens), while  $99\%$  answers are completely covered, the performance degrades as the segment center moves away from the answer center, and the segment contains fewer relevant contexts. When the distance reaches 96, more than half of the predicted spans are incomplete.
In the experiments, we evaluate the proposed  $\mathrm{RCM}^1$  on three MRC datasets: CoQA, QuAC, and TriviaQA. Experimental results demonstrate that RCM leads to consistent performance gains on these benchmarks. Furthermore, it also generates segments that are more likely to cover the entire answer spans and provide richer contextual information around the ground truth answers.
The primary contributions of this work are:
- We propose a chunking mechanism for machine reading comprehension to let a model learn to chunk lengthy documents in a more flexible way via reinforcement learning.
- We also apply recurrence to allow information transfer between segments so that the model can have knowledge beyond the current segment when selecting answers.
- We have performed extensive experiments on three machine reading comprehension datasets: CoQA, QuAC, and TriviaQA. Our approach outperforms two state-of-the-art BERT-based models on different datasets.
The proposed recurrent chunking mechanisms (RCM) are built upon the pre-trained BERT models. We will briefly introduce the basic model in Section 2.1, and then the RCM approach in Section 2.2 and 2.3. More details of our model in training and testing are presented in Sections 2.4 and 2.5.
2.1 Baseline Model
Pre-trained BERT model has been shown to achieve new state-of-the-art performance on many MRC datasets (Devlin et al., 2019). Here, we introduce this basic BERT model, which is used as our baseline. As the maximum input length in BERT is restricted to be 512, a widely adopted strategy is to chunk a long document into multiple segments with a fixed stride size (i.e., 128). Following the input format of BERT, the input for each document segment starts with "CLS" token, which is followed by question tokens "Q" and document segment tokens. We use "SEP" token as a separator between the question and the segment. We also append a special "UNK" token at the end of the segment to handle unanswerable questions. If a given question is annotated as unanswerable, we mark the "UNK" token as the ground truth answer during training. Accordingly in evaluation, if "UNK" token is selected by the model from the input segment, we output the answer as "unanswerable".
Answer Extraction. Following previous work on extractive machine reading comprehension, we predict the start and the end positions of the answer span in the given document segment. BERT first generates a vector representation  $\mathbf{h}_{c,i}$  for each  $i$ -th
<sup>1</sup>The code is available at https://github.com/HongyuGong/RCM-Question-Answering.git.
6752
Figure 2: BERT generates representations for each input sequence, and recurrence accumulates information over segments. Based on these representations, the answer extractor extracts answers from the current segment, and the policy network takes chunking action and moves to the next segment. Chunking scorer scores each segment by estimating its likelihood of containing an answer and selects answers among predictions from multiple segments.
token in the  $c$ -th segment. Given  $\mathbf{h}_{c,i}$ , the model scores each token in terms of its likelihood of being the start token of the answer span.
$$
l _ {c, i} ^ {\text {s t a r t}} = \mathbf {w} _ {s} ^ {T} \mathbf {h} _ {c, i}, \tag {1}
$$
where  $\mathbf{w}_s$  is the model parameter. The probability  $p_{c,i}^{\mathrm{start}}$  that the answer starts at the  $i$ -th token is computed by applying the softmax to  $l_{c,i}^{\mathrm{start}}$ .
$$
p _ {c, i} ^ {\text {s t a r t}} = \operatorname {s o f t m a x} \left(l _ {c, i} ^ {\text {s t a r t}}\right) \tag {2}
$$
Likewise, the model scores how likely the answer ends at the  $j$ -th token in segment  $c$  using
$$
l _ {c, j} ^ {\text {e n d}} = \mathbf {w} _ {e} ^ {T} \mathbf {h} _ {c, j}, \tag {3}
$$
where  $\mathbf{w}_{\mathrm{e}}$  is the model parameter. The probability of the  $j$ -th token being the end of the answer (denoted as  $p_{c,j}^{\mathrm{end}}$ ) is calculated in a similar manner as Eq. (2).
Answer Ensemble. The baseline model adopts a max-pooling approach to ensemble candidate answers from multiple segments. The answer with the highest probability is selected.
2.2 Recurrent Mechanisms
The baseline model makes the answer prediction for each document segment independently, which may cause incomparable answer scores across segments due to the lack of document-level information. We propose to use a recurrent layer to propagate the information across different segments and a chunking scorer model to estimate the probability that a segment contains the answer.
For an input sequence containing the segment  $c$ , BERT's representation for its first token "CLS" is taken as the local representation  $v_{c}$  of the segment. The segment representation is further enriched with the representations of previously generated segments via recurrence. We denote the enriched segment representation as  $\tilde{\mathbf{v}}_c$ :
$$
\tilde {\mathbf {v}} _ {c} = f \left(\mathbf {v} _ {c}, \tilde {\mathbf {v}} _ {c - 1}\right), \tag {4}
$$
where  $f(\cdot)$  is the recurrent function. We consider two recurrent mechanisms here: gated recurrence and Long Short Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997) recurrence.
Gated recurrence is simply a weighted sum of its inputs:
$$
f _ {\text {g a t e d}} \left(\mathbf {v} _ {c}, \tilde {\mathbf {v}} _ {c - 1}\right) = \alpha \mathbf {v} _ {c} + \beta \tilde {\mathbf {v}} _ {c - 1}, \tag {5}
$$
where  $\alpha$  and  $\beta$  are coefficients depending on the inputs. We have  $\alpha, \beta = \mathrm{softmax}(\mathbf{w}_r^T[\mathbf{v}_c, \tilde{\mathbf{v}}_{c-1}])$ , where  $\mathbf{w}_r$  is a model parameter.
The LSTM recurrence, which uses LSTM unit as the recurrence function, takes  $\mathbf{v}_c$  as the current input and  $\tilde{\mathbf{v}}_{c - 1}$  as the previous hidden state.
$$
f _ {\mathrm {L S T M}} \left(\mathbf {v} _ {c}, \tilde {\mathbf {v}} _ {c - 1}\right) = \operatorname {L S T M} \left(\mathbf {v} _ {c}, \tilde {\mathbf {v}} _ {c - 1}\right). \tag {6}
$$
Chunking Scorer. Given the enriched segment representation  $\tilde{\mathbf{v}}_c$  as input, the chunking scorer produces an scalar  $q_{c}$  by:
$$
q _ {c} = \sigma \left(\mathbf {W} _ {c} \tilde {\mathbf {v}} _ {c} + \mathbf {b} _ {c}\right), \tag {7}
$$
where  $\mathbf{W}_c$  and  $\mathbf{b}_c$  are model parameters, and  $\sigma(\cdot)$  is the sigmoid function. The scalar  $q_c$  is an estimation of the probability that an answer is included in
6753
segment  $c$ . Then, the chunking scorer uses  $q_{c}$  to further refine the likelihood of the candidate answers from different segments (see Sections 2.4 and 2.5 for more details on this part of chunking scorer).
2.3 Learning to Chunk
The baseline approach divides a long document into multiple segments with a fixed stride size, from left to right. We will present an approach that could allow the model to choose the stride size flexibly by itself when reading the document. Our motivation, as mentioned in Section 1, is to prevent the answer span from being too close to the segment boundary and covering incomplete answers.
We formulate the problem of learning-to-chunk under the framework of reinforcement learning. We define the state  $s$  of the model to be the segments that a model has processed up to the current segment  $c$ , i.e.,  $s = \{1,2,\dots,c\}$ . The action  $a$  is the stride size and direction (forward or backward) the model chooses to move to the next document segment. We define the action space  $A$  as a set of strides, e.g.,  $A = \{-16,16,32\}$ , where 32 indicates moving forward with stride size 32 and -16 indicates moving backward with stride size 16. In this work, we represent the state  $s$  with the enriched segment representation  $\tilde{\mathbf{v}}_c$ .
Chunking Policy. The chunking policy gives the probability  $p^{\mathrm{act}}(a \mid s)$  of taking an action  $a$  at the current state  $s$ , which is modeled by a one-layer feedforward neural network:
$$
p ^ {\mathrm {a c t}} (a \mid s) = \operatorname {s o f t m a x} \left(\mathbf {W} _ {a} \tilde {\mathbf {v}} _ {c} + \mathbf {b} _ {a}\right), \tag {8}
$$
where  $\mathbf{W}_a$  and  $\mathbf{b}_a$  are trainable parameters.
Fig. 2 gives an overview of the proposed recurrent chunking mechanisms built upon the BERT model: the chunking policy network takes the enriched segment representation as the input to generate the chunking action, which decides the next segment to be processed.
2.4 Training
In the training phase of the recurrent chunking mechanisms, the stride actions of moving to the next segment are sampled according to the probability given by the chunking policy (Sutton and Barto, 2018). Our model generates a sequence of document segments for each question. We train the answer extractor and chunking scorer network with supervised learning, and we train the chunking policy network via reinforcement learning.
Supervised Learning for Answer Extraction. Just as the baseline model, we train the answer extraction network via supervised learning. Given a question, the answer extractor classifies whether a word from a document segment is the start or the end of the answer. The cross-entropy loss can be computed given the ground-truth answer and the predictions of the answer extractor. Suppose that the  $i^{*}$ -th and  $j^{*}$ -th tokens are the answer start and end, respectively. The training objective to minimize the following cross-entropy loss,  $L_{\mathrm{ans}}$ :
$$
L _ {\text {a n s}} = - \sum_ {c} \log p _ {c, i *} ^ {\text {s t a r t}} - \sum_ {c} \log p _ {c, j *} ^ {\text {e n d}}, \tag {9}
$$
Supervised Learning for Chunking Scorer. A binary variable  $y_{c}$  indicates whether the segment  $c$  contains an answer or not. Chunking scorer estimates the probability  $q_{c}$  that the segment contains an answer. Similarly, the chunking scorer network can be trained in a supervised manner by minimizing the cross-entropy loss,  $L_{\mathrm{cs}}$ :
$$
L _ {\mathrm {c s}} = - \sum_ {c} y _ {c} \log q _ {c} - \sum_ {c} \left(1 - y _ {c}\right) \log \left(1 - q _ {c}\right), \tag {10}
$$
where the chunking score  $q_{c}$  is given in Eq. (7).
Reinforcement Learning for Chunking Policy. Since the selection of the stride actions is a sequential decision-making process, it is natural to train the chunking policy via reinforcement learning.
First of all, the accumulated reward for taking action  $a$  at state  $s$  is denoted as  $R(s,a)$ , which is derived in a recursive manner:
$$
R (s, a) = q _ {c} r _ {c} + \left(1 - q _ {c}\right) R \left(s ^ {\prime}, a ^ {\prime}\right), \tag {11}
$$
where  $q_{c}$  is the probability that segment  $c$  contains an answer as given in Eq. (7), and  $(s', a')$  denotes the next state-action pair. The value of  $r_{c}$  indicates the probability of the correct answer being extracted from the current segment  $c$ . The mathematical definition of  $r_{c}$  is given as:
$$
r _ {c} = \left\{ \begin{array}{l l} p _ {c, i ^ {*}} ^ {\text {s t a r t}} \cdot p _ {c, j ^ {*}} ^ {\text {e n d}}, & \text {i f a n s w e r i n c l u d e d}, \\ 0, & \text {e l s e .} \end{array} \right. \tag {12}
$$
The first term in Eq. (11) is the reward of the answer being correctly extracted from the current segment. The answer is included in the current segment  $c$  with probability  $q_{c}$ , and thus the first term is weighted by  $q_{c}$  in reward  $R(s, a)$ . The second term in Eq. (11) indicates that  $R(s, a)$  also
Dataset	Train	Validation
Question #	Avg tokens #	Max token #	Question #	Avg tokens #	Max token #
CoQA	108,647	352	1,323	7,983	341	1,037
QuAC	83,568	516	2,310	7,354	576	2,146
TriviaQA (wiki)	61,888	2,622	5,839	7,993	2,630	6,690
Table 1: Statistics of the CoQA, QuAC and TriviaQA datasets. We report the number of sub-tokens generated by the BERT tokenizer.
relies on the accumulated reward  $R(s', a')$  of the next state when the answer is not available in the current segment.
The chunking policy network can be trained by maximizing the expected accumulated reward (as shown in Eq. (13)) through the policy gradient algorithm (Williams, 1992; Sutton et al., 2000; Gong et al., 2019).
$$
J = \mathbb {E} _ {p ^ {\mathrm {a c t}} (a \mid s)} [ R (s, a) ]. \tag {13}
$$
To be consistent with the notations in answer extraction and chunking scorer modules, we denote the loss function of chunking policy as  $L_{\mathrm{cp}}$ , which is the negative expected accumulated reward  $J$  in Eq. (13):  $L_{\mathrm{cp}} = -J$ . Thus, the stochastic gradient of  $L_{\mathrm{cp}}$  over a mini-batch of data  $\mathcal{B}$  is given by:
$$
\nabla L _ {\mathrm {c p}} = - \sum_ {(s, a) \in \mathcal {B}} \nabla \log p ^ {\mathrm {a c t}} (a \mid s) R (s, a), \tag {14}
$$
where  $p^{\mathrm{act}}(a \mid s)$  is the chunking policy in Eq. (8). Training procedure. The overall training loss  $L$  is an sum of all three losses:  $L = L_{\mathrm{ans}} + L_{\mathrm{cs}} + L_{\mathrm{cp}}$ . In addition, we initialize the bottom representation layers with a pre-trained BERT model and initialize other model parameters randomly. We use the Adam optimizer with peak learning rate  $3 \times 10^{-5}$  and a linear warming-up schedule.
2.5 Testing
In the testing phase, the model starts from the beginning of the document as its first segment. Later on in state  $s$ , the model takes the best stride action  $a^{*}$  according to the chunking policy:
$$
a ^ {*} = \underset {a \in A} {\operatorname {a r g m a x}} p ^ {\text {a c t}} (a \mid s) \tag {15}
$$
After the stride action  $a^*$  is taken, a new segment is taken from the given document, and so on until the maximum number of segments  $C$  is reached. Now for a document segment  $c$ , we score its candidate answer spanning from the  $i$ -th to the  $j$ -th token by  $p_{i,j,c}^{\mathrm{A}}$ :
$$
p _ {i, j, c} ^ {\mathrm {A}} = p _ {c, i} ^ {\text {s t a r t}} \cdot p _ {c, j} ^ {\text {e n d}} \cdot q _ {c}. \tag {16}
$$
The best answer span  $(\bar{i},\bar{j})$  across multiple segments can be obtained by selecting the one with the highest score  $p_{i,j,c}^{\mathrm{A}}$ .
$$
\bar {i}, \bar {j} = \underset {i \leq j, 1 \leq c \leq C} {\operatorname {a r g m a x}} p _ {i, j, c} ^ {\mathrm {A}}, \tag {17}
$$
where dynamic programming is used to find  $(\bar{i},\bar{j})$  efficiently in linear time.
We use three MRC datasets, CoQA (Reddy et al., 2018), QuAC (Choi et al., 2018) and TriviaQA (Joshi et al., 2017)) in our experiments.
(1) CoQA. Answers in the CoQA dataset can be abstractive texts written by annotators. It is reported that an extractive MRC approach can achieve an upper bound as high as  $97.8\%$  in F1 score (Yatskar, 2019). Therefore, We preprocess the CoQA training data and select a text span from the document as the extractive answer that achieves the highest F1 score compared with the given ground truth.
(2) QuAC. All the answers in the QuAC dataset are text spans, which are highlighted by annotators in the given document.
(3) TriviaQA. TriviaQA is a large-scale MRC dataset, containing data from Wikipedia and Web domains. We use its Wikipedia subset in this work. It is reported to be challenging in its variability between questions and documents as well as its requirement of cross-sentence reasoning. Documents in TriviaQA contain more than 2,000 words on average, which is suitable for evaluating the capability of a model to deal with long documents.
The dataset statistics are summarized in Table 1, including the data sizes, the average and maximum number of sub-tokens in documents.
3.2 Baselines and Evaluation Metric
Baselines. We have two strong baselines based on the pre-trained BERT, which has achieved state-of-the-art performance in a wide range of NLP tasks
6755
Dataset	CoQA	QuAC
Max sequence length	192	256	384	512	192	256	384	512
BERT-Large (Devlin et al., 2019)	72.8	76.2	81.0	81.4	34.5	50.6	56.7	61.5
Sent-Selector (with previous questions)	54.5	63.8	75.3	79.4	33.9	38.8	47.6	55.4
Sent-Selector (only current questions)	57.5	66.5	76.5	79.5	34.3	39.1	47.6	56.4
BERT-RCM
- Gated recurrence (no RL chunking)	74.5	78.6	81.0	81.4	48.8	51.4	56.2	61.4
- Gated recurrence	76.0	79.2	81.3	81.8	51.6	55.2	59.9	62.0
- LSTM recurrence (no RL chunking)	74.1	78.5	81.0	81.3	49.2	51.5	56.4	61.6
- LSTM recurrence	75.4	79.5	81.3	81.8	53.9	55.6	60.4	61.8
Table 2: Comparison of F1 scores (%) achieved by different algorithms.
including machine reading comprehension.
(1) BERT-LARGE MODEL. It achieves competitive performance on extractive MRC tasks such as SQuAD (Rajpurkar et al., 2016, 2018). It adopts a simple sliding window chunking policy – moving to the next document segment with a fixed stride size from left to right. We also analyze the performance of the Large BERT model with different stride sizes in training and testing (see Section 4.1 for details). The best performance is obtained by setting stride size as 64 in CoQA and QuAC, and 128 in TriviaQA.
(2) SENTENCE SELECTOR. Given a question, the sentence selector chooses a subset of sentences that are likely to contain an answer. The selected sentences are then concatenated and fed to the BERT-Large model for answer extraction. For conversational datasets CoQA and QuAC, since a question is correlated with its previous questions within the same conversation, we apply the sentence selector to select sentences based on the current question alone or the concatenation of the previous questions and the current question. We only use the current question as the input to the sentence selector for TriviaQA, which does not involve any conversational history. The sentence selector we used in experiments is released by Htut et al. (2018).
Evaluation Metric. The main evaluation metric is macro-average word-level F1 score. We compare each prediction with the reference answer. Precision is defined by the percentage of predicted answer tokens that appear in the reference answer, and recall is the percentage of reference answer tokens captured in the prediction. F1 score is the harmonic mean of the precision and recall. When multiple reference answers are provided, the maximum F1 score is used for evaluation.
3.3 Results on CoQA and QuAC
We first perform experiments on two conversational MRC datasets, CoQA and QuAC.
Setting. We perform a set of experiments with different maximum sequence lengths of 192, 256, 384, and 512. Our model fixes the number of segments read from a document for each question. It generates 4, 3, 3, and 2 segments under the length limit of 192, 256, 384, and 512, respectively.