Understanding Attention Mechanisms
Authors: Bingyuan Liu
Yogesh Balaji
Lingzhou Xue
Martin Renqiang Min
Venue:  (2020)
Keywords: Attention
deep learning
sample complexity
self-attention

Abstract: Attention mechanisms have advanced the state of the art in several machine learning tasks. Despite significant empirical gains, there is a lack of theoretical analyses on understanding their effectiveness. In this paper, we address this problem by studying the landscape of population and empirical loss functions of attention-based neural networks. Our results show that, under mild assumptions, every local minimum of a two-layer global attention model has low prediction error, and attention models require lower sample complexity than models not employing attention. We then extend our analyses to the popular self-attention model, proving that they deliver consistent predictions with a more expressive class of functions. Additionally, our theoretical results provide several guidelines for designing attention mechanisms. Our findings are validated with satisfactory experimental results on MNIST and IMDB reviews dataset.

Under review as a conference paper at ICLR 2020
UNDERSTANDING ATTENTION MECHANISMS
Anonymous authors
Paper under double-blind review
ABSTRACT
Attention mechanisms have advanced the state the art in several machine learning tasks. Despite significant empirical gains, there is lack of theoretical analyses on understanding their effectiveness. In this paper, we address this problem by studying the landscape of population and empirical loss functions of attention- based neural networks. Our results show that, under mild assumptions, every local minimum of a two-layer attention model has low prediction error, and attention models require lower sample complexity than models not employing attention. Additionally for the popular self-attention, our theoretical results provide several guidelines for designing attention mechanisms. Our findings are validated with satisfactory experimental results on MNIST and IMDB reviews dataset.
INTRODUCTION
Significant research in machine learning has focused on designing network architectures for superior performance, faster convergence and better generalization. Attention mechanisms are one such design choice that is widely used in many natural language processing and computer vision tasks. Inspired by human cognition, attention mechanisms advocate focusing on the relevant regions of input data to solve a desired task rather than ingesting the entire input.
Several variants of attention mechanisms have been proposed, and they have advanced the state of the art in machine translation (Bahdanau et al. 2014: Luong al. 2015 Vaswani et al. 2017). image captioning (Xu et al. 2015). video captioning (Pu et al. 2018). visual question answering (Lu et al. 2016), generative modeling (Zhang et al. 2018). etc. In computer vision, spatial/ spatio-tempon
attention masks are employed to focus only on the relevant regions of images/ video frames for the underlying downstream task (Mnih al. 2014). In natural language tasks, where input-output pairs are sequential data, attention mechanisms focus on the most relevant elements in the input sequence to predict each symbol of the output sequence. Hidden state representations of a recurrent neural network are typically used to compute these attention masks. The most popular implementation of this paradigm is self-attention (Vaswani et al., 2017). which uses correlation among the elements of the input sequence to learn attention mask.
Substantial empirical evidence demonstrating the effectiveness of attention mechanisms motivates us study the problem from theoretical lens. In this work, attempt understand the loss landscape of neural networks employing attention. Analyzing the loss landscape and optimization of neural networks is an open area of research, and is a challenging problem even for two-layer neural
networks (Poggio & Liao| 2017]] Rister & Rubin 2017 Soudry & Hoffer 2018 Zhou & Feng 2017]
Mei al. 2018b KSotanolkotabii al. [017]] Ge et al. 2017 Nguyen Hein 201 Arora et al.
2018). Convergence of gradient descent for two layer neural networks has been studied in (Allen -hu
al. 2019 Mei al.) 20186; Du et al. 2019). Ge al. (2017) shows that there is no bad local minima for two-layer neural nets under a specific loss landscape design. Unfortunately, these results cannot directly be applied to attention mechanisms, as attention modifies the network structure and introduces additional parameters which are jointly optimized with the model. To the best of our knowledge, our work presents the first theoretical analysis on antention-basee models.
Our main result shows that, under some mild conditions, every stationary point of attention models achieve a low prediction error. We perform an asymptotic analysis where we show that expected prediction on eтror goes to
co. We also show that attention models achieve lower sample complexity than the models not employing attention. We then discuss how the result can be extended to recurrent attention and multi layer cases, and discuss the effect of regularization. In addition,

---

