Title: DOUBLY-NORMALIZED ATTENTION

Abstract: Models based on the Transformer architecture have achieved better accuracy than models based on competing architectures. A unique feature of the Transformer is its universal application of a self-attention mechanism, which allows for free information flow at arbitrary distances. In this paper, we provide two alternative views of the attention mechanism: one from the probabilistic view via the Gaussian mixture model, the other from the optimization view via optimal transport. Following these insights, we propose a new attention scheme that requires normalization on both the upper and lower layers, called the doubly-normalized attention scheme. We analyze the properties of both the original and the new attention schemes, and find that the doubly-normalized attention mechanism directly mitigates two unwanted effects: it resolves the explaining-away effect and alleviates mode collapse. We conduct empirical studies that quantify numerical advantages for the doubly-normalized attention model, as well as for a hybrid model that dynamically combines both attention schemes to achieve improved performance on several well-known benchmarks.

1. INTRODUCTION:

The Transformer architecture (Vaswani et al., 2017) has been successfully used to improve state-ofthe-art 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; Goodman et al., 2019) , dialog (Mazaré et al., 2018; Cheng et al., 2019) , image captioninig (Sharma et al., 2018; Zhao et al., 2019) , and visual question answering (Yu et al., 2019b; Tan & 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 (such as visual question answering).

Despite the empirical success of the self-attention 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 current attention formulation from both a probabilistic view via the Gaussian mixture model and from an optimization view via optimal transport. First, we reveal the mathematical connection between the Transformer attention and the Gaussian mixture model. In particular, we show that 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.

A formulation in which the upper layer acts as the generated data while the lower layer acts as the Gaussian centers aligns well with the purpose of a decoder mechanism, where a final upper layer ultimately generates the outputs. However, this design does not correspond well to the purpose of an encoder mechanism, in which the bottom layer takes input data and the upper layer encodes the data. To address this mismatch, we describe in this paper a new attention mechanism, in which the role of the upper and lower layers in the GMM formulation are reversed: the lower layer represents the generated data, while the upper layer represents the Gaussian centers. It is worth noting that a similar design was also applied in the Capsule Networks algorithm (Hinton et al., 2018) .

The Maximum Likelihood Estimate (MLE) solution of the reversed GMM model leads to a very similar attention update as the original, except for the attention weight normalization. The original attention scheme only normalizes the attention weights once, over the lower layer of every upper-layer neuron. By contrast, the doubly-normalized attention mechanism requires a two-step attention weight normalization: the first normalizes over the upper layer for each lower-layer neuron, and the second normalizes over the lower layer for each upper-layer neuron. In the rest of this paper, we denote the original, lower normalized attention scheme as LNAS, and the doubly-normalized attention scheme as DNAS.

In addition, our analysis shows that the two normalization schemes can also be connected as optimizing two closely related constrained objective functions, with one additional normalization constraint for DNAS. We also show that DNAS updates correspond to one iteration of the Sinkhorn algorithm in optimal transport (Peyré & Cuturi, 2019) .

We mathematically analyze the statistical properties of these two attention schemes, finding that the DNAS proposal possesses two advantages over LNAS. First, it avoids the "explaining away" effect that the LNAS scheme suffers from, in which information present in the input is filtered out too early and cannot be recovered at a later stage when needed. Second, it alleviates the mode-collapsing phenomenon that also plagues the LNAS formulation, which fails to accommodate multiple modes and concentrates its entire mass in a single mode.

We also formulate a hybrid scheme, HNAS, that dynamically combines both attention schemes, and can provide a handle on the preference betewen LNAS and DNAS, as resulting from the optimization algorithm. We perform empirical studies that quantify the numerical advantages of the doublynormalized attention model, as well as the ones for the hybrid model. We obtain clear numerical improvements using the HNAS formulation over several well-known benchmarks, including a new state-of-the-art result on the Gigaword bechmark for text summarization.

2. ATTENTION FROM THE PROBABILISTIC PERSPECTIVE:

In this section, we review the Transformer self-attention mechanism and analyze how it relates to the Gaussian Mixture Model formulation from the probabilistic perspective.

The Transformer attention mechanism involves two layers of neurons. In what follows, we denote the lower-layer neurons as x j , and the upper-layer neurons as y i . The self-attention mechanism first transforms the input features to a query and a key by applying the transformations q i = Q x i and k j = K x j , where Q and K are trainable transformation matrices. The value of an upper-layer neuron y i is then computed as the weighted sum over the lower-layer neurons x j , using the attention update equation,

EQUATION

which is then followed by an additional output-value transformation.

The Transformer attention updates from Eq. ( 1) can be regarded as the solution of the most likely data generation of a Gaussian mixture model (GMM). To make this connection clear, let us consider the log-likelihood function of a GMM. We denote the Gaussian cluster centers as k j , the priors of the clusters as α j satisfying j α j = 1, and the generated data 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:

EQUATION

To find the most likely data generated by the Gaussian mixture model, we take the derivative of q i ,

∂ ∂ q i i log p(q i ) = ∂ ∂ q i log p(q i ) = j α j N (q i |k j , 1) ∂ log N (q i |kj ,1) ∂ q i j α j N (q i |k j , 1)

.

We introduce the abbreviation

EQUATION

where j π ij = 1. We approximate the cluster priors as α j ∝ exp( 1 2 k j k j ), and therefore

π ij = exp(q i k j ) j exp(q i k j )

.

Also, since

∂ log N (q i |kj ,1) ∂ q i = k j -q i , then, ∂ ∂ q i i log p(q i ) = j π ij (k j -q i ) = j π ij k j -q i .

Let ∂ ∂ q i log p(q i ) = 0, then we have:

EQUATION

If we compare Eq. ( 4) with Eq. ( 1), the two equations are equivalent modulo a transformation matrix K, which can be incorporated in the additional output-value transformation that follows.

3. DOUBLY-NORMALIZED ATTENTION:

As we have shown, in the original Transformer self-attention mechanism LNAS, the lower layer corresponds to the Gaussian centers, while the upper layer corresponds to the data generated from these centers. This flow of information aligns well with the goal of a decoding mechanism, in which the upper layer generates the decoded outputs. However, for an encoding mechanism, the information flows in the other direction: data comes from the lower layer, while the role of the upper layer is to represent the signal from the lower-layer. As such, the current design of the Transformer self-attention mechanism, in which the LNAS scheme is applied universally for both the encoder and the decoder networks, does not align well with the information flow of the encoding process.

To resolve this discrepancy, we propose a new attention mechanism, DNAS, resulting from reversing the role of the upper and lower layers in the GMM. Under this new scheme, the lower layer neurons k j are treated as data, and the upper layer neurons q i are the Gaussian centers. The log-likelihood function becomes:

EQUATION

with priors β i satisfying i β i = 1. We take the gradient with respect to q i ,

∂ ∂ q i j log p(k j ) = j β i ∂ ∂ q i N (k j |q i , 1) i β i N (k j |q i , 1) = j β i N (k j |q i , 1) ∂ ∂ q i log N (k j |q i , 1) i β i N (k j |q i , 1)

.

We introduce the abbreviation

EQUATION

where i ξ ji = 1. If we approximate β i ∝ exp( 1 2 q i q i ), then

EQUATION

)

