Reinforced Mnemonic Reader for Machine Reading Comprehension

Introduction

Teaching machines to comprehend a given context paragraph and answer corresponding questions is one of the long-term goals of natural language processing and artificial intelligence. Figure 1 gives an example of the machine reading comprehension (MRC) task. Benefiting from the rapid development of deep learning techniques Goodfellow et al. (2016) and large-scale benchmark datasets Hermann et al. (2015); Hill et al. (2016); Rajpurkar et al. (2016), end-to-end neural networks have achieved promising results on this task Wang et al. (2017); Seo et al. (2017); Xiong et al. (2017a); Huang et al. (2017).

Despite of the advancements, we argue that there still exists two limitations:

To capture complex interactions between the context and the question, a variety of neural attention Dzmitry Bahdanau (2015), such as bi-attention Seo et al. (2017), coattention Xiong et al. (2017b), are proposed in a single-round alignment architecture. In order to fully compose complete information of the inputs, multi-round alignment architectures that compute attentions repeatedly have been proposed Huang et al. (2017); Xiong et al. (2017a). However, in these approaches, the current attention is unaware of which parts of the context and question have been focused in earlier attentions, which results in two distinct but related issues, where multiple attentions 1) focuses on same texts, leading to attention redundancy and 2) fail to focus on some salient parts of the input, causing attention deficiency.

To train the model, standard maximum-likelihood method is used for predicting exactly-matched (EM) answer spans Wang and Jiang (2017). Recently, reinforcement learning algorithm, which measures the reward as word overlap between the predicted answer and the groung truth, is introduced to optimize towards the F1 metric instead of EM metric Xiong et al. (2017a). Specifically, an estimated baseline is utilized to normalize the reward and reduce variances. However, the convergence can be suppressed when the baseline is better than the reward. This is harmful if the inferior reward is partially overlapped with the ground truth, as the normalized objective will discourage the prediction of ground truth positions. We refer to this case as the convergence suppression problem.

To address the first problem, we present a reattention mechanism that temporally memorizes past attentions and uses them to refine current attentions in a multi-round alignment architecture. The computation is based on the fact that two words should share similar semantics if their attentions about same texts are highly overlapped, and be less similar vice versa. Therefore, the reattention can be more concentrated if past attentions focus on same parts of the input, or be relatively more distracted so as to focus on new regions if past attentions are not overlapped at all.

As for the second problem, we extend the traditional training method with a novel approach called dynamic-critical reinforcement learning. Unlike the traditional reinforcement learning algorithm where the reward and baseline are statically sampled, our approach dynamically decides the reward and the baseline according to two sampling strategies, namely random inference and greedy inference. The result with higher score is always set to be the reward while the other is the baseline. In this way, the normalized reward is ensured to be always positive so that no convergence suppression will be made.

All of the above innovations are integrated into a new end-to-end neural architecture called Reinforced Mnemonic Reader in Figure 3. We conducted extensive experiments on both the SQuAD Rajpurkar et al. (2016) dataset and two adversarial SQuAD datasets Jia and Liang (2017) to evaluate the proposed model. On SQuAD, our single model obtains an exact match (EM) score of 79.5% and F1 score of 86.6%, while our ensemble model further boosts the result to 82.3% and 88.5% respectively. On adversarial SQuAD, our model surpasses existing approahces by more than 6% on both AddSent and AddOneSent datasets.

MRC with Reattention

For the MRC tasks, a question $Q$ and a context $C$ are given, our goal is to predict an answer $A$, which has different forms according to the specific task. In the SQuAD dataset Rajpurkar et al. (2016), the answer $A$ is constrained as a segment of text in the context $C$, nerual networks are designed to model the probability distribution $p(A|C,Q)$.

2 Alignment Architecture for MRC

Among all state-of-the-art works for MRC, one of the key factors is the alignment architecture. That is, given the hidden representations of question and context, we align each context word with the entire question using attention mechanisms, and enhance the context representation with the attentive question information. A detailed comparison of different alignment architectures is shown in Table 1.

