Stochastic Answer Networks for Machine Reading Comprehension

Abstract: We propose a simple yet robust stochastic answer network (SAN) that simulates multi-step reasoning in machine reading comprehension. Compared to previous work such as ReasoNet which used reinforcement learning to determine the number of steps, the unique feature is the use of a kind of stochastic prediction dropout on the answer module (final layer) of the neural network during the training. We show that this simple trick improves robustness and achieves results competitive to the state-of-the-art on the Stanford Question Answering Dataset (SQuAD), the Adversarial SQuAD, and the Microsoft MAchine Reading COmprehension Dataset (MS MARCO).

Machine reading comprehension (MRC) is a challenging task: the goal is to have machines read a text passage and then answer any question about the passage. This task is an useful benchmark to demonstrate natural language understanding, and also has important applications in e.g. conversational agents and customer service support. It has been hypothesized that difficult MRC problems require some form of multi-step synthesis and reasoning. For instance, the following example from the MRC dataset SQuAD BIBREF0 illustrates the need for synthesis of information across sentences and multiple steps of reasoning:
 $Q$ : What collection does the V&A Theator & Performance galleries hold?
 $P$ : The V&A Theator & Performance galleries opened in March 2009. ... They hold the UK's biggest national collection of material about live performance.
To infer the answer (the underlined portion of the passage $P$ ), the model needs to first perform coreference resolution so that it knows “They” refers “V&A Theator”, then extract the subspan in the direct object corresponding to the answer.
This kind of iterative process can be viewed as a form of multi-step reasoning. Several recent MRC models have embraced this kind of multi-step strategy, where predictions are generated after making multiple passes through the same text and integrating intermediate information in the process. The first models employed a predetermined fixed number of steps BIBREF1 , BIBREF2 , BIBREF3 , BIBREF4 . Later, shen2016reasonet proposed using reinforcement learning to dynamically determine the number of steps based on the complexity of the question. Further, shen2017empirical empirically showed that dynamic multi-step reasoning outperforms fixed multi-step reasoning, which in turn outperforms single-step reasoning on two distinct MRC datasets (SQuAD and MS MARCO).
In this work, we derive an alternative multi-step reasoning neural network for MRC. During training, we fix the number of reasoning steps, but perform stochastic dropout on the answer module (final layer predictions). During decoding, we generate answers based on the average of predictions in all steps, rather than the final step. We call this a stochastic answer network (SAN) because the stochastic dropout is applied to the answer module; albeit simple, this technique significantly improves the robustness and overall accuracy of the model. Intuitively this works because while the model successively refines its prediction over multiple steps, each step is still trained to generate the same answer; we are performing a kind of stochastic ensemble over the model's successive prediction refinements. Stochastic prediction dropout is illustrated in Figure 1 .
The machine reading comprehension (MRC) task as defined here involves a question $Q=\lbrace q_0, q_1, ..., q_{m-1}\rbrace $ and a passage $P=\lbrace p_0, p_1, ..., p_{n-1}\rbrace $ and aims to find an answer span $A=\lbrace a_{start}, a_{end}\rbrace $ in $P$ . We assume that the answer exists in the passage $P$ as a contiguous text string. Here, $m$ and $n$ denote the number of tokens in $Q$ and $P$ , respectively. The learning algorithm for reading comprehension is to learn a function $f(Q, P) \rightarrow A$ . The training data is a set of the query, passage and answer tuples $P=\lbrace p_0, p_1, ..., p_{n-1}\rbrace $0 .
We now describe our model from the ground up. The main contribution of this work is the answer module, but in order to understand what goes into this module, we will start by describing how $Q$ and $P$ are processed by the lower layers. Note the lower layers also have some novel variations that are not used in previous work. As shown in Figure 2 , our model contains four different layers to capture different concept of representations. The detailed description of our model is provided as follows.
Lexicon Encoding Layer. The purpose of the first layer is to extract information from $Q$ and $P$ at the word level and normalize for lexical variants. A typical technique to obtain lexicon embedding is concatenation of its word embedding with other linguistic embedding such as those derived from Part-Of-Speech (POS) tags. For word embeddings, we use the pre-trained 300-dimensional GloVe vectors BIBREF5 for the both $Q$ and $P$ . Following chen2017reading, we use three additional types of linguistic features for each token $p_i$ in the passage $P$ :
In summary, each token $p_i$ in the passage is represented as a 600-dimensional vector and each token $q_j$ is represented as a 300-dimensional vector.
Due to different dimensions for the passages and questions, in the next layer two different bidirectional LSTM (BiLSTM) BIBREF6 may be required to encode the contextual information. This, however, introduces a large number of parameters. To prevent this, we employ an idea inspired by BIBREF7 : use two separate two-layer position-wise Feed-Forward Networks (FFN), $FFN(x)=W_2 ReLU(W_1 x +b_1) + b_2$ , to map both the passage and question lexical encodings into the same number of dimensions. Note that this FFN has fewer parameters compared to a BiLSTM. Thus, we obtain the final lexicon embeddings for the tokens in $Q$ as a matrix $E^q \in \mathbb {R}^{d \times m}$ and tokens in $P$ as $E^p\in \mathbb {R}^{d \times n}$ .
Contextual Encoding Layer. Both passage and question use a shared two-layers BiLSTM as the contextual encoding layer, which projects the lexicon embeddings to contextual embeddings. We concatenate a pre-trained 600-dimensional CoVe vectors BIBREF8 trained on German-English machine translation dataset, with the aforementioned lexicon embeddings as the final input of the contextual encoding layer, and also with the output of the first contextual encoding layer as the input of its second encoding layer. To reduce the parameter size, we use a maxout layer BIBREF9 at each BiLSTM layer to shrink its dimension. By a concatenation of the outputs of two BiLSTM layers, we obtain $H^q\in \mathbb {R}^{2d \times m}$ as representation of $Q$ and $H^p\in \mathbb {R}^{2d \times n}$ as representation of $P$ , where $d$ is the hidden size of the BiLSTM.
Memory Generation Layer. In the memory generation layer, We construct the working memory, a summary of information from both $Q$ and $P$ . First, a dot-product attention is adopted like in BIBREF7 to measure the similarity between the tokens in $Q$ and $P$ . Instead of using a scalar to normalize the scores as in BIBREF7 , we use one layer network to transform the contextual information of both $Q$ and $P$ :
$$C=dropout(f_{attention}(\hat{H}^q, \hat{H}^p)) \in \mathbb {R}^{m \times n}\\$$   (Eq. 8) 
 $C$ is an attention matrix. Note that $\hat{H^q}$ and $\hat{H^p}$ is transformed from $H^q$ and $H^p$ by one layer neural network $ReLU(W_3x)$ , respectively. Next, we gather all the information on passages by a simple concatenation of its contextual information $H^p$ and its question-aware representation $H^q \cdot C$ : 