Let ∂ ∂ q i j log p(k j ) = 0, we have 0 = j ξ ji (q i -k j ), which results in:

EQUATION

Comparing ( 4) and ( 8), we notice that the only difference between the two updates is the normalization process of the attention weights. Under the Transformer attention scheme LNAS from Eq. ( 3), its weights π ij are normalized just once over the lower layer neurons j. For DNAS in Eq. ( 6), the attention weights are computed in two steps: first, an upper-layer normalization is applied over the upper-layer neurons i according to Eq. ( 7), for every lower-layer neuron j; second, a lower-layer normalization is applied over the lower-layer neurons j according to Eq. ( 8), for every upperlayer neuron i. The name doubly-normalized attention scheme, DNAS, is chosen to reflect this normalization mechanism. It is worth noting that DNAS is closely related to the EM routing algorithm in the capsule networks (Hinton et al., 2018) , which is also derived from GMM. Due to space limitations, this connection is discussed in more detail in Appendix A.

4. ATTENTION FROM THE OPTIMIZATION PERSPECTIVE:

The Transformer LNAS scheme can also be understood from an optimization perspective. Consider the following constrained optimization problem that characterizes π ij ,

EQUATION

Introducing the Lagrange multipliers λ i , this formulation is equivalent to optimizing the Lagrangian, whose gradient with respect to π ij gives