where $E_{ij}$ indicates the similarity between $i$-th question word and $j$-th context word, and $f$ is a scalar function. Different methods are proposed to normalize the matrix, resulting in variants of attention such as bi-attentionSeo et al. (2017) and coattention Xiong et al. (2017b). The attention is then used to attend the question and form a question-aware context representation $H=\{h_{j}\}_{j=1}^{m}$.

where $\mathds{1}_{\{\cdot\}}$ is an indicator function ensuring that the context word is not aligned with itself. Finally, the attentive information can be integrated to form a self-aware context representation $Z=\{z_{j}\}_{j=1}^{m}$, which is used to predict the answer.

We refer to the above process as a single-round alignment architecture. Such architecture, however, is limited in its capability to capture complex interactions among question and context. Therefore, recent works build multi-round alignment architectures by stacking several identical aligning layers Huang et al. (2017); Xiong et al. (2017a). More specifically, let $V^{t}=\{v_{i}^{t}\}_{i=1}^{n}$ and $U^{t}=\{u_{j}^{t}\}_{j=1}^{m}$ denote the hidden representations of question and context in $t$-th layer, and $H^{t}=\{h_{j}^{t}\}_{j=1}^{m}$ is the corresponding question-aware context representation. Then the two similarity matrices can be computed as

3 Reattention Mechanism

To address these problems, we propose to temporally memorize past attentions and explicitly use them to refine current attentions. The intuition is that two words should be correlated if their attentions about same texts are highly overlapped, and be less related vice versa. For example, in Figure 2, suppose that we have access to previous attentions, and then we can compute their dot product to obtain a “similarity of attention”. In this case, the similarity of word pair (team, Broncos) is higher than (team, Panthers).

Therefore, we define the computation of reattention as follows. Let $E^{t-1}$ and $B^{t-1}$ denote the past similarity matrices that are temporally memorized. The refined similarity matrix $E^{t}$ ($t>1$) is computed as

Dynamic-critical Reinforcement Learning

In the extractive MRC task, the model distribution $p(A|C,Q;\theta)$ can be divided into two steps: first predicting the start position $i$ and then the end position $j$ as

where $\theta$ represents all trainable parameters.

The standard maximum-likelihood (ML) training method is to maximize the log probabilities of the ground truth answer positions Wang and Jiang (2017)

where $y_{k}^{1}$ and $y_{k}^{2}$ are the answer span for the $k$-th example, and we denote $p_{1}(i|C,Q;\theta)$ and $p_{2}(j|i,C,Q;\theta)$ as $p_{1}(i)$ and $p_{2}(j|i)$ respectively for abbreviation.

Recently, reinforcement learning (RL), with the task reward measured as word overlap between predicted answer and groung truth, is introduced to MRC Xiong et al. (2017a). A baseline $b$, which is obtained by running greedy inference with the current model, is used to normalize the reward and reduce variances. Such approach is known as the self-critical sequence training (SCST) Rennie et al. (2016), which is first used in image caption. More specifically, let $R(A^{s},A^{*})$ denote the F1 score between a sampled answer $A^{s}$ and the ground truth $A^{*}$. The training objective is to minimize the negative expected reward by

where we abbreviate the model distribution $p(A|C,Q;\theta)$ as $p_{\theta}(A)$, and the reward function $R(A^{s},A^{*})$ as $R(A^{s})$. $\hat{A}$ is obtained by greedily maximizing the model distribution:

The expected gradient $\nabla_{\theta}\mathcal{L}_{SCST}(\theta)$ can be computed according to the REINFORCE algorithm Sutton and Barto (1998) as

where the gradient can be approxiamated using a single Monte-Carlo sample $A^{s}$ derived from $p_{\theta}$.

