Linear attention is (maybe) all you need (to understand Transformer optimization)
Authors: Kwangjun Ahn
Xiang Cheng
Minhak Song
Chulhee Yun
Ali Jadbabaie
Suvrit Sra
Venue: ICLR 2024 poster (2024)
Keywords: Transformer
optimization
adam
clipping
heavy-tailed noise
directional smoothness
Primary area: optimization
TL;DR: Shallow linearized transformer exhibits same training difficulties as a real transformer, hence opens up the exciting prospect of having an analyzable proxy to reality.

Abstract: Transformer training is notoriously difficult, requiring a careful design of optimizers and use of various heuristics. We make progress towards understanding the subtleties of training Transformers by carefully studying a simple yet canonical linearized *shallow* Transformer model. Specifically, we train linear Transformers to solve regression tasks, inspired by J. von Oswald et al. (ICML 2023), and K. Ahn et al. (NeurIPS 2023). Most importantly, we observe that our proposed linearized models can reproduce several prominent aspects of Transformer training dynamics. Consequently, the results obtained in this paper suggest that a simple linearized Transformer model could actually be a valuable, realistic abstraction for understanding Transformer optimization.

Published as a conference paper at ICLR 2024
LINEAR ATTENTION IS (MAYBE) ALL YOU NEED (TO UNDERSTAND TRANSFORMER OPTIMIZATION)
Kwangjun Ahn" MIT EECS/LIDS kjahn@mit edu
Xiang Cheng
Minhak Song* KAIST ISysE/Math
MIT LIDS
chengxêmit edu
minlaksongkkaist.ac.kkk
Chulhee Yun
Ali Jadbabaie MIT CEE/LIDS
Suvrit Sra MIT EECS/LIDS suvrit 8mit eau
KAIST AI chulhee, chulhee.yunlkaiie yunkkaist ac.kr
j jacbabai eetaaammmiii..mm odu
ABSTRACT
Transformer training notoriously difficult, requiring a careful design optimiz- ers and use of various heuristics. We make progress towards understanding the subtlcties of training Transformers by carefully studying a simple yet canonical train linear Transform-
linearized shallow Transformer model. Specifically, ers to solve regression tasks, inspired by J. von Oswald et al. (ICML 2023), and K. Ahn et al. (NeurIPS 2023). Most importantly, we observe that our proposed linearized models can reproduce several prominent aspects of Transformer training dynamics. Consequently, the results obtained in this paper suggest that simple linearized Transformer model could actually be a valuable, realistic abstraction for understanding Transformer optimization.
INTRODUCTION
Transformer architectures (Vaswani et al., 2017) (henceforth, referred to as Transformers) have shown impressive performance in various applications (Devlin et al., 2019; Bubeck al., 2023). However, training Transformers is notoriously difficult and laborious; see, e.g., observations given by Liu et al. (2020) well scaling laws (Kaplan al., 2020). In particular, training Transformers requires carefully designed optimizers as well as use of various heuristics. For instance, as illustrated in Figure 1, stochastic gradient descent (SGD)-th the workhorse of most deep learning optimization problems-faite to train Transformers effectively. This failure is in contrast to the success of SGD when applied train convolutional neural networks (CNNs) on vision tasks.
Several recent papers propose number of different explanations to why Transformer optimization so difficult. There general consensus in the literature that the loss landscape of Transformers has number of distinctive features that significantly differ from standard optimization theory assumptions. Most notably, it is empirically verified through various experiments that stochastic gradient noise is heavy-tailed and non- non-Gause GansGan (Zhang et al., 2020b; Kunstner et al., 2023) and the loss landscape is significantly ill-conditioned (Zhang et al., 2020a: Jiang et al., 2022; Pan and Li, 2023). In particular, standard assumptions are incapable of dealing with and explaining these observations, and as a result, Transformer optimization has become more of an art than science.
A major obstacle in understanding Transformer optimization is that flllefledged Transformers extremely complicated to model. One can probe the Transformer's 's properties by measuring quantities, such gradient norm or smoothness, but it is much harder to parse the inner-layer workings, and to satisfactorily answer questions such as: why does the loss landscape have such features, or how do algorithms like Adam perform better than SGD Transformer training?
Therefore, having an appropriate mathematical abstraction is necessary for progress in understanding Transformer optimization-an an abstraction that is as simple possible, while still being able to capture the essence of Transformer optimization. The main message of this paper is that distinctive features of Transformer training also arise a far simpler setting: the linear attention model, without
* Equull contribution, alphahetical order.

---