∂L(π ij , λ i ) ∂π ij = D(q i , k j ) + 1 + log π ij + λ i ,

and leads to the same solution as (3) when D(q i , k j ) := -q i k j .

The doubly-normalized attention scheme DNAS can be derived from a very similar constrained optimization except that an additional normalization constraint is added to the lower layer neurons:

EQUATION

)

The above objective function is well-known in the optimal transport literature, and the classical iterative algorithm for finding the solution is called the Sinkhorn algorithm (Peyré & Cuturi, 2019) . This algorithm uses the initial condition π 0 ij = exp(-D(q i , k j )), and iterates

ξ t ji = π t-1 ij i π t-1 ij , π t ij = ξ t ji j ξ t ji .

If we write D(q i , k j ) := -q i k j then the doubly-normalized attention weights in Eq. ( 8) correspond exactly to the updates of the Sinkhorn algorithm for one iteration.

Comparing the two constrained optimization problems in ( 9) and ( 10), the only difference is that the constraint on the attention-weight sum for the lower neurons is removed in (9). The removal of the constraint allows solutions in which a lower-layer neuron j has an arbitrary contribution to the upper layer. This is the reason for which the LNAS scheme suffers from the so-called "explaining-away" effect, as we discuss in the next section.

5. PROPERTIES OF DOUBLY-NORMALIZED ATTENTION:



5.1. DOUBLY-NORMALIZED ATTENTION AVOIDS EXPLAINING AWAY:

Under the original Transformer normalization scheme, LNAS, the mechanism allows the nodes in the higher layer to not attend to some of the nodes in the lower layer. In particular, for a certain lower-layer neuron j, the sum of its attention assignments to the upper layer is i π ij , where i denotes the location of the indices over the upper-layer neurons. Since the only restriction under LNAS is j π ij = 1, the summation i π ij can be as low as 0, which means the contribution of the lower neuron j is effectively filtered out in the upper layer. This is the very definition of the "explaining-away" phenomenon, which can potentially have negative effects for an encoder: if early encoder layers filter out the contribution of certain inputs that could be useful at later stages in the encoder, that information can no longer be recovered. In contrast to LNAS, the DNAS scheme avoids completely the "explaining-away" phenomenon. The following lemma formalizes this property by showing that each lower-layer neuron gets to contribute with a weight assignment of at least 1/L , where L is the layer size.

Lemma 1 For any lower-layer neuron j, the sum of the doubly-normalized attention weights over the upper layer neurons i π ij = i ξji j ξji is lower bounded by 1/L. The proof of the lemma is provided in Appendix B. To illustrate the difference between the two attention schemes and how different they behave in practice with respect to the "explainingaway" phenomenon, we use the Visual Question Answering setup described in more detail in Sec. 7.1. Fig. 1 shows the minimum attention weight sum values, min j i π ij , achieved by both LNAS and DNAS during training. As predicted by our analysis, the LNAS scheme has the minimum attention weight sum close to 0 (meaning that it explains-away at least some of its inputs), while the DNAS scheme maintains a minimum attention weight sum above its lower-bound value of 1/L = 0.01.

5.2. DOUBLY-NORMALIZED ATTENTION ALLEVIATES MODE COLLAPSE:

Another shortcoming of the LNAS scheme is its tendency to collapse modes * . To study the modecollapsing effect, we analyze the speed of two clusters approaching each other in a 1D scenario using the two attention schemes. Let us assume we have two sets of data, one containing N 0 data points centered at value a, and another containing 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 † .

Due to space limitations, the details of the derivations are provided in Appendix C and only the main results are presented here. If r := N 0 /N 1 , after one LNAS update, the new center distance is:

EQUATION

For DNAS, the new center distance is

EQUATION

)

where q = r+s rs+1 . To better understand the difference between the values of Eq. ( 11) and Eq. ( 12), we plot them on the y-axis against that of r = N 0 /N 1 on the x-axis, for several different a values, see Fig. 2 . 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 LNAS, with the gap getting larger as the cluster sizes get more unbalanced (r = 1).

The mode collapse effect is even more obvious in multi-layer attention. In Appendix C, we show that when the two clusters are unbalanced (Fig. 6 ), the LNAS collapses to a single cluster after 4 steps, while the DNAS maintains two separate clusters. 

6. HYBRID ATTENTION:

In practice, we can also combine the LNAS and DNAS formulations. We investigated doing so with a trainable variable u ∈ [0, 1] that controls the contribution of the attention weights of the two normalization schemes:

π H ij = u π D ij + (1 -u)π L ij ,

where π D denotes the DNAS weights and π L denotes the LNAS weights. We call this combination form to be the hybrid normalized attention scheme (HNAS). The HNAS is attractive because it allows the model to learn, at different layers l, which of the two normalization schemes achieves better results; the u l parameter is trained jointly with the other parameters to improve the representation power of the model and better fit the data. As a side effect, this approach also allows one to visualize how the values of the u 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 Sec. 7.

7. NUMERICAL EXPERIMENTS:



7.1. VISUAL QUESTION ANSWERING:

The goal of a visual question answering model is to provide an answer to a natural language question relevant to the contents of a given image. For our first experimental setup, we use a model similar to the one proposed in (Yu et al., 2019a) , which uses one attention layer to combine multi-view features, i.e., visual features produced by different image processing modules.

Experiment Setup. We conduct experiments on the most commonly used VQA benchmark dataset, VQA-v2 (Goyal et al., 2017) . Our core VQA model uses as a backbone the Pythia architecture (Jiang et al., 2018) . Aside from the backbone network, a crucial factor in the performance of a good VQA system is its visual feature extraction. Currently, virtually all high-performing model use boundingbox visual features extracted by object-detector models trained on the Visual Genome Dataset (Krishna et al., 2017) . Additionally, Yu et al. (2019a) showed that it is beneficial to use bounding-box features from multiple object detectors. In our experiments, we use three object detection models, where each detector generates 100 bounding-box features. All three object detection models are trained over the Visual Genome dataset. The difference between detection models is in the backbone network: 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 VQA model 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 multi-view features using a one-layer attention mechanism, 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 and values. The resulting output feature is a weighted sum of the features according to the attention weights. With this attention scheme, the original 300 bounding-box features from the three object-detection models are transformed into 100 features, which are then fed into a Pythia model. We experiment with two versions of this scheme: one with an attention mechanism using lower-normalized attention (LNAS) and one using doubly-normalized attention (DNAS).

Results and analysis. The results are summarized in Table 1 .

Method:

Test-dev Test-std 100-boxes Pythia (Jiang et al., 2018) 68 Confirming the findings from (Yu et al., 2019a) , we see that using visual features from three object detectors improves performance over using the one from a single object detector (+0.46 on Test-dev). Furthermore, using an attention mechanism over the 3x100 boxes further improves the accuracy over the 3x100-boxes baseline, but the DNAS mechanism achieves a better utilization of the signal provided by the three object detectors compared to the LNAS mechanism (+0.56 on Test-dev and +0.51 on Test-std). Moreover, using the HNAS formulation (Sec. 6) allows us to both visually and empirically confirm the superiority of the DNAS mechanism: as we plot the hybrid weight u in Fig. 3 , the hybrid weight rapidly converges to 1.0 in this scenario.

7.2. HEADLINE GENERATION:

The goal of a headline generation model is to abstractively (as opposed to extractively) generate a short, headline-like summary given a text document. Similar to recent work that achieves state-of-theart results on the task (Goodman et al., 2019) , we use an encoder-decoder Transformer architecture with 12 layers of attention.

Experiment Setup. The standard benchmark for the headline generation task is the Gigaword dataset (Graff & Cieri, 2003) , which 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 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.

The backbone model is a Transformer encoder-decoder containing 12 layers, each with a hidden size of 768 and 12 attention heads. We use an Adam optimizer (Kingma & Ba, 2015) and a learning rate of 2e -5 . We truncate (or pad) the input and output sequences to a fixed number of word-piece positions, namely 128 encoder positions and 64 decoder positions, to accommodate hardware and model-architecture limitations. Similar to (Goodman et al., 2019) , we run experiments starting from both random initializations and BERT-checkpoint initializations.

Results Analysis. The results are summarized in Table 2 . Our Transformer LNAS model reproduces the configuration from (Goodman et al., 2019) 

7.3. 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. BERT (Devlin et al., 2019) established itself as a high-performing contextual representation model. We report here experiments done with the base setting for BERT: a Transformer network with 12 layers of attention, the hidden and embedding size set to 768, and 12 attention heads. Following recent work (Joshi et al., 2019) , we use ngram masking (n ≤ 3), with the length of each n-gram mask randomly selected.