However, a sampled answer is discouraged by the objective when it is worse than the baseline. This is harmful if the answer is partially overlapped with ground truth, since the normalized objective would discourage the prediction of ground truth positions. For example, in Figure 1, suppose that $A^{s}$ is champion Denver Broncos and $\hat{A}$ is Denver Broncos. Although the former is an acceptable answer, the normalized reward would be negative and the prediction for end position would be suppressed, thus hindering the convergence. We refer to this case as the convergence suppression problem.

Here, we consider both random inference and greedy inference as two different sampling strategies: the first one encourages exploration while the latter one is for exploitationIn practice we found that a better approximation can be made by considering a top-$K$ answer list, where $\hat{A}$ is the best result and $A^{s}$ is sampled from the rest of the list.. Therefore, we approximate the expected gradient by dynamically set the reward and baseline based on the F1 scores of both $A^{s}$ and $\hat{A}$. The one with higher score is set as reward, while the other is baseline. We call this approach as dynamic-critical reinforcement learning (DCRL)

Notice that the normalized reward is constantly positive so that superior answers are always encouraged. Besides, when the score of random inference is higher than the greedy one, DCRL is equivalent to SCST. Thus, Eq. 3 is a special case of Eq. 3.

Following Xiong et al. (2017a) and Kendall et al. (2017), we combine ML and DCRL objectives using homoscedastic uncertainty as task-dependent weightings so as to stabilize the RL training as

where $\sigma_{a}$ and $\sigma_{b}$ are trainable parameters.

End-to-end Architecture

Based on previous innovations, we introduce an end-to-end architecture called Reinforced Mnemonic Reader, which is shown in Figure 3. It consists of three main components: 1) an encoder builds contextual representations for question and context jointly; 2) an iterative aligner performs multi-round alignments between question and context with the reattention mechanism; 3) an answer pointer predicts the answer span sequentially. Beblow we give more details of each component.

Encoder. Let $W^{Q}=\{w_{i}^{q}\}_{i=1}^{n}$ and $W^{C}=\{w_{j}^{c}\}_{j=1}^{m}$ denote the word sequences of the question and context respectively. The encoder firstly converts each word to an input vector. We utilize the 100-dim GloVe embedding Pennington et al. (2014) and 1024-dim ELMo embedding Peters et al. (2018). Besides, a character-level embedding is obtained by encoding the character sequence with a bi-directional long short-term memory network (BiLSTM) Hochreiter and Schmidhuber (1997), where two last hidden states are concatenated to form the embedding. In addition, we use binary feature of exact match, POS embedding and NER embedding for both question and context, as suggested in Chen et al. (2017). Together the inputs $X^{Q}=\{x_{i}^{q}\}_{i=1}^{n}$ and $X^{C}=\{x_{j}^{c}\}_{j=1}^{m}$ are obtained.

To model each word with its contextual information, a weight-shared BiLSTM is utilized to perform the encoding

Iterative Aligner. The iterative aligner contains a stack of three aligning blocks. Each block consists of three modules: 1) an interactive alignment to attend the question into the context; 2) a self alignment to attend the context against itself; 3) an evidence collection to model the context representation with a BiLSTM. The reattention mechanism is utilized between two blocks, where past attentions are temporally memorizes to help modulating current attentions. Below we first describe a single block in details, which is shown in Figure 4, and then introduce the entire architecture.

Finally, a BiLSTM is used to perform the evidence collection, which outputs the fully-aware context vectors $R=[r_{1},...,r_{m}]$ with $Z$ as its inputs.

Multi-round Alignments with Reattention. To enhance the ability of capturing complex interactions among inputs, we stack two more aligning blocks with the reattention mechanism as follows

Experiments

We mainly focus on the SQuAD dataset Rajpurkar et al. (2016) to train and evaluate our model. SQuAD is a machine comprehension dataset, totally containing more than $100,000$ questions manually annotated by crowdsourcing workers on a set of $536$ Wikipedia articles. In addition, we also test our model on two adversarial SQuAD datasets Jia and Liang (2017), namely AddSent and AddOneSent. In both adversarial datasets, a confusing sentence with a wrong answer is appended at the end of the context in order to fool the model.