Published as a conference paper at ICLR 2024
nonlinear activations and feedforward networks, being precisely the sought abstraction. We verify that training this model on a www-dmeeniionl linear regression task displays all the distinctive features that have been observed on the full Transformer, suggesting that surprisingly simple model can serve as a testbed for rigorous understanding of Transformer optimization.
Main contributions. We summarize our main contributions as follows:
We propose the problem of training shallow linear Transformer model on random linear regression model for understanding Transformer optimization. We verify that this model reproduces all the optimization features and phenomena that have been previously reported for full Transformers. We leverage the simplicity of our model to look deeper into how these features arise, by changing settings (e.g., data distribution, the number of layers). Our results reveal that the unique features from previous work get more pronounced in our linear Transformer setting when the data distribution becomes more heavy-tailed, or the number of layers increases.
We expect that such a simple abstraction has great value not only for theoretical research but also for development of optimization methods for Transformers. However, these directions are out-of-scope of this work, and left for future work. As preliminary, we first survey the previous works that seek to characterize and understand the Transformer optimization landscape.
DISTINCTIVE FEATURES OF TRANSFORMER OPTIMIZATION
Numerous recent papers have identified a number of distinctive features of the Transformer optimiza- tion problem, which set it apart from commonly studied optimization objectives, even other neural networks such as CNNs. As shown in Figure 1, one of the most striking features is the following:
Adaptive methods like Adam are significantly better than SGD!
(Adam>
This is in stark contrast with the training of other neural networks (e.g., convolutional neural networks) for which several works have shown that the values of adaptive methods are marginal (Wilson et al., 2017). This phenomenon sparked the interest of the optimization community in investigating the main causes, and subsequently, recent works (Zhang et al., 2020b; Kunstner et al., 2023; Jiang et al., 2022; Pan and Li, 2023) have identified various "unique" features of Transformer optimization.
SGD(+m)
Adam(+m)
SGD( m)
Adam( m)
Training loss
10
101-
8
10
10
10
10:3
10
10-9
10-5
2
10
10
Epoch
100
Epoch
100
0
Epoch
100
0
Epoch
40
0
Epoch
5
(a) CNNS MNIST and CIFAR-10
(b) Transformers PTB, WikiText2, and SQuAD
Figure 1: Adaptive optimization methods like Adam are much more effective than SGD for training Transform- This experimental result is taken from (Kunstner al., 2023, Figure 1). (+m) denotes "with momentum".
In this section, we discuss them one by one in detail, building preliminaries for our main results. In order to discuss each feature, we first give a whirlwind tour on some background in optimization.
2.1 WHIRLWIND TOUR OF (CONVEX) OPTIMIZATION THEORY
For symmetric matrix M, we denote by Xmaax (M) and Ami (M) the largest and smallest eigenvalue of M, and by ||M]| the spectral norm of M. For simplicity, we assume the training loss function f twice differentiable. We introduce the following standard concepts in the optimization literature.
Lipschitzness. We say is G-Lipschitz if /VIll, S G. Smoothness. We say f is L-smooth if / < L. Strong convexity. We say f is ju-strongly convex if Amin (V2f) μ. Condition number. The (local) condition number (z) is defined as X= ((x))//mmnn (V2 provided that Amin 2 fx)) >
Bounded stochastic gradient noise. In most SGD analyses, it is assumed that the stochastic gradient ((r) satisfies the bounded variance property: E [[g())) 2 o
2

---

Published as a conference paper at ICLR 2024
Transformers (in practice)
Shallow linear Transformers (see Subsection 3.1 and Table 1)
1. Gap between Adam vs. SGD (Zhang al. 2020b; Kunstner et al. 2023; Jiang al. 2022; Fan and LL, 202312
(GD
Adam(+m)
(GD(-m)
Adam(-m)
101
10-1
10"
20-3
24:0
Epoch
100
Epach
Epoch
Iteration
105co
Iteration
Iteration
13000
Figure 2: For Transformer optimization, adaptive methods like Adam are strictly better than SGD. (+m)
denotes "with momentum" and (-m) denotes without momentum. Our plots only show the momentum variants of SGD and Adam as they perform better in all cases.
Left 3 plots: Full Transformers, from (Kunstner et al., 2023, Figure 1). Right plots: Shallow linear Transformers (see Settings 1, and 3 from Table 1).
2. Hcavy-tailed stochastic gradient noise (Zhang 20206; Kunsiner al,, 2023);
PTB
Wik-
SQUAD
2004
100
20000
100
50-
250
5000
10000
2500
Gradient
0.7 0.8 0.9
Gradient
1.0
1.2
2.5
Gradient
5.0
Gradient
Gredient
Gradient
Figure 3: The stochastic gradient noise is heavy-tailed for Transformer optimization. The top-right comer of each plot is the quantile-quantile (q-q) plot between the histogram (y-axis) and its best fit Gaussian (æ-axis). The q-q plot above y = line toward right. showing its heavy-tailedness.
Left 3 plots: Full Transformers, from (Kunstner et al., 2023, Figure 1).
Right 3 plots: Shallow linear Transformers (see Settings 1, 2. and 3 from Table 1).
3. Robust condition number of the landscape (Jiang al., 2023)2
Iteration 750
Iteration 1250
Iteration
Iteration 1250
Layer#
Pe Aa
15
1.65 (0.65)
2,01 (1.00)
Setting
1.76 (0.40)
1.58 (0.41)
17
1.91 (0.53)
1.43 (0.63)
Setting 2
3.14 (0.97)
5.98 (2.86)
22
3.54 (1.21)
2.28 (1.18)
Setting 3
9.57 (13.3)
6.53 (3.55)
Figure 4: The comparison of the robust condition number (see Subsection 2.3) between SGD and
Adam for Transformer optimization. Numbers in parentheses show standard deviation. Left table: Full Transformers, from (Jiang et al., 2022, Table 1). Right table: Shallow linear Transformers, see Table 1.
4. Directional smoothness gap between SGD v.s Adam (Zhang et al., 2020s; Pan and Li, 2023)2
deem
Figure 5: log(directional smooth- ness) against iteration (see Subsec- tion 2.4) for shallow linear Trans- formers (see Settings 1, from Table
2030
The concepts defined above are typically of great importance the theory convex optimization, as the convergence rate of gradient-based optimizers (e.g., gradient descent) typically depend these quantities. For instance, the convergence rate of gradient descent gets better as the Lipschitzness or smoothness constant gets smaller, the condition number gets smaller (Bubeck, 2015; Nesterov

---

Published as a conference paper at ICLR 2024
MNIST
CIFAR-10
PTB
Wex--
SQuAD
15-
200
100-
Counts
10-
40
100
250
0.1
0.2 Gradiant error
25
30
0.7 0.8 0.9
1.0
1.2
2.5
5.0
Gradient error
Gradient erorr
Gradient error
Gradient error
Figure 6: The heavy-tail stochastic gradient noise for Transformers. Under the same setting as Figure 1, Kunstner ct al. (2023) plot the stochastic gradient noise at the initialization. The top-right corner of cach plot is the quantile-quattll (9-9) plot between the histogram (y-axis) and its best fit Gaussian (æ-axis). Notice that the stochastic gradient noise for the convolutional neural networks vision tasks (MNIST, CIFAR-10) is much less heavy-tailed than the Transformers on NLP tasks. We will revisit this plot Figure 10.
et al., 2018). Building on these concepts, we now discuss the previous studies on Transformer optimization. Several recent works have connected the difficulties of training Transformers to the unconventional features arising from the loss landscape of Transformer optimization.
2.2 HEAVY-TAILED GRADIENT NOISE (ZHANG ET AL., 2020B; KUNSTNER ET AL., 2023)
In (Zhang et al., 2020b) (entitled Why are adaptive methods good for attention models?), was observed that the stochastic gradient is typically more heavy-tailed for Transformer optimization than other neural network optimization. In particular, they make case that this is opposed to the standard bounded variance condition for SGD analysis see Figure 3 and Figure 6. They posit that this phenomenon might one of the main reasons behind the phenomenon (Adam>SGID) they also theoretically show that adaptive step sizes in the form of gradient clipping is required for convergence.
noteworthy follow-up work by Kunstner al. (2023) reveal that the heavy-tailed stochastic noise may not explain the full picture. In particular, they compare the full-batch versions (hence no stochastic noise), and notice the phenomenon (Adam>SGD) still hold. Since there no stochastic noise in this setting, the explanation based on heavy-tailed noise does not apply here.
2.3 ILL-CONDITIONED LANDSCAPE (JIANG ET AL., 2022)
In another inspiring work (Jiang et al., 2022), the authors seek to understand the difficulty of Transformer optimization through the lens of condition numbers. In particular, they consider a "robust" condition number defined as ROPT == Amax (V2)), f)/Auau (V=)) and here the reason for Xmedian instead of Amia handle degenerate Hessians. They observe that during Transformer optimization, non- adaptive optimizers like SGD tend to have larger robust condition number than adaptive optimizers like Adam; they posit that this phenomenon Figure 4. Jiang et al. (2022) also report that this gap is not there when training convolutational neural networks on image classification tasks, and suggest that this phenomenon may be rooted in unique features of the Transformer which are missing in other popular neural networks.
one of the main reasons for (Adam>SGD) see
2.4 DIRECTIONAL SMOOTHNEES (PAN AND LI, 2023)
In a follow up work by Pan and Li (2023) (entitled Toward understanding why Adam converges faster than SGD for Transformers). the authors again corroborate (Adam>SGD). In addition, they further observe in (Pan and Li, 2023, Figure 6) that proper gradient clipping techniques further accelerate optimization. In order to understand this phenomenon, they propose an explanation based on "directional smoothnesss" along the iterates Iy. More formally, they consider the following Taylor expansion along the iterates: for == || zell-
((++11) - f(x) = (x:) (#+1 -x.) + (at+1 2.) VV2 (2:)()( - xx.) + O(1²),
and define the directional smoothness as (x++1-2,) T 4 ((x))((+))1)) met1 = In particular, based on the above calculations, one can infer that smaller directional smoothness implies better optimization
In fact, in their paper, they instead consider the maximum diagonal entry of the Hessian divided by the median diagonal entry approximation of this quantity.
4

---

Published as a conference paper at ICLR 2024
f f(s) becomes smaller. They claim that the directional smoothness holds the key to understanding (Adam>SGD) (as well as Transformer optimization in general). They also verify that adaptive optimizers tend to have smaller directional smoothness values, and employing gradient clipping further reduces the directional smoothness. Once again, Pan and Li (2023) hypothesize that this feature is unique to Transformers, they observe that adaptive algorithms demonstrate worse directional smoothness than SGD for, c.g., ResNet training.
2.5 GENERALIZED SMOOTHNESS (ZHANG AL.,, 2020A)
We discuss
noteworthy work (Zhang et al., 2020a) that identifies another unconventional feature. We note that the main motivation of (Zhang et al., 2020a) was not about understanding (Adam>>GD)) they also observe their proposed feature in some other non-Transformer networks such as ResNets. The main observation made by (Zhang et al., 2020a) is that the standard smoothness assumption not suitable for neural network training. Instead, they observe that the spectral norm of Hessian typically grows with the norm of gradient at the current iterate (see Figure 16). Based on this observation, the authors define the following generalized smoothness:
Definition We say f is (Lo, Im)-shooth if [²)(()] s
Lo + L1 IVf(())] When L, this condition recovers the standard smoothness condition.
A coordinate-wise version of Definition was considered in (Crawshaw et al., 2022). Under Definition 1, they demonstrate that non-adaptive SGD needs more iterations to converge than an adaptive method based on the global clipping of gradients.
Thus far, we have seen several features identified in the previous works that set Transformer opti- mization apart from other neural network optimizations. In the next section, propose simple yet canonical Transformer model that exhibits all these features.
LINEAR SHALLOW TRANSFORMERS HAVE THE SAME LOSS LANDSCAPE AS PRACTICAL DEEP TRANSFORMERS
In this section, we show that a simple yet canonical Transformer model exhibits all the features in Section 2. Specifically, the optimization problem to be solved the training of linear Transformers on random instances of linear regression, model recently proposed for understanding of in-context learning (Garg et al., 2022; Akyürek et al., 2022; von Oswald al., 2023; Ahn et al., 2023b; Zhang et al., 2023; Mahankali et al., 2023).
3.1 LINEAR TRANSFORMER ON LINEAR REGRESSION
Data distribution. The data distribution can be thought of as the random instances of linear regression. Concretely, for 1, 2 ... n 1, let (+)) Rd be drawn i.i.d. from a distribution Dx. We then
draw w. Dw and then generate the scalar responses y [(x(1), wx)...., (x((), (e/)),,.)] ) RR.. Now
the input of the data set consists of these linear regression examples:
(1))
x22)
(())
x(++1)]
Input matrix: Zo
R((+1)((+11)
(1)
(22)
(m)
The goal is to predict the missing (n+1), detail below.
Optimization objective. Let (F W) ((d+1)x(++1) - R denote the prediction of the I--layered
linear Transformer with parameters W. Our optimization objective given by
f (W) = Eizo...) ( (TFFL(Z: W) - w. w.Ix(x1))) x(x+t1)
In words, we train the linear Transformer to predict (n++1) using TF (Zu; W); we will formally define the linear Transformer architecture below. This objective the center of study in number of recent empirical and theoretical works on understanding Transformers (von Oswald al., 2023; Ahn et al., 2023b; Zhang al., 2023; Mahankali et al., 2023).
5

---

Published as a conference paper at ICLR 2024
Linear Transformer (self-attention) architecture. We will now present the neural network archi- tecture that will be used throughout this paper. Given matrices P. Q E IR (dd++)) we define the linear self-attention architecture as
Attne (2) P QZ) where M ==
In 0 0
0 E (R(+1)x((+1)
(1)
Finally, for a positive integer L, we define an L-layer linear Transformer TFL as a stack of L linear attention units. Specifically, let the output of the 1 layer attention, Zz, be recursively defined as
1 71
Ze+1 Ze + Ptttn, (Ze) for l = 0, 0.1....,4 L-1
Then we define TF 11) The reason for the minus sign is be consistent with (von Oswald al., 2023; Ahn et al., 2023b), where such choice was motivated by theoretical considerations.
i.e., the (d + 1. + 1) entry of ZL.
We emphasize here that the linear attention unit, defined in (1), differs from the standard attention unit in (Vaswani et al., 2017): our architecture does not have feedforward networks, and we use single matrix Q to represent the product of key, query matrices. More importantly, we remove the softmax activation outside Z QZ. There two key reasons for choice:
2.5
linear
softmax
0.0-
-2.5-
The linear attention unit is much better suited to the task of lin- regression. For instance, (von Oswald al., 2023, Appendix A.9) demonstrates that the performance of softmax Transformer with twice many heads matches that of linear Transformers; in other words, we need two softmax attention heads to recover the perfor- mance of single linear head. In Figure 7, show that linear attention performs significantly better than standard attention with softmax.
-5.0-
0
10000 Iteration
5000
Figure 7: log(loss) against iter- ation. Comparison between lin- ear attention and softmax atten- tion for the 3-layer Transformers. Note that the loss of linear Trans- former decreases much faster.
2. Our goal in this paper is to find the simplest abstraction which representative of the Transformer's optimization landscape. As we will see in Subsection 3.2, the loss landscape of the linear Trans- former well approximates that of the actual Transformer, even without the softmax activation, feedforward networks, and other components of standard Transformers.
We also note that the key-query matrix parametrized by a single matrix Q, which another difference relative to standard Transformers. We make such parametrization for simplicity, and in the left plot of Figure 8, we verify that the loss plot for the standard parametrization largely similar to ours. We also remark that the lack of softmax may result in different Iearned attention scores. In particular, it may lead to denser attention scores than the attention scores for softmax Transformers (Oymak et al., 2023; Li et al., 2023a;b). On the other hand, the sparsity of learned attention scores depends on the data distribution; for instance, we observe that orthogonal covariates (as in (Huang al., 2023)) lead to sparser attention scores for both linear and softmax Transformers.
3.2
LINEAR TRANSFORMERS AS A FRUITFUL ABSTRACTION
Setting 1 (Ahn et al., 2023b)
Setting 2 (fewer covariates)
Setting (heavy-tailed covariates)
(d 5)
#contexts n
20
20
distribution of (x)
NO, Ia)
N00, Ia) (0,
To..1.0 Unif(Sa-1) (O, Ta)
distribution of w,
N(O, Ia)
Table 1: Settings for (the right-side plots of) Figures 2, 3, 4, and 16.
Setting for the experiments. Having established the framework in Subsection 3.1. we now describe details of our experiments. Our base-setup is the 3-layer linear Transformer, with 5-dimensional covariates, i.e. (I 3, d 5). This is the minimally complex setting that still recovers all of the

---

Published as a conference paper at ICLR 2024
discussed features of full Transformers. Transformers with larger L or d are qualitatively similar to the (L = 3, d = 5) setting, and we provide such an example in the right plot of Figure 8.
Our "default" setup is Setting 1 of Table I, where the context consists of 20 context demonstrations; each context covariate is sampled from the standard Gaussian, i.e., x:(() ~ NO, Ia). and we draw w, ~ N (0, Ia). This consistent with previous works (Garg et al., 2022; Akyürek et al.. 2022; von Oswald et al., 2023; Ahn al., 2023b). In order to see the effect of nonlinearity in data distribution, we conduct an additional set of experiments for a nonlinear regression where the covariates are distorted by multilayer perceptron (MLP) with nonlinear activations; see Appendix B for details.
In order to understand the effect of context length, we also consider the setting when context length this is Setting of Table 1.
instead;
Adam
SCDM
D
10-1
S-
Adiam
SCOM
Finally, to investigate the effect of heavy-tailed covariates on various aspects of the loss landscape, we consider Set- ting 3 in Table 1, where draw each i; instead uniformly from the unit sphere, and then scale it by the square root of heavy-tailed Gamma random variable with shape param- eter k = 0.1 and scale parameter 0 = 10. Furthermore, in Subsection 4.1, we study the effect of heavy-tailedness of the covariates in more detail.
10-2
-10-
-15
10-3
soino
10000
solon Iteration
10000
Iteration
Figure 8: Left: The case when Transformer is parameterized by separate Q,K K (query, key) matrices, instead of a single matrix as in (1). The setting is the same as Setting I in Table 1. Right: The setting of 8-layer linear Transformer with covariate dimension d 20 and context length 60.
For each different setting, we pick the best learning rate from grid search over 10 different choices. We choose the momentum parameter 0.9 for SGD, and 31 = 32 =
0.9 for Adam. We also employ the (global) gradient clipping where the thresholds chosen to be 1 for all settings (i.e., the clipped gradient direction is the same as the non-clipped direction). All the experiments are run over 6 different random seeds. See Appendix A for details.
Discussion of results. Below we provide detailed discussion of the results.
1. Gap between SGD and Adam. In Figure 2 (right), we plot the training loss for the three settings in Table 1. Notice that we observe the phenomenon (Adam SGD) over three different settings, to different extents. These loss behaviors resemble those of the practical Transformer optimization (left plots of Figure 2).
2. Heavy-tailed stochastic noise. In Figure 3 (right), following (Zhang al., 2020b; Kunstner et al., 2023), we plot the stochastic gradient noise at the initialization. Notice the similarity between the left plots and the right plots, showing that the shallow linear Transformers also exhibit the heavy-tailed stochastic gradient noise phenomenon.
3. Condition number of the landscape. Following (Jiang et al., 2022), we measure the "robust" condition numbers of different optimizers along the trajectory. Figure 4 shows that the condition numbers of adaptive methods are lower than those of SGD,, similar to (Jiang et al., 2022).
4. Directional smoothness. As observed by previous works (Zhang et al., 2020a;b; Pan and Li, 2023), in our experiments, we also observe that Adam has better directional smoothness than SGD, which correlates with the speed-up of Adam over SGD. We present this in Figure 5.
5. Generalized smoothness. As discussed in Subsection 2.5, the generalized smoothness condition of Zhang et al. (2020a) might not be a unique feature to Transformer optimization. Neverthe- less, interestingly, we also observe such a phenomenon (to a certain extent) in shallow linear Transformer optimization as shown in the right plots of Figure 16.
In this section, we have seen that simple linear Transformers described in Subsection 3.1 suffice to recover all the main features identified in previous works (Section 2). In the next section, we take advantage of the concreteness and simplicity of our linear Transformer to explore and understand the role of heavy-tailedness in data distribution and depth of the network.
UNDERSTANDING FEATURES BASED ON LINEAR TRANSFORMERS
The main advantage of our toy linear Transformer comes from its simplicity and concreteness. In particular, thanks to the concreteness of the setting, one can conduct various "controlled" experiments understand the features observed in Subsection 3.2. Recall that the data set used in experiments
7

---

Published as a conference paper at ICLR 2024
consists of nothing but random linear regression instances. This data set is far simpler and more concrete than the language modeling data sets (c.g., Wikipedia texts, question&answering) of the previous works discussed in Section 2.
We first take advantage the concreteness of our data distribution, and look deeper into how the main distinctive features of Transformer optimization arise. We first investigate how the Theavy-tailedness. of the data distribution affects the extent the features from Section 2.
4.1 EFFECT OF DATA DISTRIBUTION
Given that observe the "leavy-tailedness" of stochastic gradient noise, perhaps natural question ask the following:
Spherical 2:(0's
Heavy-tailed (000s
1. Comparing SGD v.s Adam:
O. Does the "heavy-tailedness" of data distri- bution exacerbate the features Section 2?
SGD(
Aaami ml
Adem S6DM
10
10-t
Settings. In order to investigate the above ques- tion, we consider the following distributions for the covariates (:()'ss of linear regression for (L = 3, d = 5, N = 20):
10-2
2012
5000
10flo
5000
10000
Figure 9: Plot of log(loss) against iteration for SGD and Adam.
Spherical covariates. We sample (f)'s uni- formly at random from the unit sphere Sd-1
2. Stochasticc gradient noise:
Heavy-tailed covariates. We first sample (s''ss uniformly at random from the unit sphere gd-1, and then multiply each covariate by dom scale drawn i.i.d from a heavy-tailed distri- bution, specifically the square root of Gamma random variable from o Note that k = 2.5 and 0
Dersity
Density
6.1n
00
5.0
Gradient
Gradient error
Figure 10: Comparing distribution of stochastic
precisely corresponds to the case where ,²()) ~ (0, Is). In experiments, use k = 0.1 and = 10 to make the distribution more heavy-tailed, while keeping the variance the same.
gradient noise at the initialization
3. Robust condition number:
RIGT
Discussion. We now discuss the experimental results one by one:
Adam
SCDM
10a
SGDN
12f
Adam
5000
10000
10000
In Figure 10, we see that "heavy-tailed"-nee of covariates is reflected in the "heavy-tailed" ness of the stochastic gradient. Notably, the con- trast between the two plots in Figure 10 reminds us the contrast we see between CNNs and Transformers in Figure 6.
Iteration
Iteration
Figure 11: Comparing the robust condition num- ber from Jiang al. (2022)
In Figure 11, it appears that there is some correlation between the gap in robust condition number, and the "heavy-tailed"-ee tain----eildd -ness of the data distribution, with heavier tails leading to larger gaps.
Finally, Figure 9 shows how the optimization speed of SGD and Adam vary with the heavy- tailedness of covariates. First, given spherical (light-tailed) covariates, both SGD and Adam converge much faster than Gamma-scaled (heavy-tailed, covariates. On the other hand, the relative gup between the speed of Adam and SGD does not seem to improve noticeably under light-tailed noise.
Together, Figure 9 and Figure 10 show that the relationship between heavy-tailed gradient noise and optimization speed may be a little more complicated than suggested in (Zhang et al., 2020b). Specifically, adaptivity seems be equally beneficial regardless of the heavy-tailedness of the gradient noise. Instead, these two plots seem to align more with the message in (Kunstner et al., 2023) that noise may not be the sole contributor of (Adam>SGD).
We next take advantage of the concreteness of our model, and investigate the effect the number of layers the optimization.
8

---

Published as a conference paper at ICLR 2024
L = 2
L = 4
L 6
L = 8
1. Loss against time
SGD(+m)
Adam(+m)
Adam SGOM
0.0
56DM
Adam
0-
Adam 5GDM
0-
Adam SGDM
-5.0
-10
-10-
-7.5
5000
10000
-10.0
10000
-15
-20-
20000
Iteration
lteration
Iteration
Iteration
Figure 12: Comparison of log (loss) between SGD and Adam for different number of layers.
2. Stochastic gradient noise:
oa
1.00
Density
Denatty
Denuty
a Dens Rt
0.25
0.2
oz
0.00
on
Gradient error
Gradient enror
Gradient error
Gradient error
Figure 13: Comparing the stochastic gradient noise for differrnt number of layers.
3. Robust condition number:
Adam SGOM
ROFT
102
10f
10"
10f
50.00
1,0000
5000
5000
10000
15000
Figure 14: Comparing the robust condition number for different number of layers.
4.2
EFFECT OF MORE LAYERS
We investigate the effect of the number of layers L on the optimization. Specifically,
Q. Will deeper linear Transformer exacerbate the features in Section 2?
Settings. In order to investigate the above question, consider repeating the experiments in Setting 1 of Table for the number of layers (2, 6,8.
Discussion. We present the experimental results one by one:
As one can see from Figure 12, the gap in loss between adaptive methods and SGD become more and more pronounced we increase the number of layers.
On other hand, the absolute value of the loss decreases with increasing depth, for both SGD and Adam, which makes sense considering the larger capacity of deeper models.
In Figure 13, we see that the stochastic gradient noise for the case = 6, are more heavy-tailed than the case of T. 2, 4.
Lastly, we observe in Figure 14 that the gap in the robust condition number of SGD and Adam is more pronounced in deeper models (L 4, 8) than the shallow model (L=2).,
CONCLUSION
The complexity of modern neural networks, especially Transformers, often eludes precise mathemati- cal understanding, and hence calls for such "physics-style" approaches (c.f. Zhang et (2022); Ahn et al. (2023a); Abernethy et al. (2023); Allen-Zhu and Li (2023); Li et al. (2023b); Dai et al. (2023)) based on simplified models. This work presents a concrete addition to this viewpoint, and it builds valuable, realistic proxy for understanding Transformers. However, our findings currently lack solid theoretical foundation, and our linear regression setting may not fully capture the features the language data utilized in Transformer optimization. We hope that work will serve the stepping stone for building a more precise theory of Transformer optimization, as well as contributing to the development of efficient training methods for Transformers.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work stems from group project at MIT; we thank the collaborators in the group, Hadi Daneshmand, Haochuan Li, Zakaria Mhammedi, Swati Padmanabhan, Amirhossein Reisizadeh, and William Wang for their time and intriguing discussions.
Kwangjun Ahn and Ali Jadbabaie were supported by the ONR grant EN00014-23-1-2222 and MIT- IBM Watson as well as Vannevar Bush fellowship from Office of the Secretary of Defense. Xiang Cheng and Suvrit Sra acknowledge support from NSF CCF-21 12665 (TILOS Al Research Institute) and an NSF CAREER award (1846088). Minhak Song and Chulhce Yun were supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant (No. 2019-0- 00075, Artificial Intelligence Graduate School Program (KAIST)) funded by the Korea government (MSIT), two National Research Foundation of Korea (NRF) grants (No. RS-2023-0021 1352) funded by the Korea government (MSIT), and grant funded by Samsung Electronics Co., Ltd.
REFERENCES
Jacob Abernethy, Alekh Agarwal, Teodor V Marinov, and Manfred K Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arX1v:2305.17040, 2023. 9
Kwangjun Ahn, Séhastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via the "edge of stability". NeurlPS 2023 (arXiv:22/2.07466), 2023a.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurlPS 2023 (arXiv:2300.00227), 2023b. 5, 6,
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. International Conference on Learning Representations, 2022. 5,
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.13673 2023. 9
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231 2015. 3
Sébastien Bubeck, Varun Chandrasekaran. Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint 2333 2023.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in Neural Information Processing Systems, 35:9955-9968. 2022.
Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-aware minimization. NeurIPS 2023 (arXiv:2305. 0..3305..22777 15287), 2023.
J Devlin, MW Chang, K Lee, and K Toutanova. Bert: Pre-training of deep bidirectional transform- ers for language understanding in: Proceedings of the 2019 conference of the north american chapter of the association for computational linguistics, 4171-4186.. acl. ACL. DOI: https://doi. prg/10.1853//1/ (19):1423, 2019.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30599 2022. 5.7
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint 10.0029, 2023. 6
10

---

Published as a conference paper at ICLR 2024
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural network geometry? arXiv preprint arXiv:2211.02254, 2022. 2, 4,,3,4,7 8
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprini arXiv:2001.066 008361, 2020.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be, In International Conference Learning Representations (ICLR) (arXiv:2304. 133600) 2023. 1, 2, 3, 7,8
Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015, 2023a. 6
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: To- wards a mechanistic understanding. International Conference on Machine Learning (ICML)
(arXiv:2303.04255; 2023b. 6,
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, pages 5747-5763. Association for Computational Linguisties (ACL), 2020.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023. 5
Yurii Nesterov et al. Lectures convex optimization, volume 137. Springer, 2018. 3
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning arXiv preprint arXiv:2306.0344. 2023.
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transformers. arXiv preprint arXiv:2306.0022 2023. 1, 5,
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 2017. 1, 6
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR, 2023. 5, 6,
Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, pages 4148-4158, 2017. 2
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. International Conference on Learning Representations (ICLR), 2020a. 1, 3,
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15393, 2020b. 1. 2, 3, 4, 7.8
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprini arXiv:2306.0990 2023. 5
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint 220.:220000430. 2022. 9
11

---

Published as a conference paper at ICLR 2024
Appendix
A Hyperparameters for the experiments
12
Additional experiments for nonlinear regression
12
Additional plots
13
A HYPERPARAMETERS FOR THE EXPERIMENTS
In this section, we summarize the choice hyperparameters for Subsection 3.2 and Section 4. We choose the momentum parameter 0.9 for SGD, and 31 32 0.9 for Adam. We also employ the (global) gradient clipping where the thresholds are chosen to be for all settings (i.e.. the clipped gradient direction the same the non-clipped direction). The choice of learning rates is summarized in the following table for (1) Setting from Table 1, (2) Setting 2 from Table 1, (3) Setting 3 from Table (4) Spherical covariates setting of Subsection 4.1, (5) Heavy-tailed covariates setting of Subsection 4.1, (6) L 2 setting of Subsection 4.2, (7) L = 4 setting of Subsection 4.2,
(8) L 6 setting Subsection 4.2, and (9)
setting of Subsection 4.2
Irs of
(1)
(2)
(3)
(4)
(5)
(6)
(7)
(8)
(9)
SGDM
0.02
0.01
0.02
5
0.02
0.1
0.05
0.05
0.05
Adam
0.005
0.02
0.02
0.1
0.02
0.1
0.05
0.05
0.02
Table 2: The choice of learning rates for experiments.
B ADDITIONAL EXPERIMENTS FOR NONLINEAR REGRESSION
2.5
Adam
250-
SGDM
smoothness
101
SGDM
log(loss)
Density
200
log(R OPT medd 6-
Adam
2.5
150
10"
-5.0-
100
4-
-1--
7.5
50
10000
0.005 0.110 0.015
5000
10000
10000
Iteration
Gradient error
Iteration
Iteration
Figure 15: The results for the nonlinear regression where the covariates distorted by ReLU network.
In this section, we consider the case of nonlinear regression, where the covariates (''s of the linear regression are distorted by a multilayer perceptron (MLP). Let us describe the setting:
Analogous to the Setting of Table ie.. N = 20, 5, ()) ~ N(0 Ia). and w, ~ N( Ta). On the other hand, to generate the responses y ( we first fix a randomly generated one-hidden- layer multilayer perceptron (MLP) with ReLU activation that we denote by MLP R + R with 5 hidden neurons and consider (() (w., (t)) In particular, use the code nn Sm..Squuentaa ial (nn. Linear (5, 5) nn. ReLU () nn. .n. Linear (5, 5) (where nn is the torch. nn in PyTorc for generating the random ReLU network MLP.
In order to cope with the MLP, in our linear Transformer architecture, we add an additional ReLU MLP layer with 15 hidden neurons before the linear Transformer blocks.
For the choice of learning rates, the optimal learning rates for this setting is 0.01 for Adam and 0.05 from Figure 15, we get similar plots the case of linear regression.
for SGD. As
12

---

Published as a conference paper at ICLR 2024
C
ADDITIONAL PLOTS
133
139
339
Halgradeen norm)
Ingigred norm)
kggggrde 10
logig nodlent noem)
Figure 16: The plot of log(1 V f(x)) II) against (sgosoothhssss) Following (Zhang al., 2020a), measure the directional smoothness instead of [ We observe similar trends with || V (av))l2. Left plot: LSTM from (Zhang et al., 2020a, Figure 1). Right plots: Shallow linear Transformers trained with Adam, see Settings 2, 3 in Table
13


Reviews:
[{"review_id": "HKu0Jpj3O5", "reviewer": "Reviewer_GNbE", "summary": "This paper studies the property of the shallow linear Transformer model. The evaluated phenomenon includes comparing Adam and SGD, heavy-tailed gradient noise, condition number, smoothness, etc. The experimental results show that the linear Transformer model reproduces the phenomena that have been observed for full Transformers. The results also show that more heavy-tailed data distribution and more layers can enhance the conclusion.\n\n--------------------------------------------------------\n\n**After rebuttal**: Thank you for your response. I am sorry for the late reply. I am satisfied with the feedback and revisions. I will keep my rating as 6 and support you. A minor point and suggestion: Please clarify what the four figures in Fig. 15 are to be compared within the main body. Later, I realize you are mentioning some figures on page 3. It is not a big issue anyway.", "questions": "1. One conclusion may not be obvious for linear Transformers compared with softmax Transformers. That is, the attention weights are more concentrated (even sparsely) on some key features/tokens after training, which is observed in several existing works [Li et al., 2023a, Li et al., 2023b, Oymak et al., 2023] for softmax Transformers. Can the authors provide a comparison empirically on this? Also, I think it is better to cover such a discussion in the revision.\n\nOymak et al., 2023, \"On the Role of Attention in Prompt-tuning. \"\nLi et al., 2023a, \"A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity.\"\nLi et al., 2023b, \"How do transformers learn topic structure: Towards a mechanistic understanding.\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work is novel in terms of new experiments on linear transformers from many aspects of evaluation. The paper is well-written and clear to follow. The studied problem is interesting to the community, from my understanding. The designed experiments are concise to support the conclusion.", "weaknesses": "1. This work lacks theoretical understanding or explanation after the experiments in Section 3.2, 4.1, 4.2. \n2. The limitation is not discussed. One thing that should be emphasized in the introduction or the abstract is that the experiments are linear regression, which is a good fit for linear Transformers.", "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": "", "extra_scores": {}}, {"review_id": "cr6PdyUnzk", "reviewer": "Reviewer_bjjH", "summary": "This paper identifies \"use shallow linear Transformers for regression task\" as a simple setup that serves as an effective mathematical abstraction for studying Transformer optimization.\n- The linear Transformer has attention (inner product) and residual connection, without softmax and MLP. For attention, a single $Q$ matrix is used to replace the product $W_Q^\\top W_K$.\n- The aim is to find the _simplest_ setup as a surrogate for studying optimization: this paper uses a 3-layer network with 5-dim covariate, on the linear regression task.", "questions": "- What are some phenomena that are not able to be captured in this simplified setup?\n    - e.g. using the $Q$ parameterization rather than $W_Q^\\top W_K$ might lead to different implicit regularization effects.\n    - There might be effects about architectural choices, e.g. consequences and limitations from the softmax.\n    - In general, it would be helpful to better clarify what is consider in-scope and what is not.\n- Could you give some examples of new discovery or actionable advice from the simple setup?\n- The current setup is entirely linear; if the data is non-linear, would there be similar optimization characteristics?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The simple setup is able to reproduce several characteristics of the training dynamics, which are not captured by canonical optimization theory. These include heavy-tailed gradient noise, robust condition number (larger for SGD than for Adam), as well as better directional and generalized smoothness from Adam.\n- The synthetic setup allows for more precise control.\n    - The paper varies the number of layers, data distribution, and context length.\n    - It's affordable to grid search over 10 learning rates.", "weaknesses": "- The main claim is that the simplified setup can be used as a proxy for larger scale systems. However, some connections are drawn by having \"qualitatively similar\" results.\n  - For example it's known that higher-order moments of the gradient noises matter, which cannot be captured by examining plots visually.\n  - e.g. How to judge whether the noises are similar based on Fig 3, or that 8-layer is qualitatively similar to 3-layer based on Fig 8?\n- The current empirical results are for confirming existing findings, but it's unclear what new discoveries can be enabled by this simplified setup.", "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": "", "extra_scores": {}}, {"review_id": "S42L4mY7tb", "reviewer": "Reviewer_dN4V", "summary": "This paper offers empirical evidence highlighting parallels between shallow linearized transformer models and the dynamics observed during practical Transformer training. These properties span various well-known aspects: the superiority of adaptive methods over SGD, the presence of heavy-tailed stochastic gradient noise, a robust condition number characterizing the optimization landscape, and the directional smoothness gap. The findings advocate for the potential of simple linear transformers as an insightful and accessible model to elucidate the mystery of Transformer optimization trajectories.", "questions": "- Are there other possible tasks that can be expressed by linear transformers? Do those tasks on simple linear transformers also exhibit those behaviors?\n- If we add other linearities (such as softmax) to the shallow neural networks, can you also show those properties along the training trajectory?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The presented work is overall clear and of notable quality. Its innovative aspect is the unexpected empirical evidence indicating that, under certain conditions, shallow linear transformers can closely resemble conventional attention models, revealing a significant similarity between the simple model and its real-world analogs. \n\nGiven the historical challenges associated with analyzing attention mechanisms involving softmax or other nonlinearities, it is commendable to identify such marked resemblances under broad observations using a simplified framework. This study illuminates the theoretical analysis for optimization on linear transformers, which makes a valuable contribution to the literature.", "weaknesses": "The paper is somewhat limited in its data distribution and optimization objective setting to support its conclusion. The data distribution discussed in this setting focuses on in-context linear regression tasks—a largely statistical setting, rather than one that closely resembles language. This might be attributed to the limited expressivity of linear transformers. \n\nI fully understand that the authors are trying to *find the simplest abstraction* as a representative of the transformer’s landscape. However, since the landscape is likely heavily dependent on the data distribution, it would be more convincing to see a variety of data models that exhibit the same phenomenon in simple models. Furthermore, it's difficult to assert that the linearized model is sufficient to understand Transformer optimization, especially since the properties verified in this paper represent only a portion of the optimization narrative.", "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": "", "extra_scores": {}}, {"review_id": "QQAc62zhnn", "reviewer": "Reviewer_gJVN", "summary": "This paper aims at deriving a simplified model for theoretical study of Transformers' optimization. The authors study a pure-attention without softmax Transformer variant (which is called linear attention) on a synthetic regression task. By conducting extensive empirical study, the authors claim that the proposed linear attention Transformer exhibits similar optimization properties that have been observed in the training of vanilla Transformers. Thus, this paper concludes that the proposed linear attention Transformer can be a realistic proxy for understanding Transformer optimization.", "questions": "The question here is just for clarification.\n\n- What is the box in the upper right of the left 3 plots in Fig. 3?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow.\n- Sec. 2 offers a concise and informative overview of existing study on Transformers' optimization. It's helpful for readers to digest the context of this study.\n- The empirical study and comparisons look solid and convincing. I'm looking forward to optimization theories based on the proposed simplified model.", "weaknesses": "- The proposed simplified model does not have feed forward network (FFN), which is yet another difference from standard Transformers. The authors should explicitly point this out (e.g., in the introduction and the formulation of model architecture) in the paper.\n\n- I strongly recommend the authors to summarize the detailed training recipes (e.g., hyperparameters) in a table and release the code for calculating relevant quantities (e.g., gradient error, robust condition numbers) to facilitate reproducibility and follow-up research.\n\n- In the beginning of Sec. 3.1, $x^{(i)}$ is defined as a row vector. But in the input matrix it seems that one has to interpret $x^{(i)}$ as a column matrix. Moreover, the matrix shape here is very confusing. The input of $F$ is $(n+1)\\times (d+1)$, while $Z_0$ is $(d+1)\\times (n+1)$. I suggest following the \"_Attention Is All You Need_\" paper and using shape $(n+1)\\times (d+1)$ consistently. The authors should fix the notation and formulas here.\n\n**Minor issues.** \n\n- Sec. 2.1: _The features defined above are typically of ..._ $\\to$ _The concepts defined above are typically of ..._\n\n- I recommend the authors to capitalize the first letter of \"Transformers\".", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)