We use the BOOKCORPUS (Zhu et al., 2015) and English Wikipedia (Devlin et al., 2019) to pretrain two contextual representation models, one using LNAS and one using HNAS.

We evaluate the resulting representations by using them as a starting point to finetune for the SQuAD task (Rajpurkar et al., 2018) and the MNLI task (Williams et al., 2018) .

Results Analysis. The results are summarized in Table 3 . The HNAS experimental condition improves over the LNAS condition on all tasks considered. Aside from the numerical improvements when finetuning on the task, we also inspect what happens to the hybrid weight u during HNAS pretraining. In Fig. 4 , we plot the hybrid weight for all 12 layers and find that they are always larger than 0.5, meaning that the DNAS method is preferred during the pretraining optimization. In agreement with the GMM view, the plot shows that the LNAS method starts to have more weight for the layers that are closer to the output layer, approaching a point where LNAS and DNAS are similarly useful. This suggests that the resulting network parameters encode their language representations by making use of the advantages of the DNAS method, positively contributing to the resulting empirical advantage of the HNAS-driven language representations.

8. DISCUSSION:

The attention mechanism of the Transformer, here called LNAS, can be framed under both a probabilistic perspective and an optimization perspective. The new attention scheme we propose based on these perspectives, called DNAS, compensates for certain shortcomings of the LNAS scheme. Together, these two attention normalization schemes can be combined into a hybrid one, HNAS, which we show here to be empirically superior to the original normalization mechanism. Further research needs to be conducted to thoroughly analyze the interplay of the attention scheme with other Transformer components, such as the residual layer.

A CONNECTION TO CAPSULE NETWORKS WITH EM ROUTING:

DNAS is closely 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. ( 5); the new pose matrix µ j in (Hinton et al., 2018) is similar to q i in Eq. ( 5). The difference is that in DNAS, the q i is only updated for one iteration, while the pose matrix µ j is updated until it converges. There is no variance σ 2 i and β i estimation in DNAS, as we assume σ 2 i = 1 and β i ∝ exp( 1 2 q i q i ). One interesting question is whether DNAS could potentially perform better with more iterations for the updates in Eq. ( 6) and Eq. ( 8), as well as updates for σ 2 i and β i . In our experiments, we observed that adding more update iterations increases computational time but does not improve the performance significantly. Furthermore, trying to estimate σ 2 i tends to hurt the empirical performance of the algorithm.

B PROOF OF LEMMA 1:

Lemma 1 For any lower-layer neuron j, the sum of the doubly-normalized attention weights over the upper layer neurons i ξji j ξji is lower bounded by 1/L.

Proof Since i ξ ji = 1, i ξ ji j ξ ji ≥ i ξ ji max i ( j ξ ji ) = i ξ ji max i ( j ξ ji ) = 1 max i ( j ξ ji ) ≥ 1 j max i (ξ ji ) ≥ 1 L

C DETAILS ABOUT THE MODE COLLAPSING ANALYSIS:

To study the mode-collapsing effect, we analyze the speed of two clusters approaching each other in a 1D scenario using the LNAS and DNAS attention schemes. Let us assume we have two sets of data, one containing N 0 data points centered at value a, and another containing 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.

Applying Eq. ( 4) for the LNAS scheme, the new center distance of the lower-normalized attention scheme are:

c L 0 = N 0 t N 0 t + N 1 s - N 1 s N 0 t + N 1 s a = N 0 t -N 1 s N 0 t + N 1 s a c L 1 = N 0 t N 0 t + N 1 s - N 1 s N 0 t + N 1 s a = N 0 s -N 1 t N 1 t + N 0 s a

and the distance between the two updated centers is:

c L 0 -c L 1 = 2N 0 N 1 (t 2 -s 2 )a (N 1 t + N 0 s)(N 0 t + N 1 s) .

Since we have that t = 1, defining r = N 0 /N 1 then gives

c L 0 -c L 1 =

2r(1 -s 2 )a (1 + rs)(r + s) .

(13) 

:

* Note that most multi-layer attention models such as the Transformer avoid such collapsing effect by adding a residual layer after attention.† The attention weights are computed with a Gaussian. But the same result holds with dot product attention, where the inter-attention weight is s = exp( -a, a ) = exp(-a 2 ) and 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.

Figures and Tables:

Figure 1: The minimum of attention weight sum in LNAS and DNAS.

