Analyzing Attention Mechanisms through Lens of Sample Complexity and Loss Landscape
Authors: Bingyuan Liu
Yogesh Balaji
Lingzhou Xue
Martin Renqiang Min
Venue:  (2021)
Keywords: Attention mechanisms
deep learning
sample complexity
self-attention

Abstract: Attention mechanisms have advanced state-of-the-art deep learning models in many machine learning tasks. Despite significant empirical gains, there is a lack of theoretical analyses on their effectiveness. In this paper, we address this problem by studying the sample complexity and loss landscape of attention-based neural networks. Our results show that, under mild assumptions, every local minimum of the attention model has low prediction error, and attention models require lower sample complexity than models without attention. Besides revealing why popular self-attention works, our theoretical results also provide guidelines for designing future attention models. Experiments on various datasets validate our theoretical findings.

Under review as a conference paper at ICLR 2021
ANALYZING ATTENTION MECHANISMS THROUGH LENS OF SAMPLE COMPLEXITY AND LOSS LAND- SCAPE
Anonymous authors Paper under double-blind review
ABSTRACT
Attention mechanisms have advanced statc-of-the-art deep learning models for many machine learning tasks, Despite significant empirical gains, there is lack of theoretical analyses on their effectiveness. In this paper, we address this problem by studying the sample complexity and loss landscape attention-based neural networks. Our results show that, under mild assumptions, every local minimum of the attention model has low prediction error, and attention models require lower sample complexity than models without attention. Besides revealing why popular self-attention works, our theoretical results also provide guidelines for designing future attention models. Experiments on various datasets validate our theoretical findings.
INTRODUCTION
Significant research 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 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 et al. 2015 Vaswani et al., 2017). image captioning (Xu al. 2015), video captioning (Pu al. 2018), 8), visual question answering (Zhou et al., 2015 Lu et al., 2016), generative modeling (Zhang et al. 2018 etc. In computer vision, spatial//patiooemmpraal attention masks are employed to focus only on relevant regions of images/video frames for underlying downstream tasks (Mnih et 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 predict each symbol of the output sequence. Hidden state representations of recurrent neural network are typically used to compute these attention masks.
Substantial empirical evidence on the effectiveness of attention mechanisms motivates study the problem from a theoretical lens. To this end, it important to understand the loss landscape and optimization of neural networks with attention. Analyzing the loss landscape of neural networks an active ongoing research area, and it can be challenging even for two-layer neural networks (Poggio
Liao 2017: Rister & Rubin 2017 Soudry & Hoffer 2018 Zhou & Feng 2017 Mei al. 2018b Soltanolkotabi et al.) 2017 Ge et al. 2017] Nguyen & Hein 2017a Arora et al.) 2 8). Convergence gradient descent for two-layer neural networks has been studied
Allen-Zhu et al.)
(2019): Mei et al. (201 8b); Du et al. (2019). Ge et al. (2017) shows that there is no bad local minima two-layer neural nets under a specific loss landscape design. These works reveal the importance of understanding loss landscape of neural networks.
Unfortunately, these results cannot be directly applied for attention mechanisms. In attention models, the network structure is different, and the attention introduces additional parameters that are jointly optimized. To the best of our knowledge, there is no existing work analyzing the loss landscape and optimization of attention models. In this work, we present theoretical analysis of self-attention models /Vaswani et al. 2017), which uses correlations among elements of input sequence to learn an attention mask.

---