$$U^p = concat(H^p, H^qC) \in \mathbb {R}^{4d \times n}$$   (Eq. 9) 
Typically, a passage may contain hundred of tokens, making it hard to learn the long dependencies within it. Inspired by BIBREF10 , we apply a self-attended layer to rearrange the information $U^p$ as: 
$$\hat{U}^p = U^p drop_{diag}(f_{attention}(U^p, U^p)).$$   (Eq. 10) 
In other words, we first obtain an $n \times n$ attention matrix with $U^p$ onto itself, apply dropout, then multiply this matrix with $U^p$ to obtain an updated $\hat{U}^p$ . Instead of using a penalization term as in BIBREF10 , we dropout the diagonal of the similarity matrix forcing each token in the passage to align to other tokens rather than itself.
At last, the working memory is generated by using another BiLSTM based on all the information gathered: 
$$M=BiLSTM([U^p; \hat{U}^p])$$   (Eq. 11) 
where the semicolon mark $;$ indicates the vector/matrix concatenation operator.
Answer module. There is a Chinese proverb that says: “wisdom of masses exceeds that of any individual." Unlike other multi-step reasoning models, which only uses a single output either at the last step or some dynamically determined final step, our answer module employs all the outputs of multiple step reasoning. Intuitively, by applying dropout, it avoids a “step bias problem" (where models places too much emphasis one particular step's predictions) and forces the model to produce good predictions at every individual step. Further, during decoding, we reuse wisdom of masses instead of individual to achieve a better result. We call this method “stochastic prediction dropout" because dropout is being applied to the final predictive distributions.
Formally, our answer module will compute over $T$ memory steps and output the answer span. This module is a memory network and has some similarities to other multi-step reasoning networks: namely, it maintains a state vector, one state per step. At the beginning, the initial state $s_0$ is the summary of the $Q$ : $s_0=\sum _j \alpha _j H^q_{j}$ , where $\alpha _j = \frac{exp(w_4 \cdot H^q_j)}{\sum _{j^{\prime }}exp(w_4 \cdot H^q_{j^{\prime }})}$ . At time step $t$ in the range of $\lbrace 1, 2, ..., T-1\rbrace $ , the state is defined by $s_t = GRU(s_{t-1}, x_t)$ . Here, $x_t$ is computed from the previous state $s_{t-1}$ and memory $s_0$0 : $s_0$1 and $s_0$2 . Finally, a bilinear function is used to find the begin and end point of answer spans at each reasoning step $s_0$3 .
$$P_t^{begin} = softmax(s_tW_6M)$$   (Eq. 12) 
$$P_t^{end} = softmax([s_t; \sum _j P_{t,j}^{begin}M_j]W_7M).$$   (Eq. 13) 
From a pair of begin and end points, the answer string can be extracted from the passage. However, rather than output the results (start/end points) from the final step (which is fixed at $T-1$ as in Memory Networks or dynamically determined as in ReasoNet), we utilize all of the $T$ outputs by averaging the scores: 
$$P^{begin} = avg([P_0^{begin}, P_1^{begin}, ..., P_{T-1}^{begin}])$$   (Eq. 14) 
$$P^{end} = avg([P_0^{end}, P_1^{end}, ..., P_{T-1}^{end}])$$   (Eq. 15) 
Each $P_t^{begin}$ or $P_t^{end}$ is a multinomial distribution over $\lbrace 1,\ldots ,n\rbrace $ , so the average distribution is straightforward to compute.
During training, we apply stochastic dropout to before the above averaging operation. For example, as illustrated in Figure 1 , we randomly delete several steps' predictions in Equations 14 and 15 so that $P^{begin}$ might be $avg([P_1^{begin}, P_3^{begin}])$ and $P^{end}$ might be $avg([P_0^{end}, P_3^{end}, P_{4}^{end}])$ . The use of averaged predictions and dropout during training improves robustness.
Our stochastic prediction dropout is similar in motivation to the dropout introduced by BIBREF11 . The difference is that theirs is dropout at the intermediate node-level, whereas ours is dropout at the final layer-level. Dropout at the node-level prevents correlation between features. Dropout at the final layer level, where randomness is introduced to the averaging of predictions, prevents our model from relying exclusively on a particular step to generate correct output. We used a dropout rate of 0.4 in experiments.
 Dataset: We evaluate on the Stanford Question Answering Dataset (SQuAD) BIBREF0 . This contains about 23K passages and 100K questions. The passages come from approximately 500 Wikipedia articles and the questions and answers are obtained by crowdsourcing. The crowdsourced workers are asked to read a passage (a paragraph), come up with questions, then mark the answer span. All results are on the official development set, unless otherwise noted.
Two evaluation metrics are used: Exact Match (EM), which measures the percentage of span predictions that matched any one of the ground truth answer exactly, and Macro-averaged F1 score, which measures the average overlap between the prediction and the ground truth answer. Implementation details: The spaCy tool is used to tokenize the both passages and questions, and generate lemma, part-of-speech and named entity tags. We use 2-layer BiLSTM with $d=128$ hidden units for both passage and question encoding. The mini-batch size is set to 32 and Adamax BIBREF12 is used as our optimizer. The learning rate is set to 0.002 at first and decreased by half after every 10 epochs. We set the dropout rate for all the hidden units of LSTM, and the answer module output layer to 0.4. To prevent degenerate output, we ensure that at least one step in the answer module is active during training.
The main experimental question we would like to answer is whether the stochastic dropout and averaging in the answer module is an effective technique for multi-step reasoning. To do so, we fixed all lower layers and compared different architectures for the answer module:
The main results in terms of EM and F1 are shown in Table 1 . We observe that SAN achieves 76.235 EM and 84.056 F1, outperforming all other models. Standard 1-step model only achieves 75.139 EM and dynamic steps (via ReasoNet) achieves only 75.355 EM. SAN also outperforms a 5-step memory net with averaging, which implies averaging predictions is not the only thing that led to SAN's superior results; indeed, stochastic prediction dropout is an effective technique.
The K-best oracle results is shown in Figure 3 . The K-best spans are computed by ordering the spans according the their probabilities $P^{begin} \times P^{end}$ . We limit K in the range 1 to 4 and then pick the span with the best EM or F1 as oracle. SAN also outperforms the other models in terms of K-best oracle scores. Impressively, these models achieve human performance at $K=2$ for EM and $K=3$ for F1.
Finally, we compare our results with other top models in Table 2 . Note that all the results in Table 2 are taken from the published papers. We see that SAN is very competitive in both single and ensemble settings (ranked in second) despite its simplicity. Note that the best-performing model BIBREF14 used a large-scale language model as an extra contextual embedding, which gave a significant improvement (+4.3% dev F1). We expect significant improvements if we add this to SAN in future work.
We are interested in whether the proposed model is sensitive to different random initial conditions. Table 3 shows the development set scores of SAN trained from initialization with different random seeds. We observe that the SAN results are consistently strong regardless of the 10 different initializations. For example, the mean EM score is 76.131 and the lowest EM score is 75.922, both of which still outperform the 75.355 EM of the Dynamic step ReasoNet in Table 1 .
We are also interested in how sensitive are the results to the number of reasoning steps, which is a fixed hyper-parameter. Since we are using dropout, a natural question is whether we can extend the number of steps to an extremely large number. Table 4 shows the development set scores for $T=1$ to $T=10$ . We observe that there is a gradual improvement as we increase $T=1$ to $T=5$ , but after 5 steps the improvements have saturated. In fact, the EM/F1 scores drop slightly, but considering that the random initialization results in Table 3 show a standard deviation of 0.142 and a spread of 0.426 (for EM), we believe that the $T=10$ result does not statistically differ from the $T=5$ result. In summary, we think it is useful to perform some approximate hyper-parameter tuning for the number of steps, but it is not necessary to find the exact optimal value.
Finally, we test SAN on two Adversarial SQuAD datasets, AddSent and AddOneSent BIBREF22 , where the passages contain auto-generated adversarial distracting sentences to fool computer systems that are developed to answer questions about the passages. For example, AddSent is constructed by adding sentences that look similar to the question, but do not actually contradict the correct answer. AddOneSent is constructed by appending a random human-approved sentence to the passage.
We evaluate the single SAN model (i.e., the one presented in Table 2 ) on both AddSent and AddOneSent. The results in Table 5 show that SAN achieves the new state-of-the-art performance and SAN's superior result is mainly attributed to the multi-step answer module, which leads to significant improvement in F1 score over the Standard 1-step answer module, i.e., +1.2 on AddSent and +0.7 on AddOneSent.
For practical deployment scenarios, prediction speed at test time is an important criterion. Therefore, one question is whether SAN can train with, e.g. $T=5$ steps but test with $T=1$ steps. Table 6 shows the results of a SAN trained on $T=5$ steps, but tested with different number of steps. As expected, the results are best when $T$ matches during training and test; however, it is important to note that small numbers of steps $T=1$ and $T=2$ nevertheless achieve strong results. For example, prediction at $T=1$ achieves 75.58, which outperforms a standard 1-step model (75.14 EM) as in Table 1 that has approximate equivalent prediction time.
The average training time per epoch is comparable: our implementation running on a GTX Titan X is 22 minutes for 5-step memory net, 30 minutes for ReasoNet, and 24 minutes for SAN. The learning curve is shown in Figure 4 . We observe that all systems improve at approximately the same rate up to 10 or 15 epochs. However, SAN continues to improve afterwards as other models start to saturate. This observation is consistent with previous works using dropout BIBREF11 . We believe that while training time per epoch is similar between SAN and other models, it is recommended to train SAN for more epochs in order to achieve gains in EM/F1.
To see whether SAN performs well on a particular type of question, we divided the development set by questions type based on their respective Wh-word, such as “who" and “where". The score breakdown by F1 is shown in Figure 5 . We observe that SAN seems to outperform other models uniformly across all types. The only exception is the Why questions, but there is too little data to derive strong conclusions.
MS MARCO BIBREF27 is a large scale real-word RC dataset which contains 100,100 (100K) queries collected from anonymized user logs from the Bing search engine. The characteristic of MS MARCO is that all the questions are real user queries and passages are extracted from real web documents. For each query, approximate 10 passages are extracted from public web documents. The answers are generated by humans. The data is partitioned into a 82,430 training, a 10,047 development and 9,650 test tuples. The evaluation metrics are BLEU BIBREF28 and ROUGE-L BIBREF29 due to its free-form text answer style. To apply the same RC model, we search for a span in MS MARCO's passages that maximizes the ROUGE-L score with the raw free-form answer. It has an upper bound of 93.45 BLEU and 93.82 ROUGE-L on the development set.
The MS MARCO dataset contains multiple passages per query. Our model as shown in Figure 2 is developed to generate answer from a single passage. Thus, we need to extend it to handle multiple passages. Following BIBREF13 , we take two steps to generate an answer to a query $Q$ from $J$ passages, $P^1, ..., P^J$ . First, we run SAN on every ( $P^j, Q$ ) pair, generating $J$ candidate answer spans, one from each passage. Then, we multiply the SAN score of each candidate answer span with its relevance score $r(P^j, Q)$ assigned by a passage ranker, and output the span with the maximum score as the answer. In our experiments, we use the passage ranker described in BIBREF30 . The ranker is trained on the same MS MARCO training data, and achieves 37.1 p@1 on the development set.
The results in Table 7 show that SAN outperforms V-Net BIBREF31 and becomes the new state of the art.
The recent big progress on MRC is largely due to the availability of the large-scale datasets BIBREF0 , BIBREF27 , BIBREF32 , BIBREF1 , since it is possible to train large end-to-end neural network models. In spite of the variety of model structures and attenion types BIBREF33 , BIBREF34 , BIBREF35 , BIBREF21 , BIBREF13 , BIBREF19 , a typical neural network MRC model first maps the symbolic representation of the documents and questions into a neural space, then search answers on top of it. We categorize these models into two groups based on the difference of the answer module: single-step and multi-step reasoning. The key difference between the two is what strategies are applied to search the final answers in the neural space.
A single-step model matches the question and document only once and produce the final answers. It is simple yet efficient and can be trained using the classical back-propagation algorithm, thus it is adopted by most systems BIBREF34 , BIBREF21 , BIBREF19 , BIBREF18 , BIBREF36 , BIBREF37 , BIBREF17 . However, since humans often solve question answering tasks by re-reading and re-digesting the document multiple times before reaching the final answers (this may be based on the complexity of the questions/documents), it is natural to devise an iterative way to find answers as multi-step reasoning.
Pioneered by BIBREF1 , BIBREF2 , BIBREF3 , BIBREF4 , who used a predetermined fixed number of reasoning steps, Shen et al shen2016reasonet, shen2017empirical showed that multi-step reasoning outperforms single-step ones and dynamic multi-step reasoning further outperforms the fixed multi-step ones on two distinct MRC datasets (SQuAD and MS MARCO). But these models have to be trained using reinforcement learning methods, e.g., policy gradient, which are tricky to implement due to the instability issue. Our model is different in that we fix the number of reasoning steps, but perform stochastic dropout to prevent step bias. Further, our model can also be trained by using the back-propagation algorithm, which is simple and yet efficient.
We introduce Stochastic Answer Networks (SAN), a simple yet robust model for machine reading comprehension. The use of stochastic dropout in training and averaging in test at the answer module leads to robust improvements on SQuAD, outperforming both fixed step memory networks and dynamic step ReasoNet. We further empirically analyze the properties of SAN in detail. The model achieves results competitive with the state-of-the-art on the SQuAD leaderboard, as well as on the Adversarial SQuAD and MS MARCO datasets. Due to the strong connection between the proposed model with memory networks and ReasoNet, we would like to delve into the theoretical link between these models and its training algorithms. Further, we also would like to explore SAN on other tasks, such as text classification and natural language inference for its generalization in the future.
We thank Pengcheng He, Yu Wang and Xinying Song for help to set up dockers. We also thank Pranav Samir Rajpurkar for help on SQuAD evaluations, and the anonymous reviewers for valuable discussions and comments. 

Figure and table captions:
Figure 1: Illustration of “stochastic prediction dropout” in the answer module during training. At each reasoning step t, the model combines memory (bottom row) with hidden states st−1 to generate a prediction (multinomial distribution). Here, there are three steps and three predictions, but one prediction is dropped and the final result is an average of the remaining distributions.
Figure 2: Architecture of the SAN for Reading Comprehension: The first layer is a lexicon encoding layer that maps words to their embeddings independently for the question (left) and the passage (right): this is a concatenation of word embeddings, POS embeddings, etc. followed by a position-wise FFN. The next layer is a context encoding layer, where a BiLSTM is used on the top of the lexicon embedding layer to obtain the context representation for both question and passage. In order to reduce the parameters, a maxout layer is applied on the output of BiLSTM. The third layer is the working memory: First we compute an alignment matrix between the question and passage using an attention mechanism, and use this to derive a question-aware passage representation. Then we concatenate this with the context representation of passage and the word embedding, and employ a self attention layer to re-arrange the information gathered. Finally, we use another LSTM to generate a working memory for the passage. At last, the fourth layer is the answer module, which is a GRU that outputs predictions at each state st.
Table 1: Main results—Comparison of different answer module architectures. Note that SAN performs best in both Exact Match and F1 metrics.
Table 2: Test performance on SQuAD. Results are sorted by Test F1.
Table 3: Robustness of SAN (5-step) on different random seeds for initialization: best and worst scores are boldfaced. Note that our official submit is trained on seed 1.
Figure 3: K-Best Oracle results
Table 4: Effect of number of steps: best and worst results are boldfaced.
Figure 5: Score breakdown by question type.
Table 5: Test performance on the adversarial SQuAD dataset in F1 score.
Table 6: Prediction on different steps T . Note that the SAN model is trained using 5 steps.
Figure 4: Learning curve measured on Dev set.
Table 7: MS MARCO devset results.

Questions about this paper:
How much performance improvements they achieve on SQuAD?

Answers:
Compared to baselines SAN (Table 1) shows  improvement of 1.096% on EM and 0.689% F1. Compared to other published SQuAD results (Table 2) SAN is ranked second.