Figure 2: Center distance values after lower-normalized attention (LNAS, blue solid curve) and doubly-normalized attention (DNAS, red dashed curve), as a function of cluster mass ratio r = N 0 /N 1 with different a values (initial distance between centers is 2a).

Figure 3: The hybrid weight heavily favors DNAS for the VQA task.

Figure 4: The hybrid weight favors DNAS in all layers of the encoder (u ≥ .5); LNAS gains more weight for closer-to-output layers.

Figure 6: Mode-collapsing behavior on unbalanced mixture of Gaussian data: LNAS collapses to one cluster after 4 steps, while DNAS maintains 2 clusters.

Test Accuracy on VQA v2.0, over Test-dev and Test-std splits.

Table content: <table><tr><td>.31</td><td>-</td></tr></table>

, and obtains 35.23 ROUGE-L F1 score with random initialization and 36.24 with BERT initialization. Consistent with the previous results, the Transformer HNAS model performs better compared to the LNAS version by +0.5 ROUGE-L F1 points with random intialization and +0.3 ROUGE-L F1 points with BERT initialization. To the best of our knowledge, the ROUGE scores of the Transformer HNAS BERT, BERT configuration establish a new state-of-art on the Gigaword headline generation benchmark. ROUGE F1 scores for headline generation on the Gigaword benchmark.

Table content: <table><tr><td>Method</td></tr></table>

Pretraining with HNAS and finetuning on SQuAD 1.1 and SQuAD 2.0 (F1/EM) and MNLI.

Table content: <table/>

Peer review (iclr2020):
This paper proposes a novel (self)-attention mechanism dubbed *doubly normalized attention*. When taking the attention between keys and query, standard attention proceeds by computing o_i = 1/Z sum_j exp(q_i * k_j) * v_j. This corresponds to normalizing the logits exp(q_i * k_j) over the keys (j). The authors point out that this can lead to an *explaining away* phenomenon: basically, the score for a particular j, exp(q_i * k_j), could be 0 for all i, which means that information coming from position *j* is basically discarded, and the authors argue this is bad. Moreover, they argue that standard attention suffers from *mode collapse* (i didn*t quite get / found clear what this effect corresponds to in standard architectures). To palliate the two problems, they propose to re-normalize the attention weights along the columns. Basically, first computing *responsibilities* (in a EM flavor) p_j = 1/Z sum_i exp(q_i * k_j) , and then computing the attention by renormalizing p_j, o_i = (1 / sum p) sum_j pi_j * v_j. The author show that this corresponds to one step of EM (or basically kmeans if q_i are cluster centers and k_i = v_i are the data points). The authors test their model on headline generation, a visual qa setting and some natural language understanding task, reporting modest gains (0.2-0.4 on SQUAD/MNLI) w.r.t. standard attention. This work has some quite interesting aspects due to the fact that it tries to interpret and generalize attention computations under a more general framework. I appreciated the theoretical work done by the authors, even if sometimes it appeared to me more complex than needed to be. However, this paper also has some weaknesses which prevent me for putting it above the acceptance bar, namely (i) lack novelty of proposed method and (ii) somewhat weak motivation. I*d be happy to increase my score if the author shed light on the problems below. 1) About the lack of novelty: 1.1) The doubly normalized attention appears to me identical to EM routing proposed in the capsule network (ICLR *18, https://openreview.net/pdf?id=HJWLfGWRb). I think the author mention this in the appendix. If it is, what*s the main contribution of the first part of the paper ? I agree that the analogue to Gaussian Mixture Models appears novel but solving GMMs are just a special applications of EM-type algorithms. 2) About weak motivation: 2.1) The authors mention the “explaining away effect” as a problem to solve. It is however unclear as to whether/how the performance of the architecture tested suffered in virtue of this phenomena. Can you justify or quantify this ? Can you attempt at actually computing how much information is discarded by looking at the attention distribution ? You could maybe link this to works analyzing *pruning heads* https://arxiv.org/abs/1905.09418 ? This could give more depth to the paper. 2.2) I didn*t quite get the section on *mode collapse* and why adding residual layers would help in solving mode collapse in standard transformers architectures (as written in a note in the paper). Could you clarify this point ? 2.3) The analogue to optimal transport felt less well-motivated and it was unclear what new understanding was gained from the reframing of self-attention as constrained optimization. I don*t know what to get out of it.