Under review as a conference paper at ICLR 2021
We summarize our work as follows. We carefully analyze attention mechanisms the loss landscape in Sections and 4. In Section 3, we show that, under mild assumptions, every stationary point of attention models achieves low generalization error. Section 4 studies other properties of attention models on the loss landscapes. After the loss landscape analyses, we discuss how our theoretical results can guide the practitioners design better attention models in Section 5. Then validate our theoretical findings with experiments on various datasets in Section 6. Section 7 includes a few concluding remarks. Proofs and more technical details are presented in the appendix.
ATTENTION MODELS
Attention mechanisms are modules that help neural networks focus only on relevant regions of input data to make predictions. To compare attention model with non-attention model, we first introduce a two-layer non-attern non-attention model as the baseline model. The network architecture consists of linear layer followed by rectified linear units (ReLU) as a non-linear activation function, and second linear layer. Denote the weights of the first layer by (1) E IRpxd the weights of the second layer by (2) E Rd, and the ReLU function by <(-). Then the response function for the input a E IR" can be written as y = (2)T ( 6((w(()) x)). We call the above function "baseline model" since does not employ any attention.
To study such mechanisms, we mainly focus on analyzing the most popular self-attention model. In this paper, consider two types of self-attention model.
For the first type of self-attention model, we consider attention weights that are determined by a function (a)): y = (2)T Φi( (w x @ (*))) where f()) a known mapping function from RR to IR, representing the attention weight of each feature with any given æ. This model is prototype version of transformer model (Vaswani et al. 2017). with pre-determined function as attention weights.
Second, we introduce more practical self-attention setup, which is the transformer model proposed in Vaswani al. (2017). To mimic the NLP task, set the input æ; = (z!...., x²)) R'xp, where a E IR', are t-dimensional vectors. Intuitively, each æ; corresponds to independent sentences for i 1,...,,4. .,, and x/'s are fixed dimensional vector embedding of each word in sentence Φ;- wº,w* _x E 12 are query and key weight matrices, and w Y xt is the value matrix. For each input æ; the key calculated as: K, (w K xx,"" E Rexd;; For 2h vector in the input, the query vector is computed as: Q = (w E Rlxdd for 2 = 1,...,P The value matrix V = wv V x, E Rds Then the attenniton w.r.t to the 2thh vector in the input æ; is computed as:
Q.KT
self(z) wei,w0 w*) = softmar
a,
(1)
Vdg
sel/(p)). This self-attention vector represents the interaction between different words in each sentence. The value vector for each word in the sentence 2 can be calculated as V,* Va
for I....,p. p. And a; self (a; (1)
sel Rdd This value vector is then passed 2-layer MLP
parameterized by (1) Rpd, xd and (2) E Rdx1, resulting in the following general model:
V w (2)) (( vec(w> x,,; ())))(())) +
where vec(-) represents the vectorization of matrix, and e, are i..d sub-Gaussial error.
3 SAMPLE COMPLEXITY ANALYSES
In this section, focus analyzing the loss landscape for the the self-attention model as introduced in Section In Section 3.I we consider the sample complexity of the model with known attention weight function fx:) In Section 3.2 we consider transformer self-attention model, in which the attention weight function is also need to be learnt. Section 3.3 discusses the sample complexity result for multi-layer self-attention model.
To avoid the non don-differentiabee point of ReLU é, we use the softplus activation function e (x). log(1 e702) Note @To converges to ReL. - (Glorot al. 2011). All i... Pro (x: To
2

---

Under review as a conference paper at ICLR 2021
theoretical results derived with 6ro (z) holds for any arbitrarily large To. For ease of notation, still use (zx) to denote the softplus function.
In this paper, we focus the regression task which minimizes the following loss function:L = E( [/(x)) - yll.. where h(x)) is the defined baselinc/attention models. Our theory can be easily extended to classification tasks well.
3.1 ASYMPTOTIC PROPERTY OF SELF-ATTENTION MODEL WITH KNOWN WEIGIIT FUNCTION
We start with a self-attention model with known attention function f(.) The objective function is:
min 2n
w (2)T ( (0(2(1 It ① f(xx))) y.)²
(3)
1=1
where [lf(2)]l2 1, as the sum one attention weights.
Before proceeding, we introduce three major assumptions for our analysis: (Al): Model output can be specified by a two-layer neural network with attention model structure; (A2): The Attention weights are mainly contributed by top several masks; (A3): Hidden layers of the network are sufficiently overparameterized such that they have sufficient expressiveness to predict y.
For space consideration, we provide the detailed justification of these assumptions and explain why they can be intuitively summarized as (A1)-(A3) in the Appendix (A.1). Specifically, (A1) supported by previous works. For (A2) and (A3), we verify them both theoretically and empirically in Appendix Section A.1 and in our experiments Section 6 In short, (A2) naturally holds since softmax function in attention model helps us to obtain concentrated attention weights instead of evenly spread out weights; (A3) naturally holds for overparameterized neural networks. More details can be found in Appendix Section A.1. In the following, we present their mathematical presentations.
Let (((((1) x @ f (m))) represents the d-dimensional random vector of hidden units with network weights ( For any vector 2 and 22, we denote var(2) (22) as the covariance matrix of residual
2 after taking linear regression on 22. The explicit form can be represented as var(21 (2)) =
Eaya Σ2ga Σz Σ2rzg where Σn represents the cross-covariance matrix of vector 2 and
22. The explicit derivation can be found on page 176 of Izenman (2013). In our setting, 21 = o( (1) x @ fx)),, and 22 = (((w(()) (1), a @ fax)).
(A1) There exists set of parameters (w ()) w()))) such that y, (2) (w ()) xz O f(x)))+ 6; for i 1, nn.,, where e,'s follow the sub-Gaussian distribution subG((0, C3) and a; 1 G
(A2) For any æ;, we order the attention weights f fxx:) in descending order as (a))) (a))))))
There exist positive integer < so < p and 0 T such that largest so attention weights satisfies [I ll 2 T.
f.. = (f.
(A3) For 7 > 0, Amax (Var(V( ((( x. @ (()))))( x @ f(x)))) o(72) n 2
log( L (pd+ + d) and n oc for constant k which satisfies 550 + TVP
Other than these three major assumptions, we also assume the regularity assumption (B1): the model samples x,'s 's are i.i.d. with maximum bound; network weights (1) and (2( have bounded Eg norms, denoted by C and C2, respectively; and the output y centralized. The regularity assumption (B1) is standard in the literature, which will be justified in the Appendix (A.2). Given these assumptions, we show that the sample complexity bound of attention model:
Theorem 1. (Sample complexity for attention model) Under (AI) and regularity assumption (BI), for any 0 and 80 s.I. k S /80 + T /P, given the sumple size
50 + ( 217
log( (V50 + +T/)) (pd + d)
7
where CRCr. with probability converging to I, any stationary point (1)(1), (((2)) of the objective function satisfies that: E( ((ū(1)) x @ f()))) E(y (z)) $ 72.
3

---

Under review as a conference paper at ICLR 2021
The proof sketch and complete proof are provided in Appendix Section E. To explicitly compare the sample complexity bound of our attention model with non-attention model, we provide the following corollary for sample complexity bound of baseline model. We just need to set so 0 and 1 in Theorem 1, fixing all attention weights to be equal, then the model is the baseline one as follows.
Corollary 1. (Sample complexity for non-attention model) Under (AI)-(3) and regularity as- sumption (BI), for non-attention model with [(x,) 1/ for all features æz at given the sample size:
p²12 Pd 72 log( (pd d)
With probability converging to I, any stationary point (16(1, w(()) of the objective function satisfies that: E( ((w(1), x 3 ff f()))) [(yx))) 2 2²
where
C2C2.
There is an important message from the comparison of Theorem 1 and Corollary 1. The corresponding sample complexity bound from Corollary This fact can be observed as follows. We can see the effect of concentration in attention models from Theorem 1. In Theorem I, the sample complexity bound is proportional to (v5g + T T When the attention is sufficiently concentrated, it means that we have much lower order of so p and T - 0, then our sample complexity will be significantly reduced.
higher order than the sample complexity Theorem I.
Also from the sample complexity bound of Theorem up log term, prediction error y is proportional to n -1/2, which is the optimal rate of convergence in regression. In Imaizumi & Fukumizu 2018) and similar works, they showed the generalization error convergence rate is O(n where 0 < <<< 1/2. These facts all imply that the bound is tight in aspect of sample size compared with existing works.
3.2
ASYMPTOTIC PROPERTY OF CELE-ATTENTIO MODEL WITH UNKNOWN WEIGHT FUNCTION
In this section, we extend our analyses to the transformer self-attention model introduced in Section
Theorem implies that if the self-attention mask can be precisely computed with f(x), we can derive its sample complexity bound. However, the function f(.) is not necessarily known, and needs be learned in real- applications. Therefore, the transformer self-attention setup is more desired in real-world settings. It provides concrete model to learn the parameter of attention weight function. Denoting w (1) 22)), w w w))
the two-layer self-attention model can be formulated as:
min w
w (2)2' (((w0)) vee(( "xaa self
y1)²
(4)
2n
We now introduce necessary assumptions.
(A4) For each column a." we order the entries of a. self in descending order as a
self
self
k(p)
There exist positive integer 0 < so < and 0 < t < 1, such that the largest so attention weights self ((leading) = {a {a self
sels self satisfies ||a (leading) (1) ...., (((0))
1l2 2 1 T.
(A5) Denote Φ = dl(w (1), ver(w "a,,
sel sel })), and Φ (((w(1), vec(w ())) For any a,,0
7 > we have Amax (V Var(( ( x ( = o( as n 2 ddd log( (k+ vpt) and n + co.
where similar to (Al), (w (1)x w (2)x V* a self correspond to the true parameter set such that = w (2)TT of (1)
(1), vee(w" V "naa eelf+))) + 65-
(5)
(A4) and (A5) are parallel to (A2) and (A3), assuming the attention weights are focused on 50
items, and a sufficient expressive power of +. Furthermore, sparse transformer Child et al. (2019).
sparsemax attention model Martins Martins & Astudillo (2066) and local attention model [odel
al. (2015)
can be regraded as a special case the assumption with r = 0, in which we only consider attention
weights between partial locations instead of all locations. We further verify this assumption in
4

---

Under review as a conference paper at ICLR 2021
experiments that 8o p and 0. Other regularity assumptions feature and parameter space
similar with the ones in Theorem 1. We denote is as (B2): The l2 norm bound for w((1), w(() an and R are, respectively, C1, C2, Cs. and Cs. and the & column norm of un is Cs- They are presented in Appendix (A.3). Now provide sample complexity bound for transformer self-attention models: Theorem 2. (Sample complexity for transformer self-attention model) Under (A4), (A5) and regu- larity assumptions (B2), any Y >> O and So s.t k S so, given the sample size:
(/50 + TVD) + (
((v50 + TVD]² + vpt
2
log(
??
C3 + Co. with probability tending to 1, any stationary point (ib(1, w((2, w ))) of the objective function satisfies that: E( ((()))o(((()))) veo(w" a.a (())e())) - E(yl=))? S²2.
where ?/e
CPO2C& and th
There are two important messages from Theorem First, theorem shows that, with the help of self-attention, we can achieve consistent predictions under more expressive class of models (equation 5). which is analyzed in Yun et al. (2019). Non-attention model does not have consistency for this class of models. To train non-attention model on data with self-attention structure, more layers and larger network parameter size required to reduce such bias. Second, we help of concentrating the attention in designing transformer model. Similar to Theorem 1, we see the sample complexity bound is proportional to (vsi + + +T -v))". Therefore, a properly small so and can significantly reduce the sample complexity of self-attention models. Later in our experiment, we show that it is exactly what is happening in the real transformer model. Our theorem also answers why the self-attention design with softmax function can effectively helps us achieve better prediction results. What's more, it also explains the effectiveness of sparse design in attention model. Sparse attention is one special case of our concentration condition with T = O. There has been work showing that sparse attention weights can significantly reduce computational cost and improve the performance, and is verified in sparse transformer, local attention model and sparsemax attention models et al. Luong et al., 2015; Martins & Astudillo, 2016).
the
2019
3.3
MULTI-LLAYER SELFFATTEENION MODELS AND RECURRENT ATTENTION MODELS
Theorems and 2 can be extended to multi-layer neural nets. Due to page limit, the rigorous definition, notations, assumptions and statements of Theorem are deferred Appendix Section A.4. Here we provide plain statement of it as follows.
Theorem 3. Given the overparumeterized and regularity condition, for any given generalization error level gamma, with high probability, multi-layer self-attention model can achieve generalization error smaller than 7/. given the sufficient large sample size.
Our analyses can be also extended to recurrent attention models, following the recurrent attention setup in Luong et al. (2015). The analysis recurrent attention model is deferred to Appendix Section B due to page limits.
NON-LINEARITY, FLATNESS OF MINIMA AND SMALL SAMPLE SIZE
In this section, we further investigate several additional properties how attention mechanisms improve the landscape of neural networks and keep the nice properties of baseline models in aspects of reducing unnecessary non-linear regions, sharpness of local minimum and it doesn't 'tt affect the loss landscape in small sample case.
4.1 ON THE NUMBER OF LINEAR REGIONS
We first study how attention mechanisms affect the number of linear regions (Montufar al.) 2014] in a wide two-layer neural network with attention of known attention weight function, when the number of hidden units is larger than the number of non-zero weights in f(x). This result shows how the connetticooccnttaation of attention weights effects the non- coneaiiyyy of loss landscape.
Su. which is the sparsity of the attention mask matrix, and the number of units in the hidden layer n, >> Sp- Then the maximal number of linear regions of the function by two-layer attention model with ReL.U activation function, is lower bounded by m 190
Theorem 4. Assume [=)
5

---

Under review as a conference paper at ICLR 2021
In general, the bound of attention model is smaller than the of baseline model. The corresponding plots on bounds can be found in appendix Section C.I. The result implies that, when appropriate attention mechanism is used, the reduction of number of linear regions leads to simpler landscape, yet the approximation error remains small.
4.2 ON PLATNESS/SHARPPESS OF MINIMA
Many recent works, such as Keskar et al. 2016), argue that flatter local minimum tends generalize well. However, in 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 of Hessian matrix can also be arbitrarily large. We show that this fact also holds for the self-attention mechanism, if no & norm bound on parameter (1i(1) (((2) is imposed. Here we introduce the definition of c-flatness in Hochreiter & Schmidhuber [197].
Definition 1. Given > 0, minimum 0, and loss L, C(L, 0, is the largest connected set containing 0 such that v E C(L,0,e) L(0') L(0) c, and its volume is called the e-flateess.
In the following Theorem, we analyze the flatness of stationary point for self-attention model.
Theorem 5. Consider the two-layer ReLU neural network with self-attention mechanism as stated in Section 3.2: Ys w(2)xT Φf (w(() vve(w" x,a; 2,a;e )t), and a minimum 0
(ub(1), T6(2), w. w² " satisfying that ui + 0 for i (1), V.Q, K. For any c > C(L, 0,e) e) has infinite volume, and for any M O, we can find stationary point such that the largest eigenvalue of V L() larger than M.
Theorem 5 iindicates that property on flatness of minima is maintained when attention mechanism is applied. Furthermore, & norm bound helps remove sharp minima which are bad in generalization. also coincides with theoretical and empirical result that flat minimum expected to generalize better in general (Keskar et al., 2016). We also dicuss the loss landscape attention model under small sample size. The results are deferred Appendix Section C.3.
GUIDANCE ON IMPROVING THE ARCHITECTURE OF ATTENTION MODELS
In this section, we provide insights into future attention model design through our analyses.
Our analyses suggest proper regularization is helpful in training an attention model. We can that l2 norm bound Ci. C2 play an important role in sample complexity bound. It implies that an & and & regularization on network weights (!) and (2) are effective in reducing the sample complexity. In Theorem 5, we also find that imposing constraints and regularization on network weights help remove sharp minima and keep flat minima with good generalization.
Regularization:
From the discussion of Theorem 1, Theorem and Corollary we conclude that proper concentration design with small so and can significantly reduce sample complexities. Our analysis show that the soft-max design concentrate attention on limited number of entries, which help reduce the sample complexity. In different problems, can further concentrate the attention by adjusting the temperature of the softmax function. The smaller the temperature, more concentrated the attention weights are.
Concentration on attention:
Overparameterization in query/key weights From our analyses, we can see that as dg. the dimen- sion of key and query matrices, increases, the sample complexity will not increase significantly. It indicates that we can obtain high expressive power attention model through overparameterizztion in query and key matrices increase expressiveness, without hurting sample complexity.
EXPERIMENTS
In this section, we validate our theoretical findings empirically. This section is divided into the following parts: (1) Verification of concentration assumption (A2) and (A4) on Portuguese to English translation task. (2) Ablation study the IMDB reviews dataset, showing the effectiveness of attention, regularization, and cenntiionconcentraatinn on self-attention and recurrent attention models. (3) Experiments constructed noisy-MNIST dataset using self-attention models.

---

Under review as a conference paper at ICLR 2021
Experiment 1: Concentration of top attentions
To verify that our key assumptions/obserrations (A2) and (A4) hold in the real-world task, we investigate the distribution of attention weights in the transformer model on the tasks of translating Portuguese to English as proposed in Vaswani et al. (017)) Then we randomly select the trained attention weights vector for 100000 different words in training samples. For each attention weight vector with sentence length p;, we calculate the Eg norm of largest [VPi] number of attention weights. We find the top LyP:] weights on average contribute 90.9% of attention weights, very close 1. The histogram is plotted as follows Figure T We see that most of sum weights are close to 1. The result indicates that the largest LyPe] attention weights contribute to most of the attention weights. It indicates that our assumption (A2) and (A4) hold with reasonably small < p and such as
80 = [Vp]-
Experiment 2: Ablation study on self-attention and recurrent attention on the IMDB reviews dataset:
Outtrbb 201 of top stoeesss wnggtt
We consider the problem of sentiment classification on an IMDB
reviews dataset (Maas et al. 2011). The task is to classify the senti- ment of a sentence as either positive or negative one. We zero-pad all our sentences to make their length equal to the sentence length 130. For every input word, we train their embeddings with random initialization (of dimension 100) which is then passed to the neural network. Hence, the dimension of input is 130 100. We con- sider baseline 2-layer MLP, and then consider adding self-attention, weight regularization and tempered softmax function into the model to verify our theoretical analysis and corresponding guidance in Section 5.
2009-
1500
05
07
85
Figure 1: Distribution of lz norm of top L/Pe] of attention weights
Baseline model: To train the baseline model, we first flatten the input one large vector of dimension 130 x 100 and pass a 1-hidden layer MLP with h hidden units. The model is trained using binary cross entropy loss,
Self-attention model: For self self-attention model, the dimensions of query, key and value matrices are w E IR x 100 w E R 100x 100
wy E IR 100x 130 respectively. We first compute the attention mask a; a per equation I (afe
geis 1300 Using this attention mask, the attended feature then computed as fau = w x,a,d The feature vector faut E R 100x 130, We then flatten this attended feature and then pass it through 1-hidden layer MLP with h hidden units.
Regularization: We impose a 10-4 & regularization on both wT and u72.
Tempered Softmax: We calculate the attention weights as a; fol (x w 20
softmaat 5.Q:K"" K", We multiply the inner product by 5(or temperature as 1/5). Thus, the softmax
Vdg
operator pushes the small attention weights close to zero while retaining all large attention weights. In this way, it achieves higher level of concentration of attention weights, comparing with the standard softmax function based attention.
Optimization: All models were initialized randomly with Xavier initialization. Binary cross-entropy loss was used to train the models. All models were trained using Adam optimizer with a learning rate 10-3
To test the sample complexity, vary the number of training samples in each experiment, train all models and compute the performance on the test set. We varied the fraction of training samples from Ik to 10k of the training data. Each experiment was repeated for 10 replications, and mean and standard deviation was reported in Table [
Table 1: Testing accuracy of self-attention ablation study on the IMDB reviews dataset.
Training sample
T0W
Buseline
0.892 (0.015)
0.309 (0.016)
0.821 (0.015)
0.542 (0.006
Buscine+rnaaaaizatin
0.737 (0.009)
0.823 (0.010)
0.837 (0.005)
0.841 (0.005)
Self Atlention
0.806 (0.007)
0.821 (0.012)
0.832
0.861 (0.004)
Self Anemlioneeeuuaaiiiatii
(0.11)
0,83 (0.011)
0.853000022
0.865
Self Ateetliontemmpene
0.798 (0.010)
03300 (0.010)
0.840000000
0.864 (0.01
Self Attentonn rennttionregulliiiaaaae sagglaarzaaiinneeeeeeee
0.817 (0.010)
0.843 (0.010)
0.856 (0.013)
0.868 (0.067)
7

---

Under review as a conference paper at ICLR 2021
Recurrent attention experiments: Beyond self-attention, also provide ablation study recurrent attention model, verifying the superiority of recurrent attention model, as our analysis in the Appendix B.
(1) RNN baseline:
The baseline here a bi-directional RNN attention model with LSTM cell size 32, and put the final output of RNN into 2-hidden layer MLP with hidden units 64 and 20 separately to return the final prediction. (2) Recurrent attention: For recurrent attention, we design the structure the same as Luong al. 2015). (3) Regularization: We impose the & regularization of 10-4 on network weights in the attention layer. Here we reported the test accuracy results in the following Table 2
Table 2: Testing accuracy of recurrent attention ablation study on the IMDB reviews dataset.
Training sample size
1k
3k
5k
10k
RNN Baseline
0.776 (0.023)
0.804 (0.009)
0.823 (0.006)
0.829 (0.015)
RNN eeeeiiiieeeellaizaationn
0.777 (0.012)
0.806 (0.010)
0.814 (0.008)
0.832 (0.007)
Recurrent Attention
0.789 (0.010)
0,821 (0.007)
0,820 (0.009)
0.835 (0.008)
Recurrent Atcanion-reculaiiztion
0.792 (0.010)
0.817 (0.009)
0.834 (0.009)
0.862 (0.002)
From the experiment results, significant effect of regularization and softmax temperature in self-attention model, and we also see regularization also helps in the recurrent attention model. They all coincide with our theoretical findings.
Experiment 3: Self-attention on Noisy-MNIST dataset
To prove the applicability of our analyses, further verify the effectiveness of attention on image classification task. We construct a noisy MNIST dataset based on the original MNIST dataset. For each original 28 28 image, we separate a 56 56 image into 16 square grids, and put the whole digit image randomly into 4 neighboring grids. Finally we generate all other grids with uniform random variables. Examples of our generated dataset are provided in Figure
We consider a standard 2-layer Convolu- tional neural networks with fillter size 64, kernel size 3 and max pooling size 2x2, following hidden fully-connected layer with size 128. (2) Self-attention-CN CNN We fit 2-layer convolutional network with filter size 64, kernel size * with a max pooling size 2 * 2, obtaining 1600-dimensional embedding for each of 16 square grids. Then treat these embedding of 16 grids as "16 words embedding" in sentence, fitting it into equation 2
We consider three models here. (1) CNN:
9
3
(3) Equal weight self-attention-CNN The model structure is the same as model 2, with fixing the attention weights between all grids are the same. Model (3) designed as alter- native baseline guaranteeing that the improvement of model is from the attention block instead of the model structure. The testing accuracy of three models are reported in Table 3. Model 1 has 5558k number of parameters, and model 2 has only 670k number of parameters. This fact also guarantees that the superiority of our model is not from training a bigger model.
experiment
Figure
Noisy MNIST dataset
Table 3: Testing accuracy self-attention on the Noisy MNIST dataset
Sample
CNN
54
Self-ameetionnCNN
0.970/0
(00)
(000)
Equal weight self attention CNN
05000000001
0.92/000411
0.00000022 (25))
From the table, we see that Attention-CNN model achieves almost perfect testing accuracy as the original MNIST dataset. Although equal-attention CNN has the same expressiveness power, its performance is much worse than Coention-CNN model. Again, It shows the usefulness of concentrating self-attention weights properly in the task.
7 CONCLUSIONS
In this paper, we study the loss landscape of neural networks on attention models, and show that attention mechanisms help reduce the sample complexity and achieve consistent predictions in the
8

---

Under review as a conference paper at ICLR 2021
large sample regime. Besides theoretical analyses of loss landscape, empirical studies validate theoretical findings. Based on our analyses, we discuss how regularization, concentration on attention, and overparameteriization in attention weight matrices can further improve the attention model.
REFERENCES
Zeyuan Allen- -he,, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in neural information processing systems, pp. 6155-6166, 2019.
Zeyuan Allen-Zhu. Yuanzhi Li, and Zhao Song. A Convergence Theory for Deep Learning via PaeerParameteiization.. 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üssun, Stockholm, Sweden, pp. 244-253, 2018.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. CoRR, 0bs//4409.4477 2014.
Andrew R Barron and Jason M Klusowski. Approximation and estimation for high-dimensional deep learning networks. arXiv preprini arXiv: /809 000900, 2018.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
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 of 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 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 2017.
design. CoRR,
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.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Masaaki Imaizumi and Kenji Fukumizu. Deep neural networks learn non-smooth functions effectively. arXiv preprint arXiv: 1.X2::8802.0444. 2018
Alan Julian Izenman. Multivariate regression. In Modern Multivariate Statistical Techniques, pp. 159-194. Springer, 2013.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch batch training for deep learning: Generalization gap and sharp minima. arXiv preprint urXiv:1600 1609.04836, 2016.
Jiasen Lu, Jianwei Yang, Dhruv Batra, and Devi Parikh. Hierarchical question-image co-attention for visual question answering. NIPS, pp. 289-297, 2016.
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to iitention-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.

---

Under review as a conference paper at ICLR 2021
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 Anmual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150. Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL http:
7/ www.aaclwebe
Andre Martins and Ramon Astudillo. From softmax to sparsemax: sparse model of attention and multi-label classification. In International Conference on Machine Learning, pp. 1614-1623, 2016.
Song Mei, Yu Bai, Andrea Montanari, al. The landscape of empirical risk for nonconvex losses. The Annals of Statistics, 46(6A):2747-2777 2a.
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, 11663)) 2018b. ISSN 0027-8424. doi: 10. 10.1073/pnas... 180657911
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.
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: 7704. 08045, 2017b.
Tomaso A. Poggio and Qianli Liao. Theory II: landscape of the empirical risk in deep learning. CoRR, abs/1700.00983. 2017.
Yunchen Pu, Martin Renqiang Min. Zhe Gan, and Lawrence Carin. Adaptive feature abstraction for translating video to text. Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 7284-720 72911 2018.
Blaine Rister and Daniel L. Rubin. Piecewise convexity of artificial neural networks. Neural Networks, 94:34-4,, 2017.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D. Lee. Theoretical insights into the optimization landscape of werr-parameterized shallow neural networks. CoRR, aws//107.00 04926, 2017.
Daniel Soudry and Yair Carmon. No bad local minima: Data independent training error guarantees for multilayer neural networks. arXiv preprint arXiv:1605. 08361, 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. In Advances in Neural Information Processing Systems, pp. 5998-6008, 2017.
Roman Vershynin. Introduction to 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. 2048-2084 2057, Lille, France, 07-09 Jul 2015. PMLR.
10

---

Under review as a conference paper at ICLR 2021
Chulhee Yun, Srinadh Bhojamapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of Sequence-to-sequenee functions? arXiv preprint arXi::1912.10077, 2019.
Han Zhang, Ian J. Goodfellow, Dimitris N. Metaxas, and Augustus Odena. Self-attention generative adversarial networks. CoRR, (bs//805.08311 2018.
Bolei Zhou, Yuandong Tian, Sainbayar Sukhbaatar, Arthur Szlam, and Rob Fergus. Simple baseline for visual question answering. CoRR, abs/1512.02167,, 2015. URLlnttp://arraiv.or///bb.
1512202160
Pan Zhou and Jiashi Feng. The landscape of deep learning algorithms. CoRK, abs// 705.07038. 2017.

---

Under review as a conference paper at ICLR 2021
APPENDIX
In this appendix, Section A presents detailed theorem assumptions and their justifications; Section B presents the theorem of sample complexity bounds for recurrent attention models; Section C presents details results in Section 4 of main paper. Section D presents key lemma of sample complexity bound on attention model with fixed attention masks for all samples; Finally, we provide proofs in section E.
A THEOREM ASSUMPTIONS
In this section, we provide more rigorous presentations on the standard assumptions of Theorems 1-3 and further justify these assumptions.
A.1 JUSTIFICATION OF ASSUMPTIONS
Here we justify (A1) to (A3) for Theorem 1, and explain why they can be intuitively summarized as in the main paper. (A1) assumes the expectation of output y can be specified by two-layer neural network with attention model structure. It has been studied that general bounded functions with a Fourier representation [-1, 1] can be well approximated by the defined two-layer network (Barron & Klusowski] 2018). Specific to attention models. Yun et al. (2019) prove the strong expressiveness of transformer type self-attention model give sufficient large network. Therefore very mild to assume attention model can identify the mean of output y.
(A2) assumes that the attention mask is mostly concentrated on the largest Sy attention weights. This assumption is naturally satisfied due to the softmax function in the self-attention equation. Sofumax function makes attention weight proportion to the exponential of the inner product of query and key vectors. This fact significantly enlarges the difference between the inner products, and makes the attention weights do not evenly spread out over all other entries, but only a few of them. Empirically, we verify that this assumption is satisfied in real-worke transformer translation task in our experiments. Please check experiment in Section 6 for more details on the distribution of attention weights in Portuguese to English translation. It shows that this assumption is well satisfied.
(A3) assumes that (o....(((w() y when sample size sufficiently large. In an overparameterized network with large d we know (é...(((w() O ff (x)))) have up to 2d number of linear regions over n samples. And (A2) assumes these linear regions space all directions of (((ww))) 3, 3 f(x)) up to o(y) term. It says that given the 2d linear regions and strong expressiveness of (ói...((w") x @ f f()))) }. the residual vector is a o(>) term with respect to all directions in Ed. This assumption is parallel to the full column rank condition in (Nguyen & Hein, 2017b), where we essentially assume in overparameterized network, the linear combination spans all directions in RR. AAllenZhu et al. (2019) also shows the fact in their Lemma B. and Corollary B.2 that in overparameterized networks, in every small region of parameter space, there exists set of parameters with good prediction. And it leads to a good landscape in each neighborhood. These results all indicate that it is reasonable to assume sufficient expressiveness of (1)(((()) a @ a)) in overparameterized networks as n + oo. What's more, with required large sample size, we can also straightforwardly evaluate this assumption by checking whether these (((w(1), x O a)) spread out the whole space and their linear combinations have good estimation on y. All these results guarantee that (A3) can be achieved overparameterized networks. We also empirically validate (A3) by computing the largest eigenvalue of such conditional covariance in (A3). We generate a two-layer network under self-attention model (4) with n = 500, d = 256, p = 2142, d, = 100. We choose the set of parameter same with our experiment setup in Section 6, and hidden layers are sufficient overparameterized with p = 2142. Then we compute the empirical conditional covariance Amax (Var(ób (w x, @ f()))) (w x @ f()))) = 3.13 * 10-8, This result indicates that the largest eigenvalue term Amare (V Var(o(( (r(((((())))) x, @ ((=))))(((w)) O f(x))) small enough to be assumes o(>²) as in (A3).
@ }())) obtain sufficient expressiveness to predict the output
12

---

Under review as a conference paper at ICLR 2021
A.2 REGULARITY ASSUMPTIONS OF THEOREM
In section 3, we described the regularity assumptions of Theorem 1 in words. Here we provide rigorous mathematical representations of these assumptions in the following (A1.1) to (A1.3).
(A1.1) my are i.idd and || so < Cz for any i ,.
(A1.2) There exist C1. C2. Cy such that Amax Ex) Cx2 1/ <i and 1/
C2 for any w0 w2(2) S.
(A1.3) ((w (w a f(x)))) E(v)= 0.
(A1.1) to (A1.3) are standard assumptions for parameter and feature space. (A1.1) and (A1.2) require upper bounds on the imput 2; and Eg bound for network weights. It is a standard assumption in landscape analysis (Mei et al. 201 8a b), and also it is crucial to remove sharp minima which may not generalize well (Keskar et al, 2016 Dinh al. (017). discussions after Theorem 4), These assumptions can be achieved through regularization.
(A1.3) can always be achieved by making the data centered.
A.3 ASSUMPTIONS OF THEOREM? 2
Similarly, for transformer-type self-attention, we also provide similar assumptions (A2.1) and (A2.2)
as follows:
(A2.1): (A1.2) holds with Eg = (ov(((((u())) veelw" a. (el Further there exist
Cs.Co and C7 such that || r S Cz. ||w Ir S Co. || V I|r S Cr, and ||
max Σ w I Cs. And || K ll Co for ,.
(A2.2) There exists a set of parameters (a", w((),, w(())) such that Y/
(2)))xi''(( (1) vec(wV a; (())(())))) es where aels is calculated by (I); 6
subG(0, Cf) for = 1, n, with o; 1L e,. And the output is centered as (A6).
These assumptions are all parallel to (A3) to (A6).
A.4 DETAILS ABOUT THEOREM 3
We consider D-layer network with self-attention structure. We denote the 1tth self-attention layer follows (a = w=*0(((" w (a where æ is the output of (A 1)th self-attention layer, with w Rdext, z Btxde-1, a E Rdi-ixde-1, w E Rd and =x Rd Rdd. xqa, zelf is calculated in the same way with two-layer self-attention net- works. Then we have the final output h(x) w (Poo(("), (eeo(, vee(ww xg D- ( where a D- (9D-1(-..99 (z)), and Dt IRlxd. D Rd In this way, the network caleu- lates self-attention D times and finally produce the final prediction. is worth mentioning that, to obtain scalar prediction in regression model, we flatten the value matrix of the last layer as in the two-layer model. We still denote u = (o " V D- x aa)lf)))) - E(y|æ). Then the necessary assumptions parallel to (A2) are as follows:
xy
(A6) There exists integer l and r such that E (1..... D} and r E {1,2), such that cov( (h(w"r) u) 2 ey²) for some constant and such that [Vh(w* ll S ck-
This Theorem also requires mild regularity conditions as follows:
(A3.1) All weights th for k = 1,...,D D and j = 1,2 satisfy || ||2 S C1o. And we assume the prediction is centered, i.e. F(u)
(A3.2) There exists
set of parameters (w) w²()) such that y Po Pi, vec( w V x D-1 ( + e, as defined, where E ~ subG(0, C for 1, 2, ., with æ, 1 6. And the output y is centered as (A6).
13

---

Under review as a conference paper at ICLR 2021
The following theorem provides sample complexity bound of multi-layer self-attention model: Theorem Under (A6) and regularity assumptions (A3.1) and (A3.2), drelf is the total number of parameters in all value, query, key matrices. Then for any Y 0. given the sample size: n Z log( S*t (dsett + where og is the Lipschitz constant of (h(w">). With probability tending to any stationary point (mr(), u0(2) vi0 .ū*)) satisfies that: E(((()) [B(y)=)))² - S
Because multi-layer self self-attention models include a large parameter set with complicated gradients, the assumptions are not as intuitive as the two-layer model. But the main assumptions are parallel, such that in an overparameterized network, Vh(w) spans almost all directions, and some of them are correlated with the bias term u. The results show that mainly the size of network parameters determines the sample complexity bound. It implies that efficient architecture design is critical in reducing the sample complexity. For multi-layer cases, the regular assumptions are stated as following (A3.1) and (A3.2).
B RECURRENT ATTENTION MODEL
In this section, consider analyzing the sample complexity bound for representative recurrent attention framework in Bahdanau et al. (2014). In the recurrent attention network, we follow the setting in self-attention model, such that æ; = (a ..... =²") E Rixp, corresponding words with t-dimensional embedding. And æ is the population version of æ;, Then the generative model can be represented as:
y/ w (2)T (w
(2)T
Σ [axx))=))) +
Analogous to NLP setting, a(x;) is a unknown function mapping æ; to t-dimensional vector, where axz;), represents the effect of the jth word in the sentence for point i. Then following the RNN setup in Bahdanau et a (2014), using data features themselves as their annotations, then for time
stamp k 1.... T. The recurrent attention model estimates for 4thh time stamp a* (x:) as follows:
8 = f(sk-11 Ck- -1); Ckj = score( (8 (1 x
p
0kk =
Σ Ckj
Ckj
Ch Σ
anjæ!
1=1
(k) w (2)T (1) ,cx)))
where score(-) : is the score function representing how well the inputs around position j and the output at position match. It can either be a dot product MLP. y; And we denote a(xx), = aT as the attention mask for the final time stamp. And f(-) is the function to update sk- Suppose the parameter set inside these two functions are w, and wy with number of parameters as d, and d accordingly. Here we show that when these two functions are expressive enough, recurrent attention networks also have sample complexity bound parallel to self-attention models. Here we introduce necessary assumptions.
(k) is the prediction in htth time stamp.
(A7) When T is sufficiently large, the output can be predicted by the two-layer network with independent sub-Gaussian error with variance o², i.e, there exists a set of pa- rameters (w(1), (2))) such that 3/4 (2)x (((w())+ E a(xx) + es where
6, subG(O, C3) for 2. with æ; E
(A8) Suppose (A2) holds when we substitute a ④ f(x) with E a(xx))ar
(A9) We assume l|w,l|2 S Css and ||wlll Cg.
(A7) (A9) are parallel the assumptions in the self-attention case. They be justified similar to them, which is discussed in Section of Appendix. When T is large, recurrent attention models can represent wide class of attentions weights. Thus assume yi be expressed by such recurrent attention models.
Now can provide the following sample complexity bound.
14

---

Under review as a conference paper at ICLR 2021
Theorem 6. Under (A7) (A9). for any O, suppose
GF n 7?
CC (td + d da + dy)
log(
where 1l C2C2 such that if there exist stationary poini(s), then with probability tending any stationary point (wo(1, 2()), ..y wa) satisfies the following prediction error bound:
P ()) $22
E (
==1
Theorem provides a sample complexity bound for recurrent attention networks. The bound holds under expressiveness assumption (A7). It also shows trade-off of expressiveness and sample complexity. When the number parameters is too large, the sample complexity will be too large; When the number of the parameters is small, we don't have enough expressiveness consistency. If f-) and a--) are properly selected, they will be sufficiently expressive to obtain good stationary points, and also the number of parameters du and dy will not be too large. In this way, an ideal sample complexity bound to these good stationary points can be achieved as Theorem 7 says. However, with an wer-complicated design in these functions, the sample complexity bound will be large; With an over-simple design, such good stationary points don't It exist. It is parallel to trade-off between approximation error and estimation error learning theory. The theory implies a good design of the recurrent structure will help achieve an optimal sample complexity in recurrent attention model.
achieve
C APPENDIX FOR SECTION 4
C.I DISCUSSION OF THEOREM
For illustration, the linear region bounds under different sparsity levels are plotted in Figure 1. The top red line for baseline model with p 100, and other lines are bounds for attention model with different sparsity level so-
In a sufficiently wide network, 1 80 is much smaller than [ Then, ( (2x))² holds as long as n1 exp( since _ is close to p when 80 is relatively small, the 8-80 result still holds when 71 is larger than the order of p. For illustration, the linear region bounds under different sparsity levels are plotted in Figure 1. The top red line for baseline model with p 100, and other lines are bounds for attention model with different sparsity level the bound of attention model is
p p- log #a), ). Given p-og log S P__ p-80
p-80 log p.
p-
450
Figure 3: Number of linear regions in log scale v.s. sparsity
80- In general, smaller than the one of baseline model. The result implies that, when appropriate attention mechanism is used, the reduction of number of linear regions leads to simpler landscape, yet the approximation error remains small. We interpret that attention mechanisms help reduce unnecessary non-linearity inside the landscape.
C.2 DISCUSSION OF THFOREM 5
Theorem Sareemm@iiiie inddicates that property on flatness minima is maintained when attention mechanism is applied, and there exist good sharp minima, coinciding with the observation in Dinh al. (2017).
However, there is no guarantee that all sharp minima are good in generalization. Revisiting our analysis in Section 3 the restriction on the parameter space helps remove these sharp minima. Specifically, provide upper bounds the & norm of (w( (2)). These constraints restrict the parameter space and remove all sharp minima which we construct in the proof Theorem [
In these constructed sharp minima, (% 02 goes infinity, and e, norm bound guarantee that
15

---

Under review as a conference paper at ICLR 2021
it cannot happen. In these constructed sharp minima, &2 norm bound guarantees that it cannot happen. Practically, &g norm bounds can be achieved through a proper la regularization, which will be discussed in Section 5. It also coincides with our theoretical and empirical result that flat minimum are expected to generalize better in general (Keskar et al. 2016).
C.3 SMALL SAMPLE SIZE
In this section, study the local minimum of wide neural networks in small sample regime. Nguuyen & Hein 2017b) proves that a two-layer neural network model can always achieve perfect empirical estimation error when the sample size is small. Here, we extend this result for self-attention model.
Theorem 7. For self-attention model in Sec 3.2 if frank(l( (wi(), vee( a,
a. ()) M)).2,,,))))) = n.
self
x) of object function (4) is a global minima. Then every stationary point (wo(1, rx(2) u. "0,
rank(o(((()) C a))==1......) = n. is a mild assumption in a wide network with over- choose the number of units d to be larger than n, a) 1=12...... holds with measure zero. In other words, almost
parameterization. We can see that long the linear dependence of (wx(1) F. ④ 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 2017b), where they assume the number of units in some layer is larger than the sample size. When the sample size is smaller than the number of units in the network, this theorem holds for the network without attention. It has been proved by Nguyen & Hein (2017b) and Soudry & Carmon (2016) under different conditions.
D LEMMA: MODELS WITH A FIXED ATTENTION MASK
In this section, we introduced attention model with fixed attention mask, as fundamental building block of attention mechanisms. In the fixed-mask attention model, we consider a dataset D = 1.V {0,, Ji=1 D2 E IR, Vi E I. where the output y4 depends only certain regions of input æ,, i.c., y/ = f 3 æ;), where a" is an unknown fixed attention mask, and f*(.) is the ground-truth function that is used generate the dataset and the vector a* [0, 1j. The set of entries {aş la # 0}
corresponds to the relevant region of the input, while the complementary set {a; la = 0} corresponds the irrelevant region.
The following Lemma 1 on the sample complexity of attention model with fixed attention mask, is not only important building block of the proof of Theorem 1 but also provides helpful insights for understanding attention mechanisms and revealing the main idea of our proof. With the sparsity structure of attention mask a. attention mechanisms constrain parameters in a smaller space, thus they reduce the variability of the empirical landscape, and also reduce the covering number of parameter space. These results lead lower sample complexity compared with the baseline model not employing attention. Similar to Corollary result, it is straightforward to calculate the sample complexity bound for the baseline moo employing attention). To achieve the same error bound, we substitute so with p in the bound, and this results in much larger sample complexity bound.
The attention model we use can written as:
f(x) = w (2)T ((w (1) @ a))
(6)
The assumptions for analyzing this model is the same as the assumptions of Theorem 1(A1 to A3 and A to AI.3), where we substitute f(x) ) by a in all assumptions with ||allo S sg. For clarity, we state them here explicitly:
(A.L.I) There exists set of parameters (w(1)+ (())) such that y: (2)) oi (1)())) @ a)) 6; for i = 1,2,..n, where G's follow the sub-Gaussian distribution subG(0, C2) and x; Le E;
(A.L.2) ||allo 8o with at most 8o non-zero weights.
(A.L.3) For any 7 > Amax (V (Var(o(( (w(()) x, ① a*)) (o( (1) x @ a))) = o(72) as fo log( (pd + d) and n + xo.
(A.L.4) x; i.i.d and ||ax|loc < C2 for any = 1.
16

---

Under review as a conference paper at ICLR 2021
(A.L.5) There exist C1, C2. Cx such that Amax (Ex) Cx. /20 < C1 and 1/12
C2 for any w11), wll) E S.
(A.L.6) E(w (2) (ww a C a))) = BByy) = 0.
Lemma 1. Under (A.L.) to (A.L6), then for any 7 > 0 and so such that l S O(sa), suppose
n S log(
307 (pd p
where - CCC2 Then with probability tending to 1, any stationary point (ā, , ()(2)) of the objective function (4) satisfies the following prediction error bound:
E[ E( (2)T 6((0011) @ a)) -(ylx))) S 72
(7)
E PROOFS
In this section, first provide a proof sketch for Theorem for readers understand the high-lelvel idea. Then we provide proofs of all our theorems. Lemma 1, Theorem 1, Theorem 2, Theorem 3 and Theorem are proved in similar manners. For Theorem 1, Theorem 2 and Theorem 3, we omit the exact same part of the proof as Lemma 1.
E.1 PROOF SKETCH OF THEOREM
We target at proving that under our assumptions and required sample size, all local minimum must have prediction ability as good global minimum up to y2, The proof can be divided into two major steps. First, we show that for all parameter sets with bad prediction, their ((w(2), It @ a) term must be correlated with the bias F(ylæ;; - (( ) (wr(1) 2; @ a). It leads to large magnitude of population gradient with respect to (2) Second, we construct an c-cover of parameter sets (w (2)) to show that sample gradients converge to population gradient, such that sample gradient with respect to (2) also away from zero. Thus these parameter sets with bad prediction cannot be local minimum. And we conclude that all local minimum must have prediction as good as global minimum up to ()²²). The complete proof is provided in the Appendix Section
Proof Sketch:
E.2 PROOF OF LEMMA
Proof. As described in the proof intuition, the proof divided into two steps. First we study the lower bound of population risk gradient [EE.,v(V R, (2)) (t)))(()),, In step we study the convergence of Rt, (wn(2) [|) to the population gradient [| ( (V Rt (w)() )l22. We further separate these two steps into three lemmas.
Lemma 1.1 proves the first step, we study the landscape of population risk, showing that with high probability, we know the population risk with respect to x and y || |E (V. (wx())) lz is large; Lemma 1.2 and Lemma 1.3 prove the second step. Specifically, in part (b), consider the convergence of population risk only taking expectation on y. Ey, R2 (un())) )). In part (c), finally we consider the convergence of empirical risk gradient V R,, (w, (2) to (VR,(,
We introduce necessary notations beforchand. To emphasize the role of x and y separately, here we denote (1,(1), (2) a) EE,, (R,, (1) random, and V R(w) as corresponding derivatives. And we denote E,, (V (1) w (2) (2)))a)))))))) Er.y (V R,, (1) (,(2) a)), which is the expectation of the empirical loss function with expectation to both æ and y. the proof, we may use o(7) for vector/matrix case. In these cases, it means that every element in vector/matrix is o(7).
which is the expectation of the empirical
loss gradient with respect to y. treating
Lemma 1.1 Under the assumption of Lemma 1, when equation 7 is violated, with probability going to that, the population risk gradient with respect to (2) satisfies that:
|rr, (V Pe, (2) 0(r)
17

---

Under review as a conference paper at ICLR 2021
Proof: We denote:
( , 3 aa)) E(y/æ)
(8)
and u, as the version with the specified sample index. Then the derivatives of population risk with the expectation to y can be presented follows:
1
(2(w
( (0, Oa))
7
1
VR(w())
ui (x; ① (a) (2) ② ((w()
2, o a)))
i=1
VR(a)=
1 " Σ
us (2x; O (w (w (2) oo (, a)))))
0,,00
By (A.L.I), we know that B(y((x))) = w(2) (( (1) x, 2 a*)). Therefore when (a, w(), w(2)) = (()), w(2))), all the u,'s are zero, and all the gradients with expectations to æ and y are zero. Thus for any true set of parameter (a", w('), w(2))), they have zero gradient expectations automatically. And the key of our proof is showing that with high probability, any
parameter (a, ((),, w()) cannot be stationary point if B( (1()()) o (w)(), c ā)) [(ylx)) (3) 7²
because their gradients w.r. to w(22 must be bounded away from zero.
In the following section, we prove Lemma 1.1 by showing that E(u²) 2 72, we must have [[E( ((w())l2 2 O(7). We prove it by contradiction.
In this proof, we denote r cov(E(y)a)) ((w xxa)))) E Rd, and the covariance matrix for o (1) @ a)) as:
EE coo(di((w x 1 2 a)), (, x xxaa))
If we have |[ (2))) la o(7),
||cou(u, (( m4 a))l(2 o(7);
and plug it into cov((w (2) (2)))) ((((w(())
a)) (1 x 0 aa))) + o(7), we have:
(e (2) - r|12 = | (2))
# a)) F(ylar) ((w 2 a)))ll2
cou o (w xù a))ll ((7)
Since we know the true mean y can be specified by a true set of parameter (wo()))), (2) a²), and we denote that the covariance matrix for o( (1) a²)) as:
ΣS = cov(ói((u))
x
By
xx0a)))
And we denote the cross-covariance matrix of vector oi ((w (1)+ xOa*)) and vector (( as:
EE = cor(o:((w o a 25 ((w x a)))
Then with zero expectation on E(ylx) and (2 ((w 2 a)) we have:
E(u²)= var(E(ylxx) - w (2)T (( aa)
vur(EEylaa) + var (2)T aw (2) x; & a)) - 2c0r (zov(E(y)æ), (217 (2 ( xx
(2) (((w(( (1)*
+ car(w 2)T ((w x, Ra)
2cor(w (2) (((w
fl x a))
(2)+T Σ w (2)a 2w
+ ESr )2)
w (2)+T Σ (2)x 2w (2)-T (2).TTE Σ. E. (2)- + w (2)-T (2: (2- +c)0(7)
w (2)+T E E.ET E.)w (2) + 2
18

---

Under review as a conference paper at ICLR 2021
where 2/ E E. is exactly the residual of covariance matrix after taking regression ((... x f(x)))), as we defined in (A.L. 176 of Izenman 2013)
w (())+T (E - -.EE Σe )2) is a quadratic form of that conditional variance matrix. By (A.L.5),
know its largest eigenvalue Amax ES Σ.=g" Ee) o()). Finally we obtain that:
E(u²) (E(( ))12 cxx((²2)
I w (2 * Amax (2) E.E0 Ee) + cxo(?²)
= =(
By contradiction, we conclude that if B(u²) 2 72, we must have ([ 2 0()). Here we finish the proof of Lemma 1.1.
Lemma 1.2 Under the assumption of Lemma I, when equation 7 is violated, the risk gradient of w with expectation to y satisfies:
Ey (V (w (2) )) 2 O(7)
proof: Lemma we have shown that population level, parameter sets with bad prediction has large magnitude of E(V ((v()) In Lemma 1.2, we use e-cover technique to bound the gap
between (()w(()) (2) (recall that V RR( ([ww (2) is short notation of (V,R,,, (w)())) and E(VR(w (2)).
Our parameters (1)) (2) and a are inside the Eg balls (O, Boxd, C2) and (9 22) where (e r2) represents a p-dimensional l with center and radius r. By Lemma 5.2 in Vershynin (2010), the e-covering number Negg-Neee for these three balls are upper bounded by:
Ne. S (3C1/e) Ne, S (3C2/e)nd, No. S (3s
Then we know the joint cc-coveingg number for the union of all three parameters Na. satisfies that Ns. < Ne, Ne, For the case of notation, we denote = (a, w()) 2())). Let θ = {0,,... 0N, } be corresponding cover with N3. elements. Following the e-covering we construct can always find O such that for any feasible 0, there exists j E [N] such that max([|w/ [1) - wl2) w!l ( - (1) [|a(] a all2) S e. In this proof, we use (5)
for (w(1), w(2)) a) separately,
parenthesis subscription (j) represent the jth element in the cover, to distinguish from other subscriptions.
By triangle inequality, we have:
() 2 (V 12) ([
VR(W (2)
Therefore, we only need bound both parts on the r.h.s of equation 9 We start with the first term (2) ||V R(w () lle To achieve this, we first bound the gradient term (1)(((()) O a). For any fixed parameter set, we calculate:
| (Calló(w G a))l2))(lo((u)) a e aa)
Here we denote 1.etivee as the & norm of (1))) on active features, i.c. on the feature such that its attention weight a not zero, and from the sparsity condition we know that there are most so such nonzero elements in x @ a. Combining with the &g bound of (2) we apply Cauchy-Schward inequality:
[|
(1) (1) ② a)ll2 S max(læ O @al (1) Jaactive C2
It implies:
1i (1)
()) x. aa =
(10)
From Lemma 1.1. we know there exists a constant e such that |E, (V 2 cy for some constant We denote g2 of C2C.. Then apply Hoeffding bound the & norm of
19

---

Under review as a conference paper at ICLR 2021
()) E. (VR(W ()) ), with the upper bound of it O(o2). Denoting VR(w
(2)
as the gradient with
respect to jth parameter set in c-cover for j E
P P(IVR((W))) E. (VR z 67
S exp( n
3
Then by union bound, over all N. elements in O
P(jj E [N.] VV 21 207 2 1 exp(p-1r
(1) la term. Here we use u, to represent the prediction error for nd instant with respect to parameter (a, w(1), w( and use W:() to represent the term with respect to the parameter from jth element in e-cover set. By triangle inequality, we have:
Secondly we analyze (R(w(()) V R
vi(w (3) I
2 n2
(u:6((w0))
2e ④ a)) ( (1) w; @ a())))ll
2 (1) 2
[u - (10)))(((00 (1) m. a))]l
mo a a)) (( 3 as o a()))))laa
i=1
$ xx C2E
We choose c = 30d Co C2 and plug back above results into equation then at least with probability
1 O(N, exp(exppp n
have (VR(w(()))l T. Therefore choose:
log( S0C, (pd + p d)
ey
such that N: exp(- 2272 = o1). Finally we can conclude that with probability - on(),, for any (a, (1) (2)) such that (E (2) (( x o a)) [yyl=))) 2 7, we have [| > ey
Lemma 1.3 Under the assumption of Lemma I, when equation iis violated, with probability going that, the empirical risk gradient with respect to (2 satisfies that:
(w (1) 2 0(7)
Proof: So far in Lemma 1.2, have shown that for population risk with respect to y. with high probability, all the parameter sets with poor prediction expectation, i.e E(u²) 2 O(72), their population risk gradient with expectation y must be away from zero. Now we move forward to show that empirical risk R, (2) converges to [[V In aspect of w??), they can be represented as:
V (w) vR(w (2))
co((w()) x. Oa))
With (A.L.1), know that e sO, CJ), thus - vn) by C.L.T, combining
the bound for d((w(!), , Dy O a)) we have derived in the proof of Lemma 1.2, with sample size n Spp log( SoC, (pd + p + d), conclude that with probability - (1)) oy
(w) vR(w (2)
)12
Cy
(11)
(12)
Recalling part (a), under the first case that w.h.p (( 2 S for any parameter (a, (1) () with () ((w(())) X @ a)) Elyl X)ll 2 7. Combining this with (1). we conclude that for any positive constant y > with required sample size, with high probability that ||V R., (w 22) > 0, thus they cannot be stationary solution for our loss function.
20

---

Under review as a conference paper at ICLR 2021
In other words, under our assumptions, all the stationary points (ā, w()) ((() in our programming satisfy the prediction error upper bound rate ry w.h.p, when sample size:
shn n 2 c²q]
S09 )pd p d)
log(
67
E.3 PROOF OF THEOREM
be proved following the same manner as Lemma I, substituting a by f()). only changing two bounds. Here we specify these two different bounds with Lemma 1. First difference is the bound of (1o(((()), æ, O f(x)))) in Lemma 1.2, and the second is the e-cover number Lemma 1.2.
Proof. Theorem
By assumption (A2), we denote f())deed weights such that (xf 122 T. And ( as the other 8p attention weights and it satisfies sævbll
(co) } as the top so leading attention
f [lf(z) + ( = 1. We denote Vioad and Wlead as the features and network weights corresponding to leading attention weights, and ssub.
(1)
w enb as the features and network weights corresponding other attention weights.
(1)
Parallel Lemma 1.2, we just need to adjust the bound of of ((w f(x))) term as:
( @ f(x)))l22 = (lo(( Blead 3 f(x) + (w0 bub: æsub @ f()) f(x) Read) 1l2 + ( (1)
S V2(llo((w))) (wofod
lead Blead
suhs Bsub o (x))
The two terms in the last inequality can be bounded by Cauchy-Schwarz inequalll separately: [ lead Blaad (c)
(1)
s l/00 112 |ateed f( (x))laddl2
< || n0 (1) ||2 [ [ ( x lead 1/2 S C1 / * C. * -1 -) s Va,CiCr
And
| (w (1)
sub- Deub 2 f(x) I |!l(() 12 | Deub 2 f(x)
[|w sullarsut subll2 - C1 +7 FA
Combining both inequalities, we have:
S (1 m 5 + +VDDCG
(13)
Further we have
| $ (Callo((wo (2))))) x, f())) (V50 +
(14)
Second, in this case, since fæx is not optimized together, but caleulated from instead. We don't have to consider the e-cover number for in the maximum operator. Therefore the new c-cover number for w(!) and (2) are upper bounded as:
N₂, < (3C1/e)4, I (3C2/c)pd
Substituting the new e-cover bound to the theorem, obtain the final sample complexity bound:
2 (vbu + T T/pp'
logí 807 (pd d)
22
&
21

---

Under review as a conference paper at ICLR 2021
E.4 PROOF OF THEOREM 2
Proof. Theorem 2 can be proved following the same manner as Lemma 1, substituting a by a add only changing several bounds. Here we specify the differences with Lemma
Finally, we obtained a new 6-covering bound accordingly:
No s (3C//e)", No, I (3C2/c)puud Ne S (3Cs/6)"d,, Na S (3Co/c) ( NNo S $3307/e)t
where N = Ne. Then, the new v; terms parallel to Lemma 1.2 is:
8; = uu0((1ll (w(1), veelw> 2; ((a
where u4 = (w(()))(( vec(w' w æ. E(u(z;)) Under assumptions, using the same argument Lemma 1.1 and Theorem under overparameterized network, we have that, there exists constant e such that either [|R(V ((())( 2 ey-
In the case when (E(V(( 2 cy. Here we derive the new bound of || with respect to æ. We have:
11:113 (C2llo((w vecl(( x; sel/ ",aa"")))))
11l))) veo(w> x;a; sel 1))
Then we derive the bound of ( vee( æ,a; 2e la (leadd llz 2 - T, following the same steps derived in equation 13 define Lead and self top so leading attention weights and other following attention weights, Then we have:
self |)) Using our assumption (A4): self
(o vee(w" z,a; 2,,; 12))2 s w |12lll w I1all2,, @ a; a
S 220, (atead @ alead ll2 Doub @ a self ll2))
S v2(v00 TVDJCCCs
Finally we obtain:
|| & (Callo((wo)) veo(w" (1) vec(w" -
Denoting 5² =
parallel to Theorem 3, applying Hoeffding bound and union
bound:
P(jj [N., 1u5l22 - 2g1) s N. cxp(-n- E²
207,
o272
Second, we bound [[ V (2) )o)1l term The attention weight gap is bounded as:
1Q: Qigy KT lmaxx [1(x;) (llet)(((o))) (ww0(TTTTW))))) (2t)
(0)))((u)) (w ()) w 20
S <t(Cs Co)
With this bound, we know a has lower bounded. Recall the softmax function for vector / is defined as:
sel Lipschitz function under assumption (A9) such that [Q:K7]
exp(13)
softmar( Σ. exp(33,)
It is Lipschitz continuous on B when D exp(B:) has lower bound, which is satisfied by (A2.1).
Thus the bound can be derived as:
K.K. softmax
2i0) KT (0)
| self(z)
self(z) a.6)
112 S softmax
)
die
Vdl.
I Vpt(Cs C..
22

---

Under review as a conference paper at ICLR 2021
With these bounds, have:
|| VRR(w 2) (R(w())
"6 vec(w" x,aad(1)) - ( vedww" x,a Mll
c "(o)) veclw" a;a;"!) ))l2
Σ " w x, unyz vec(w])) 1)2)))
((( T WD ( Ca 2 (1v50 / C (C, Co))
Recall that (4()) and a:G)) are corresponding to the jth element in the epsilon cover. Denote & ((vRo + T /D)² + CC, (Cg + Ca)), and choose - and combine the above results. Then at least with probability (N, exp( dag )) we have [ |l 3 Therefore we can choose n 2 2oa log( (aCG (pd, + d + 2pd,), such ca
then N, exp(-n - o(). Thus with this required sample complexity, we have | ((
E(VR(w())) 2 z 2cm 3
Finally conclude that with high probability, any parameter (wo(), wx(2) (0) (R) w()) with E(u²) 2 7 have || [( 3. Then following the same empirical risk convergence argument as Lemma 1.2, with high probability they cannot be stationary point. We conclude the sample complexity bound as:
/ppt2pd,ddii 72
log( ((V50 TVD)² vpt)nam
2
E.5 PROOF OF THEOREM
Proof. Under assumption (A3.1), we know all input features and weights are bounded. Therefore know ( Ill is a Lipschitz continuous function on all parameters, and we denote its Lipschitz
constant Cj. For =",,, we can derive that:
VR(w
= Σ u vh(w 1m1
Under (A6), if have ((((x) E(y 2 22 then:
| R
= [|ooo(Vh(( u))l Z 0(7)
Then similar with Theorem 1 and 2, we construct an e-cover over all parameters 0 = (w=+, ==2 w 2 w²k), and denote it (O....., 0v) } such that for any feasible parame- ter, there exist j E [N] such that the maximum & distance to f is smaller than e. By calculating the number parameters in all matrices in O, we have
1 log(N.) = Oddeet + Edd (dx + do)gk)
e
Denoting ( as the gradient with respect to jth parameter set in e-cover for E (l,..., N.):
-E, 2 ?) s exp( -pp--n G2
ck
23

---

Under review as a conference paper at ICLR 2021
By union bound, we have:
P E [Nc. VV 207, x No exp(-xppp
e22
C
Secondly we analyze [R((">)) V (i(w (3) )l1 term. As we have shown that the gradient is Lipschitz continuous, thus we have:
[VR((=*>) VR(W cke
C7_ then at least with probability 1-0(N, exp(-n - we have /| R (( ( > 3ex*
We choose e
3 Therefore can choose 2 log( (dder) + D(d dodx)), such that N₂ exp( +
o11). Finally can conclude that with probability 1 1-0,(1) (1), for any (a, w(1), wl() such that a)) [yy))))² 2 % have [R () > eT Then following the convergence of empirical risk procedure of Lemma 1.2, we show with probability going to such that | Rm (w(2))ll > and all parameters with prediction error O(7) cannot be stationary point as long log( Ek (d D(d + d,gx)). Thus complete the proof.
E( (((())((())))))))))))
E.6 PROOF OF THEOREM 4
Proof. First with [2) so, know all the inputs æ; with corresponding a; 0, will be inactive in the network. We can omit all these inactive inputs. Then we split nı units into 8o group, with % number of units in each group, and discard the leftover units. So different groups correspond to Su active inputs with non-zero attention weight.
Inside each group, for example in jth group, denoting = 1. we choose the input weights and biases for = 1,2, 1,,,,.-,0 as:
h (a) max w,
h2 = max(0, 2w 1}.
ha(a) max (0, 2w, - (q-
here we assign wy to be row vector with jth variable equal to 1 and all other entries to be 0. And in the second layer, we choose w!?) = (wy, ws), where Wa = 11, 1,... (-1)9+1), corresponding to hn to hy each group. Then the designed network has linear regions inside each group, giving by the intervals:
(-oo, 0), (0, 1), (1,
lg - 1, 00)
Each of these intervals has a subset that is mapped by wzh(æ) onto the interval (0,1) Montufar et al. (2114) Therefore the total number of linear regions is lower bounded by 21 |
E.7 PROOF THEOREMS
Proof. Here we define an a scale transformation such that:
Ta (1) (2)) (1. ( (2)
And all the value,query and key matrices remain the same. Then we know the Jacobian determinant for T. = (pd.-1jd Since pd,,d 2 d, as we assign a - oc, such that the Jacobian determinant goes infinity, and the volume of C'L, 0, e) goes infinity.
For the Hessian matrix, we still assume a positive diagonal element 8 in w()) Simi- larly we have the Frobenius norm ( of
L²)( =
[
a
I aI
0
a a-'
aI
is lower bounded by 28. When we choose sufficient small a. have the biggest eigenvalue of 2 L2L(Ton,.o (0)) is larger than any constant M. Therefore there exists a stationary point such that the operator norm for Hessian is arbitrary large.
24

---

Under review as a conference paper at ICLR 2021
E.8 PROOF OF THEOREM 6
Proof. First, we obtained new e-covering bound for the parameter set (wl2), w(l) w. wa):
No < (3C1/e)", No S (3C2/e))",M, Nog s (3Cs/e) No. S (3Co/e))
And N₂ S Ne, Similar to Theorem we denote
(w (w (2)2 o((w!))
Σ a((x))a)) -((y)))
(15)
Then the derivatives of population risk with expectation y be presented follows:
1 n
vR(w(()
uo (1 a(x))a)?
And also, the vy term is: uup((u)) (1)
) alx; )a?)) Following (A15) and same procedure as Lemma 1.1, we have [B(VR(w( 12 2 cy. we have new bound of ||b|| with respect upper bounded by o² CCC follows for the case when [[E(VR( )) cy. Then following the same approach as Theorem
with normalized attention weight. Same argument
and 2, we obtain the sample complexity bound:
o² 22g2
log( 1C1C2C&CgC,, )d (d+ + td + +d + da)
ni
67
E.9 PROOF THEOREM
Proof.
Here we only consider the derivatives with respect to wn(1) and (2(2), they can be presented as follows:
n
V (2) = Σ
( 2za;;e!1)
n
(R)(w())
u wvec(w" z.a e.aeell (w) o o ((w(1), veclw> x z,aeel)) ))
By assumption, (onk(((((w()) veo(w" w x.a )==.,... n. Solving the linear system, we must have u; 0 for any = 1, ..., to satisfy that V (2) = O. Thus we know that the loss is exactly zero inside sample. Thus must be global minimum.
F FIXED ATTENTION MASK EXPERIMENTS ON NOISYMNIST
In this additional experiments, another variant of MNIST dataset is constructed where the images of digits from the MNIST dataset are embedded noise, as shown in Figure (a). Then we compare a two-layer MLP baseline model (in blue) with global (fixed) attention model proposed in Supplemental (in red). Fig (b) shows plot of epochs v.s. test loss while Fig (c) shows a plot of epochs test accuracy. We observe that the global attention model obtains lower loss values and higher test accuracy with fewer epochs. Again it verifies the superiority of attention weights are from the effect of concentrating attention weights.
25

---

Under review as a conference paper at ICLR 2021
6
2
(a)
0.20
12.3
10.0
0.006
0.04
(b)
(c)
26


Reviews:
[{"review_id": "uTYJi7YX-Xc", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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 provides theoretical analysis of the sample complexity and the loss landscape of attention mechanism. Under several assumptions, the paper claims that attention models have lower sample complexity than non-attention alternatives. Besides, it also shows every local minimum of the attention model has low prediction error. The theoretical findings of the paper can further provide guidelines for designing attention models. A few empirical studies are conducted to confirm the theoretical findings.\n\nThe theoretical results rely on a few assumptions, which I am afraid are quite strong.\n\n1. (A1) assumes the data is generated by the attention model plus sub-Gaussian noise, which seems very strong to me. Since your data is obtained from an attention model, it is natural that an attention model can fit the data better with less data than a non-attention alternative. In other words, the sample complexity comparison (Theorem 1 vs Corollary 1) is not fair to me.\n\n2. (A3) looks weird to me, and I don’t see any explanation in the paper. What does this complex assumption mean? Is it just for conquering a step in the proof?\n\n3. Why is the assumption (B1) hidden in the text but not listed in parallel with other assumptions?\n\n4. Again, I have difficulty in understanding the complex (A5). \n\nMy feeling is that the author makes strong and hard-to-interpret assumptions, in order to apply the classical learning theory to show the results. \u0010The findings are not surprising nor inspiring that the guidance are just common practices in real applications.\n\nWhen you talk about the loss landscape, is it possible to draw some figures to compare the loss surface of attention models with that of non-attention models? An example would be in this paper: https://papers.nips.cc/paper/7875-visualizing-the-loss-landscape-of-neural-nets.pdf\n\nTypo: The Attention weights are mainly contribute —> contributed.", "extra_scores": {}, "extra_text": {}}, {"review_id": "gGnrlQtALdG", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "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:\nThis paper studies attention mechanisms.The main results of the paper are upper bounds on sample complexities of attention models and the baseline non-attention model. The paper also further investigates several additional properties and provides empirical studies.  \n\nConcerns:\n(1) Without going into the details, the main problem of the paper is that it only gives upper bounds on the sample complexity. When comparing the upper bound of two algorithms, the comparison does not show which one is better. To show one algorithm is better than the other, you need to prove that one’s upper bound is smaller than the other’s lower bound on sample complexity.\n(2) The paper proves the result only for a specific non-attention model. Without further justification, it is very unclear why this is enough to represent other non-attention models.\n\nReasons for score:\nI vote for rejection. The theoretical results are not sufficient to support the claims made in the paper about attention models. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "VDNGLUB506p", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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:\n\nThis paper analyzes the sample complexity and properties of loss landscape for self-attention neural networks. Specifically, under the assumptions about teacher-student setting, \"sparsity\" of the attention mask, enough over-parametrized student model, and regularity assumptions (all the weights are bounded, etc.), in a regression task with square loss, the authors proved that a two-layer MLP with attention enjoys a smaller upper bound on the sample complexity compare to the baseline without attention. Using very similar techniques and proof procedures with stronger assumptions, this result can be extended to multi-layer or recurrent attention models. Experiments were done on NLP and CV tasks with shallow networks/small datasets to validate the theoretical results.\n\nMain assumptions:\n\nThe assumptions for different theorems are slightly different, but they are somewhat parallel and can be summarized into the following categories. Each theorem requires one assumption from each category.\n\n1. Teacher-student setting: There exists a two-layer teacher model with attention that the ground-truth labels are the outputs of the teacher network plus i.i.d. sub-gaussian noise.\n\n2. The largest few (less than input dimension) weights in attention mask has a constant share of the norm of the entire mask, i.e., the weights of the attention mask are sparse/highly concentrated.\n\n3. When the sample size is large enough (tend to infinity), the student model with learned weights is significantly close to the teacher model in terms of covariance matrix of residual w.r.t. linear regression.\n\n4. Regularity assumptions, e.g., the norms of all the weights and samples are bounded.\n\nMain results:\n\n1. For two-layer neural networks, under all the assumptions mentioned above, the authors proved that the model with a fixed attention mask has a smaller sample complexity upper bound compare to that without an attention mask. By using the same procedure and techniques with stronger assumptions, the result can be extended to attention with a normal self-attention mask, multi-layer self-attention models, and recurrent attention models. The proofs basically come from the same lemma (Lemma 1 in appendix) and can be considered as different corollaries of Lemma 1 under different sets of assumptions.\n\n2. Experiments were done to validate the theoretical results. Specifically, the authors did experiments on Portuguese to English translation, IMDB reviews sentiment classification, and noisy MNIST with small models, e.g., two-layer conv nets or fully-connected nets with/without attention. The experiments showed that the weights of the attention mask are concentrated, and attention in these tasks has the effect of reducing sample complexity and improving performance.\n\n3. The authors also showed that the sparse attention model has a smaller lower bound on the maximal number of linear regions and suffers from the scaling issue when considering the flatness of minima.\n\nPros:\n\n1. The problem analyzed in this paper, i.e., understanding the underlying reasons why attention works, is an important problem, and this paper provides some theoretical together with experimental evidence for this.\n\n2. This paper is written in a clear and well-organized way so it is easy to follow. The experimental settings are explained in detail, and the experiments are replicated multiple times.\n\n3. The theoretical proof and experimental methods appear to be correct and reasonable.\n\nCons:\n\n1. The authors seem to use attention somewhat like a black box. The weight matrices, e.g., queries, keys, values, are learned during the training process, but the authors just assume that they have some nice properties as stated in Assumption A3. This makes the attention mechanism in the proof essentially the same as the known fixed function f. However, I think one of the core parts of the attention mechanism and part of the reason why it works is the learning process of the query, key, value matrices, which is ignored by the authors in this paper.\n\n2. Similar to 1, the authors also hide other properties of the loss landscape in assumption A3. This assumption is not only about over-parametrization because it is a universal property of all possible weights along the training trajectory or on the entire landscape. I believe that over-parametrization can be shown to prove the existence of such weights, but this assumption is much stronger than the over-parametrization itself because another important thing is how the training algorithms find such points. This is also part of the reason why the authors can have strong results that hold for *any stationary point*.\n\n3. The comparisons in this paper are about the upper bounds or lower bounds, but there is a lack of evidence about the tightness of these bounds. For instance, for the comparison between the sample complexity for models with/without attention, the bounds seem to be very large and a bit loose since empirically the networks can achieve small error with much smaller sample sizes. Therefore, it is somewhat unfair to do a direct comparison between the upper bounds. Besides, the loss landscape of models with and without attention may have different $\\gamma$, which is not addressed by the authors in the paper.\n\n4. The logic for this flatness part is kind of confusing to me. The scaling issue is to show that flatness alone cannot fully explain the generalization performance, but not having that issue does not mean the model can generalize better. Besides, the bound on the model weights are the assumptions, not conclusions, of this paper. The authors did not rigorously prove that their model does not suffer from scaling issues. \n\n5. For the noisy-MNIST experiments, the authors were claiming that effect of attention is not mainly due to the model structure by showing that an equal weight attention model performs poorly. However, the authors did not decouple the performance gain achieved by attention in the concentration of the weights and the learning process of the mask. I wonder whether the model will perform well if a fixed concentrated instead of equal weight mask is applied. If the concentration of the mask weights is very important, it should improve the performance of the model by a lot.\n\nRecommendation:\n\nI tend to vote for rejecting this paper. My biggest concern about this paper is that this paper considered some parts of the problem as black boxes and directly have some assumption about them (in other words, the assumptions made by the authors, especially A3, are too strong and hide a lot of things like implicit dependencies on network width, attention mechanism and model structure in parameter $\\gamma$), but some of the things skipped by the authors may be the core part of the problem, making the analysis provided in this paper a bit superficial and the results somewhat limited.\n\nSupporting arguments for recommendation:\n\nSee Cons, especially points 1 and 2 there.\n\nQuestions for the authors:\n\n1. Please address the cons mentioned above.\n\nAdditional feedback:\n\nThere are some notation confusions. For instance, $\\tau$ is used as both the parameter for the softplus activation function (last paragraph of page 2) and in assumption A2. Besides, for the paragraph before assumption (A1) and assumptions (A1) - (A3), the authors are using both $x\\odot f(x)$ and $x_i\\odot f(x)$.", "extra_scores": {}, "extra_text": {}}, {"review_id": "5zRWMh0lBp", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "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\n\nThe paper presents a theoretical analysis of self-attention mechanisms in neural networks by considering a two layer neural network with known and learned attention mechanisms. The paper shows that in the non-agnostic case, i.e. assuming that data is generated from a function in the function class of two-layered neural networks with self-attention, then every local minima of the squared loss has low expected generalization error. The theoretical analysis reveals that using neural networks with self-attention can have lower sample complexity than standard 2-layered neural networks when the data is in fact generated by a two-layered neural network with concentrated self-attention weights. The paper justifies some of the assumptions through experiments on IMDB reviews data set and also shows both theoretically and empirically that weight regularization can improve generalization.\n\n# Strength\n\n1. While attention mechanisms have had tremendous empirical success in NLP, little is known about their theoretical properties. Towards that end the paper takes a first stab at theoretically analyzing attention mechanisms by considering the expected generalization error of local minima of the loss functions.\n2. Through some toy experiments the paper reveals that weight regularization can improve generalization of two-layered NNs with self attention.\n\n# Weakness\n\n1. The paper goes for too many results at the cost of readability. Lots of results (for recurrent self-attention) are deferred to Appendix and even formal statements of Theorem are not in the main paper. \n2. While analyzing two-layered NNs is acceptable for making progress, the main problem is the non-agnostic setting where the authors assume that data is being generated by a two-layered NN with concentrated self-attention. Existing results for two-layered NNs (e.g. Allen-Zhu et al. 2017) are for the agnostic case. This is a significant shortcoming of the current work.\n3. Better sample complexity of attention model comes from the assumption that the true generative model has concentrated self-attention weights, i.e. most of the attention weights are concentrated in a few attention masks. So fitting a model that exploits this assumption, , i.e. a model that belongs to the same class as the true function, would have better sample complexity.\n4. Apart from the non-agnostic assumption, the paper also makes a number of very unrealistic assumptions on the weights (Assumption A3, A1.3 and A4) of the model class. All these assumptions make assumptions on not only the true weights (the weights of the NN generating data) but on the weights of all NNs in the model class. So it is not clear if the model class is even realizable.\n5. There are also many problems with experiments which are detailed below.\n\n# Questions to address in rebuttal\n\n1. For which w's does Assumption A3 (maximum eigenvalue of covariance matrix of residual), A1.3 (zero mean), and A4 (query, key, value) need to hold? If this needs to hold for all w's (and not just the true w's) then this is a very strong assumption. Is the model class realizable? \n2. \"Assumption A1 (non-agnostic setting) is supported by previous works\" which works? Needs reference. Experiment 1 shows that the learned self attention weights are concentrated not that the true data is being generated from such a model.\n3. What is k in Assumption A3 ?\n4. Seems like $1/\\sqrt{n}$ rate comes from the assumption that maximum eigenvalue is $o(\\gamma^2)$ for some $\\gamma > 0$ ? Why use $\\gamma^2$ ? Why not just $\\gamma ?$ \n5. Not clear where assumption A3 is used in the proof. Apparently A3 is used in proof of Theorem 3 but A3 is also needed in proof of Theorem 1. It is stated in proof of Theorem 3 that \"by Assumption A3.1 we know all input features and weights are bounded.\" but that is not what A3 states.\n6. The theoretical claim is that l2 norm of weights determines generalization error but experiments use l1 norm. Why?\n7. Analysis is for regression tasks but experiments are on classification.\n8. Why noisy MNIST ? Why not on the original MNIST ? It is obvious that paying attention on the digit blocks and ignoring the uniform noise blocks will help.\n\n# Minor comments\n\n1. we first introduced a two-layer non-attention model ⇒ we first introduce a two-layer non-attention model.\n2. we consider attention weights are determined ⇒ we consider attention weights that are determined.\n3. There are also a number of other grammatical mistakes that should be fixed to improve readability.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The three reviewers seem to reach a consensus that the assumptions made in the paper are too strong and hard to interpret. In particular, R1&R2 made the comments that the generative model for the data by itself uses attention, which seems to be make the comparison unfair. The authors seem to argue that the attention model is still expressive enough, which in the AC's opinion could be true but does not justify the use of the generative model when comparing the sample complexity of the two methods. The reviewers also pointed out a few other limitations of the paper. The AC mostly agrees with the reviewers' points (though perhaps with one or two exceptions.) In summary, I think the assumption of using an attention model itself seems to be a big enough issue that makes the paper not read for publication at ICLR. 
Decision: Reject