{'fraction_non_alphanumeric': 0.03792988127752076, 'fraction_numerical': 0.01948051948051948, 'mean_word_length': 4.445068285280728, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 1, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Attention that does not Explain Away

Models based on the Transformer architecture have achieved better accuracy than the ones based on competing architectures for a large set of tasks. A unique feature of the Transformer is its universal application of a self-attention mechanism, which allows for free information flow at arbitrary distances. Following a probabilistic view of the attention via the Gaussian mixture model, we find empirical evidence that the Transformer attention tends to"explain away"certain input neurons. To compensate for this, we propose a doubly-normalized attention scheme that is simple to implement and provides theoretical guarantees for avoiding the"explaining away"effect without introducing significant computational or memory cost. Empirically, we show that the new attention schemes result in improved performance on several well-known benchmarks.


Introduction
The Transformer architecture (Vaswani et al., 2017) has been successfully used to improve state-of-theart performance in a variety of machine learning tasks, such as machine translation (Vaswani et al., 2017;Dehghani et al., 2019), language modeling (Devlin et al., 2019;Yang et al., 2019), summarization (Cohan et al., 2018;, dialog (Mazaré et al., 2018;Cheng et al., 2019), image captioning (Sharma et al., 2018;Zhao et al., 2019), and visual question answering (Yu et al., 2019b;Tan and Bansal, 2019). One of the most important components of the Transformer architecture is its self-attention mechanism, applied universally to both the encoder and the decoder components. This attention mechanism allows for information to freely flow between inputs at arbitrary distances, which is intuitively appealing for modeling natural language or tasks that need to model cross-modal relationships between their inputs.
Despite the empirical success of the selfattention mechanism, little formal work has been done to analyze its statistical properties and relate it to previously known classical models. Better understanding its properties can lead to insights into what it does and does not do well. This in turn can lead to improvements to the attention mechanism and ultimately to a better-performing Transformer network.
In this paper, we closely study the Transformer attention formulation from a probabilistic view via the Gaussian mixture model. If we consider the Transformer model as a stack of layers with data flowing from lower to upper layers, then the output neurons (from the upper layer) of an attention unit can be regarded as the most likely data generated by a Gaussian mixture model (GMM), while the input neurons (from the lower layer) of the attention unit act as the Gaussian centers.
Our insight here is that this Transformer attention scheme has an "explaining away" effect, which means that the information present in certain lower layer neurons may be filtered out completely. This is because for a GMM, not all Gaussian centers (lower layer neurons) are required to contribute in generating output data (upper layer neurons). The information of the centers that do not generate data is lost after observing the data. This "explainingaway" effect is related to the one in the directed graphical model, in the sense that the existence of the few contributed lower neurons "explain away" the other muted lower neurons on generating upper neurons.
In order to compensate for this, we describe an alternative probabilistic model for attention, in which the role of the upper and lower layers in the GMM formulation are reversed. This new attention scheme requires all the generated data (lower layer neurons) to be explained by at least one Gaussian center (upper layer neurons). Therefore, it guaran-tees the preservation of information for all lower layer neurons, as we prove in this paper.
The MLE equation of the reversed GMM model leads to a simple attention update that is similar to the original one, except for the attention weight normalization. The original Transformer attention scheme only normalizes the attention weights once for every upper-layer neuron. By contrast, our new attention mechanism requires a two-step attention weight normalization procedure: the first normalizes each lower-layer neuron, and the second normalizes each upper-layer neuron. In the rest of this paper, we denote the original, upper normalized attention scheme as UNAS, and the new doublynormalized attention scheme as DNAS.
We also show that DNAS updates correspond exactly to one iteration of the Sinkhorn algorithm (Peyré and Cuturi, 2019) in a constrained optimization problem. As a result, iterating DNAS until convergence results in a doublystochastic attention matrix where the attention weights of all upper and lower neurons are normalized. We also showed that UNAS can be formulated in a similar constrained optimization problem, except that the optimization problem of UNAS does not have the constraint which presents "explaining away" compared to DNAS.
Mathematically, we also formalize the concept of "explaining away" of a lower neuron by using the sum of its attention weights. We prove that the attention weights sum of the lower neurons of DNAS are lower bounded by 1/(sequence length), therefore completely avoid the "explaining away" effect of UNAS.
Last but not least, we formulate a hybrid attention scheme, HNAS, that dynamically combines both attention schemes, and can provide a handle on a task-based preference between UNAS and DNAS, as resulting from the learning algorithm. We perform empirical studies and obtain clear numerical improvements using DNAS and HNAS formulation in several wellknown benchmarks, with minor computational overhead and negligible increase of model size.

Transformer Attention and Gaussian
Mixture Models In this section, we review the Transformer selfattention mechanism and analyze how it relates to the Gaussian Mixture Model. Assuming a sequence of length S, we first focus on the Transformer single-headed attention formulation involving two layers of neurons: the lowerlayer neurons are the input representations denoted as x j at position j ∈ {1, . . . , S}, and the upperlayer neurons are the output representations denoted as y i at position i ∈ {1, . . . , S}. We assume both x j and y i are 1-d tensors of the same size D.
The self-attention mechanism first transforms the input representations x j to queries and keys by applying q j = Q x j and k j = K x j , where Q and K are trainable transformation matrices of size D × D. The value of an upper-layer neuron y i is computed as the weighted sum over the lower-layer neurons x j followed by the value transformation V of size D × D, .
Since in this formulation the attention weights π ij are normalized for every upper layer neuron i over the lower layer neurons j, we refer to this attention scheme as upper-normalized attention, UNAS.

Relation to GMM
The UNAS scheme (1) relates to a Gaussian mixture model (GMM) in the following way. Let us use k j to denote the positions of the Gaussian cluster centers, and the cluster priors denoted as α j , satisfying j α j = 1. The generated data position is denoted as q i . If we assume the variance of the Gaussian distributions to be equal to 1 * , then the log-likelihood of the GMM is: We can compute the optimal q i by taking the derivative of q i and solve the following equation, .
If we assume the cluster priors † as α j ∝ exp( 1 2 k j k j ), we have . (2) Using the fact that j π ij = 1 and we obtain a fixedpoint equation: If we compare Eq. (3) with Eq. (1), the Gaussian cluster centers k j play exactly the same role as the key representation k j of the lower-layer neurons in Eq. (1). The data position q i in Eq.
(2) plays the same role as the query representation q i in Eq. (1). By iterating the fixed-point equation (3) for one iteration, the new data position q new i = j π ij k j corresponds to the upper layer neuron y i in Eq. (1) after applying the transformation V K −1 .
Note that computing the most-likely data positions q i given the Gaussian centers is non-standard for probabilistic inference. A more natural way would be the MLE estimation for the Gaussian centers given the data. That is exactly what doublynormalized attention corresponds to, as we will discuss in the next section.

Multi-head attention
The multi-head (H heads) attention can be derived similarly. The lower neurons x j are pro- where V h is the value transformation matrix of size D H × D. Similar to (1), (4) corresponds to † The cluster prior αj favors the neurons with larger |kj|, which intuitively are the ones carrying more information. a GMM followed by value transformations ‡ . Hheads attention corresponds to H GMMs followed by value transformations. The final output is a concatenation of all H heads: y i = concat(y h i ).

Doubly-normalized Attention
As we have shown, in the original UNAS scheme, the lower layer neuron representations correspond to the Gaussian centers, while the upper layer neuron representations correspond to the data generated from these centers. The maximization with respect to the data positions is unnatural. In addition, the formulation has an "explaining away" effect, because for a GMM, not all Gaussian centers (lower layer neurons) are required to contribute in generating output data (upper layer neurons). As a result, the information of the centers that do not generate data is completely lost. For tasks such as summarization, "explaining away" may be acceptable, while for other tasks such as visual question answering and language modeling, the attention mechanism may benefit from a more "conservative" formulation, with the upper layer preserving the neural information at all positions.
To this end, we propose to reverse the role of the upper and lower layers in the GMM, so that all the generated data (lower layer neurons) will be explained by at least one Gaussian center (upper layer neurons). This results in a new doubly-normalized attention scheme (DNAS) (the derivation will be given shortly): .
Comparing (1) with (5), the only difference between the two is the normalization process of the attention weights. The DNAS scheme applies two normalization steps: first for each lower layer neuron j and then for each upper layer neuron i.

Relation to GMM
We present here the derivation of (5) from a GMM. When we reverse the role of the upper and lower layers, we use q i to denote the Gaussian centers ‡ Some special treatments are needed to handle the value transformation since K h is no longer square matrices. See the details in the Appendix. and k j as the data generated by GMM. The loglikelihood function of the GMM is: where the priors β i satisfy i β i = 1. We take the gradient with respect to q i , At optimum ∂ ∂ q i j log p(k j ) = 0, we have 0 = j ξ ij (q i − k j ), and therefore the fixed-point equation is, By iterating the fixed-point equation (8) for one iteration and assuming β i ∝ exp( 1 2 q i q i ), then the new center position q new i = j ξ ij j ξ ij k j is equivalent to the upper layer neuron y i of Eq. (5), modulo a transformation matrix V K −1 .
Similar to the UNAS, it is also straightforward to extend the above derivations to the multi-head (H-heads) DNAS scheme, where it would be H GMMs followed by value transformations.

Relation to Double Stochasticity
It should be emphasized that our doublynormalized attention is not doubly-stochastic (where the columns and rows of the attention matrix π ij all sum to 1). After applying DNAS, the attention weights of the lower layer neurons are not normalized, since the upper layer normalization in the second step of DNAS denormalizes the lower layer. However, as we show in the following, doubly-stochastic attention can be achieved by applying the two normalization steps for multiple iterations until convergence.
Consider the following constrained optimization problem that characterizes π ij , This problem is well-known in the optimal transport literature. The classical iterative algorithm for finding the solution is called the Sinkhorn algorithm (Peyré and Cuturi, 2019), which uses the initial condition π 0 ij = exp(−D(q i , k j )) and iterates If we write D(q i , k j ) := − q i k j then the doublynormalized attention weights computed in Eq. (5) correspond exactly to the updates (10) of the Sinkhorn algorithm for one iteration. If more iterations are applied, the attention weights will eventually satisfy both constraints in (9), and become doubly-stochastic. One question is whether DNAS could perform better with more iterations for the updates in Eq. (10). Empirically, we find that adding more update iterations increases computational time but does not improve performance. Interestingly, the attention weights of the original UNAS scheme can be obtained from a very similar constrained optimization except that the normalization constraint on the lower layer neurons j is removed: Introducing the Lagrange multipliers λ i , this formulation is equivalent to optimizing the Lagrangian, whose gradient with respect to π ij gives and leads to the same attention weights as in Eq. (1) Comparing the two constrained optimization problems in (11) and (9), the removal of the constraint in (11) allows solutions in which a lowerlayer neuron j has an arbitrary contribution to the upper layer, causing the "explaining-away" effect.

Relation to Capsule Networks
It is also worth noting that DNAS is related to the EM routing algorithm in the capsule networks (Hinton et al., 2018). In particular, the vote matrix V ij in (Hinton et al., 2018) is similar to k j in Eq. (6); the new pose matrix µ j in (Hinton et al., 2018) is similar to q i in Eq. (6). However, unlike Capsu-leNet, there is no variance σ 2 i and β i estimation in DNAS, as we find that estimating variance σ 2 i significantly hurts the empirical performance of the DNAS algorithm. In addition, we only iterate the fixed-point equation (8) for one iteration, as more iterations are computationally expensive and does not improve the performance.

Doubly-Normalized Attention Avoids Explaining Away
In this section, we formalize the definition of "explaining-away" and compare UNAS and DNAS theoretically and empirically with respect to the "explaining-away" phenomenon.
Definition 1 In an attention unit, a lower-layer neuron j is considered -"explained away", if the sum of the attention weights over the upper layer neurons i π ij is less than .
We consider to be some small value (fixed at 10 −8 in the rest of this paper). For the original Transformer UNAS, the only constraint in (11) is j π ij = 1. It does not require all lower layer neurons to be attended by the upper layer. Therefore, for a certain lower-layer neuron j, the total attention weights to the upper layer i π ij can be as low as 0 so that it is -"explained away".
In contrast, the DNAS scheme attempts to optimize the objective with both lower and upper layer normalization constraints (9) by one iteration of the Sinkhorn algorithm. It turns out that this is sufficient to avoid the "explaining-away" phenomenon. The following theorem formalizes this fact by showing that each lower-layer neuron contributes with a total attention weight of at least 1/S, where S is the sequence length.
Theorem 2 For any lower-layer neuron j, the sum of the doubly-normalized attention weights over the upper layer neurons i π ij = i ξ ij j ξ ij is lower bounded by 1/S.
We illustrate the difference between the two attention schemes, and how different they behave in practice with respect to the "explaining-away" phenomenon, using the multi-view attention model (with a single-layer, single-head attention) described in the VQA experiments later. Fig. 1 shows the histogram distribution of log e ( i π ij ) between UNAS and DNAS. As the graph indicates, a large proportion of the UNAS attention weights-sum is -"explained-away" (log e values < −20), meaning that the information of only a few of the lower neurons are passed to the upper layer. In contrast, DNAS preserves more information from all lower layer neurons, as indicated by their weights-sum log values (> − log e S, where S = 100).
Finally, we would like to emphasize that DNAS does not work against attention sparsity. It allows the attention map π ij = 0 between any pairs of neurons. What it forbids is the 0 total "contribution" of any lower neuron j: i π ij = 0. Therefore, our method is compatible with existing faster sparse attention structures such as (Parmar et al., 2018).

Hybrid Attention
Since the formulations of UNAS and DNAS result in attention mechanisms with quite different prop-erties, it is beneficial to combine them together. A direct way to do so is by using trainable variables u h l ∈ [0, 1] that control the contribution of the attention weights (for layer l and head h) of the two normalization schemes (we use u here to simplify the notation): where π D denotes the DNAS weights and π U denotes the UNAS weights. We call this combination form the hybrid normalized attention scheme, HNAS. HNAS allows the model to learn, at different layers l and different heads h, which of the two normalization schemes fits the data better, for a given task. Each u h l parameter is trained jointly with the other parameters to improve the representation power of the model and better fit the data. Moreover, this approach also allows one to visualize how the values of the u h l parameters change as the model is training, and therefore provides direct evidence of how much and where the different normalization schemes lead to better training performance. We provide examples of such visualizations in the experiments.

Computational Cost of DNAS and HNAS
The pseudo-code of the (multi-headed) UNAS, DNAS and HNAS is summarized in Algorithm 1. Note that for notational clarity, we wrote multihead operations in a for-loop over different heads h ∈ {1, . . . , H}. However, an efficient implementation should use single tensor products across all heads, similar to the original Transformer method.
We can see that the additional computational cost of the DNAS scheme compared to the original Transformer's UNAS scheme is the two normalizations in Step-4 as opposed to one in Step-3.   In a vision-and-language multimodal system (e.g., Visual Question Answering), a crucial factor in the performance is the quality of the visual features. A good example is the work of (Yu et al., 2019a), where they show that it is beneficial to use visual features produced by different image processing modules (multi-view). They combine these visual features using an attention layer over the boundingbox features derived from multiple object detectors (Fig. 2). Experiment Setup. Our experimental setup is similar to the one proposed in (Yu et al., 2019a). We conduct experiments on the VQA benchmark dataset, VQA-v2 (Goyal et al., 2017). Our core VQA model uses as a backbone the Pythia architecture (Jiang et al., 2018). We used three object detection models, where each detector generates 100 bounding-box features. All three object detection models are trained over the Visual Genome dataset (Krishna et al., 2017), but use different backbone networks: the first uses a ResNet-101 network (He et al., 2016), the second a ResNet-200 network, and the third an Inception-ResNetV2 network (Szegedy et al., 2016). Multi-view features can be used in a straightforward manner by concatenating them all together before feeding them into the Pythia model; we call this approach the 3x100-boxes baseline. The proposal from (Yu et al., 2019a) combines the multiview features using a one-layer attention model as follows: one object-detector model is designated as primary, and its corresponding features are used as queries (after transformation); the second and third object detection models are designated as secondary, and their corresponding features are used to obtain keys (see Figure 2). The resulting output feature is a weighted sum of the features according to the attention weights. More details about the mutliview attention model and the experiment hyperparameter settings are provided in the Appendix. We use a single-layer and single-head attention model and experiment with two versions of the attention scheme: UNAS and DNAS. Results Analysis. The results are summarized in Table 1. Confirming the findings from (Yu et al., 2019a), using an attention mechanism (UNAS) over the 3x100 boxes improves the accuracy over the 3x100-boxes no-attn baseline, but the DNAS mechanism achieves a better utilization of the signal provided by the three object detectors compared to the UNAS mechanism. Moreover, HNAS allows us to visually confirm the superiority of the DNAS mechanism for the VQA task: as we plot the hybrid weight u from Eq.(12) in Fig. 3, it rapidly converges to 1.0, meaning that the model learns to heavily favor DNAS over UNAS for combining multi-view features. Combining the findings in Fig. 1, we believe that UNAS performs worse because it -"explains-away" too many box features in this stage, while DNAS preserves information from all bounding boxes.

Language Representation Learning
The goal of language representation learning is to pretrain textual representations that are useful for solving natural language understanding (NLU) tasks like entailment or question answering. Experiment Setup. We use the BERT (Devlin et al., 2019) setting for our language representation learning setup: a Transformer network with 24 layers of attention, the hidden and embedding size set to 1024, and 16 attention heads.
Our experiment is based on the ALBERT platform (Lan et al., 2019) § . We use the BOOKCOR-PUS (Zhu et al., 2015) and English Wikipedia (Devlin et al., 2019) to pretrain three contextual representation models, using UNAS, DNAS, and HNAS respectively. Each pretraining uses a batch size of 4096 and a LAMB optimizer with learning rate 0.00176 for 125k steps on the Cloud TPU V3 with 64 TPUs. We evaluate the resulting representations by using them as a starting point to finetune for a number of representative NLU tasks (Rajpurkar et al., 2018;Williams et al., 2018). Due Method Test-dev Test-std 10-100-boxes Pythia (Jiang et al., 2018) 66.91 -3x100-boxes (no-attn baseline) 68.79 69.22 3x100-boxes UNAS (attn baseline) 69.14 69.50 3x100-boxes DNAS 69.70 70.01   to space limitation, more experimental details are provided in the Appendix. Results Analysis. Each fine-tuning experiment is done 5 times, and the mean number and their standard error are reported. The main results are summarized in Table 2 and more detailed results are available in the Appendix. Overall, the network parameters encode their language representations by making use of DNAS, resulting in the empirical advantage of the DNAS and HNAS based models over the UNAS based models on most tasks considered. Aside from the numerical improvements when finetuning on the task, we also inspect what happens to the hybrid weight u of Eq.(12) during HNAS pretraining. In Fig. 4, we plot the hybrid weights (averaged over all heads of each layer) for all 24 layers and find that they are always larger than 0.5, meaning that the DNAS method is preferred for pretraining (masked-LM & sentenceordering) tasks. The UNAS method has more weight for higher layers, meaning that "explaining away" is more allowable when it is closer to the output.

Headline Generation
We also present empirical results on a summarization task. As already mentioned, summarization aligns well with the tendency of UNAS of "explaining away" unimportant information.
Experiment Setup. We use the Gigaword dataset (Graff and Cieri, 2003), which is a standard benchmark for headline generation. We preprocess this dataset as in (Rush et al., 2015), and further tokenize the words into word-pieces (Devlin et al., 2019), which results in a vocabulary size of 30,522 word-piece types. We use a 10k dataset for validation, and the standard 2k test set (Rush et al., 2015) as the evaluation test.
Our model and training hyperparameters are adapted from . The transformer contains 12 layers, each with a hidden size of 768 and 12 attention heads. We keep the attention mechanism in the decoder as UNAS, and compare the DNAS and HNAS with UNAS as the encoder attention mechanism. Our training uses a batch size of 512 and an Adam optimizer (Kingma and Ba, 2015) with learning rate of 2e −5 for 500k steps. The training is done on Cloud TPU V3 with 16 TPUs for each job.
Results Analysis. Each experiment is run 5 times, and the mean number and standard error are reported in Table 3. We also plot the averaged hybrid weights for all layers in Fig. 5 which shows that the HNAS model favors UNAS, especially in the top and bottom layers of the encoder. Nevertheless, DNAS still makes a positive contribution in the middle layers, which allows the model based on HNAS to perform better compared to the UNAS- Hybrid weight Figure 5: The hybrid weights favor UNAS in the encoder of headline generation, because the task requires filtering unimportant information. However, the ROUGE scores of DNAS is higher than UNAS.
based one. Somewhat surprisingly, DNAS alone performs competitvely: all of its ROUGE scores are higher than the ones of UNAS and are close to the ones of HNAS. This indicates that complete "explaining away" by UNAS is unnecessary for filtering unimportant information. DNAS provides a conservative alternative which achieves better generation performance.

Conclusion
The formulation of the attention mechanism of the Transformer, here called UNAS, leads to "explaining away" effects in which the information of certain input neurons is completely ignored. Our new DNAS scheme compensates for UNAS's weaknesses by avoiding "explaining away", as we show both theoretically and empirically. Empirically, we show DNAS and a hybrid HNAS to be superior to the original attention mechanism, at the cost of minor computational overhead.

A Multi-head attention and GMM
In multi-head attention, the lower neurons x j are projected into H heads with different q h j = Q h x j and k h j = K h x j where Q h and K h are transformation matrices of size D H × D. This yields H outputs y h i , If we follow the same idea as in single-head attention, the corresponding GMM becomes, In order to convert q h,new i to y h i , one difficulty is that K h is a down-projection matrix. Therefore, the inversion of K h K h does not exist. In order to avoid the problem, one can use the same key transformation for all headsK h = K which is

One can show that, ifq
Therefore, the corresponding GMM becomes

B Experiment Details about the Multi-view Attention Model for VQA
Dataset and evaluation The VQA-v2 (Goyal et al., 2017) dataset contains a training set (with 80k images and 444k QA pairs), a validation set (with 40k images and 214k QA pairs), and test set (with 80k images and 448k QA pairs). For each question, there are 10 answers provided by 10 different human annotators. Following the same setting as Pythia (Jiang et al., 2018), we augment the train set with a part of validation set (train + val2train) and use the remaining data in validation set for validation (minival). The test set is split into test-dev and test-std, and the evaluation can only be conducted online. Same as other work on VQA, we report a robust accuracy metric as the average score over 9 subsets of the groundtruth 10 answers, where each score is computed as follows: Acc(ans) = min{(#human that said ans)/3, 1}.
Detailed Model Descriptions Our VQA model uses as a backbone the Pythia architecture (Jiang et al., 2018). In order to combine the 100 features from each of the three object detection models, we use a one-layer attention mechanism as in (Yu et al., 2019a). The features from one objectdetector model is used as the primary feature. The features of the second and third object detection models are designated as secondary features. In order to obtain keys and queries, we apply transformation on the secondary and primary features, so that k S1 However, we find that it is better to directly use the features as the values without transformation. For the primary view, the output value of the i-th feature is y P i = x P i . For each secondary view, the feature is computed as y S1 i = j π S1 ij x S1 j For the UNAS scheme, π S1 ij = exp(q i k S1 j ) j exp(q i k S1 j ) .
The final output feature integrates the 100 features from different views via an element-wise summation, followed by layer normalization, Hyperparameters During the hyperparameter tuning process, we train on training set only and manually tune our hyperparameter based on the accuracy on the validation set. We use the same model hyperparameters as the Pythia model. Our image feature dimension is 2048 and the query and key transformation matrices are of size 2048×2048.
For the attention layer, we experiment with multiple number of heads including 1, 2, 4, and 8, and we find the single head attention gives the best performance. We also did a grid search on the dropout probability in attention layer from 0.05 to 0.4, and set it to 0.1 after the search. The hybrid attention weight is initalized to be 0.5. For optimization, we use Adam optimizer with learning rate 10 −4 , and use batch size 192. We train the model for 500, 000 steps. The training was done on 4 Cloud TPUs. The total training time is approximately 38 hours for each model. The validation performance on the minival dataset is reported in Table 4.

C.1 Downstream Evaluation Tasks
SQuAD SQuAD is an extractive question answering dataset built from Wikipedia. The answers are segments from the context paragraphs and the task is to predict answer spans. We evaluate our models on two versions of SQuAD: v1.1 and v2.0. SQuAD v1.1 has 100,000 human-annotated question/answer pairs. SQuAD v2.0 additionally introduced 50,000 unanswerable questions. For SQuAD v1.1, we use the same training procedure as BERT, whereas for SQuAD v2.0, models are jointly trained with a span extraction loss and an additional classifier for predicting answerability (Yang et al., 2019;Liu et al., 2019). We report the results on the development set.
RACE RACE is a large-scale dataset for multichoice reading comprehension, collected from English examinations in China with nearly 100,000 questions. Each instance in RACE has 4 candidate answers. Following prior work (Yang et al., 2019;Liu et al., 2019), we use the concatenation of the passage, question, and each candidate answer as the input to models. Then, we use the representations from the "[CLS]" token for predicting the probability of each answer. The dataset consists of two domains: middle school and high school. We train our models on both domains and report accuracies on the development set.
GLUE GLUE (Williams et al., 2018)  There is about 20% overhead which is much higher than our theoretical estimation. This is because our BERT pretraining used 64 TPUs that are highly efficient for parallelizing large matmul ops. As a result, the runtime of two consecutive normalization steps of smaller tensors could be longer than a single-step matmul of a much larger tensor. We expect the relative overhead to be smaller with other types of processing units. Hyperparameters for downstream tasks are shown in Table 5. These hyperparameters were copied from  which were adapted from (Liu et al., 2019), (Devlin et al., 2019), and (Yang et al., 2019). We used the ADAM optimizer for fine-tuning as in .

D Experimental Details about Headline Generation
The Gigaword dataset (Graff and Cieri, 2003) consists of about 4M article, headline pairs. We pre-process this dataset as in (Rush et al., 2015), which results in an average article length of 31.4 words, and an average headline length of 8.5 words. We further tokenize the words into wordpieces (Devlin et al., 2019), which results in a vocabulary size of 30,522 word-piece types. We use a 10k dataset for validation, and the standard 2k test set (Rush et al., 2015) as the evaluation test.
Method minival 3x100-boxes (no-attn baseline) 68.26 3x100-boxes UNAS (attn baseline) 68.34 3x100-boxes HNAS 68.99   Our backbone Transformer model is adapted from ) that contains 12 layers, each with a hidden size of 768 and 12 attention heads. The total number of model parameters is about 108M. We truncate (or pad) the input and output sequences to a fixed number of wordpiece positions, namely 128 encoder positions and 64 decoder positions, to accommodate hardware and model-architecture limitations. The hybrid attention weight is initalized to be 0.1, because the headline generation task favors UNAS to "explain away" unimportant neurons. We use an Adam optimizer (Kingma and Ba, 2015) and a learning rate of 2e −5 for 500 steps. The training was done on Cloud TPU V3 with 16 TPUs for each job. The total training time is approximately 16.5 hours for DNAS/HNAS and 16 hours for UNAS.
The ROUGE-L score on the validation set is 45.75 for HNAS, and 45.63 for UNAS.

E Doubly-normalized Attention
Alleviates Mode Collapse Attention model tends to collapse modes. In particular, the data at different positions tend to move closer to each other after attention. We illustrate the collapsing effect in a 2-D example in Fig. 6, where two separated clusters of data converge to a single point after only 4 steps of UNAS (left). Most multilayer attention models such as the Transformer try to avoid such collapsing effect by adding a residual layer, which pulls the data back to its original position.
To compare the mode-collapsing effect of UNAS and DNAS analytically, we study a 1-D toy example which contains two clusters. One cluster contains N 0 data points centered at value a, and the other contains N 1 data points centered at value −a. The distance between the two centers is 2a. Assuming the relative distance between the data points within each set is negligible compared to 2a, the unnormalized attention weights between one center and the data from the other set is s = exp(−(2a) 2 /2) = exp(−2a 2 ), and the weights between one center and the data within that set is t = exp(0) = 1 ¶ . We compare the center dis- tance between the two data clusters after applying the UNAS and DNAS self-attention updates. Applying Eq. (3) for the UNAS scheme, the new center distance of the upper-normalized attention scheme are: and the distance between the two updated centers is: 2 )a (N 1 t + N 0 s)(N 0 t + N 1 s) .
the intra-attention weight is t = exp( a, a ) = exp(a 2 ). The ratio s/t = exp(−2a 2 ) is identical to the Gaussian case.
By contrast, if we apply the Eq. (8) updates for the DNAS scheme, the new center position of the doubly-normalized attention scheme are: and the distance between the two updated centers is: =2N 1 a{ t(N 0 t + N 1 s) N 0 s(N 0 s + N 1 t) + N 1 t(N 0 t + N 1 s) − s(N 0 t + N 1 s) N 0 t(N 0 s + N 1 t) + N 1 s(N 0 t + N 1 s) }.
We plot the values of Eq. (13) and Eq. (14) on the y-axis against that of r = N 0 /N 1 on the x-axis, for several different a values, see Fig. 7. We see that in both cases the distance between the two centers decays after the attention updates. However, the center distance of DNAS always upper bounds the one of UNAS, with the gap getting larger as the cluster sizes get more unbalanced (r = 1). The above result holds for the 2-D example in Fig. 6 as well, where the UNAS collapses to a single cluster after 4 steps (left) while the DNAS maintains two separate clusters (right). The mode collapse effect is even more obvious in multi-layer attention. In Fig. 8, when the two clusters are balanced (both clusters contain 225 data points), both normalization schemes yield similar results. However, when the two clusters are unbalanced (the red cluster contains 500 points and the blue one contains 50) (Fig. 9), UNAS collapses to a single cluster after 4 steps, while the DNAS maintains two separate clusters.