Under review as a conference paper at ICLR 2020
show how attention further helps improve the loss landscape by studying three properties: number of linear regions, flatness of local minima and small sample size training. We validate our theoretical results with experiments MNIST and IMDB reviews dataset.
2 ATTENTION MODELS
Attention mechanisms are modules that help neural networks focus only the relevant regions of input data to make predictions. To study such behavior, we analyze different types of attention models. We start with naive global attention model. Then we analyze the most popular self-attention model and discuss the extension to recurrent attention in Appendix. In the naive global attention model, we consider dataset D { yi f1=1+ 3 IRP, yi I, where the output yi depends only on certain regions of input æ;, i.e., y/i = l"(a" @ æi), where a" is an attention mask, and (*)() is the ground -rouud-truth function that is used to generate the dataset and the vector a* E [0, 1]" The set of entries {a, la # 0} corresponds to the relevant region of the input, while the complementary set {ağ |a; 0} corresponds to the irrelevant region.
We consider a two-layer ReLU-based neural network to approximate the function f". The network architecture consists of a linear layer followed by rectified linear units (ReLU) as a non- -ineaiittitty and a second linear layer. Denote the weights of the first layer by w(1), the weights of the second layer by 2x(2, and the ReLU function by ((-). Then the response function for the input x can be written as f(x) (2yT Φ( (w(()), x))). We call the above function "baseline model" since it does not employ any attention. To incorporate attention, we introduce the attention mask a as additional neural network parameters. The attention model we use can be written as:
f(x) w (2)T (((w (1) x a))
(1)
In this paper, we focus on the regression task which minimizes the following loss function:L E(m.))=Dl (f) - v12.. While we present analysis on the regression task, our theory can easily be extended to classification tasks well.
After a thorough analysis of global attention, we analyze a more practical self-attention setup, which comes from the transformer model proposed in Vaswani al. (2017). The input æ; (at....,,ff x Rexp, where æi. are t-dimensional vectors. Each æ; corresponds to independent sentences for i = L.... ...... and x 's are the fixed dimensional vector embedding of each word in sentence æ. R E ig, are the query and key matrices, and w/ Rd.xt is the value matrix. For each input 2;, the key is calculated the query vector is computed as: Q = (w xwx= E R gixd for z = 1,...,p. The value matrix V = wv V x, E Xp, Then the self-attention w.r.t to the th vector in the input æ; is computed as:
K, (w=xx;))) E Roxd,; For 2thh vector in the input,
Q.K7 VTg
self(s) (x w wx = softmar
a;
for 2 I.....p. .. And a; gots (a (a (e))(()))
sels())). This self-attention vector represents the interaction between different words in each sentence. The value vector for each word in the sentence self(z) Rdd. This value vector then passed a 2-layer ML.
af can be calculated as V.² = Va Vagel(a)
parameterized by (1) E Rdd.. and (2) E Rdd11 resulting in the following generative model:
(2)xT (1) vec(w" (ooewW,,aaee a;a,t"/)) +
(12))o(((l (L)
(3)
where vec(-) represents the vectorization of a matrix. We also discuss the extension to recurrent attention model and multi-layer self-attention model in Appendix due to the page limit.
Note that naive global attention not widely used in practice, because here we assume the attention mask is globally fixed for all points, but not function of input. In real world application, the attention weights depend on the input, such as the self-attention framework we just introduced. Despite the limitation of naive global attention model, it is fundamental building block of attention models, and needed to be analyzed first for two following reasons.
First, this fixed attention shares the core idea of attention: There is specific intrinsic structure in to different input
data, and this intrinsic structure requires that we should assign different weights features accordingly. And this weight assigning strategy should be learned from data. In naive
2

---

Under review as a conference paper at ICLR 2020
global attention, attention weights a are parameters themselves; In the self-attention model, the attention weights a depend on a function parameterized by value/key/query matrices, and we learn these matrices attention parameters. And for both global and self atff-attention, we jointly learn the network weights(w' ( 2) and attention parameters(a in global attention, value/key/query matrices in self-attention) at the same time. Therefore this naive global attention is a good starting point for analyzing attention mechanisms.
Second, we can gain helpful insights on attention by analyzing the global attention case. The number of non-zero elements of in global attention represents both the size of attention parameters and the sparsity level of attention. In the standard self-attention model, size of attention parameters is determined by the size of value/key/querry matrices. And the sparsity level is how many words we allow one word to attend to. By studying the effect of this quantity, we can have a better understanding of how the sparsity and parameter size of attention affect the model performance and sample complexity. The detailed discussions can be found in Section 3.
3
LOss LANDSCAPE ANALYSES ON ATTENTION MODELS
In Section 3.1
we analyze the loss landscape for the the naive global attention model, which is (2)T ((w(() xz 2 a)) y2)²
defined as
1
72
min (1) (2) JaES 2n1
where w(1) RR w((2 a R", and S is the parameter space of a. Section 3.2 extends the loss landscape analysis to the self-attention model with a = f(x). Section 3.3 discusses the extension to recurrent attention model. Section (.4tioon3.4444 discusees the effect of regularization and how the result can be extended to multi-layer networks.
To approximate the non-differentiable ReLU o, we use the softplus activation function (>(x), i.e.,
1 (x) = log(1 [ + +
Note Φt converges to ReLU as r Glorot et al. 2011). Theoretical results with (x:) will hold for any arbitrarily large For ease of notation, we still use ((x) to denote the softplus function.
3.1 ASYMPTOTIC PROPERTY OF NAIVE GLOBAL ATTENTION MODEL
We first analyze the asymptotic prediction error of local minimum with large sample size. Let the covariance matrix of ((w(1), x @ a) be (Zz)is = cou(o:((w x O (y x O a))) and Mxxx) = x. (2)() @ o' (LTT x @ a)), where I. represents 1tth feature in æ, and og is the attention mask for Z.A. & ((T()) x @ a) is the first order derivative with respect to the value in Φ(-) and it belongs to Rdd. Let u (w (1 x @ a)) - R(yla:)). Before proceeding, we introduce several necessary assumptions:
(A1) x; are i.i.d with ||ar <. for 1,2,
(A2) There exist C1, C2 such that ||w("|| < C1 and (| < C2 for any w()), (2) in S.
(A3) The output y can be specified by the two-layer neural network up to an independent sub- Gaussian error with variance o², i.e., there exists a set of parameter (a" w(()), (()))), such that
V: (n ((((())) 2n; @ a*) a²)) ))++ + 6,, where e; ~ subG(0, C+) for i=1,2.... with æ; Il 6,, And [[ (1)+ n. @ a*) a))))l12 S C4
(A4) ||allo < so such that so S p, which represents the sparsity of the attention model, and 0 S a, < 1 for any i= 1, .p. (A5) Amin (E)) Co. and when the estimation of is inaccurate, i.e, E(I ((ll@((())) ()) "x a) o' (1, @ a)))l2) 2 0()), then there exists feature s. and for the 4hh element of he( it satisfies sd(he (x.+)) = O(1) and cor(u, he (x.x.)) O(1) for w in the form of Σ r + o(7).
(2) ()))))(((())))))))) (1) a))) E(y) o(7/ /=0)- (A6) E(
3

---

Under review as a conference paper at ICLR 2020
(A7) The sum of the weights ||a|| 1.
These assumptions are mild when the dimension of x: is large. The justification of these assumptions are provided in the appendices. Given these assumptions. we show in Theorem that the sample complexity required to converge to a good minimum is reduced with attention mechanism.
Theorem 1. Under (Al) to (A6), for any >> O, suppose
$00 72
log( 5077 (pd + p
T
where 1l - CrC Then with probability tending to 1, any stationary point (ä, ( objective fanction satisfies the following prediction error bound:
((()) of the
[(w (2)T 6((20)) (1) a) Blyle)) 72
Remark: The sample complexity bound of Theorem provides helpful insight for understanding attention mechanisms. With the sparsity structure of attention mask a, attention mechanisms constrain the parameters in a smaller space, thus reducing the variance and the covering number. This leads to lower sample complexity compared to the baseline model not employing attention. It is straightforward to calculate the sample complexity bound for the baseline model. To achieve the same error bound, we substitute So with in the bound, and this results in much larger value. When 8o is fixed, we can see up to log term, prediction error is proportion to n n-1/2, which is the optimal rate of convergence in regression. This shows that the bound is tight in this aspect.
Next, extend Theorem to the attention model with an additional sum-to-one constraint (A7). The discussion of the following Corollary provided in the appendix.
Corollary 1. Under (AI) to (A7), for any y > O, suppose n Z SGGV" log( ( (pd + p + d), where
" CzC., Then with probability lending any stationary point (ā, (((1) (((2)) of the objective function (4) satisfies the following prediction error bound:TE (ú (2) (((((1)) x a)) E(y)==))² x y
3.2 ASYMPTOTIC PROPERTY OF SELF-ATTENTION MODEL
In this section, we extend our previous analysis to the self-attention model. In self-attention, the attention mask is more fixed globally, but instead a function of the input. We begin by analyzing self-attention model with a known attention function f(-), in which the weight a is not optimized together with w. Similar bound can be derived for the following model:
min (1 30 ES 2n
Σ (wwll (w (2)T (lww)) x, O f(x))) - -1)²
(5)
Proposition 1. Under (AI) to (A6), suppose that a = f(x). For any 7 > O, given the sample size n 2 C00 log( sall (pd + d), where 1j = C.C2Ca, with probability converging to any
7/3
stationary point (1) to (2) of the objective function equation 5 Ssatisfies that: (2) (((1((1)))))))))))))))))))) xQ a)) - - )² S 22
Proposition implies that if the self-attention mask can precisely computed, global attention results can be extended to self afffattention ones. However, the function f(-) is not necessary known, and needs to be learnt in real world applications. Therefore the self-attention setup as we introduced in section 2 is more desired in real-world setting. Denoting w (w()), (2)) w w w w)), the two-layer self-attention model can be estimated by:
n
Σ
(2)T (2TTT) (((w()) vec(w" æ, (() y1)
min t0 2n
(6)
We now introduce necessary assumptions for analyzing self-attention model.
(A8) There exist Cs and C7 such that | S Cs. w K ||r S Co. 11/0 1/r Cr.
4

---

Under review as a conference paper at ICLR 2020
(A9) The output y can be predicted by the two-layer network (3) with an independent sub- Gaussin error with variance o², i.e, there exists a set of parameters (a*, w()),, wl(²)) such that yi = w (2)xx ((ww!))) vec(w" a,a; (()()((())))))) ))) + C, where a is calculated by 2 G ~ subG(O, Cf) for i i= 1,,, with æ; 1 6;.
(A10) We assume (A5) and (A6) holds, substituting at with veolw> a. (el and h (x..)) (vee(w) V (,,aa(f)) (ux(2)))))) C o' (w(t), a4 3 a)), where (vee(w) (ww ) is the k-th element of the value matrices.
The assumption (A9) states that self-attention model can correctly predict the conditional mean (z.). |(ys/z,). Note that (A9) encompasses the models used in practice such as the transformers. (A10) is parallel to (A5) and (A6). Under these assumptions, we can obtain its sample complexity as given by following theorem:
expressive class of models than (A3), which includes
Theorem 2. Under (AI), (A2), and (A8) to (A10), for any > 0, given the sample size:
C/ 2
log( CsCoCTM
(pd, + d 2dgt + dut)
72
?
where 17 = C2C with probability tending to 1, any stationary point (ii(1, (2)), u i ;x of the objective function 6 satisfies that: (2( o0 (((ū(1), vec(w" x;a; B()l=))) S 72
shows that with the help of self-attention, we can achieve consistent prediction under more expressive class of models (assumption (A9)) which considers the interactions between vectors in data. It is worth pointing out that both global attention model and baseline model do not have consistency for the class of models beyond the ones stated in (A3). In other words, consistent prediction on the data distribution generated from equation 3 using baseline and global attention models requires introducing larger parameter space, for example, using more layers of network or more units each layer. Self-attention model, on the other hand, achieves the more accurate estimation by constraining the parameter space and input space. And parallel to the sparsity level so in Theorem 1, a proper choice of value/query/key matrices can help reduce sample complexity. If a sparse attention(i.e. one word should attend all words, but only some relevant words), the sample complexity can also be further reduced similar with Theorem 1.
Remark:
Theorem
3.3 EXTENSION TO RECURRENT ATTENTION NETWORK
Sample complexity analysis can be extended to recurrent neural networks. This is included in the appendix. The key messages from our analysis include: (1) A good design of recurrent framework can help the network converge to good stationary point with small sample complexity, and (2) An arbitrarily complex framework increases the sample complexity. Since in real world, the optimal recurrent framework is unknown, careful design choice has to be made for obtaining good sample complexity.
3.4 DISCUSSION: REGULARIZATION AND BEYOND 2 LAYERS
So far our analyses provide some theoretical justification on how attention mechanisms help learn superior models. Furthermore, our analysis also suggests proper regularization is helpful in training an attention model. An L regularization on attention weights and &2 regularization on network weights are effective in reducing the sample complexity. We also find that imposing constraints and regularization on network weights can help remove sharp minima, and keep flat minima with good generalization. Detailed discussions on regularization are provided in Appendix Section D.
Also, Theorem 1 and 2 can be extended to multi-layer attention network, under the assumptions parallel to (A8), (A9) and (A10): There exists a correct multi-layer self-attention network can specify the model, and the bias and gradient with respect to network weights are not uncorrelated. Under these assumptions, we provide sample complexity bound for multi-layer self-attention models. And all the discussions and insights are applied to multi-layer models. Explicit assumptions, discussions and theorems are provided in Appendix Section D.2. We avoid multi- setting in main context because it leads to comp-ccmlle derivations and assumption justifications, and will distract readers from main idea of the paper. We believe two-layer models are representative enough to provide theoretical evidence on why attention reduces sample complexity.
5

---

Under review as a conference paper at ICLR 2020
4 ON IMPROVING THE LANDSCAPE STRUCTURE
In this section, we further investigate three additional properties on how attention mechanisms improve the landscape of neural networks. First, we show that in global attention model, attention mechanisms reduce unnecessary number of linear regions and maintain a low approximation error; Second, we show that flatness properties of minima are retained when attention mechanisms are used for both global attention and self-attention. Furthermore, our analysis indicates that for attention models, smaller sample size suffices converge to good minima which generalize well in prediction. Finally, we show that the perfect in-sample prediction on small sample size is also achieved in attention networks for both global attention and self-attention.
4.1 ON THE NUMBER OF LINEAR REGIONS
We first study how attention mechanisms affect the number of linear regions (Montufar al. 2014) in wide two-layer network, where the number of units in the hidden layer is larger than the sparsity of the attention mask matrix.
Theorem 3. Assume ||allo= so. which is the sparsity of the mask matrix, and the number of units in the hidden layer n, > Sg- Then the maximal number of linear regions of the function by two-layer fully connected neural network with ReLU activation function, is lower bounded by 50 100
The theorem implies that when appro- priate attention mechanism is used, the number simpler
Remark:
of linear regions reduces leading landscape, yet the approximation error remains small. This leads to lower sample complexity for achieving desired prediction error. More detailed discussion can be found in the appen- dices. The result of Theorem also applies to the self-attention with different attention spar- sity(i.e. allowing how many words allow one word to attend to).
Figure 1: Number of linear regions in log scale v.s. sparsity
4.2 ON FLATNESS/SHARPNESS OF MINIMA
Many recent works, such as Keskar et al. (2016). argue that flatter local minima tend to generalize well. However, in a recent study, Dinh et al.) (2017) observes that by scale transformation, the minima which are observationally equivalent, can be arbitrarily sharp, and the operator norm a Hessian matrix can also be arbitrarily large. We will show that this fact also holds for the global attention mechanism, if no constraint parameter (ā, (1), (2)) ) is imposed. Here we introduce the definition of e-flatness in Hochreiter Schmidhuber (1997).
Definition Given > O, minimum 0, and loss L, C(L, the largest connected set containing O such that VO' E CCL 0). L(0')= S L( + e, and its volume is called the e-flatness.
In the following Theorem, we analyze the flatness of stationary point for both naive global and self-attention model.
Theorem 4. (a) Consider the two-layer ReLU neural network with naive global attention in Section 3.1:
yi = wo (2)+T (((ww()) (1)+ x0, + aa)))
and a minimum 0 - (ä, (1))) (2) satisfying that à + O. (6(1) * (, + O. For any c >> O, C(L,0, 0,e) has an infinite volume, and for any M > O, we can find stationary point such that the largest eigenvalue of L(0) is larger than M;
(b) Consider the two-layer ReLU neural network with self-attention mechanism stated in Section 3.2:
y1 (2)-T d((ww(1)) vec(w> a,aa (r(aa

---

Under review as a conference paper at ICLR 2020
and a minimum 0 (wb(1, x
w wº, ) satisfying that w * for (1), (2), V, Q. K. For any E > 0, C 8 e) has an infinite volume, and for any M >> 0, we can find stationary point such that the largest eigenvalue of L(0) is larger than M.
Theorem indicates that property on flatness of minima is maintained when attention mechanism is applied, and there exist good sharp minima, coinciding with the observation Dinh et al. (2017). However, there is guarantee that all sharp minima are good in generalization. Revisiting our analysis in Section 3 the restriction on the parameter space (see (A2)) help remove these sharp minima. Specifically, (A2) provides an upper bound on the magnitude of (a, w("), 2x22) and (A5) bounds the magnitude of )((w((), x C a)) from below. These constraints control the parameter space and remove all sharp minima generated in Theorem which on or 022 goes to infinity. The &2 bounds in (A2) can be achieved through a proper & regularization (See Section 3.4..
4.3 ON SMALL SAMPLE SIZE
We conclude by studying the local minima of wide neural networks in small sample regime. (Nguyen & Hein| 2017b) proved that two-layer neural network model can always achieve perfect empirical estimation error when the same size is small. Here, we extend this result for global and self-attention model. The discussion of Theorem 5 is deferred to the appendices.
Theorem 5. (a) For naive global attention model, if rank(ol (w(()
② a) 11)..)) n. Then
every stationary point (ā, (()), ( of the objective function (4)
is a global minimum;
(b) For self attention model, if rank(ol (io(1), vec(w> a;a; stationary point (wo(1) (0(2),10) ,0,,00 wx of object function a global minimum.
self ) n. Then every
EXPERIMENTS
5.1
SAMPLE COMPLEXITY
5.1.1
GLOBAL ATTENTION MODEL
Theorem I proves that attention models require lower sample complexity than baseline models, i.e. attention models require fewer samples to achieve the same test error as baseline models. This result is validated empirically in this experiment. To mimic the assumptions of Theorem []] we consider ground truth two-layer neural network G* (æ) is formed using random weights as inputs. The network G maps the input vector æ E 256 to 10-dimensional output. The input vector æ is randomly sampled such that each element is drawn i.i.d from (00, 1). An attention mask a* is then constructed with & randomly chosen elements as and the rest as The ground -oound-truth labels are generated from y G (x a*).
Table I: Experiments on sample complexity for global attention model. np denotes the number of trainable parameters in each model. All results are averaged over 5 runs
Number of training samples
Test loss:
Test loss:
Test loss: regularized attention (My 34442)
baseline (Mp 34186)
attention (np 34442)
10000
2.1063
0.5484
0.0143
14000
0.4109
0,0382
0.0107
16000
0.1811
0.0211
0.0100
18000
0.1072
0,0163
0.0122
20000
0.0769
0.0101
0.0098
50000
0.0511
0.0060
0.0072
To test the sample complexity, we generate multiple datasets, each containing 10k, 14k, 16k, 18k, 20k and 50k unique samples respectively using the scheme mentioned the previous section. A common test set of 5000 samples is created to evaluate each of the models. A regression model then trained on each these datasets. All models are trained with SGD optimizer with a fixed learning rate of 10-3 Table 1 reportt test emrrr for baseline and attention models 400k iterations
7

---

Under review as a conference paper at ICLR 2020
the number of training samples vary. We observe that attention models need fewer training samples than baseline models to achieve a desired error. For instance, to attain the desired error 0.07, attention models need 14000 samples, whereas baseline models need 20000 samples. We would like to point out that improvements obtained by attention models is not because of increase in model parameters. As shown in Table T the number of parameters in baseline and attention models are comparable. Hence, the performance gain is solely due to the attention mechanism.
Regularization: In section 3.4, discuss how regularizing attention vector helps obtain better attention model. To empirically validate this claim, we train a model with L1 regularization on the attention vector. Same experimental setting as Section 5.1 is used, regression model is trained using a two-layer neural network. We add the L1 penalty on the attention vector to the objective: Treg = E; las]- The results are also shown in Table [ We observe that models trained with I, regularization achieves better sample complexity than its unregularized counterpart. We also observe the faster convergence when the models were regularized, as shown the appendices.
5.1.2 SELF-ATTENTION MODEL.
We extend the sample complexity experiments self-attention model discussed in Section 3.2 Since our model is tailored towards natural language tasks, we consider the problem of sentiment classification on IMDB reviews dataset (Maas et al.) 201I). Note that our analysis needs fixed length sentences which hardly holds true in any NLP dataset. So, we zero-pad all our sentences to make their length equal the maximum sentence length in the dataset (2142 for IMDB reviews). For every input word, we first obtain their corresponding pre-trained GloVE embeddings (E IR 100) which is then passed the neural network. As baseline model, we flatten the input to one large vector of dimension 2142 x 100 and pass it to 1-hidden layer MLP with 256 hidden units. For self-attention model, we use 8 E IR. 10 100 u E TR 1Q100x100 w V E TR. 10100x1112
Once the attended features computed per equation 2 is passed to 1-hidden layer MLP with 256 hidden units as the baseline model. All models were trained using Adam optimizer with learning rate 10 10-3 This was the setting that gave the best performance among optimizer and learning rate configurations comparison of sample complexity of baseline model and the self-attention model provided in Table We clearly observe that self-attention model requires low sample complexity to achieve the same error as the baseline model. To test if improvements are obtained in attention model due to increase in model parameters, ran the baseline model with twice the number of parameters the self-attention model. Even with large parameter size, baseline model performs poorly compared to self-attention models.
tried.
Table 2: Experiments on sample complexity for self attention model. up denotes the number of parameters used in each model. All results are averaged over runs.
Number training samples
Test accuracy (in %)
Baseline
Baseline
Self-attention
(np = 57335913)
(n2
111744625)
(np 57365913)
875
63.38
64.38
70.52
1750
64.48
65.21
83.10
3500
63.51
64.32
86.14
5250
71.92
69.51
87.96
7000
75.49
76.32
87.10
8750
78.64
78.66
88.58
13125
79.97
79.72
88.98
17500
80.85
80.52
88.59
5.2 CONVERGENCE PLOT
This experiment studies the convergence of the empirical risk of the baseline and attention models. modified MNIST dataset called NoisyMNIST is constructed where the images of digits from the MNIST dataset is embedded in noise as shown in Panel (a) of Figure 2 We consider the classification task to predict the labels of the digit in each image. Since the ground truth label depends only on certain regions of input, MNiST mimics the data generating process consider in this paper.
8

---

Under review as a conference paper at ICLR 2020
6.00
200
(a)
(b)
(c)
Figure Visualization of NoisyMNIST (a) and convergence plots for baseline and attention models on NoisyMNIST, including (b) plot of test loss and (c) plot of test accuracy over iterations.
For baseline model, we train a two-layer neural network with 128 hidden units using stochastic gradient descent. For attention models, the input tensor is multiplied element-wise with a learned attention mask a as in equation]] | The attended input is then passed to the two-layer network. The convergence plots for baseline and attention models are plotted in Figure 2 We observe that the attention model converges faster, and to better minimum than the baseline model. Similar behavior observed for different learning rate and scale configurations shown in the appendices.
5.3 HESSIAN PLOOT
We study the Hessian matrix of the loss surface to validate the loss land- scape of attention models. The same classification setup the previous ex- periment is considered. The follow- ing two-layer neural network architec- ture employed: 576 -> 16 10. The Hessian matrix of loss landscape about the computed minimum was, spectively, computed for the baseline and attention models, and their top sorted eigenvalues are plotted in Fig- ure Baseline models exhibit higher eigenvalues than attention models, so the loss landscape of attention mod- els are flatter than the baseline mod- els. Since flat landscapes lead better generalization, models with attention generalize better than models without attention as shown Scction 5.2
Spesvola
300
Pigur inbio
Figure 3: Top 100 eigenvalues of the Hessian matrix for baseline and attention models.
CONCLUSIONS
In this paper, we study the loss landscape of two-layer neural networks on global and self attention models, and show that attention mechanisms help reduce the sample complexity and achieve consistent predictions in the large sample regime. Additionally, by analyzing the number of linear regions, the loss landscape under small sample regime, and flatness of local minima, we demonstrate that attention mechanisms produce well behaved loss landscape that leads to a good minima. Extensive empirical studies on NNIST dataset and IMDB reviews dataset validate theoretical findings.
REFERENCES
Zeyuan Allen-Zhu Yuanzhi Li, and Zhao Song. A Convergence Theory for Deep Learning via Over- Parameterization. Proceedings of the 36th International Conference on Machine Learning, ICML '19, 2019.

---

Under review as a conference paper at ICLR 2020
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparamettriizatonn. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssan Stockholm, Sweden, pp. 244-253, 2018.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. CoRR, abs/1409.00 0473, 2014.
Andrew Barron and Jason M Klusowski. Approximation and estimation for high-dimensional deep learning networks. arXiv preprint arXiv: 0009,03990 2018.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 Proceedings of Machine Learning Research, pp. 1019-1028, 2017.
Simon S Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Proceedings of the 36th International Conference on Machine Learning, ICML 11, 2019.
Rong Ge, Jason D. Lee, and Tengyu Ma. Learning one-hidden- layer neural networks with landscape
design. CoRR, abs/171 ftos//7111.00 1.05511, 2017.
Xavier Glorot, Antoine Bordes, and Yoshua Bengio. Deep sparse rectifier neural networks. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp. 315-323, 2011.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On latge-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.0480 2016.
Jiasen Lu, Jianwei Yang. Dhruv Batra, and Devi Parikh, Hicrarchical question- image co-attention for visual question answering. In NIPS, pp. 289-297, 2016.
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention- based neural machine translation. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 1421. Association for Computational Linguistics, 2015.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-1500 Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL http
7/ h/www.aclwett www.aaclweb.org/anaee www.etthhhoooooooooooeooee
Song Mei, Yu Bai, Andrea Montanari, et al. The landscape of empirical risk for nonconvex losses. The Annals of Statistics, 46(6A):2747-2774, 2018a.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 11533))E7665--77771 2018b. ISSN 0027-8424. doi: 1ttp://ppaa.18679911111
Volodymyr Mnih, Nicolas Heess, Alex Graves, and koray kavukcuoglu. Recurrent models of visual attention. In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems 27, pp. 2204-2212. Curran Associates, Inc., 2014.
Guido F Montufar, Razvan Pascanu, Kyunghyun Cho, and Yoshua Bengio. On the number of linear regions of deep neural networks. In Advances in neural information processing systems, pp. 2924-2932, 2014.
10

---

Under review as a conference paper at ICLR 2020
Quynh Nguyen and Matthias Hein. The loss surface of deep and wide neural networks. In Proceedings of the 34th International Conference on Machine Learning (ICML 2017), Sydney, NSW, Australia, 6-11 August 2017, pp. 2603-2612. 2017a.
Quynh Nguyen and Matthias Hein. The loss surface of deep and wide neural networks, arXiv preprint arXiv:/704.08044 2017b.
Tomaso A. Poggio and Qianli Liao. Theory II: landscape of the empirical risk in deep learning. CoRR, /e//1703.09333 2017.
Yunchen Pu, Martin Renqiang Min, Zhe Gan, and Lawrence Carin. Adaptive feature abstraction for translating video to text. In Proceedings of the "hirty-Seconde AAAI Conference Artificial Intelligence, (AAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 7284-7291, 2018.
Blaine Rister and Daniel L. Rubin. Piecewise convexity of artificial neural networks. Neural Networks, 94:34-45, 2017.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D. Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. CoRR, Abs//707.0493 2017.
Daniel Soudry and Yair Carmon. No bad local minima: Data independent training error guarantees for multilayer neural networks. arXiv preprint
2016.
Daniel Soudry and Elad Hoffer. Exponentially vanishing sub-optimal local minima in multilayer neural networks, 2018
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 5998-6008, 2017.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint 1rXiv:1011.3021 2010.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. Lille, France, 07-09 Jul 2015. PMLR.
Han Zhang, Ian J. Goodfellow, Dimitris Metaxas, and Augustus Odena. Self-attention generative adversarial networks. CoRR, htss/8805.0831 2018.
Bolei Zhou, Yuandong Tian, Sainbayar Sukhbaatar, Arthur Szlam, and Rob Fergus. Simple baseline for visual question answering. arXiv preprint arXiv: 2015.
Pan Zhou and Jiashi Feng. The landscape of deep learning algorithms. CoRR, abs/1 705.07038, 2017.

---

Under review as a conference paper at ICLR 2020
Appendices: Understanding Attention Mechanisms
In appendices, Section A presents the extensions of our analyses recurrent attention model; Section B provides detailed justification of our assumptions: Section C discusses the implications of theorem results in more detail(Theorem discusses the regularization effect in attention networks and potential extensions beyond 2 layers and to CNN/RNN. Finally We provide proofs and additional experiment results in section and separately.
Corollary Theorem and Theorem in order): Section D
A EXTENSION TO RECURRENT ATTENTION MODEL.
Here we consider analyzing the representative recurrent attention framework in Bahdanau et al. 2014)) In the recurrent attention network, still have each data point æ; (at , at)) RE corresponding to words with (-dimensional embedding. Then the generative model can be represented as:
y1 = (2)T (w
Σ aax;))a²)) +
j=1
Analogous to NLP setting, a(x;) is unknown function mapping æ; to t-dimensional vector, where aı;;); represents the effect of the jth word in the sentence for point i. To simplify the model, we use data features themselves as their annotation, then for time stamp k 1,..., The recurrent attention model estimates ax;; as follows:
8k f(sk-1, --11: 1): Ckj score(sk-1.
0kj Σj=1 Ekj
Ekj
Ch Σ
akjo
j=1
yk w (2)T (1) .c)))
where score(-) is the scoring function representing how well the inputs around position j and the output at position match. can be dot product or MLP. And f()) is the function to update 8k- Suppose the parameter inside these two functions are w, and wy with number of parameters as da and df accordingly. Here we show that when these two functions are expressive enough, recurrent attention network will also have sample complexity bound parallel previous sections. Here we introduce necessary assumptions.
(A11)
The output y can be predicted by the two-layer network with an independent aab-Gaussian error with variance o², i.e, there exists set of parameters (w) (2))) such that yi (2)))TT (1n())), Ej= a())) e,, where 6 ~ s0,G((, Cf) for 1, ...m, with I; G.
(A12) Suppose (A5) holds when we substitute O a with Σj=1 a(x;)aa? and h,(x,) a a(xx)) )al, (2)) @ ((w()) 2 @ a))).
(A13) We assume |/all/2 S Cs and ||wy|ll S Cg.
(A11) and (A13) are parallel to (A1) to (A5) in global attention case. They can be justified similar as them, which is discussed in Section B of appendices.
Now we can provide following sample complexity bound extended from previous sections.
Theorem 6. Under (A2))))))( to (A13), there exists sufficient large T, for any y > 0, suppose
GF 72
((dC log( nCsS09 (td + d + du + d
n n>
where 1 = C1 C2C such that if there exist stationary point(s), then with probability tending to any stationary point (1) 2)) (
, w.) satisfies the following prediction error bound:
E(w (2) (((())
Σ (a(x))=)))) -
E(yl=))² 2
12

---

Under review as a conference paper at ICLR 2020
This theorem provide sample complexity bound for recurrent attention network. It holds when such good stationary point' exists. It also shows a trade-off between a complicated recurrent attention network and the sample complexity bound. If f-) and a(-) are properly selected, they will be sufficient expressive to obtain good stationary points, and also the number of parameters du and dz will not be too large. In this way. an ideal sample complexity bound to these good stationary points can be achieved as theorem says. However, with a over complicated design in these functions, the sample complexity bound will be large; With a over simple design, such good stationary points don't 't exist. It is parallel to a trade-off between approximation error and estimation error in learning theory. The theory implies a good design of recurrent structure will help achieve an optimal sample complexity in recurrent attention model.
Remark:
B JUSTIFICATION OF ASSUMPTIONS (A1)-(A6)
In this section, we discuss the rationality of Assumptions (AI) to (A6). Note that (Al) to (A6) are required to prove the main result in Theorem and they have also been studied by Keskar et al. (2016); Dinh et al. (2177); Mei et al. [2018]] b): Nguyen & Hein (0077)). In what follows, we will show in details that (A6) are all reasonable assumptions.
First of all, (A1) and (A2) require upper bounds on the input æ; and &g bound for network weights. It is a standard assumption in landscape analysis (Mei et al.] 20188 b), and also it is crucial to remove sharp minima which may not generalize well (Keskar et al.] 2016 Dinh et al., (See remark after Theorem
4 These assumptions can be achieved through regularization.
(A3) requires that this two-layer network are rich enough specify the condition mean Ey, (z,). It has been studied that general bounded functions with Fourier representation on [-1, 1] can be well approximated by the defined two-layer network(Barron & Klusowski 2018).
(A4) requires a sparse structure on a"; otherwise, the model would be equivalent to the baseline model, simply just choose attention masks all being 1. And ||a|lce S requires the attention weight ranges from 0 to 1.
(A5) is technical condition for our analysis of stationary point. (A5) includes two parts. Both of them hold naturally when dimensionality is large. Firstly, the lower eigenvalue bound assumes o( ((w)(1) I; @ a)) is not degenerated. If eigenvalue assumption is violated, the model is equivalent to a network with fewer number of units in hidden layers, and we can study this equivalent degenerated one instead; This assumption also guarantees us to remove sharp minima(Keskar et al., 2016)] Dinh et al. 2017). same as (A1) and (A2). Secondly, (A5) assumes when o( (((w(), x. O a) is not well estimated, there exists an "active feature'. The correlation between this feature and bias cannot be
cancelled by the direction of a specific linear combination of of (((w(1) x; O a)). Intuitively, this
assumption says that the correlation between (x(xx)) and E(y(x.+) cannot be fully explained by a fixed linear combination if there is some systematic bias in oi (( x O a)). Since there is systematic bias, it is reasonable to assume this systematic bias cannot be uncorrelated with all the directions spanned by h4 (a:). This correlation assumption between h:(a: (mx.y) and parallel to the full column rank condition in Nguuen& & Hein 2017b). Considering all active terms in he (z..),,, they span a larger space comparing to o0 (w æt O a)), considering dl is a fixed dimension. Thus it is a natural assumption if we have reasonable large dimensionality and the model doesn't 'tt degenerate. What's more, with required large sample size, we can also straight forwardly evaluate this assumption by checking empirical correlation, and avoid this type of bad minima through random initialization and proper gradient descent type algorithm(Allo Allen -Zu et al. 2019). And it also will not affect the key structure of our proof.
Specifically, for the sd(ha = O(1) part, since hy(x..)) = (w (2) o ((u az O ()))),
where o' ((w(1), Ii O a)); (1)) with positive probability otherwise the network always predicts zero, therefore we have sd sd(h.zz.)) = O(1) as long as sd(2.)) = =(1) and a = O(1.. ). For the correlation assumption cor(u. (u, h.(xx))) = 0(1), if this assumption doesn't ||cor(u, ((x)))) o(1) for all k = 1, By linear combination, for any vector 2 = (zi,..., zx) with bounded l2 norm, we have cor(u, z (z) O (ww)" oo ((w()) x a)) o(7) for all t = I,...,d, d. We (x @ a)((w( (2) o o (1 x @ a)) can be represented as the combination of term 2T(m @ a) (100 @ of ((w(1) xr a)). By the arbitrary choice of
know for any ((()) E Rpxd the term (0
13

---

Under review as a conference paper at ICLR 2020
2, it means that all the directions of x are almost not correlated with u. When p and So
high-dimensional, (x (z @ a)( (2) @ ø ((w(l) a @ a)) spans all the directions as p -> oc, thus
there must be some direction is correlated with u. Therefore it is reasonable to assume O(1) correlation occurs.
(A6) assumes that we fit model with expectation close to truth. We can always achieve this by centralization.
With these assumptions, we just make sure we remove unnecessary minima, such that we can concentrate on analyzing the behaviour of good stationary points of our interest. Assumptions have no influence the main idea of our following theorem: showing the required sample complexity to approach a good minimum is reduced with attention mechanism.
The assumptions are also related to the question that when attention mechanism should be applied network. (Zhou et al. 2015) shows that attention sometimes can be badly used in certain cases, which shares the similar philosophy with analyses. For example, theoretical analyses depend on the assumption (A3), that is, the model can be correctly specified when attention mechanism applies. This can be violated when all the variables included in the model. In this case, the model with attention will be inconsistent. It indicates that in this case when assumptions are violated, neural networks can achieve the precise estimation only through the over- pareraeteettiiiiooonnn of (
Remark:
useful and they all need to be
C DISCUSSIONS ON THEOREM RESULTS
C.I THEOREM 1: COMPARISON OF SAMPLE COMPLEXITY WITH BASELINE MODEL
When we compare sample complexity of attention model versus baseline model, we say the key difference will be that so is substituted with p. which can be larger diverging dimensionality. Constants C1 and C2 can be regarded remain same when we compare attention model with the baseline for following reasons: (1) We assume same generative model, and the network size is the same, thus the optimal weight is known to be same. To make sure trained network weight is on the same scale with the optimal one, it is fair to keep there & bound constant same. (2) In this framework, we study the effect when diverges to Do as n + co. In this aspect, we don't expect the weight norm also diverges, since diverging weight leads to overfitting. By imposing lg regularizations on weights, we can always control the upper bound of l2 norm. Therefore it is reasonable to assume its norm bounded by a sufficient large constant. By imposing f2 regularization on weights, we can always control the upper bound of & norm by this sufficient large constant. (3) Even with overfitting, C2 in baselines are expected to be even larger due to the overfitting effect(Ch.7 Goodfellow et al. (2016). It will further explain why sample complexity of baseline is even larger comparing with fixing C1 and C2. In the experiment result, we also observe that the &g norm of weights from baseline larger or equal to the attention network.
C.2 COROLLARY 1: SUM-TO-ONE GLOBAL ATTENTION
Given the assumption (A7), we do not need 8p on the sample complexity bound for the sum-to-one global attention model. However, if we rescale æ properly, the result will be parallel to Theorem result.
C.3 THEOREM 3: NUMBER LINEAR REGIONS
In sufficiently wide network,
N1 s0 long n1 2 exp( plog log s0) p-80
is much smaller than [ p. Then, (#)))0 <()² ) holds as
Given plogg-su so log
p-#p < - -=,_oog log... p. since _P_ is close to p when p-s
is relatively small, the result still holds when n1 is larger than the order of p.
For illustration, the bounds are plotted in Figure in paper. The red line is for baseline model with p 100, and others attention model with different sparsity level sp. In general we can the bound for attention model is smaller than that for baseline model. The implication is that, when proper attention mechanism is applied, the approximation
remains small, and we
14

---

Under review as a conference paper at ICLR 2020
simpler landscape structure with less number of linear regions to reduce the estimation error. This is the why we can achieve specifie prediction error rate with a smaller sample complexity.
C.4 THEOREM 5: SMALL SAMPLE SIZE RESULT
n is mild assumption in a wide network with over prrr-aaammteeiztionn We can see that as long as we choose the number of units d to be larger than n, the linear dependence of (wo(1, Dy @ 1,,.1, holds with measure zero. In other words, almost surely this matrix has full column rank n. Thus after the nonlinear activation, The full column rank still holds almost surely. This assumption is similar to the condition in Theorem 3.8 of Nguyen & Hein (017b), where the number of units in some layer is larger than the sample size. When the sample size smaller than the number of units in the network, this theorem holds for the network without attention. has been proved by Nguyen & Hein (2017b) and and (Sudryy & Carmon (2016) under different conditions.
In assumptions, rank(o((1)) at C ())..1.....))
D REGULARIZATION EFFECT AND MULTI-LAYER EXTENSIONS
Our analyses provide some theoretical justification how attention mechanisms help learn superior models. The benefit mainly comes from that the attention weight a shrinks the whole parameter space, while this space is still large enough to capture all the necessary information. Thus the gradient and Hessian are more controllable this space, and the landscape of loss function behaves better compared to the baseline model. It will be shown again in Theorem 4 in the following section. This fact holds for both global and self-attention model. Our theory further validates that as long as the attention masks are learnt well, the performance is expected to improve. And this effect can be more significant when the weight is sparse. We also find that imposing constraints and regularization on network weights
also help remove sharp minima, and keep flat minima with good prediction. These discussions can be found in D.1.
The main idea of Theorem and 2 can also be extended multi-layer networks under the assumption that bias and gradients of weights are not uncorrelated. They can be found in D.2.
D.1 REGULARIZATION
Our analyses indicate that it worth considering approaches to control the bounds for the network weight matrices and pursue more accurate estimation of attention weighted input space. Motivated by this message, we suggest two possible regularization methods, which help to improve currently used attention mechanisms.
f Regularization on w
It is known that in optimization, & regularization on w is equivalent to specific & bound for w. For our two-layer model, if we choose regularization level properly, the & bound for w()), (2) also w0,w"* and w in self-attention model) will match our assumption. In practice, a proper regularization should be as large as possible, as long as a nice in sampee prediction is still achieved. This is important for attention mechanisms to keep prediction power while improving the landscape behavior of loss function.
f Regularization on achieve sparse attention weight
As we discussed, regularization can help achieve sparsity and more precise estimation. For our 2-layer model, imposing & regularization will not affect the analysis since we didn't 't use the gradient w.r.t a in the analysis. It means the same theoretical guarantee holds for the regularized network. At the same time will be more sparse, practically lead to more precise estimation, and more interpretable result.
This idea can also be adapted to self-attention model. In self-attention, we can add regularization a softmaz( QK") With this regularization, only part of the value matrices proceeds to the decoding procedure. In real world application such as in transformer, the sparsity in a; corresponds to the case when there some words in sentence should not attend each other. In this case, a sparse a self help us only focus attention between useful words in sentence, thus improve predition.
Vddy
15

---

Under review as a conference paper at ICLR 2020
It is worth mentioning that, through our theorem, can tell why & loss are more helpful than lg loss under sparsity assumption. Both & loss and &g loss will control the magnitude of attention mask, however, & loss can also help control sparsity level Sg, and we can see that the sample complexity bound is proportion to sh- Although E loss on attention mask can also help reduce the sample complexity, its effect will only be reflected in the log term, therefore it is less effective. Similarly in Theorem 2, if we assume the attention to be sparse(i.e. one word should not attend all words, but only some relevant words), then & loss can help further reduce the sample complexity.
Regularization one experimental approach for estimating sparse and precise and w. If some other methods can achieve this target, they are also the right directions to reduce sample complexity and improve current attention structures.
D.2 BEYOND TWO LAYERS
Here discuss how results of Theorem and 2 can be extended to multi- neural nets.
First we consider a D-layer network with naive global attention structure: f(x) af())
(1)))((((D-1))... of (ww(), x @ a))). Denoting vf(=") Fuity + under the assumption that at least one gradient term (((=") for k = 1, ...... D has correlation rate ((1) with bias u(this assumption is parallel to (A5)), we can show the expectation term EE ( (V)((()) Z O(z). Then using the similar e-covering and uniform convergence technique as Theorem goes, we can still show attention mechanism leads smaller e-covering number and tighter Hoeffding bound, thus leading to a smaller sample complexity bound comparing to baseline model.
Then we consider a D-layer network with self-attention structure. We denote the htth self-attention layer follows 9s. ( (x--)) = w"*0(((u (ww",, w x x== aaddf))), where x is the output of (k- - 1)th self-attention layer, with w E Rd,, x--1 E RRxxd 1, anls Rd 1xd, 1, w E Rd. xax and = E Rd
calcualted in the same way with two-layer self-attention network. Then we have the final output f(x) = w DD of (wD1, vec(w" x ==p-la==/))), where x (gD-1(-- (-.. g (æ)), and D: R2x D E RRx In this way, the network calculate self-attention D times and finally produce the final prediction. It is worth mentioning that, To obtain a scalar prediction in regression model, we flatten the value matrix of the last layer as same as the two-layer model. We still denote (w (epp=(((w'" (w -veew" V x (y)eff)) E(ylz)). Then the necessary assumptions parallel to (A2), (A9) and (A10) are as follows.
(A14) All weights uh for = I...., and j = 1,2 satisfy ||00* ||2 S Cio. And we assume the prediction is centered, i.e. E(u) 0.
(A15) The output y can be predicted by the D-layer self-attention network with an in- dependent sub-Gaussian error with variance o², i.e, there exists set of parameters (w( w (2))) such that y (0"o(((")) vvecw" x: ( + as defined, where 6 su, Cf) for i = 1, with aq L 44-
(A16) There exists interger & and r such that h E (1,...,D) D} and {1,2), such that ( ), u) O(1 and sd(Vf(w = 0(1).
Then we can have the following theorem of sample complexity bound of multi-layer self-attention model:
Theorem 7. Under (AI) and (A14) to (A16), we assume the weight term -h satisfies (A16) with | (1o )ll S Ck. dsetf is the total number of parameters in all value, query, key matrices. Then for any > 0, given the sample size:
Ck n R log( (dwats + (dx + dg)gs))
where n = C2C.C with probability tending to 1, any stationary point (z()) 20, 2 k) of the objective function (6) satisfies that: ((f((z)) Eyy 2z)) $ n2
16

---

corresponding deriva- Under review as a conference paper at ICLR 2020
Remark: Because multi-layer self-attention models include large parameter set with complicated gradients, the assumptions are not as intuitive as the two-layer model. But the main assumptions parallel, such that the bias u, cannot be uncorrelated with all possible directions. And this assumption is reasonable considering the high-dimensionality nature of networks. The extension of this multi-layer model is omitted in the main paper since it leads to over-complicated derivations and complicated assumption discussion and will distract readers from main ideas of the paper. We believe the two-layer attention model is representative enough to provide theoretical evidence on why attention reduces sample complexity.
D.3 CNN AND RNN
Our analyses are based on fully connected network, and may not be able to apply to more involved task such as CNN and RNN directly. But we believe the key message of our analyses also provide insights in analyzing CNN/RNN with attention: The attention mechanisms can help effectively shrink the parameter space, thus reducing most of the noise and unnecessary variability in training. Thus the stationary solutions are more likely to generalize well. That's why start with analyzing the naive global attention model and self-attention models, which can help inspire the analysis of these complicated tasks. More detailed analysis/experimennt with CNN/RNN important future work. To this point, our experiments aim at validating the theoretical analyses and explaining why attention works in general.
E PROOFS
E.
PROOF OF THEOREM
Proof. The proof divided into two parts. Firstly, we study the landscape of population risk in part (a), then we evaluate the convergence of empirical risk to the population risk in part (b).
a. Landscape of population risk
We introduce necessary notations beforehand. To emphasize the role of æ and y separately, here denote (1)(1), (2) a) Ey (w(()), ((2), a)),
w()), which is the expectation of the empiri- cal loss gradient with respect to y. treating as random, and V VR(ww tives. And we denote E,, (V R(w(()) w2( a)) Ez,g (Rn (w(1), w()) a)), which is the expectation of the empirical loss function with expectation to both x and y. In our analysis, first we will study E,, (VR(W()) (()),, (,,2) a)). then we analyze Ry (w(1), (2( a). The motivation of using R,, Rw."",,""" (1), (²) a) comes from that it simplifies the part (b) analysis of empirical risk conver- gence, since the randomness of has been included the population risk analysis, making advantage that the noise is independent from predictors. In the proof, we may use o(7) for vector/matrix case. In these cases, means that every element in vector/matrix is o(7). In the proof, we will regard d as an arbitrary large fixed value, not diverging with n.
We denote
u = (w d( (w x e a)) E(y|z))
and u; as the version with specified sample index. Then the derivatives of population risk with expection to y can be presented as follows:
1
()
Σ
u
(w x, Ca))
n
iml
(2)
Σ n
1
x2
(x (w o (1((((()) o a)))
1
VR(a)
(, 3 (IT o ((w x, o a))))
17

---

Under review as a conference paper at ICLR 2020
I, a.)). Therefore when (a, w( 2( (a", w(()) (u(t),,((()))) w(())), all the u, are zero, and all the gradients expectations are zero. Thus for any true set of parameter (a", w(1), w(())) ). they have zero gradient expectation automatically. And the key of our proof is showing that with high probability, any parameter (a, w(t), ((2) cannot be stationary point if (((((())((((())
By (A3), know that E(ys[a,) = (2) ( (((w())>
@ a}} [()())))) 2 O(?²), because their gradients w.r.t to w)(2 or w(() must be bounded away from zero.
By our assumption (A2) and (A4), our parameters wL))w m (2) and a are inside the l balls (0, C)),Bxx((,, C2) and (0, sg). By Lemma 5.2 in Vershynin (2010), we know the e-covering number Ne, Ney for these three balls are upper bounded by:
No. S (3C1/c)² No, S (3C2/c)m), Ne, S (38g/c)²
Then we know 3e-covering number for the union of all three parameters Nae satisfies that Nsc S Ne, Neg.Nee For the case of notation, we denote 0 = (a, w(1), w(2)), Let θ = (01,... NN.1 be a corresponding cover with N3., elements. Then we can always find O such that for any feasible 0, there exists j E [N] such that ()
all2) S e. In this proof, we use parenthesis subscription (j) to represent elements in the cover, to distinguish it from other subscriptions.
(1) - w0 [[ (2) w l|2 ||ay)
By triangle inequality, we have for 1,2: [[ 2 |V (1) () - IVV v
(8)
And in the following section, we prove that if E(u2) 2 O(+²). we must have | (
O(7) or there exists htth variable of æ, such that the hth column of V R(w satisfies
|[E (VR(( 1122 O(7). Here the subscription le without parenthesis corresponds to h feature.
First consider (((w(2)). By (A4), there are most so nonzero elements in + a, and denote || llctive as the é norm on the elements with corresponding non-zero attention mask. Thus by inequality between norms, we have:
((Ilo((w 2z ① ()2) S a O a almaxx//0 (201)) 11 ll,cciive I (oCC,C==((VS
Then we derive
| (VR (2) ll = |cco((u, (((10()) x a))) + (() x a)))l
> |,cov(u, di (1(1(()) 2 a))) 0l// veol
Therefore if ||coo(u, of (((w()) xz O aa))lll 2 0(7), we already have ((VR(( 2 0()). Then we consider the case when |,oo(u, ob (((w0()) n. @ a)))]l2 o(7). In this case, we denote r (oy(E())))), (((w(()) 2 O a))) Rd, and denote the covariance matrix for (1(w(1) x o a)) as:
(Ezli cowl x O a)), ((w x (a)))
Then plugging into con( ( ((((()) x x @)a)) ((( a ()))) = r + o(7), using subtraction and addition, we have:
Eg = =r
With (A5), covariance matrix Eg is invertible with smallest eigenvalue lower bounded, we have (2) = E01 r + o(7).
Next we argue the following term by contradiction:
2 a)))l2) ((llo((w()) x @ a)) xz,
0())
(9)
where 8² ('((w()) x @ a)) represents the function corresponding to the true parameter set
(++ w() ). If we have equation 9 violated, we know:
Edn EØg
(4) i x B(d((w) () 21x0 ())
- ((
() x x a))) a a @)))))
((w(" x a))o"' (1 x x00 E(d ((0) x o a) ((( (,, xx)a)))))
s ((( x o a)); - (10 x @ a)ill2 + _( x a)); - -10((0)) o(7)
x 0 a),1k2)
18

---

Under review as a conference paper at ICLR 2020
where use E(llo( (w))) x C a))](2) ((1@"(Ww x @ (0))]l2)) ((ll(((w<)) o a)) o ((w(), x c 12) = (1), and we know o² ((w(", @ a}) corresponding to true parameter set is finite. Since we have derived w(²))+ = Ez r, with lower bounded eigenvalue assumption, we can derive:
Ex- 6r
Es (2² -xz!) E (I - EEg. + (()))==))) -
(1)))) = 0(7)
We know we have ||cou(u, ((((0(1) x; O a) = o(7), thus w = o ' + +(7)) Therefore we conclude w(2) w(2)+ o(7). Plugging back the formula of u in equation 12 we have:
=
EBu²) = B(Iw( ((0 ()) x O a) - -2 2 (1v( (1) a ② a))
S (((w(()) (2) ((w x R a) - -2 o{ (w (1) 3 a))
+ E([w (2) (( n. O a)) ( (( a. 1 a = o(72)
Here we conclude the contradiction. When condition equation 9 is violated, we derive RE(u²) o(72). Therefore if [E(=2) 2 O(7²), we must have (( 3 a)) o² (ww x, 2 a)ll2) 2 0()).
Now we ready to study V R(w(()) Recall in assumption (A5), we denote ha (2)) x. (w) 1 o' ((w 2z @ a))), where x. represents hthh feature in æ, and ag is the attention weight for T.k. And we have proved that E (((((((()) Ig @ (a)) x² ((w((()) x; @ a}ll)) 0(7). Thus by (A5), we can find k, such that (x.)) O(1) and cor (u, he (x)) = O(1). Then we have (cov(u, hx(xx)))l sd( (x.,().)))))))))) cor (u, hy (z.+)) We know
sd(u) = v E((") [[x) -(7))
Therefore we have:
[V ( 2
= ||cov(u, (x.,))))) + E(u) = O(7)
Here we conclude that we can always find l: E (1..... p}. such that [E( @1) 1)l 2 0(7).
With this conclusion, we move to bound the gradient term v; = un (x @ ax (w @ o' ((w()) a O a))). With any fixed parameter set, we can calculate:
a, @ a)) and 2ik
|0:11 & (Callo(( x, O ()))))))) x, a
3
|| ll S (C a a) 1)) =
- C0C
where we use again that there are at most so nonzero elements in a + a:
(1) x O l2) < max(1æ a|H|w" (1))")))1,att
||1,active VSDC2C1
From the last section, we know there exists constant e such that ||E, (V ||E.,, ((w M12 2 ey for some constant c. Suppose |[,, (V (() 2 er, and we denote = 2. And we know || is bounded as we derived. Therefore we can use Hoeffding bound on the & norm in direction of E, (V (R(w())) since know the variance this direction is smaller than 0². Denoting ( as the gradient with respect to 9th parameter set 6-cover for j E N1):
cy or
> ") s expl 222 σ²
P P( E. (VR(
3
By union bound, we have:
RR 12 2 207 207) S N exp
222 σ²
P(j E [N.]
V R 3 la term. Here we use u, to represent the prediction error for 8th instant with respect to parameter (a, (1) (2))), and use uz(y) represent the term
Secondly we analyze (2)
19

---

Under review as a conference paper at ICLR 2020
with respect to the parameter from jth element in c-cover set. By triangle inequality, we have:
V VR ( V
2
(us(((u) (w (1) 2i a)) (( (1) ma o aلل))))ll
1=1
2 (1) Σ
[n - (u-))o))) x 2 +11 Σ (o n. a)) (( a @ a())))
2 So C2Cee
We choose c C3 and plug back above results to equation g then least with prob-
ability 1 - (N,, exp(- - )), we have (V )| -. Therefore can choose ž Gop log( C.CCC (pd + d), such that N, exp( -n n - o(1). Finally we can conclude that with probability 1 1- (1), for foraaorrryy.. any (a, w()), 2( such that (( (2) (20()) xoā)) - E(yl=))² 2 % we have [
3.
If | ( ()) 0 Applying the same technique, we can show that with probability 1 on (1), we have [R |l) l|l 3
Convergence of empirical risk
So far, we have shown that for population risk with respect to y, with high probability, all the parameter sets with poor prediction in expectation, i.e (((()())((((()))) X @ a)) [yy| )=)) 2 O(>²), their population risk gradient with expectation y must be away from zero. Now we move forward to show that empirical risk will converge to the popular risk, i.e. VR,,, (w) + RR(w). Thus these parameter sets cannot have zero empirical gradient. aspect of three parameter sets, they can be represented as:
a
(w( - (2)
col(w ,2,00))
i=1
n 7
VRn(w (B(w
(x, a) o ( (x,()x,, o
(a) Vl(a)=
(x o (w (w)((((()))) o o ((w(1), x a))))
With (A3), we know that 6, subG(o, C2), thus 1 -E 4 = ( by C.L.T, combin- ing the bound for (((w())) x @ a)) we have derived in last section, with sample size n 2 log( enCrCaCr (pd + + d), conclude that with probability 1 - (1,
|| [v, (R(w
(10)
(v (1)) v( (3)
1|1 5
(11)
Recalling part (a), under the first case that w.h.p V V (2) - for any parameter (a, (2) (2)) with ( ((0())) c a)) - E(yl X)l|2 2 7. Combining this with (10), we can conclude that for any positive constant > O, with required sample size, with high probability that ||V Rt, (2) > 0, thus they cannot be stationary solution for our loss function. As we stated, we are in another case, we have w.h.p [[ (1) 2| 2 3 w.h.p RR, (2) || 0.
we can use same techniques to show
In other words, under our assumptions, all the stationary points (ā, (1)), (2) in our programming satisfy the prediction error upper bound rate y w.h.p.
20

---

Under review as a conference paper at ICLR 2020
E.2 PROOF OF COROLLARY 1
Proof. To extend the result from Theorem to corollary 1, we simply substitute the l2 norm bound || @ all2, from SrC, to C2 since ||a||2 ||a||| 1. All the other parts keep the same. Thus the only difference is that we remove 89 in the bound comparing with Theorem 3.
E.3 PROOF OF PROPOSITION
Proof. This proposition is direct result of Theorem 3. Since the assumptions for a still hold, all the bounds apply. The only different is that since a is not optimized together, we don't 't have to consider the e-cover number for a in the maximum operator. This leads to a slightly tighter sample complexity bound in corollary 1 comparing with theorem 3.
E.4 PROOF OF THEOREM 2
Proof. Similar with Theorem we obtained new e-covering bound:
No. I (3C1/e)", Ne, S (3C2/e) Nee S (3Cs/e)*d,, No. S (3)o/e) ,No S (307/e)Md,
where N 11 Ne: And also, the new v, and Zik terms are:
v: = o( vee(w" xxaae(")))
2 ww a;a;;")) ) (2) ④ Φ ((w(()) ((w()) w v. ())
where u, = (1T (2) a((1 (wo()) veelw> x;a )). Under assumptions, using the same agrument with Theorem 3, we can show there exists a constant e such that either [|E(VR(w())]l: 2 ey or there exist , such that [IB( (2))) c7.
In the case when |E ( z cy. we have new bound of |||| with respect æ is upper bounded by o² PPP considering a is normalized by softmax function for each vector set. Parallel Theorem 3, by hoeffding bound and union bound, we have:
20y, s N. exp(-n (1 0272
P(2j [N.]. 115l12 S
3
d2
Secondly we bound (2) - V ( 112 term by subtraction and addition:
[
V (R(
u 2/0((w(1)) veclw" V x, x, (())(=)))) un vec(w" x, (())))(()))))
i=1
iu (e ((w))((()())) veelwV a (=aa)((()))
Σ (ol(ww w a,aaee((s)) woo)d((wo")) a self(=) (c((()))))))) 12)))
S C
recalling that (i() and (s()) are corresponding to the 7th epsilon cover. We choose e 3070
and combine the above results. Then at least with probability 1 - exp-nn }). have [[
||2 > 3.4 Therefore we can choose n 2 अيح log( PC,C&C,C.CCG. (pd,d + d 2pdg), such then N. exp(-n
Thus with this required sample complexity, we have (V)(w(2)) E( {()] < 3 In the same way, we can in the case when show (VR (1) (( ()) )l Bey
o(1).
20y
Finally we can conclude that with high probability, any parameter (a, (1),, w ()) with (( vec(w" x.a ())) - E(y|=))² 2 7. we have [[ > Th. or ))
(1) )|2 > z. Then following the same empirical risk convergence argument, we show [ k
that with high probability they cannot be stationary point.
21

---

Under review as a conference paper at ICLR 2020
E.5 PROOF OF THEOREM 3
Proof. First with ||a|lo sg. know all the inputs æ, with corresponding a, 0, will be inactive in the network. We can omit all these inactive inputs. Then we split n1 units into so group. with !u number of units in each group, and discard the leftover units. so different groups correspond to so active inputs with non-zero attention weight.
Inside each group, for example in jth group, denoting q = + we choose the input weights and biases for = 1,2, ... as:
h (a) max w,z)
h2(x) = max 2w, 1},
hg (x) max 2w, 1)3
here we assign u to row vector with j variable equal to 1 and other entries be And in the second layer, we choose w(22) = (w3,... wa). where Wy (1, 1, (-1)*+1), corresponding to h1 to hy in each group. Then the designed network has q linear regions inside each group. giving by the intervals:
(-o0, 0), (0, 11, (1, 1,0
Each of these intervals has a subset that is mapped by wyh(x) onto the interval (0,1). Montufar et al. (2014) Therefore the total number of linear regions is lower bounded by u1 80
80
E.6 PROOF OF THEOREM 4
Proof. Here we define an (a1, a12) scale transformation such that:
Toiigg (a,, (2) ( (2) + (ox) (2) ( (a
Then know the jacobian determinant for Tar,au is at a2 pd-d Let such that Boe (r, 0) is in C(L, 0,e) and has empty intersection with (a, wn(1), (2))) such that the jacobian determinant goes to infinity, and the volume of C(L, 0, e) goes to infinity.
Since pd > d, we assign 02 - oc,
For the Hessian matrix, without loss of generality, we assume there a positive diagonal element 5 > 0 in Therefore the Frobenius norm V L (0))]l of
2 L a1 (0))
I
a I
v
V2L(O)
2
a2
(A102)
(a102)
lower bounded by aj 28. Further we apply the fact that the biggest eigenvalue of symmetric
matrix X larger than [|]],, and pick a < V of M
then we have the biggest eigenvalue of
L (0)) is larger than M. Therefore there exists stationary point such that the operator norm for Hessian is arbitrary large. Thus we finish proving part (a).
Then we consider part (b). Similar with part (a), we define an a scale transformation such that:
Ta (w) (2)) (1)
And all the value and key matrices remain the same. Then we know the jacobian determinant for To (rd,-11) Since pd.d 2 d, as assign - oc, such that the jacobian determinant goes to infinity, and the volume C(L, 0, c) goes to infinity.
22

---

Under review as a conference paper at ICLR 2020
For the Hessian matrix, still assume positive diagonal element 8 larly we have the Frobenius norm ||V (0)) of
in w(!)
Simi-
L(T,
V
I
0
0
I 0
0
0
aT
0
V2
(()))
a
0
0
I
0
I
lower bounded by a: -28. When we choose sufficient small a, we have the biggest eigenvalue of 2 L(T((a (0)) is larger than any constant M. Therefore there exists stationary point such that the operator norm for Hessian is arbitrary large.
E.7 PROOF OF THEOREM 5
Proof. For global attention in part (a), We start with calculating the gradient of the empirical loss function V (l where Ro (wr(1), w(() a) = 2n Σ E of ((wx(1) a. 2 a)) - ys)²
Denoting (w(2)7 (2)T 6((20l)
a 3 a)) y). The derivatives can be presented as follows:
V VRn (w
( xa))
VRn
(a a) o0 1 (ww nz ( xx
n
(a) (1 @ (w (TTTT (w (2) o (ww a)))
By assumption, rank((p((u 23 O ) = thus solving the linear system, we must have M4 0 for any i = 1,2, 1,2, ..., to satisfy that RR, (2) 0. Thus we know that the loss is exactly zero inside sample. Thus it must be a global minimum.
Part (b) can be proved by substituting a to a in part (a). Here we only consider the deriva- tives with respect to (1) and w(22, they can be presented as follows:
n u (w)l veolw" (oo(wWWxre axaaaae//)
VRn (w)
(Rn
72
(Rn(w
(aaeell)) (2) o o' (lw(l), ur (vec(w" z.a.el/
(wew))"eeew"" veclw" nedlwaxaaet rea.eel))² ))T
By assumption, rank(o(( wweewwwww 2,a; self )=) n, thus solving the linear system, must have u, 0 for any 1, 2, ..., satisfy that (2 =0. 0. Thus know that the loss is exactly zero inside sample. Thus it must be global minimum.
E.8 PROOF OF THEOREM 6
Proof. First, we obtained new e-covering bound for the parameter set (w(?), ())) wy, wa):
Nr, S (3C1/e)*, N., S (3C2/e)*d Ne. s (3Ca/e) No, S (3Co/e))
And N, < II 1 N., Similar to Theorem 1, we denote
(w (2)T (1(w(1) {a(z))a)) E(ylz))
(12)
and u; as the version with specified sample index. Then the derivatives of population risk with expection to y can be presented follows:
23

---

Under review as a conference paper at ICLR 2020
vR(w(2))
uu 24 3 a))
11
1 Σ
m a(x))a) (10 o o2(((1 Σ {a(xx));))))
VR(
And also, the new Us and Zik terms are:
P v, = u Σ a(<,))²)))
1=1
Zuk u,a(<;x) (w o の ((w))) Σ (((x)))))))
In the case when | u( bounded by o² CC C the case when |EE(V R( )) le 2
cy. we have bound of ||el1 with respect is upper with normalized attention weight. Same argument follows for
C
Then follow the same approach as Theorem and 2, we
obtain the sample complexity bound:
o² n1 2 2272
10, 1C )d ld + dd + da)
log(
E.9 PROOF OF THEOREM 7
Proof. Under assumption (A1) and (A14), we know all input features and weights are bounded. Therefore know || fw""()) is Lipschitz continuous function, and we denote its Lipschitz constant Lg. For w",, we can derive that:
VR(w n
u, V f(w
Under (A14) to (A16), if we have B(f(x) E(ylxx)) S 22 then:
IE(VR(w *r 1)1l2 Z sd ( (l2))))122a)0))))) cor (VS((")), 2 0(7)
Then similar with Theorem 1 and 2, we construct an e-cover over all parameters = (w"1, ==2 w 2 wx), and denote it ( ON) } such that for any feasible parame- ter, there exist j E [N] such that the maximum &g distance to 0 is smaller than e. By calculating the number of parameters in all matrices in 0, we have
(N.) = e
1 O(deer +
Σ (dx + do)gk)
Denoting V (( ) as the gradient with respect to jth parameter set in e-cover for j E
822
P(IVR(W (V
+ S exp(-n
3
By union bound, have:
207 S N exp(
2222 (n
P(j [NJ.
24

---

Under review as a conference paper at ICLR 2020
Secondly analyze [VR(w*r) - -V VR(W()))l2 term. As we have shown that the gradient is Lipschitz continuous, thus we have:
k (dx + dr)
[V - VR(WG)))le S (seld +
1=1
3 then at least with probability 1- (N (xp(exxp(110 =2 )) we have (V (w)()))) 112
We choose
Cy Therefore we can choose Z log( Ck y (d k (dx + dg)gs). such that N₂ exp(
3
o(1). Finally we can conclude that with probability - on (1), for any (a, w()) w(2) such that
E[ (1(o() (1()(1) a @ a)) [[y)=))) 2 7, we have || [2( (((((22) (2) ]|| > S. Then following the
convergence of empirical risk procedure of Theorem 1, show with probability going to such that
| R., (2)2))ll2 > O and all parameters with prediction enror (()) cannot be stationary point as long
as Z log( 1og(*))doont (ddnts + E/= (dl + dz)gs). Thus we complete the proof.
F ADDITIONAL NUMERICAL EXPERIMENTS
F.1 CONVERGENCE PLOTS FOR CLASSTFICATION TASKS ON NOISY-MNIST DATASET
We present additional results on convergence of 2- layer neural networks on classification task involving Noisy-MNIST dataset (Section 5.1 of main paper). This dataset is formed by embedding an X digit image in a noisy image of size 48 x 48. We show convergence results for three settings: $ (5, 8, 15). each setting, the learning rate of models are varied. The plots are shown in Figures and 6 We observe that in every setting, attention models converge faster than the baseline models not employing attention.
F.2 CONVERGENCE PLOTS FOR REGRESSION TASK
In this section, we present convergence plots for sample complexity experiments discussed in Section 5.3 of the main paper. Regression task is considered in this experiment. Convergence plots of baseline model, attention model and regularized attention models are plotted for various sizes of the datsset N, shown in Figure 7 We observe that for all N₂ values, regualrized attention model converge fastest followed by attention model which is then followed by baseline model. So, we conclude that in addition achieving improved sample complexity, attention models converge faster than models not employing attention. Also, regularization helps speed up the model convergence.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. Convergence Theory for Deep Learning via Preemarraettiizationn. In Proceedings of the 36th International Conference on Machine Learning, ICML '19 2019.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization.. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, pp. 244-253, 2018
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. CoRR, abss/1409.044 2014.
Andrew R Barron and Jason M Klusowski. Approximation and estimation for high-dimensional deep learning networks. arXiv preprint arXiv: 030900, 2018.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 Proceedings of Machine Learning Research, pp. 1019-1028, 2017.
25

---

Under review as a conference paper at ICLR 2020
0.0050
0.0040
Accuracy
0.0030
0.0020
10
Epoch
Beration
Learning rate: 0.05
baseine Attention model
0.005
0.003
0.002
baseine
Attention model
10
20
10
40
50
Epoch
teration
Learning rate: 0.01
baseline Attention model
0.006
baselne
Attention model
0.005
54-
Accuracy
00
8.002
8.001
10
Epoch
Learning rate: 0.001
Figure 4: Convergence plots for classification tasks: Image size: 48 48, Digit patch size 15
26

---

Under review as a conference paper at ICLR 2020
0010
Attention
92.5
ACCuracy (6 96
90.0
07.5
0.000
0.004
80.0
0.002
10
Epoch
Learning rate: 0.01
95.0
baselne
baseine
Attention model
Attention model
90.0
0.008
Accurit (8h
87.5
02.5
0.004
77.5
10
20
30
40
50
0
30
50
Epoch
teration
Learning rate: 0.005
0.035
baseline
Attention model
0.25
Accuracy
60
0.10
baseline
Attention model
0.005
Epoch
Learning rate: 0.001
Figure 5: Convergence plots for classification tasks: Image size: 48 x 48, Digit patch size
27

---

Under review as a conference paper at ICLR 2020
90
0.00
0018
0.04
0.00
60
0.006
10
20
40
Epoch
Beration
Learning rate: 0.01
baseine
Attention model
80
70-
0.020
60-
0.10
baseine
Attention model
0.005
10
20
0
10
50
Epoch
Beration
Learning rate: 0.005
80
0.035
baseline
Attention model
0.300
0.05
0.000
0.15
baseline
8010
Attention model
10
Epoch
Learning rate: 0.001
Figure 6: Convergence plots for classification tasks: Image size: 48 x 48, Digit patch size
28

---

Under review as a conference paper at ICLR 2020
35
baseline
baseline
3.5
Atte Attnnbion model
Attention model
20tt-siinn Detized
Attention model rggilariedd
50
2.5
2.0
105
2.8-
1.0-
0.5
0.5
0.0
13000
J000
N 10000
N, 14000
baseline
baseline
Actention model
Attention model
Altantian model reguiarized
Altention modal regularized
2.5
2.0
2.0-
1255
1.5
3.5
2.-
05-
0.5
0.0
0
3003
10000
1500D
20000
21000
30300
35000
40000
5000
30000
15000
20000
25000
30000
35000
40000
Bteiition
Elratiio
N, = 16000
N. 18000
baseline
300
busefiie Attention model
2.5
Allantion model regulatized
2.5-
regularirnt
2.0
2.0
125
1.5
15
1.0
5000
10000
15000
2000C
23000
N0000
40000
5000 1000c
15000
20000
25000
30000 35003 85000
Ittrationn
berstion
N 20000
N. N₂ 50000
Figure 7: Convergence plots for regression task varying the number of samples N.. "Baseline" indicates 2- layer NN not using attention, "attention" denotes attention models, and "attention regualrized denotes attention models trained with Ln regularization
29

---

Under review as a conference paper at ICLR 2020
Simon S Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima deep neural networks. In Proceedings of the 36th International Conference on Machine Learning, ICML 19, 2019.
Rong Ge, Jason D. Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. CoRR, abs/1711 11111111.000000 1.00501. 2017.
Xavier Glorot, Antoine Bordes, and Yoshua Bengio. Deep sparse rectifier neural networks. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp- 315-323, 2011.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 1609004836 2016.
Jiasen Lu, Jianwei Yang, Dhruv Batra, and Devi Parikh. Hierarchical question-imag co-attention for visual question answering. In NIPS, pp. 289-20 297, 2016.
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention based neural machine translation. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 1412-1421. Association for Computational Linguistics, 2015,
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150 150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL http:
7/ 7/www, www.aaclweb.or.. wtw.:e/.ooggaathoooooooo//ooo///
Song Mei, Yu Bai, Andrea Montanari, et al. The landscape of empirical risk for nonconvex losses. The Annals of Statistics, 4666A)):747-27774 20188 2018a.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7655-7767 2018b. ISSN 0027-8424. doi: 10.1073/pnas.. 10177///paas111059991555 1806579115.
Volodymyr Mnih, Nicolas Heess, Alex Graves, and koray kavukcuoglu. Recurrent models of visual attention. In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems 27, pp. 2204-2212. Curran Associates, Inc.. 2014.
Guido F Montufar, Razvan Pascanu, Kyunghyun Cho, and Yoshua Bengio. On the number of linear regions of deep neural networks. Advances in neural information processing systems, pp- 2924-2932. 2014.
Quynh Nguyen and Matthias Hein. The loss surface of deep and wide neural networks. In Proceedings of the 34th International Conference on Machine Learning (ICML 2017), Sydney, NSW, Australia, 6-11 August 2017, pp. 2603-2613 2017a
Quynh Nguyen and Matthias Hein. The loss surface of deep and wide neural networks, arXiv preprint 1rXiv:1704.08000 2017b.
Tomaso A. Poggio and Qianli Liao. Theory II: landscape of the empirical risk in deep learning. CoRR, abs/1 703.09833, 2017.
Yunchen Pu, Martin Renqiang Min, Zhe Gan, and Lawrence Carin. Adaptive feature abstraction for translating video to text. In Proceedings of the Thirty-Secomt AAAI Conference Artificial Intelligence, (AAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 7284-7291, 2018
Blaine Rister and Daniel L. Rubin. Piecewise convexity of artificial neural networks. Neural Networks, 94:34 45, 2017.
30

---

Under review as a conference paper at ICLR 2020
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D. Lee. Thcoretical insights into the optimization landscape over-paramenne wer-ppraaeterized shallow neural networks. CoRR, abs/ toos/1777.0496,, 2017.
Daniel Soudry and Yair Carmon. No bad local minima: Data independent training error guarantees for multilayer neural networks, arXiv preprint arXiv:1605.085: 0606,, 2016.
Daniel Soudry and Elad Hoffer. Exponentially vanishing sub-optimal local minima in multilayer neural networks, 2018
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, pp. 5998 6008, 2017.
Roman Vershynin. Introduction the non-asymptotic analysis of random matrices. arXiv preprint arXiv:1011.3027, 2010.
Kelvin Xu, Jimmy Ba, Ryan Kiros. Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. Lille, France, 07-09 Jul 2015. PMLR.
Han Zhang, Ian J. Goodfellow, Dimitris N. Metaxas, and Augustus Odena. Self-attention generative adversarial networks. CoRR, Abs/1805.08316 2018.
Bolei Zhou, Yuandong Tian, Sainbayar Sukhbaatar, Arthur Szlam, and Rob Fergus. Simple baseline for visual question answering. arXiv preprint arXiv:1512.02166; 2015.
Pan Zhou and Jiashi Feng. The landscape deep learning algorithms. CoRR, abs/1 705.07038, 2017.
31


Reviews:
[{"review_id": "rkxwfYlG9r", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the loss landscape of two-layer neural networks on global- and self-attention models.  It shows that attention helps reduce sample complexity.  I went through the all the theorem part of the paper, but only checked the intuition and did not dig into the detailed proof.   I am less familiar with reviewing papers that are theorem-oriented, which is hard for me to justify the contributions of these bounds.  My concerns and suggestions are mainly focused on the empirical part, which I think authors need to improve.  \n\nFirst, the authors need to provide more detail about the baseline model, especially for the experiment on 5.1.1.  Also, please list the number of parameters for different models.  And please justify whether the reported results (sample complexity in table 1 & 2) are affected by the difference of number parameters if they are not the same.  The authors also need to clarify the meaning of the \"number of training samples\" in the tables.  Does it mean the number of unique training examples seen by the model or the total number of training samples?  If the authors mean the former one, then, what is the training procedure?  What are optimizers used between different models?  Are these will affect the test loss difference.  How does the author make sure the comparisons are fair?  Also, in the main text of section 5.1.1, the authors said they construct a dataset of 200,000 examples.  How are the 200,000 samples distributed?  What is the size of the testing set?   Do the reported results average over multiple runs? Similarly, in section 5.1.2 authors need to clarify these questions and the experiments are fair to support their claim.  \n\nGiven the theoretical part is a contribution, the experiments at the current version do not support their claim fully.  I will consider raising my scores if authors can clarify my questions.  ", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "SyxWAKbRYr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies attention and self-attention networks from the theoretical perspective, giving the first (as far as this reviewer knows) results proving that attention networks can generalize better than non-attention baselines. This has been observed empirically before and it is very good to start the analysis of foundations of this phenomenon.\n\nThe results cover fixed-attention (as we understand both single-layer and multi-layer) and self-attention in the single layer setting. One interesting part that is not covered (and may be very hard) though is multi-layer self-attention networks. What is the equivalent of condition (A9) there? In other words: can we prove that the attention will learn correct attention masks if they exist, or do we need to assume that?\n\nOn the experimental side, the paper introduces L1 loss on the attention mask. This is (to the knowledge of this reviewer) a new idea and it would be interesting to see larger models (e.g., a Transformer on a translation task) trained with this additional loss. Does the analysis suggest L1 loss in any way, more than say L2?\n\nI am grateful to the authors for their reply. The new multi-layer theorem is impressive and I'm grateful for clarifying the L1 loss. Have the authors considered variants of hard attention as well (e.g., training with attention that attends to at most 10 or 100 elements), could that play the role of L1 loss? I stand by my recommendation to accept this paper.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HygmoapJ9H", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This paper endeavors to develop some theoretical understanding of why attention mechanisms improve the performance of neural networks. It includes some theory around the sample complexity of models which have a multiplicative mask in them, as well as a few related experiments.\n\nReview: Having some theoretical analysis of the benefits of attention is definitely needed and I commend the authors for working towards this goal. The main issue that I have with the paper is that the model initially considered (introduced in section 2) is unrelated to attention mechanisms because the attention weights \"a\" do not depend on the input in any way. If I understand correctly, the ground-truth function y_i = f*(a* . x_i) assumes there is a single globally-correct attention mask a* for all examples. This is not attention. Every attention mechanism that I am aware of depends on the input in some way (in other words, instead of a*, you'd have a function a*(x_i) which changes according to the input). Some additional discussion is needed to explain why it is interesting to study this dramatically simplified (to the point of not being relevant) model. For example, if it was explicitly described as a preliminary for the a*(x_i) case (section 3.2 onwards) it would make more sense. Similarly, the assumption ||a_0|| ≤ s_0 is not true for most attention masks in pratice. There are some hard attention mechanisms which try to induce binary masks/hard sparsity but they are extremely rarely used in practice. The combination of these two assumptions makes for what appears to me (if I'm understanding correctly) a somewhat trivial result that the \"attention mechanisms constrain the parameters in a smaller space\". Note that this criticism does not apply to section 3.2 where the attention mask is specified as a function of the input. Similarly section 4.2 relies on the \"single attention mask\" idea, which is the sole result considering generalization in the paper. Finally 5.1.1 consists of an experiment which assumes a model of the form G*(x . a*), where a* is a fixed random binary mask. Training a model to learn G* and a* has very little to do with attention - it is a learning problem where you have removed a subset of the features and are trying to determine the parameters of the function and which features were removed. The L_1 penalty on \"a\" is certainly a good way to learn to only include a sparse subset of your features, but again has very little relevance to most attention mechanisms (i.e. I know of almost no work which regularizes the attention mask in this way).\n\nOverall I think most of the results in this paper are not relevant to attention models and as such it has limited impact. I would suggest reframing the paper around a focus on what the authors call the \"self-attention\" case (which is really just normal attention), possibly using the fixed a* case as a preliminary/motivation. I also would suggest removing the discussion of regularization the attention mask as I am not aware of this being used in practice. Finally I would suggest some more direct discussion of how the results apply to models which are used in practice.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper aims to theoretically understand the the benefit of attention mechanisms. The reviewers agreed that better understanding of attention mechanisms is an important direction. However, the paper studies a weaker form of attention which does not correspond well to the attention models using in the literature. The paper should better motivate why the theoretical results for this restrained model would carry over to more realistic mechanisms.
Decision: Reject