We evaluate the Reinforced Mnemonic Reader (R.M-Reader) by running the following setting. We first train the model until convergence by optimizing Eq. 7. We then finetune this model with Eq. 11, until the F1 score on the development set no longer improves.

We use the Adam optimizer Kingma and Ba (2014) for both ML and DCRL training. The initial learning rates are $0.0008$ and $0.0001$ respectively, and are halved whenever meeting a bad iteration. The batch size is $48$ and a dropout rate Srivastava et al. (2014) of $0.3$ is used to prevent overfitting. Word embeddings remain fixed during training. For out of vocabulary words, we set the embeddings from Gaussian distributions and keep them trainable. The size of character embedding and corresponding LSTMs is $50$, the main hidden size is $100$, and the hyperparameter $\gamma$ is $3$.

2 Overall Results

We submitted our model on the hidden test set of SQuAD for evaluation. Two evaluation metrics are used: Exact Match (EM), which measures whether the predicted answer are exactly matched with the ground truth, and F1 score, which measures the degree of word overlap at token level.

As shown in Table 2, R.M-Reader achieves an EM score of 79.5% and F1 score of 86.6%. Since SQuAD is a competitve MRC benchmark, we also build an ensemble model that consists of $12$ single models with the same architecture but initialized with different parameters. Our ensemble model improves the metrics to 82.3% and 88.5% respectivelyThe results are on https://worksheets.codalab.org/worksheets/ 0xe6c23cbae5e440b8942f86641f49fd80..

Table 3 shows the performance comparison on two adversarial datasets, AddSent and AddOneSent. All models are trained on the original train set of SQuAD, and are tested on the two datasets. As we can see, R.M-Reader comfortably outperforms all previous models by more than 6% in both EM and F1 scores, indicating that our model is more robust against adversarial attacks.

3 Ablation Study

4 Effectiveness of Reattention

Table 5 shows the results. We first see that the reattention indeed help in alleviating the attention redundancy: the divergence between any two adjacent blocks has been successfully enlarged with reattention. However, we find that the improvement between the first two blocks is larger than the one of last two blocks. We conjecture that the first reattention is more accurate at measuring the similarity of word pairs by using the original encoded word representation, while the latter reattention is distracted by highly nonlinear word representations. In addition, we notice that the attention deficiency has also been moderated: the divergence betwen normalized $E^{t}$ and ${E^{t}}^{*}$ is reduced.

5 Prediction Analysis

Figure 5 compares predictions made either with dynamic-critical reinforcement learning or with self-critical sequence training. We first find that both approaches are able to obtain answers that match the query-sensitive category. For example, the first example shows that both four and two are retrieved when the questions asks for how many. Nevertheless, we observe that DCRL constantly makes more accurate prediction on answer spans, especially when SCST already points a rough boundary. In the second example, SCST takes the whole phrase after Dyrrachium as its location. The third example shows a similar phenomenon, where the SCST retrieves the phrase constantly servicing and replacing mechanical brushes as its answer. We demonstrates that this is because SCST encounters the convergence suppression problem, which impedes the prediction of ground truth answer boundaries. DCRL, however, successfully avoids such problem and thus finds the exactly correct entity.

Conclusion

We propose the Reinforced Mnemonic Reader, an enhanced attention reader with two main contributions. First, a reattention mechanism is introduced to alleviate the problems of attention redundancy and deficiency in multi-round alignment architectures. Second, a dynamic-critical reinforcement learning approach is presented to address the convergence suppression problem existed in traditional reinforcement learning methods. Our model achieves the state-of-the-art results on the SQuAD dataset, outperforming several strong competing systems. Besides, our model outperforms existing approaches by more than 6% on two adversarial SQuAD datasets. We believe that both reattention and DCRL are general approaches, and can be applied to other NLP task such as natural language inference. Our future work is to study the compatibility of our proposed methods.

Acknowledgments

This research work is supported by National Basic Research Program of China under Grant No. 2014CB340303. In addition, we thank Pranav Rajpurkar for help in SQuAD submissions.

References