How Does Adaptive Optimization Impact Local Neural Network Geometry?
Authors: Kaiqi Jiang
Dhruv Malik
Yuanzhi Li
Venue: NeurIPS 2023 poster (2023)
Keywords: optimization
adaptive algorithms
neural networks
TL;DR: We study the difference between the local geometry of the training objective in deep learning using Adaptive algorithm and SGD.

Abstract: Adaptive optimization methods are well known to achieve superior convergence relative to vanilla gradient methods. The traditional viewpoint in optimization, particularly in convex optimization, explains this improved performance by arguing that, unlike vanilla gradient schemes, adaptive algorithms mimic the behavior of a second-order method by adapting to the *global* geometry of the loss function. We argue that in the context of neural network optimization, this traditional viewpoint is insufficient. Instead, we advocate for a *local* trajectory analysis. For iterate trajectories produced by running a generic optimization algorithm OPT, we introduce $R^{\text{OPT}}\_{\text{med}}$, a statistic that is analogous to the condition number of the loss Hessian evaluated at the iterates. Through extensive experiments on language models where adaptive algorithms converge faster than vanilla gradient methods like SGD, we show that adaptive methods such as Adam bias the trajectories towards regions where $R^{\text{Adam}}_{\text{med}}$ is small, where one might expect faster optimization. By contrast, SGD (with momentum) biases the trajectories towards regions where $R^{\text{SGD}}\_{\text{med}}$ is comparatively large. We complement these empirical observations with a theoretical result that provably demonstrates this phenomenon in the simplified setting of a two-layer linear network. We view our findings as evidence for the need of a new explanation of the success of adaptive methods, one that is different than the conventional wisdom.

How Does Adaptive Optimization Impact Local Neural Network Geometry?
Kaiqi Jiang
Dhruv Malik
Department of Electrical and Computer Engineering Princeton University
Machine Learning Department Carnegie Mellon University
Princeton, NJ 08540 jaiqijer jungiijtpinnntoon..m edu
Pittsburgh, PA 15213 dhruvm@andrew. cmu, edu
Yuanzhi
Machine Learning Department Carnegie Mellon University
Pittsburgh, PA 15213 yuanzhi1@andrez cmu edu
Abstract
Adaptive optimization methods are well known achieve superior convergence relative to vanilla gradient methods. The traditional viewpoint in optimization, particularly convex optimization, explains this improved performance by arguing that, unlike vamilla gradient schemes, adaptive algorithms mimic the behavior of a second-order method by adapting to the global geometry the loss function. We argue that in the context of neural network optimization, this traditional viewpoint is insufficient. Instead, we advocate for a local trajectory analysis. For iterate trajectories produced by running generic optimization algorithm OPT, we intro- duce SOOP, a statistic that analogous to the condition number the loss Hessian evaluated at the iterates. Through extensive experiments on language models where adaptive algorithms converge faster than vanilla gradient methods like SGD, we show that adaptive methods such as Adam bias the trajectories towards regions where RAdam is small, where one might expect faster optimization. By contrast, SGD (with momentum) biases the trajectories towards regions where RSSG is com- paratively large. We complement these empirical observations with theoretical result that provably demonstrates this phenomenon the simplified setting two-layer linear network. We view our findings evidence for the need of new explanation the success of adaptive methods, one that is different than the conventional wisdom.
Introduction
The efficient minimization a parameterized loss function is core primitive in statistics, optimiza- tion and machine learning. Gradient descent (GD), which iteratively updates parameter vector with step along the gradient of the loss function evaluated at that vector, is a simple yet canonical algorithm which has been applied efficiently solve such minimization problems with enormous success. However, in modern machine learning, and especially deep learning, one frequently en- counters problems where the loss functions are high dimensional, non-convex and mon-smoohh The optimization landscape of such problems is thus extremely challenging, and in these settings gradient descent often suffers from prohibitively high iteration complexity.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

To deal with these difficulties and improve optimization efficiency, practitioners in recent years have developed many variants of GD. One prominent class of these GD variants is the family of adaptive algorithms [DHST] T2 KB15 At a high level, adaptive methods scale the gradient with an adpatively selected preconditioning matrix, which is constructed via a moving average of past gradients. These methods are reminiscent of second order gradient descent, since they construct approximations to the Hessian of the loss functions, while remaining computationally feasible since they eschew full computation of the Hessian. A vast line of empirical work has demonstrated the superiority of adaptive methods over GD to optimize deep neural networks, especially on Natural Language Processing (NLP) tasks with transformers [VSP* 17 DCLT19].
From a theoretical perspective, adaptive methods are well understood in the traditional context of convex optimization. For instance. Duchi et al. [DSIT] show that when the loss function is convex, then the Adagrad algorithm yields regret guarantees that are provably as good as those obtained by using the best (diagonal) preconditioner in hindsight. The key mechanism that underlies this improved performance, is that the loss function has some global geometric property (such as sparsity or coordinate wise bounded Lipschitz constant), and the algorithm adapts to this global geometry by adaptively selecting learning rates for features that are more informative.
However, in non-convex optimization, and deep learning in particular, it is highly unclear whether this simple characterization is sufficient to explain the superiority of adaptive methods over GD. Indeed, for large scale neural networks, global guarantees on the geometric properties of the loss are typically vacuous. For instance, for a 20-layer feedforward neural network, if we scale up the weights in each layer by factor of 1.5, then the global Lipschitz constant of the network is scaled up by a factor of at least e.0 Hence it only makes sense to study convergence by looking at the local geometry of the loss along the trajectory of the optimization algorithm [ACH18]
SGDAN
MM ofee
Adamm
SG.D
Trahiig loss
Adam
XX0
Barataon
Barataas
Figure 1: (left) Training losses of SGD+M starting from FSGD and A (right) The 10th largest value over median in the diagonal of loss Hessian (which can be viewed as variant of R (t) defined in eq. (T)) for Adam and SGD+M. Since the full Hessian is too big. here we selected several layers and randomly sampled 200 coordinates per layer to compute.
Itratoon
Training losses of Adam and SGD+M on the sen- tence classification task described in Section 4.1
Figure 2:
Moreover, the interaction between an optimization algorithm and neural network geometry is highly complex- recent work has shown that geometric characteristies of iterates encountered during optimization is highly dependent on the choice of optimization algorithm and associated hyperpa- rameters LBD CKL+21) For instance, Cohen al. [CKL-221 demonstrate that while training neural networks with GD, the maximum eigenvalue of the Hessian evaluated at the GD iterates first increases and then plateaus at a level that inversely proportional to the step size. The viewpoint from convex optimization, where a loss function has some (potentially) non-uniform but fixed under- lying geometry that we must adapt to, is thus insufficient for neural networks, since the choice of optimization algorithm can actually interact with and influence the observed geometry significantly.
To provide another example of this interactive phenomenon, we consider the following experiment. On the same network training loss function f. we run stochastic gradient descent with momentum (SGD+M) and Adam to obtain two different trajectories. We select an iterate Z from the Adam trajectory and an iterate ZSGD from the SGD trajectory, such that f ( = f (ISGD). We then run SGD+M with the same configuration twice, once from Adam and once from SSGD- If the underlying geometry of the loss function f was truly fixed, then we would not expect significant difference in the performance of running SGD+M from either of the two iterates. However, as shown in Figure [] there is a noticeable difference in performance, and running SGD+M from I Adam achieves lower
2

---

loss than running SGD+M from TSGD. This suggests that Adam may bias the optimization trajectory towards a region which is more favorable for rapid training. This motivates the following question.
How does adaptive optimization impact the observed geometry of a neural network loss function, relative to SGD (with momentum)?
The remainder of this paper is dedicated to answering the above question. To this end, for each iterate in trajectory produced by running an optimization algorithm OPT, where the Hessian of the /th iterate is given by H() Rdxd, we define the second order statistic Roed (t) in the following fashion. For the tth iterate in the trajectory, let Roed (t) be the ratio of maximum of the absolute entries of the diagonal of H(0), the median of the absolute entries of the diagonal of H Concretely, we define
max | (1H
ROPT (t)
(1)
median (
This statistic thus measures the uniformity of the diagonal of Hessian, where a smaller value of Reed (t) implies that the Hessian has more uniform diagonal. It can also be viewed as a stable" variant of the condition number. Instead of singular values, we choose diagonal entries because adaptive methods used in practice are coordinate-wiss which can be viewed as the diagonal scaling approaches In Appendix A.9 we discuss this intuition in detail and compare Rend (t) with singular value-based metrics. As a supplementary result, in Appendix F we demonstrate that the loss Hessian approaches diagonal during training for Adam and SGD+M. There has been prior theoretical work on overparameterized neural networks showing that a smaller condition number of Hessian, Neural Tangent Kernel [[GHT8] etc. could yield to faster convergence rate for (S)GD [Z222]] As for (diagonal) adaptive methods (e.g. Adagrad), they were original designed to adapt to the nonuniform diagonal geometry. Intuitively, a smaller RorT (t), which implies more uniform diagonal geometry, could lead to faster convergence.
D
Armed with this statistic, we make the following contributions:
1. We focus on language models as on NLP tasks, adaptive algorithms show significantly faster con- vergence than SGD (with momentum). On a wide variety of neural network transformer architectures and language modeling datasets, we conduct experiments compare how Ra (1) and RSSUDM (I) evolve over time, when Adam and SGD+M are run from the same initialization and with their optimal (initial) learning rates respectively. In each case, we demonstrate that the Adam trajectory attains Rydamm (t) values that are significantly smaller than the RSGDM (t) values found by SGD+M. We show a simple example of this phenomenon in Figure rightt. This suggests that relative SGD+M, Adam biases the optimization trajectory to region where the Hessian diagonal is more uniform. We call this phenomenon the uniformity of diagonal geometry for adaptive methods. Moreover, we demonstrate that a more uniform Hessian diagonal, characterized by smaller ROOT (t), is contributing factor to faster optimization (see Section 4.3 for discussion). This suggests that a region where the Hessian diagonal is more uniform is also region that is more amenable to rapid optimization.
2. We complement our empirical results with a theoretical analysis of this phenomenon in the simplified setting of large batch Adam and SGD+M, on a two-layer linear network with d-dimensional input and hidden layer, and one dimensional output. We show that for wide range of t, Radda (t) = 1 1 o(1) but RSSIDM (f) 2(log d). Our proof reveals that Adam induces the weight matrices have low rank whose leading singular vectors have certain type of uniformity (see Section f discussion), a fact that we also observe empirically in large scale neural networks, suggesting that this may be mechanism by which adaptive methods bias trajectories to have uniformity of diagonal geometry.
2 Related work
Existing analyses of adaptive methods. The vast majority of prior theoretical work on adaptive
methods has focused on the blackbox setting [DHSI] KB15 CZT 20, RKKI8 18 WWB20 DBB U20
'Consider the case where one parameter has little impact on the loss, then the second derivative w.r.t. this
max ( (1 pi infinity. So we consider median which is more stable.
parameter is almost zero, making min (
Recall that the main theoretical bound in the original Adagrad paper [DHSI] is in terms of the diagonal scaling.
3

---

These works make minimal assumptions about the structure of the loss function, beyond (possibly) some global properties such as convexity or smoothness. These global properties (govemed by parameters such as the smoothness parameter) are assumed to hold over the entire domain. Hence this style of analysis worst case, since the resulting convergence bounds depend on polynomially on these global parameters. However, as we show in Section 3.1 in neural networks these parameters are prohibitively large. This worst case analysis is hence unlikely to explain the success of adaptive methods neural networks. By contrast, our focus is analyzing the local trajectory that is induced by running the optimization method.
ENV2I]
Existing analyses of (S)GD on neural networks. There is an extensive literature on the analysis of GD/SGD in the non-blackbox setting, e.g. overparameterized neural networks, [DZPS18] [T20 AZLLI9 AZLST9 ADH 19 [ZB22]. However, it is unclear how to translate these analyses of GDISGD, to an analysis that explains the gap between GD/SGD and adaptive methods.
Influence of algorithms on the loss geometry. In many simple convex settings, egg. linear or logistic regression and the Neural Tangent Kernel [IGHI8] the loss geometry is usually fixed and not influenced by learning algorithms. However, in neural networks the interaction between algorithms and Ioss landscapes more complicated. Lewkowycz et al. [LDD*20 20] find so-called catapult effect of initial learning rate on the training trajectory of SGD and related loss curvature. Cohen et al. [CKL demonstrate that while training neural networks with GD, the maximum eigenvalue of the Hessian evaluated at the GD iterates first increases and then plateaus at level that is inversely proportional to the step size. However, Cohen al. [CKL 21 leave open the problem of whether similar interactive phenomena occur in algorithms that are not GD, including adaptive methods.
3 Overview of main results
3.1 Issues of prior analyses on adaptive methods
As is mentioned in Section 2 existing work on adaptive algorithms has mainly focused on black- box analysis assuming some global worst-case parameters. However, these global bounds can be extremely bad in complicated deep learning models, as is discussed in Section [ see this, we initialized transformer mode@ with default initialization in Pytorch but chose a large gair and computed the smoothness parameter (denoted as 1) and the condition number (denoted as x) of loss Hessian one layer. We observed that setting the gain as a large constant (e.g. 800) results in extremely large I and x (1 2 10? and Æ 2 1)) which makes the convergence rates in prior black-box analysis vacuous.
The failure of global worst-case analysis implies that we need to focus on the local trajectory of algorithms. However, it is unclear whether when two optimization algorithms are used, they will have the same geometry in local trajectory or not. In particular, although in theory, adaptive algorithms can yield to convergence rate with better dependency on certain local geometry of the function comparing to SGD (with momentum), it could still be the case that the local geometry along the trajectory of adaptive algorithm can be much worse than that of SGD (with momentum).
That motivates us to study the local geometry, especially that obtained by adaptive methods comparing to SGD (with momentum) in the paper. Motivated by the diagonal scaling of Adagrad and Adam for neural network training, we ask the follow main question in our paper:
How does the local diagonal geometry (diagonal of the loss Hessian) along the local trajectory of adaptive algorithms compare to that of SGD (with momentum)?
3.2 Overview of the experiments
As is discussed in Section [ we consider ROPT(() (1) defined in eq. ①1 as measurement of the uniformity of the loss Hessian diagonal and conduct experiments on different NI.P tasks where adaptive methods converge faster. The detailed experimental setup will be stated in Section 4 To explore potential different patterns of different layers, we do the computation layer by layer. On
httper///yyoooooooo///oc/o/////s tutooriaarrrrtttrtaaaaao .hmml
"Thi refers to the gain parameter in some commonly used initialization functions of Pytorch, egg. torch.nn.in inw.....iit..iirrrr..mmm. uniform_()

---

wide variety of transformer architectures and language modeling datasets from the same initialization, we observe that for the vast majority of layers:
When we train the neural network using Adam, the uniformity of diagonal geometry, mea- sured by medd (t) is smaller than that when train using SGD+M from the same initializa- tion.
Table I shows typical example of Reeea (t) compared to R
(t) on a sentence classification task using BERT-small [TCLTI9] BDR21 (See Section 4. for details). We repeated the experiments for 12 times starting from the same initialization. Table shows the averaged Re (t) and Reed (t) in
Apsen (1) and their standard deviations in (e)
some randomly selected layers. We also report the averaged
the brackets
Figure 2 shows the corresponding training losses of one in these 12 experiments.
Table 1: Reham (t) and (L) in some layers, the sentence classification task (see Section 4.1
Iteration 0
Iteration 750
Iieration 1250
Layer
REAIA (1)
Heol (2)
Ried
Fasl (0)
(1)
Rea
15.7
15.7
12.76
9.65
1.45 (0.65)
11.43
14.24
0.94 (0.40)
22.63
22.63
13.17
7.41
1.92 (0.67)
10.62
9.67
1.33 (0.75)
15
9.35
9.35
80.57
53.52
1.65 (0.65)
100.65
61.80
2.01 (1.00)
82.37
82.37
405.02
223.56
1.91 (0.53)
423,28
337.32
1.43 (0.63)
18
31.32
31,32
17.07
13.24
1.43 (0.58)
18.15
13.63
1.21 (0,36)
47.13
233.72
72.67
3.54 (1.21)
158.38
93.13
2.28 (1.18)
24
31,17
31.17
17.52
17.34
1.13 (0.40)
13.51
14.23
1.05 (0.36)
To understand this phenomenon in more principled point of view, in Section B we provide formal proof of the statement in a simplified setting: large batch Adam and SGD+M on a 2-layer linear network with 1-dimensional output.
4 The uniformity of diagonal geometry
As mentioned in Section 3.2 we computed ROPT (t) defined in eq. T on different language models. In this section, we present the results of SGD+M and Adam on different architectures and datasets. In Appendix A, we present the results of other adaptive algorithms.
During training we started from the same initial weights and used the same learning rate schedule (constant or decreasing) for SGD+M and Adam. We tuned and chose the best (initial) learning rate of SGD+M. The (initial) learning rate of Adam was set as a value under which Adam converged faster than SGD+M with its best learning rate. The concrete values will be stated in later parts of this section. We used large batch sizes to make the training procedure stable. When computing Hessian, we also used large batch sizes. Due to the extremely large dimension, we did the computation on some uniformly selected coordinates, more precisely, 200 coordinates per layer.
4.1 Experiments on real datasets
Sentence classification task on BERT-small. We fine-tuned BERT-small (CCTT99 BDR21 the
IMDB dataset [MDP*11] the task is to classify whether movie reviews are positive or negative The momentum parameter in SGD was set 0.9. The two momentum parameters (3), 82) of Adam were set as (0.9, 0.999). We trained the model using linearly decreasing learning rates for 10 epochs (2500 iterations). The initial learning rates of SGD+M and Adam were 0.001 and 5e-5, respectively. As mentioned in Section 3.2 Figure and Table I show the training losses and the comparison between ReAdam (t) and /pSGbRK (t), respectively.
Translation task. We trained a Seq2Seq network that uses Transformer to solve a machine translation task Multi30k [EFSSI6]: this task is to train German to English translation model The
$ "RSCDM (t) values in Table for most layers are roughly 1.4 to 2 times (r)dam(t) in corresponding layers. In practice. it can be considered significant because might imply 1.4 to 2 times faster convergence.
16. htt.:/aaiooooottrrasermmeee.eee
5

---

momentum parameter 3 in SGD was set as 0.9. The two momentum parameters ($1, 52) of Adam were set as (0.9, 0.98). We trained the model using constant learning rates (0.03 for SGD+M and le-4 for Adam) for 60 epochs (1800 iterations). The experiments were repeated for 8 times starting from the same initialization. Figure Ileft) shows the training losses for one among them. Table 2a
shows the averaged R (t), ReedM( (1) and R (0)
(with standard deviation in the brackets) in
some randomly selected layers.
SG0-M
560
Arar
Adamn
Taaing 60
co02
65
1604
Figure 3: Training losses of Adam and SGD+M for the translation task on (left) Multi30k (see Section 4.1) and (right) Multi30k but with randomly generated targets (see weight matrix in the 27-th layer on the Section 4.2)
Figure 4: Singular values and R2 of the
translation task described in Section 4.1
Table 2: Re (t) and (S)DM(t) in some layers for the translation task. (a) on Multi30k (see Section 4.1 and (b) on Multi30k but with randomly generated targets (see Section 4.2)
(a)
(b)
19637
Lighff
(4200)
(55%%)
Epedk30
10507
125
Epeck-3
4.5
4.31
5.51
3.41
2.5 79.5
3.54
3
155 (141)
4.0
1:33
144
18
935
3.79
160
138
36.1
12.12
16.90
3M64
5
836
5.04
46.6
05.54
1.31.117)
47.33
34.15
1.41 9030
539
331
129
378
5.99
1/6
46.77
382
11.35 (237)
46.5
2.34
124
1738
5591
901.9.00
19.85
391
8315
3332
3232
46254
1130
05.73
3057
14.9
439
272
120440
11
1533
10.0
365
0.90
1180330
79.82
3.28
23.018800
913
3034
19:3
21
36.94
3636
193.1
0.89
1640000
79/40
30.80
21.8 17
30.9
3660
630
400
135
682
32
16030
18:27
1647
137
1.5
4.2
Experiments random datasets
We used the same model and momentum parameters as in the translation task described in Section 4. but generated random integers as targets. Similar to the setting real targets, the model was trained using constant learning rates (0.015 for SGD+M and 5e-5 for Adam) for 60 epochs (1800 iterations). and we repeated the experiments for 8 times starting from the same initialization. Figure right) shows the training losses for one among them. Table shows the averaged APAdam (t), R RSGDM (t) and
Ream (t)
(with standard deviation in the brackets) of the same 10 layers as in Table 2a
R RM)
4.3 Summarization of the empirical results and discussion
Overall, through extensive experiments on language models, we demonstrate that starting from the same initialization, for the vast majority of layers, the ( (1) values found by Adam are smaller than those found by SGD+M. This suggests that Adam biases the trajectory towards a region with more uniform Hessian diagonal than SGD+M. In Appendix AA we also validate this observation on the in-distribution test data.
Contribution of uniform Hessian diagonal to fast convergence. We observe that on dataset with random targets, SGD+M plateaus after about 400 steps and thus converges much slower when compared to Adam than real dataset (see Figure 3). On the other hand, the gaps of RSGDM (t) and
TTo prevent ROPT (t) from getting too large due to tiny median, added additional term 0.001 (" D the denominator of eq. (I) when computing.

---

Reeemmt) (t) are more significant on random data than on real data (see Table 2a and Table 2b as well. In Appendix A.44 we conduct another experiment where we switch from SGD Adam in the middle and compare it with the model trained by Adam from the beginning. The observation is that both the loss gap and the gap of medd (t) are gradually closed after switching (see Figure 7 and Table Hence we find a positive correlation between fast convergence and uniform Hessian diagonal (small Reee (t)). In Appendix [ we study other adaptive algorithms (Adagrad, RMSprop and AMSGrad) and get similar observation: all these adaptive methods converge faster than SGD or SGD+M and also bias the trajectory to regions with smaller m (L), suggesting the universality of our observation.
Despite the above positive correlation, it is reasonable to ask whether small R (t) indeed con- tributes to fast convergence is just byproduct of adaptive methods. To address this concern, in Appendix [.I we add a supplementary experiment similar to that in Figure [ We select two iterates Ii and 22 from two trajectories that both come from SGD+M (instead of one from Adam and one from SGD+M in Figure such that the loss f(z1) = f(22) but 22 has a smaller Read (t) than 21 We then run SGD+M with the same configuration twice, once from #1 and once from x2-- Under this setting, we get similar observation as before: rumning SGD+M from (with smaller Reme (t)) achieves faster convergence than from ø1. This suggests that the uniformity of the diagonal of loss Hessian (measured by Reed (t)) reveals some intrinsic trajectory property beyond the algorithm choice and is indeed a contributing factor to fast optimization. Appendix B.2 we theoretically prove the contribution of small R (t) to fast optimization in a simplified setting.
More discussions the trajectory difference. Considering the fact that our comparison between Reneam(L) and R (t) is conditioned on the same iteration when SGD+M has larger training loss than Adam, there is potential alternative explanation of the Hessian diagonal uniformity. That is, the global minimum has uniform Hessian, and Adam simply converges faster to it than SGD+M. To rule out this possibility, in Appendix A.3 we add a comparison of our measurements R (t) and RSCDM (t'), where t are picked such that tth Adam iterate and t' SGD+M iterate have the same training loss. The results (in Table 7 show that Riddm (t) < RSSDM (t') for most layers, thus demonstrating that the trajectories of Adam and SGD+M are truly different and that the difference is because Adam biases the local geometry (as opposed to faster convergence).
Adding regularization. People in practice usually add weight decay (equivalent to l2 regularization) to encourage better generalization ability. In Appendix [A.7 we compare SGD+M and Adam when both using small weight decay values (0.001). The results in Figure 13a and Table suggest that in this case, our observation still holds: Adam converges faster than SGD+M and in the vast majority of layers, R (t) values are smaller than Reed RSCPM(t). This reveals the robustness of our observation under weak regularization. However, under large weight decay parameters, observed cases where Adam still converged faster but Radam(t) (t) values were larger rather than smaller. In the case of strong regularization, the adaptivity of Adam requires further exploration and we hope to find new mechanisms in the future.
5 Theoretical analysis
we empirically demonstrate the uniformity of diagonal geometry. In this section, we theoretically analyze this property for large batch Adam and SGD+M on two-layer linear network with 1 I-dimensional output. Although simple, the choice of 2-layer linear networks to understand learning dynamics is common in prior works (e.g. (TG2IT). Moreover, in language transformer models when the weights are small, the softmax the key-value-query structure is near the linear regime. Then this structure might be approximated by the product of 3 linear operators, similar to a three-layer linear network. Hence the theoretical analysis of this phenomenon on linear networks would be good starting point for further understanding of more complicated language models.
In Section 4
5.1 Problem setup
Notation Let [d] = {1, (1,2,...d}} d]. We use || ||2 to denote the l2 norm of vector, and || || to denote the Frobenius norm of a matrix. Let ( be the Euclidean inner product between vectors or matrices. Let, o²) be the one-dimensional Gaussian distribution with mean μ and variance o². For scalar (vector, matrix) A which evolves over time, we use (t)) to denote its value at time t.
7

---

Let there be m data points. The data matrix is X E IR xm and the label matrix is Y E R xm" We assume that the input dataset is whitened, i.e. Arz X E Rd,xd, xd is an identity matrix.
The parameters of 2-layer linear network are given by W == (W2, Wi). Assume W; E Rd.xd,-1 for i = 1.2. We have d2 dy. do dz. We consider the square loss L(W) == 2m W2W - YII--
Denote A = 1 yYX E R2 xd=, AACHI9 show that with whitened dataset,
L(W)= = 2m I YII = L(W)+ + L(W): = 2 W -
(2)
where does not depend on W. We consider the following model with small Gaussian initialization.
Assumption (Setup), The input covariance Azz = 1 X Rda is an identity matrix. The imput and hidden layers are both of dimension d, i.e. d1 = du = d. Without loss of generality, we can assume that A is row vector (i.e. d2 = 1) whose coordinates are positive and ((1) terms of d.
Assumption 2 (Gaussian Initialization). Vi,j : ( ~ NN(0, independenly initialized with sufficiently large O.
~ NoO, +هx are w. 1
Denote and Aew as the batch versions of A and Agz. We make the following large-batch assumption. We emphasize that large batches are commonly used in NLP tasks (e.g. g. [MR 20)).
Assumption 3 (Large Batch). For the randomly selected batches, assume E[A] A. E E[ Agz-
V E [d] E (Ā. A.) < 0². E [Amm[i,s]] - (i, S 0², and o² O
poly(a)
Denote g(t) as the batch gradient at time t. The update rules of SGD+M and Adam are given by
SGD+M: (t+1) =Bult) =g((), w(++1) w(()) -nu!
Adam: M1 17 v 1- - 1
/ /(1 - -1/tt), m((11 = (1 + (1)
(3)
(1)
,((+1) = =3 + +11- (1- (2) c g(t, w(++) = -(()) 7/-- Vula +
where y is the learning rate, 3,,3,, 32 momentum parameters, and & is for numerical stability. All operations on vectors are element-wise. Here and throughout, the notation f(xx) Og((z))
(resp. (x) (x)())) [() (x)
O(g(z)) means that 3C₁, C2 > O such that f(x) S C2g(2) (resp.
f(x) 2 Cig(z). Cig(z) < f(z) < C2g(z)). Here C1, C2 may depend on B, 31, 82. We will also
( to hide factors of order log d. In our theoretical analysis, "with high probability", or "w.h.p." " for short, means that with probability least 1 poly(d))
use the notation with ~, i.. O()), ((-)
Since the weights and Hessians in different layers may have different magnitudes, we compute the ROPT (t) layer by layer. Denote RSCDM (t) (resp. RAdam (t)) as the FOOP (t) found by SGD+M (resp. Adam) w.r.t. We at time i where k 1, 1,2.
5.2 Main results
and 3 consider the weights [WSeD] } (resp. {wo00 obtained by SGD+M (resp. Adam) defined in (3).
Theorem 1. Under Assumption
W Adom 120
120
3/22 40/241
1. For any p > O, pick 0 < < IF, S O (Zrafas)) and a 2 4(p + 2). Suppose
then
there exists TsGD, TTsnn, such that w.h.p., L ( (W.Tnw.... (Taw..)) (d)), L ( (T(Tsu..))
<
and
SGD
Vt E [T
RSGOM (t) 2loo d),
k = 1,2.
'In Assumption 2 we assume Gaussian initialization. Due to the rotational invariance of Gaussian distribution. assume that all coordinates of are positive without loss of generality.
8

---

For any p > O, pick oj S O . r.a a p+1 and B2 37
Suppose
I
n2/ dis/s
Then ATddmm, 1, T Adaam such that w.h.p.. L
(W (Tmmm,))
= (d)) L W (Taim
0 ($) and
Adam
Adam
1I E [Tadam, [ndamna]
Rmeaa (t) (1)=1+0
11
k k= 1
d
Remark Theorem 1 holdd for all values of hyperparameters' (such a, in certain ranges instead of just particular values. The ranges of SGD+M and Adam overlap with each other. That means we can choose the same hyperparameters for SGD+M and Adam in the overlapped region to make fair comparison, for example, same a, a such that a 2 4(p + 2) and a I min /0/2411
1/// n²/2g² }.
An immediate corollary of this theorem below gives the difference between iterates of Adam and SGD+M that have the same loss.
Corollary Under the setup in Theorem w.h.p., for any t E [ToO,, I. and t' E (dd)
[TAdum, 1, T Tddum,2) such that I. (W3ee) L (wilt' E Adam
52 (p))
we have
med (t) = 22(log d),
B (t)) (t')=++0)
17 + di
k == 1,2.
Theorem I and Corollary [ tell us that during a long training period when the loss decreases from (d)
Ô @F) ). the diagonal of loss Hessian for Adam keeps nice uniformity: Rehea (t) + o ), k
1, 2. On the other hand, the diagonal of loss Hessian for SGD+M is less uniform. Appendix gives proof sketch of Theorem I
The detailed proof can be found in Appendix Dland]
6 Low-rank weight matrices and uniformity of leading singular vectors
The proof sketch in Appendix C highlights one crucial intuition of Theorem T After TSGD.1 (resp. 1) steps, W1 of SGD+M (resp. Adam) becomes an approximately rank-1 I matrix. Consider the left singular vector u == [u1, U2. .... ua]² which corresponds to the lcading singular value 01. We can show that the distribution of u/, uz.....u ug for Adam is more uniform than that of SGD+M. This property, we call the uniformity of the leading singular vector, is related to the uniformity of the diagonal of loss Hessian, see Appendix G for more details.
Similar low rank bias after training has been studied in prior works (e.g. GWB 17
LMZI8
CGMR20).. For more complicated models, we want check whether the weight matrices also have low rank structures and if so, whether we can still observe the uniformity of leading singular vectors. More formally, consider the weight matrix in some layer W IRmx we want check
(A) Whether W E Rm is approximately a rank & matrix with < min(m m, n). and if true,
(B) Consider the top k singular values oI..... Os and left singular vectors uI...., Uy Define u == Σ ofuu u, ① u, := [ū...... "d)" and compute R,, := max; ñ. a generalized version of median ut median ū,
in the rank-1 case. We want to see whether R., obtained by Adam is smaller than that of SGD+M.
After reviewing the weight matrices we got in different settings, we observed that (A) and (B) hold for many layers in those models. For example, on the translation task mentioned in Section 4.1 we found 12 layers which had approximately low rank structures and for 10 of them, R, values (defined in (B)) obtained by Adam were smaller than those found by SGD+M. Figure 4 Ashows the result on one typical layer. Results of more layers can be found in Appendix A.5
Remark The definition of Rg is based on the connection between diagonal of loss Hes- sian and weight matrices. Appendix G shows that for 2-layer linear network, F max; [[ 6..l
When W1 R is approximately rank k, i.c. W1 ~ Et_ ouu, de- note u, = [4,1, 142,... Uim] 1T and v, = jea Us.,... Unnl" IT, we have that for the j-th row,
median|| [addian[:::")) [.. [".]ll
The assumption B2 = 8 is only for convenience to make the proof easier but is not essential. Our result also generalize cases with other 91 and B2. See Appendix for more discussions.

---

E. of By defining ū [ū, U2 ... ū.]"
1ー=1,,v O Ug we have that |W R Vg. Although in multi-layer nonlinear neural net- works. the connection between diagonal of loss Hessian and the weight matrices is more complicated and medd2 (t) may depend on the product of many weight matrices rather than one single matrix, we still believe that this definition of R, is reasonable ratio to consider.
7 Conclusion and future work
We demonstrate that adaptive optimization methods bias the training trajectory towards a region where the diagonal of loss Hessian is more uniform, through extensive experiments on language models and theoretical analysis findings may not directly lead to an improved algorithm for practical use, they provide a new way of thinking when designing new algorithms: in contrast with the traditional view which tries to design method that performs better the bad loss geometry, our findings suggest that algorithms which implicitly avoid regions with bad geometry. There are a lot future directions along this line. For example, theoretical results the two-layer linear networks may be able to generalize to multi-layer networks. As is discussed in Section [5] the key-value query structure in the transformer models might be approximated by three-layer linear network and the analysis of multi-layer networks might provide more connection to real deep models and could be an interesting and challenging future direction. Moreover, it is also possible relax our large-batch assumption (Assumption 3 and prove similar results in the general stochastic setting.
simplified setting of two-layer linear networks. Although our
design
Acknowledgments and Disclosure of Funding
This project was supported by NSF CCF 2145703.
References
[ACH18] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In 35th International Conference Machine Learning, ICML 2018, pages 372-389. International Machine Learning Society (IMLS), 2018.
[ADH* 19] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019.
[AGCH19] Sanjeev Arora, Noah Golowich, Nadav Cohen, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
[AZLL19] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterizeed neural networks, going beyond two layers. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[AZLS19] Zeyuan Allen Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via wver-parameterizztionn. In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[BD221] Prajjwal Bhargava, Aleksandr Drozd, and Anna Rogers. Generalization in nli: Ways (not) to go beyond simple heuristics, 2021.
[MR 20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner,
10

---

Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot leamner. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[CGMR20] Hung-Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut, Gradient descent for deep matrix factorization: Dynamics and implicit bias towards low rank. arXiv preprint arXiv:2011.13771 2020.
[CKL [2] Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In 9th In- ternational Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[C] Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu. Closing the generalization gap of adaptive gradient methods in training deep neural networks. In Christian Bessiere, editor, Proceedings of the Twenty-Ninte International Joint Conference on Artificial Intelligence, IJCAI 2020, pages 3267-3275. ijeai.org, 2020.
[DBBU20] Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. arXiv preprint arXiv:2003.02395, 2020.
[DCLT19] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre- training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computa- tional Linguistics: Human Language Technologies, NAACL-IILT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pages 4171-4186. Asso- ciation for Computational Linguistics, 2019.
[DHS11] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[DZPS18] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over vverppaammeterized neural networks. In International Conference on Learning Representations, 2018.
[EFSS16] Desmond Elliott, Stella Frank, Khalil Sima'an an, and Lucia Specia. Multi30k: Multilingual english-german image descriptions. In Proceedings of the 5th Workshop on Vision and Language, pages 70-74. Association for Computational Linguisties, 2016.
[NN22]] Alina Ene. Huy L. Nguyen, and Adrian Vladu. Adaptive gradient methods for con- strained convex optimization and variational inequalities. In Thirty-Fifthh AAAI Con- ference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educa- tional Advances in Artificial Intelligence, EAA/ 2021, Virtual Event, February 2-9, 2021, pages 7314-7321 AAAI Press, 2021.
[GWB+ 17] Suriya Gunasckar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
[JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018
[J220] Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[Kaw 16] Kenji Kawaguchi. Deep learning without poor local minima. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnet,, editors, Advances in Neural Information Process- ing Systems, volume 29, Curran Associates, Inc., 2016.
11

---

[KB15] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA. USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[LBD+20] Aitor Lewkowyez, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003 022188, 2020.
[Ler19] Matthieu Lerasle. Lecture notes: Selected topics robust statistical learning theory. arXiv preprint arXiv: 10761, 2019.
[LMZI8] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over- parameterized matrix sensing and neural networks with quadratic activations. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet, editors, Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learn- ing Research, pages 2-47. PMLR, 06-09 09 Jul 2018.
[LZB22] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 2022.
[MDP+ 11] 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, pages 142-150, Portland, Oregon, USA, June 2011 Association for Computational Linguistics.
[MXBS17] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenRe- view.net, 2017.
[RKK18] Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018
[TCG21) Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10268-10278. PMLR, 18-24 Jul 2021.
[TCLT19] Iulia Ture, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read stu- dents learn better: On the importance of pre-training compact models. arXiv preprint arXiv:1908.08966 2019.
[TH+12] 12] Tijmen Tieleman, Geoffrey Hinton, et al. Lecture 6.5-rpppo::: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26-31, 2012.
[VSP+ + 17) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz. Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
[WWB20] Rachel Ward, Xiaoxia Wu, and Léon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. Journal of Machine Learning Research, 21(219):1-30, 2020.
[ZCS+22] 22] Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification on update rules. In NewrIPS, 2022.
12

---

A More experiments of the uniformity of diagonal geometry
A.1 Vanilla SGD vs. Adagrad
In this section, present the medd (t) values defined in eq. ① obtained by vanilla SGD and Adagrad on language modeling task The task is to assign probability for the likelihood of given word (or sequence of words) to follow a sequence words. We trained transformer model to solve this problem Wikitext-2 (MX BY-SA 3.0) with real targets and randomly generated targets. This model has roughly layers (not counting normalization and dropout layers)
The setup is the same in Section 3.2 decreasing) for SGD and Adagrad. We tuned and chose the best (initial) learning rate of SGD. The (initial) learning rate of Adagrad was set as value under which Adagrad converged faster than SGD with its best (initial) learning rate. We used large batch sizes to make the training procedure more stable. When computing Hessian, also used large batch sizes. Due to the extremely large dimension, we did the computation on some uniformly selected coordinates, more precisely, 200 coordinates per layer.
We used the same learning rate schedule (constant or
We tried different initialization (normal and uniform) by using different gains of the Pytorch initial- ization schedule.
A.1.1 Experiments on real datasets
Figure 5a shows the training losses Wikitext-2. Table 3 (resp. Table 4 shows the p (t) for Adagrad and SGD under uniform (resp. normal) initialization with different gains.
SGD
SGD Adagrad
Adagrad
10.5
Training loss
Training loss
.0.3
10
520
1000
1500
2000
2500
3002
3500
4032
500
1000
1500
2020
2500
3000
3502
4200
lleration
Iteration
(b)
Figure 5: Training losses of Adagrad and SGD. (a) on Wikitext-2 and (b) on Wikitext-2 but with randomly generated targets.
Table 3: RSSO (L) and AdAggaad (L) under uniform initialization with different gains
(a) Gain
Gain
Layes
AdGR (t)
Fpoch
(ellaad (t)
KeRp (4)
Fpech
Ae A (c)
A53R (t)
Fiunch 40
Relaand e
Layas!
Fpach
Addand
Fguch 30
(e (t)
Fpuch
Addland (0)
6:77
5.91
9.7
5.15
69.36
18.40
15:36
1.23
1.66
24.12
2AM
4,35
2.83
3.56
4,50
9.47
10.78
9.76
3.54
8.67
3.24
5.25
4.74
3.83
5.68
3.11
4.51
12:54
13.96
10.31
6.59
9.79
6.88
66.49
67.53
88.75
19.31
63.01
15:64
7.21
6.54
13

---

Table 4: med (t) and R Rneerrmit (t) under normal initialization with different gains
(a) Gain
Gain 0.5
Laymn
(1)
(g)
Reai (f)
40
Tayen
(1)))())
Roel (t)
20
Epuch
K (0)
5.27
12.28
9.12
14,46
5.55
3.19
10.70
15.42
2.04
3
738
7.35
9.61
3338
9.25
3.94
5.73
5.94
10.16
2.80
6,05
2.99
18.02
15.63
6.45
4.80
7.49
4.44
16,02
12.9
K90
391
8.12
4.14
9.35
15.98
16.98
2332
4.84
2.46
11.79
8.58
17.85
4.52
14.99
4.50
6.52
6.55
107.51
3.14
136.38
2.73
17.09
1273
26.00
5.66
26.91
6.73
8.19
5.30
$33.4
518
315.1
4.48
A.1.2
Experiments on random datasets
shows the training losses on Wikitext-2 but with randomly generated targets and Table shows the corresponding med (t) and paAdagrad (t) in different layers.
Figure 5b
Table 5: R (t) and R med
(1) on Wikitext-2 but with randomly generated targets
Layert
Epoch
Fuel (e) (1)
Epoch 20
(l (t)
(cSCP(t) (t)
Epoch 40
Jchlagma
10.88
10.98
9.99
18.66
9.67
22,37
13.01
3.99
3
7.45
8.52
459.71
6.09
451.16
5.11
5
7.09
7.88
103.60
353.89
184.61
190.17
18.38
4.08
18.69
2.73
7 7.81
5.40
294.68
62.72
229.25
29.76
329.12
20.59
203.70
9.57
A.2 RMSprop and AMSGrad
In this section, we present the results of RMSprop and AMSGrad and compare them with SGD+M. The experiments were conducted the translation task described in Section 4.1 We used learning rate 2.5e-5 for RMSprop, 5e-4 for AMSGrad and 0.03 for SGD+M. Both RMSprop and SGD+M used momentum parameter 0.9. The two momentum parameters (81, _33) of AMSGrad were (0.9, 0.98). Figure shows the training losses and Table shows the corresponding ROPT() (1).
P
AMSGrad
SGD+M
Training loss
600
800
1000
1800
Iteration
Figure 6: Training losses of RMSprop, AMSGrad and SGD+M on the translation task described in
Section 4.1
14

---

Table 6: R med
pRMSspo (t), Pasisaad
(t) and meal (t) in some layers, the translation task described in
Section 4.1
Epoch (1)
Epoch (2))
Layord
Re (5)
R (()
Perto (t)
Repa (1)
R 22
3,97
1668
26.17
21.19
11.36
37.11
17.83
10.85
51.94
10.22
12.31
4.10
6.98
6.12
3.94
4.95
2.92
7.58
2.29
2.58
29.41
35.72
25,86
37.81
19.89
16,00
30.88
16.24
9.97
12.67
83.06
33.63
19.51
140.99
12.22
6.72
44.77
6.98
5.37
18
8.71
2.99
9.48
3.86
2.44
4.16
3.51
2,10
2.35
8.70
5.67
6.95
8.13
3.59
5.13
6.46
2.30
2.83
28
4.44
2,42
2.64
4.67
1.85
1.81
2.63
1.46
2.13
A.3
Comparison conditioned on the same loss
In this section, we compare RSCO (t) and Rasaa conditioned on the same training loss. More precisely, we make comparison of ( (t) and RSSCDM (t'), where f, t' are picked such that ith Adam iterate and t'th SGD+M iterate have the same training loss. The details of the tasks are described in Section 4.1 Table
shows the results of Reca (t) and RSGDM (t') in some layers.
Table 7: RAdda (t) and RSGDM (t') in some layers. Dataset and task: (a) sentence classification task on BERT-small, (b) translation task on Multi30k. See Section 4.1 for detailed setup.
(b)
Losc1.22
(RA
Layon
Adde
Adde
4.45
7.44
2:54
16.77
13.69
14.14
12.71
15.17
9:56
5
31.19
27.50
44.29
21.46
57.83
19.52
18.44
7.79
51.39
45.43
5032
40:57
9
21.23
53.65
26.95
2699
44.25
26.15
208.29
464.7
315.58
313.99
13
17.77
35.80
1443
19.17
19226
18.26
12.99
15
$2.30
196.41
34.09
15.06
33.8
531
257.32
85.47
188.55
10030
197.79
139.4
21
100.43
23.06
23,45
5.12
1296
533
30
19.14
9.54
10.48
5.58
9.56
533
SGD
Adam
8-
Adam after SGD
Training loss
3-
2
1
500
1000
1500
2000
2500
Iteration
Figure 7: Training losses of SGD, Adam after SGD and Adam for the translation task
15

---

Table 8: med (t) of SGD, Adam after SGD and Adam in some layers after roughly 2160 iterations
Layer#
SGD
Adam
Adam after SGD
Adam
13
294.76
150.02
332.96
150,02
14.34
36.38
16.66
11.86
16.56
6.47
7.05
3.76
7.05
17.17
26
5.68
3.53
2.30
3.53
14.33
15.93
21.76
15.93
9.10
29
8.22
3.04
2.82
3.04
11.39
5.12
5.29
5.12
A.4 Experiments of switching from SGD to Adam
In this section we describe another learning schedule: the "Adam after SGD" schedule, where we whether the loss and R (t) can catch up with the model trained by Adam from the very beginning. Again, we used the same model as in the translation task in Section 4.1 In this section, we did not add momentum term to SGD in order to get larger gap between SGD and Adam than the case using momentum. We want to see whether this larger gap can be closed after switching to Adam in the middle.
switch from SGD Adam in the middle
As is shown in Figure 7 and Table 8. both the loss gap and the gap of ROPT (f) were closed after a period of training after switching algorithms, which provides evidence of the connection between convergence speed and uniformity of diagonal of loss Hessian.
A.5 The low rank structure
In this section, we present more results for the experiments in Section 6
We examined the weights of the model trained for the translation task in Section 4. Among roughly
30 layers, we observed that for 12 layers, at least the weight matrices obtained by Adam after training have approximately low rank structures.
Figure shows the examples of layers with or without the low rank structure.
Layer
Layer
$60 Adam
SGD
2.50
Adom
1.5
sinqular values
2
1.5
0.5
1
0
0.5
100
200
300
400
500
600
100
200
300
400
500
600
coodinate
Figure 8: Examples of singular values on (left) layers without low rank structure and (right) layers
with approximately low rank structure.
We then studied the uniformity of leading singular vectors of these 12 layers, i.e. computed R. defined in (B) of Section 6 The observation was that for 10 out of these 12 layers, R,, values of Adam were smaller those of SGD, which implies the uniformity of leading left singular vectors of Adam. Here we may also want to consider the right singular vectors U1, U.. and corresponding
mediam for Adam and SGD+M. However, on this translation task, among the 12 layers which were approximately low rank, for only 6 of them, R, of Adam were smaller, i.c. we did not observe uniformity of the leading right singular
= [01, 62 ..., 6a]" == of oo O v, and compute R2 := max; it.
16

---

vector for Adam. One possible reason is that for weight matrix, its right singular vectors are closer to the input data than left singular vectors and more easily influenced by the data, therefore may not show uniformity. Figure P shows how R2 and R, changed over time in some layers.
Layer
Layer #11
Layer
Layer
2.5g
2.5
singular values
SGD HM Adam
Shdgg values
S +M
singular volues
S +M
singular values
S D Adam
Adam
Adam
1.5
1.5
1.5
1
1
1
0.5
500
500
500
500
coodinate
coodinate
coodinate
coodinate
14
12
18-
20
16-
14
of 15
10
Adam
5GD+M
8-
Adam
SGD+M
12
Adam
S +M
S +M
Adam
10
10
50
0
50
50
epoch
epoch
epoch
egoch
10
12
10
12
of
10
of
of
of 6-
10
SGD Adam
9-
SGD Adam
SGD
SGD
Adam
8
50
0
50
50
epoch
epoch
epoch
epoch
Figure 9: Rr and R. for Adam and SGD+M in some layers
A.6 How does the (adaptive) gradient align with diagonal of loss Hessian
In this section we present the uniformity of diagonal geometry of adaptive methods from another perspective. Denote Ha as the (i. i)-th element of the loss Hessian and gi as the i-th element of the gradient. It is conjectured that when |H& is large, the corresponding lgs| is usually large as well. For adaptive methods, we can regard the update per step as the learning rate times the "adaptive gradient". Let's use Iadap,, to represent the i-th component of the adaptive gradient. Through experiments on language models, we found that [gadapt,a| for different are quite uniform and do not align with HHal as the true gradient [9:l does.
In the experiments, we first sorted |H in the ascent order: HHa. $ |H1s,2l S ... S |Hinaal (suppose H Rdxd), and then plotted the corresponding [9tal and 1addapt. for [d].
A.6.1 SGD vs. Adagrad
Here we compare SGD and Adagrad on the language modeling task on wikitext-2 described in Section A.I We observed that the figures of all layers are quite similar select one layer as an example, as is shown in Figure TO
A.6.2 SGD with momentum vs. Adam
Here we compare Adam and SGD+M on the tasks described in Section Again, we select one layer as an example for each task. Figure shows the results on the sentence classification task and Figure shows the results on the translation task.
17

---

grastem Adagrad
Adaptive gradient Acaggnd
10
10
0.5
60
100
150
200
50
100
150
200
50
100 Occeninade
150
200
apaid 240
0.05
60
100
150
200
00
100
150
200
00
100 ocordinate
150
200
Ooo diriato
coordinate
10
0.06
upodh 40
Masee
10)
0.5
200
200
200
coordinate
ccoromnate
Figure 10: How the true gradient (fl9s. 12=1) and "adaptive gradient" (1l9mtam 1k=1) align with diagonal of Hessian (L) D2==)) ). Here coordinates are sorted such that [H S ( S < Haadd (suppose E Rdxd). Experiments were conducted on the model described in Section A. This figure shows the results on the 12-th layer.
True 5rD
gradient Adam
Adaative gradient Adam
0.5
0.1
181
JE 0.05
0000
AA
50
100
150
200
0
50
100
150
200
50
100
150
200
cocrd nale
cooodinat
coordinst:
0.5
0.02
0.01
opc0
A
50
100
150
200
50
100
150
200
100 coordhatt
150
200
coord nalk
coordinata
1000
100
150
200
200
100 coordiaate
cocnd
cocndiade
Figure 11: How the true gradient (19,, 1 and "adaptive gradient" (1laadap (2 align with diagonal of Hessian ((HHay D)=z)). Here coordinates are sorted such that [H S |Hig. S A H (suppose H E Rdxd), Experiments were conducted on the sentence classification task described in Section 4.1
This figure shows the results on the 12-th layer.
A.7 Adding regularization and other tricks
In this section, add weight decay both Adam and SGD+M the translation task described in Section [ The momentum parameter in SGD was set as 0.9. The two momentum parameters set (0.9, 0.98). For both algorithms, 0.001. We trained the model using constant learning rates for 60 epochs (1800 iterations). We tuned set as 0.0001, under
(B1, (B2) of Adam
set the weight decay parameter as
and chose the best learning rate 0.03 for SGD+M. The learning rate of Adam
which Adam converged faster than SGD+M with its best learning rate 0.03. Figure 13a shows the training losses and Table 9 shows the values of RAdam (t), RSGDM (t) and selected layers.
RSOD (t)
some randomly R anl (4)
med
18

---

The grddente 560-M
True gradient
Adom Adom
Abbstve gradient Aram
x10
200
OpeCh 55
150
0
1.00
200
150
cotedirate
0300h -51
150
100
coord
coordinate
<100
<100
G00C) 56
50
100
150
200
50
100
150
200
50
100
150
200
coordinata
courd nata
coordinata
Figure 12: How the true gradient (19s,, 1==1) and "'daptiiee gradient" ( Algaaam. 1 1 align with diagonal of Hessian i.. D)==1) Here coordinates sorted such that |H an is S ... S. H Sadel (suppose H E IR IRdxd). Experiments were conducted on the translation task described in Section 4.1 This figure shows the results on the 5-th layer.
Table 9: R (t) and RSOD
(t) (both using weight decay) in some layers for the translation task.
Epach
Epech
Epoch
Layer#
Rad (1)
R (t)
Resst (6)
Praii (1)
-1 122 (1)
R Reca (t)
R
[ 722 (G)
73.09
73.09
17.65
13.11
1.35
13.38
6.28
2.13
469.88
469.88
293.48
310.85
0.94
601.68
588.12
1.02
80,78
80.78
8.22
39.65
0.21
13.65
4.85
2.81
494.27
494.27
150.14
123.79
1.21
301.89
119.53
2.53
632.10
632.10
277.18
175.34
1.58
334.48
282.88
1.18
18
55.08
55.08
6.56
4.45
1.47
23.88
4.52
5.29
549.62
549.62
257.89
44.78
5.76
515.99
53.79
9.59
107.51
107.51
8.54
3.64
2.34
53.79
3.32
16.20
13:77
13.77
4.74
2.37
2.00
15.60
2.15
7.24
30
491.62
491.62
6.91
2.66
2.60
9.60
2.02
4.77
A.8 Results on image tasks
Although in this paper we focus on language models where Adam shows significant fast convergence,
we also add results in this section on image tasks where SGD+M performs better. We trained a ResNe on CIFAR-10 dataset and compared the convergence speed and ROPT (t) of SGD+M and Adam. The momentum parameter 3 in SGD was set as 0.9. The two momentum parameters (61, 82) of Adam were set as (0.9, 0.98). The model was trained using constant learning rates for 41 epochs (2050 iterations). We tuned and chose the best learning rates for both algorithms: 0.5 for SGD+M and 0.005 for Adam. Figure 13b shows the training losses and Table 10 shows the values of RAAdam (t),
RSGDM (t) and (ESGMM()) (+)
As we can see, on this image task, Adam does not converge faster than SGD+M and in the meantime, Ryddam (t) values were longer smaller than RSGDM (t) during training. This reveals the connection between the local diagonal geometry and the convergence speed from another perspective. That is, when the diagonal of Hessian of Adam is not more uniform than SGD+M, its convergence speed is not better, either.
We borrowed the implementation here http:::///pptttochhttuu tutoriaa. rewdthedocs hatpe/cccccttttciaaaereerrededoo 1o/1////nn St//n/lattes, tutorial/ ceterraallchateerr_inee httpr/rddooitttmmeeeaeee///e 2tt.rrreeeetttt/2222222oeeoreeee wtt.222cccmmreetttccffaar0000 and replace the "layers" aray [2.2.2] with [1.1.]].
19

---

25
SGD +M +wd Adam-wrd
SGD Adam
Training loss
Training less
1.5
200
422
506
1006
1708
3400
1606
1883
500
1893
1533
2930
Iteration
Iteration
(a)
(b)
Figure 13: (a) Training losses of Adam and SGD+M for the translation task, both with weight decay. (b) Training losses of Adam and SGD+M for a ResNet trained on CIFAR-10.
Table 10: Re (1) and R (t) for ResNet CIFAR-10.
Laycsk
Epoch
Epoch
Epoch R
722 107
R00uN (1)
(20m(()
6.88
25.34
0.27
3.74
39.35
0.09
4.39
15.80
0.28
35.03
58.06
3
16.02
15.00
11.70
13.04
23.78
12.49
12.81
11.42
19.1
11.72
14.53
A.9 Comparison between Aodd
(t) and singular value-based metrics
In Section 4 through extensive experiments on language models, we demonstrate that when we train the neural network using Adam, the uniformity of diagonal geometry, measured by po (t) is smaller than that when we train using SGD+M from the same initialization, for the vast majority of layers. We are aware that people also usually consider Hessian singular values instead of diagonal entries to measure the loss geometry. Hence in this section we make a comparison between our diagonal- based metric and singular value-based metries.
First, we believe that our metric has a natural connection to the mechanism that underlies adaptive methods. Adaptive methods in practice choose coordinate-wise adaptive learning rates. From high-level perspective, this procedure can be viewed as adapting to the loss smoothness with respect to each coordinate. The smoothness of certain coordinate is measured by the second derivative with respect to this coordinate and therefore corresponds to the diagonal entries of loss Hessian. Our metric, which measures these diagonal entries, is thus fundamentally intertwined with the mechanism that underlies adaptive methods.
Next, we empirically demonstrate that our metric ROPT (t) is a reasonable proxy of singular value- (a)
as an analogy of based metrics. Define a singular value-based metric SOPT (t) == median (o,(1)
our diagonal-based metric RoOr (L), where (a,(1) jiz denotes the singular values of loss Hessian (t) E dgd at the /th iterate. We compare SOPT (t) along the trajectories of Adam and SGD+M in the translation task described in Section 4.1] Table suggests that measured by singular values,
(
Adam is also biased to region with smaller SOPT (7) than SGD+M, similar to the observation for
ROPT (t). This is expected because in Appendix F we demonstrate that the loss Hessian approaches metric and singular value-based metric give the same result also reveals the robustness of our observation to the choice of metric, demonstrating that there does exist some geometry bias of Adam towards more uniform regions even when measured by different metrics.
diagonal during training. The fact that our diagonal-basee
20

---

Finally, there is strong reason why our metric is often easier to compute empirically and analyze theoretically than singular value-based metries such as S (t).
1. From the empirical computation perspective, suppose the loss Hessian is d x d. Then computing its singular values, in general, requires computing the whole matrix with d² elements. However, metric only requires computing the d diagonal entries.
2. From the theoretical analysis perspective, in Appendix G we show that the diagonal of loss Hessian in linear networks can be connected to weight matrices by simple formulas. These straightforward formulas simplify the analysis and allow us connect our metric to the low-rank structure of weight matrices and the uniformity of their leading singular vectors (see Section
6 for more discussions). However, all these nice connections fail to hold for singular value-based metrics. The formulas of singular values are very complicated even in linear networks, making almost impossible theoretically analyze any singular value-based metrics.
Table 11: StAdam (t) and Smod (t) in some layers for the translation task.
Epoch
Epoch 30
Epoch 55
Layer#
SSGDM (t)
"ndd (1)
SSGDM (1)
(1)
(c) $2
SSGDM (4)
(C)
(03 (6)
4.53
4.53
6.08
2.87
2.12
5.92
2.58
2.30
14.64
14.64
40.01
15.38
2.60
52.28
15.05
3.47
6.91
6.91
9.84
5.06
1.94
12.25
4.22
2.90
24.12
24.12
42.02
30.89
1.36
33.20
21.54
1.54
19.07
32.41
24.84
1.30
28.83
14.23
2.03
15
47.03
47.03
69.97
11.54
6.06
42.71
7.19
5.94
18
15.96
15.96
26.03
29.73
0.88
18.46
17.94
1.03
21
31.03
31.03
25.84
7.92
3.26
19.71
7.06
2.79
24
35.42
35.42
21.31
18.08
1.18
14.62
10.33
1.41
55.38
2.77
A.10
The uniformity of diagonal geometry on in-distribution test data
In this section we compare Reddum (t) and RSGDM (t) on the an-distribution test data. The task is the translation task described in Section 4.1 Table 12 validates that on in-distribution test data, Adam is also biased to a region with smaller ROPT (t) than SGD+M, similar to what happens on the training data shown in Table 2a This is expected because of the same distribution. One other thing want to emphasize is that, in real language tasks, the dataset is typically very large and the model see each training example only once. Hence the training behavior usually implies similar in-distribution test behavior.
Table 12: AAdaa (t) and RSGDM (t) in some layers for the translation task on in in-distribution test data.
Epoch
Epech R
Epoch 55 R(dm((t)
Layer#
RstapM (1)
RM (t)
RSGIPM (t)
R SSS0 (t)
J p 127
4.39
4.39
5.80
3.06
1.89
6.79
2.81
2.42
7.90
7.90
38.01
11.71
3.24
41.40
10.21
4.06
5.77
5.77
6.00
4.61
1.30
5.53
3.20
1.73
25.09
25.09
28.81
17.17
16.67
14.85
1.12
10.24
10.24
9.13
8.63
1.06
13.78
9.09
1.52
15
79.71
79.71
77.18
13.56
5.69
37.93
9.91
3.83
18
14.78
14.78
3.94
7.15
0.55
5.42
6.04
0.90
83.25
83.25
26.04
4.79
13.11
5.57
2.36
29.91
29.91
6.89
5.42
1.27
6.51
7.16
0.91
28
22.57
22.57
5.39
3.94
1.37
6.13
2.14
2.87
21

---

B Supplementary result on the contribution of small Red (t) to fast optimization
B.1 Supplementary empirical result
As explained in Section [] Figure demonstrates that Adam biases the trajectory towards region which benefits fast optimization, and in the meantime, it has more uniform diagonal Hessian (charac- terized by smaller R (t). However, since the two trajectories in Figure come from two different algorithms (SGD+M and Adam), it is reasonable to ask whether the benefit of Adam's 's trajectory is indeed due to smaller med (t) or is due to other properties of Adam and small m (t) is just byproduct. To address this concern, we add the following experiment where we select two trajectories that both come from SGD+M and have different mod (t).
We ran SGD+M twice on the translation task described Section 4.1 with learning rates 0.01 and 0.03 and got trajectories. Then we selected an iterate æ the trajectory with learning rate 0.01 and 22 on the trajectory with learning rate 0.03 such that æ and 222 have the same training loss and starting from 21 different m (4). We then SGD+M with the same configuration twice, (characterized by and once from The observation is that ø2 has a more uniform diagonal Hessian
smaller R (t), Table 13 than 21, and in the meantime. SGD+M starting from æ2 also converges faster (see Figure 14p. Hence we demonstrate that small value of R (t) is just byproduct of adaptive algorithms (since SGD+M with different learning rates can also yield trajectories with different por (l)), and it indeed contributing factor of fast optimization.
3.15
loss starting from zn
3.1
loss starting from 21
3.5
Training loss
2.8
30
40
50
60
80
90
Iteration
Figure 14: Training losses of SGD+M starting from in and 22
Table 13: ROPT (t) values of 21 and 22 some layers med
Layer#
med
(on the iterate x1,
ROPT2 (on the iterate x²2)
3
2.89
2.73
1.06
69.73
51.44
1.36
4.39
4.81
0.91
42.93
35.89
1.20
12
7.09
6.18
1.15
97.25
68.92
1.41
3.50
3.74
0.94
21
71.69
45.07
1.59
24
7.49
5.42
1.38
28
4.85
4.00
1.21
22

---

B.2 Theoretical analysis on simple setting
In this section, we add the following lemma on simplified setting where smaller Ree (t) provably contributes to faster optimization.
Lemma 1. Consider a diagonal quadratic loss f(x) EE 2f where X 2..
Run GD with learning rate 1= b and produce {-, (0)
1120, d. We initialize the parameters
as x; (0) +1, Vi E [d]. Without loss of generality, suppose d is even. Then we have that for any log d steps, the loss fx) will shrink by half.
A p > after T 2 Rorg log median
More precisely,
(2
( +
(c!0))
S
Lemma I demonstrates that under this simplified setting, even when we just use the vanilla GD with no adaptive methods, a more uniform diagonal geometry (characterized by smaller (oPP can still lead to faster optimization.
Proof. By the update of GD with learning rate 7) = Xi
L.
we have that
Vi E [d]
(++1)
1
A
I;
(t)
(1)
A
26
(0)
2
2hrt
(0))))
S
exp
22
x1
Then after T = f median Ao |
log d steps, have that
d/2
4/2
Vi < 2
(24 s <d (cf))²=aar
(2)
x 10))
A
d
Σ
A
1
(0)
S cd (d
-=1
1=1
On the other hand,
d/2
Σ imd/2+1
Σ
(0)
2
(3)
Σ A
x )
i=d/2+1
where (1) is by definition of median and the initialization Vi E
[d] x
+1. Combining with
the fact Ei=d/24 A (z=!))² + d/2 A (1()))
-s (zx(0)) yields that Ei A a2
(T)
<
f (z)))). Hence we have
d/2
{(((T)
Σ
A
2
(7)
+
Σ
(2. 2 s ((+ (z!0) A
(e(T))²=
i=1
i=d/2+1
C Proof sketch of Theorem 1
Now give a proof sketch of Theorem which contains three major steps. The detailed proof can be found in Appendix G andd B
First we relate the diagonal of Hessian to weight matrices Wi, Ws. Under Assumption 1 denote
Wi [i.:] ] as the i-th row W1 and We == [ W22 ..., W2d). Since the input datasct is whitened, we can show that
(t) max;
| [
max;
ROPT (t)
RoeT.2(1)
(t) 2
median
median [w(") [2,
2
23

---

Next, due to the one-dimeett nne-dimensional output, we can prove that W1 converges to an approximately rank- matrix. More precisely, we have
w = ( + RT(1
W (() R_
where c(1) is a scalar, u u0),, v(
Rl E R" and R( E Rdxd Denote the i-th coordinate of (t) (0) respectively. Denote the (i.j)-th element of R(4) as R j].
( v(1), R(( as ui Di R
We have that Vi,j [d]
R
c(0
(0)
and R
1 [i
(2)
Using the rank 1 structure, can further simplify R (t) and ROST (t) by
(1)(0)
max;
R (t)
(1) 2 k 1,2.
(4)
median
u
The final step is the detailed analysis of u (1)
For SGD+M, we can prove that (1) = C(t) X2. x.)²" where C(t) E TR and X,, i E [d] are (
i.i.d. Gaussian variables. Then we have with high probability,
(llg d).
median 1
For Adam, we can try to convert signed descent (e.g. in eq. (2 and Lemma 24). The main intuition is that as long as the learning rate is small enough, the movement of the gradient will be slow and Adam will be similar to signed descent. Since the update of signed descent per step is t.
max: u4 )
we can prove that Vi E [d] u, E {+1), which gives us
1000 = 1. Substituting into eq.
median
completes the proof.
Finally, we want to emphasize that the assumption B2 = Bf is not essential. As discussed above, as long as the learning rate is sufficiently small, Adam can be approximated to signed descent regardless of the ratio between 31 82. Hence our result can generalize cases with other B1, 82, for example, 32 > B7 where Zhang et al. (2C proved the convergence of Adam to stationary points.
D Analysis of SGD+M
Note that A= m
Y AZx = m XXT. Denote g (t) Vw, Iw.(()))) 1,2. We have that - w(trr
(1) (()) (W
92 (6) (W/ w(( -A)
Let íto, ( and 0h = 1,2 be the corresponding batch versions at time Let E(1) = /h (t) ((() A, and use E A, and (w." (()). to represent the i-th coordinates of E(0), A and w: (())) respectively. By eq. 2). the update rules W1 and W2 for SGD+M are given by:
w(((-1) -w100 " g w(
(w( A) n
3 "Dg (T)
(-(t-1) w:(() 17 Σ gi-t (w( - A) w(tt Σ g 3r-TDg2"
Dg2 (T)
where
Dg1 = 91 91
(1) (1) (t)
(-(t)T Wv/"w(1 (Al) Aex) - -(A - A)) -Ax) (A -1)) W(()
Dg2
(t) = 92 92
(((!
Based on the magnitude of W2 and Wi, can intuitively divide the training procedure into 2 phases.
24

---

1. First phase: the first several iterations when W1 and W2 are "small" so that W2W1 A R
2. Second phase: later iterations when W2 cannot be ignored.
More formally, the boundary between the first and second phase is defined below.
Definition 1 (End of the first phase). The end of the first phase (denoted as Tí) defined as T inf {+ 2 0: Bi E [d]
(8)
W [. 51
1 d $
2
d
By Assumption and the assumption that Vj [d] A, > 0, A, A, @(1), at the beginning, w.h.p.. Vj E [d] (W2 w1, A, < O. During the training, each (W2W1); increases and approaches Aj. We hope that by choosing a small learning rate, when (W2W1); overshoots for some coordinate j. i.c. (W2W1); > Aj. will be close to convergence. To analyze this overshooting issue more carefully, let's 's first define the following "almost overshooting time".
Definition 2 (Almost overshooting time). For > O, denote 60 = uे + e elog V Define
T- == inf { { : 3j E [d] (w," (W/"
w('', A -
Definition 3 (Convergence time). For >> O, we define the "convergence time"
T3 == inf 2 0: ( 2 <
We can first show that after the first phase, i.e. when & = Tı, W1 will become an approximately rank-1 matrix, as described the following lemma.
Lemma 2. Under Assumption
and 3 suppose o S
1///2 d0/247-
By picking S 50 (J=). we have
that when Tú, L (W(T))) θ (d), and that
W (T(TT) u (Ti) (11),
w = R + (T)) (T.)
where c(Ti) E R, ()) (T))
(T) E Rd and RTTT)) E Rdxd, Denote the i-th coordinate (TL) (T.) R respectively, and the (i, h element of R((() as
of (T) o(Ti), (ef u, Us R((T)) (i.5). Then w.h.p.,
R (T) (i,]]
[i
11 < i.j 4 d:
di
)
1 dda-i
u ( (TU)
(( (Ti)
The following lemma tells us that this approximate rank-1 structure is preserved when Ti S 5t7 min (T,, T3).
Lemma 3. Under Assumption have that w.h.p. for Ti <t min T2, Tal.
2 and
suppose S /0/241. By picking nj S O
1/22
(~
[
W(O0) u (T) (1) + R!
w(O (t)) (T)T ROOO..
where
R [1,0] &((c))
R
AA $ iisi
u (
c(U) (Ti)
and eo is defined in Definition Moreover, when t min (Te, Ts}, I (w(t) O
The following lemma gives us more detailed description of u (T,)
Lemma 4. The a (T)) in Lemma a can be written as u (T))
= X + Y where Xz, i E [c] are
ii.d Gaussian random variables and that w.h.p. Vi E [d]
X < 0
[
d7a
Now are ready to prove the SGD+M part of Theorem
25

---

D.1 Proof of the SGD+M part of Theorem 1
we can apply Lemma (EEdd). For any > 0, by picking 0 <<<<<< 20 and a 4(p + 2), have T. (W(Taa)) = (cood) ($).
Define TSGD,1 = Ti, TSGD,2 = min (T2, By picking S 50
and 3 to conclude that (W((ssm,.)) = (dd) and L (W(Xsm)))) d 1
Moreover, when [Tscp. 1. 2l. the conditions in Lemma 31 are satisfied with Ō( Then we can apply Lemma 31 and get that
(Ti) 2
- Ō(co)
max;
"
R m (t), R med, (t) 2
1 O(co)
2 u (21)
median
By Lemma 4. u (Ti) = X + Y where w.h.p. Vi [d]
S
0
This fact yields
max, (Ti)
dt--
max; X? median X?
Vi [d]
(TL)
median
Here X, [d] are iid Gaussian random variables by Lemma
4
To prove the concentration of By setting K = N = d in this
median X?, we borrow the Proposition 12 in Chapter 2.3 of [erll proposition, have
P mediaan x² - E[x+]] 2 Var(X
S Se
Denote 0² as the variance of X,, [d]. Then E[X?] o² and Var(X?) 20 Hence
P (meddian x? 0²1 2v20² cc
That means with high probability, median X? S 2a for some C > O. By Lemma 35 in Appendix H we know that w.h.p.
max X² σ 2 d), Isedd
which gives us w.h.p.
maxisisa X? = 22log d).
median X?
Hence we have proved that 12SGDM (1), ROOM (1) 2 22llogd)
D.2 Proof of Lemma 2
In the first phase, W, W, is "small", and we write the update equations in the following way
w = w110
g T WT(T
WW/ - A) Σ 3"-rDgi (T)
r=0
= -w((0) +
W
(T)T A n
Σ
3 w w 17 Σ 3r 3'
2
000
0
0=0
=W + +n
(IT
A
gt
+
3 gt
[ - -w²!T7) A
(5)
==0
T=0
3 W
(T)
W/ (7)
W((T)
Σ
B BI-TDg
7=0
T=0
W (1)
W
(1)
A +
77
(t)
1
1
26

---

where
r = -0- (1)
A + (1 5)
Σ gh-r (W
-w( A
==0
- (1 - 3)
8 g W(fr (p)())p(() - (1 B)
3 DD (T)
==0
r=0
Similarly, we have
w W22 77
8 (// - A) w(+h = w/2h) + 1 AW (t)T (t)
+ 1- 3 2 (6)
(0)
7=0
where
(/ = -B AW + (1 - 3) Σ
gh -A (w(t)=--,("
- )
==0
- (1- - 3) Σ 3 ( - - (1- B) Σ 3 DDg2 (T)
T=0
r=0
The following lemma gives us explicit formula of w(
Lemma 5. Let \1 < X2 be the two roots of the quadratic equation x² 2x +
||A| |1
n²
0.
Pick nj
1-8 then we have that T
W = C1 + (C2 +5 X2.
W 2
C2 W W11)~ -,, (0) dg
where C
r will be specified in the proof.
Ag
We can prove that in the first phase, T5 (1) is "small". More specifically, denote its i-th coordinate as Ss and the i-th coordinate of C2 as Cz.. Then the following lemmas tell us that Vi E [4], r S
O (arlo)) ), where w.h.p. O (ar)o)) )
min |C2sl-
We first have the following bounds of 11
(0)
(t)
and
(t) 55
for i [d].
Lemma 6. Under Assumption 1
and
suppose
/0/24T and pick 17 S O ) We have w.h.p.
1//2
for all t & Tr Vi [d]
[
(t)
<
Lemma 7. Under conditions of Lemma 6 we have that w.h.p. for all t S Tú, Vi E [d]
Next we prove upper and lower bounds of |C1il and |Czil for i E [d].
Lemma 8. Under Assumption i) w.h.., Vi E [d] : (Cii <
and 3 suppose a S 2 Pick <
we have that [[Al|2"
1-8
de |C2sl S (1));
ti) C2 can be written as C2 ==
(C' + C1) where Csi [d] are i.i.d Gaussian random variables
and that w.h.p. Vi E [d]
[Cul
<
[aa
-
tii) w.h.p,, Vi E [d] |Cal Ω
1 Ω
Now we are ready to prove Lemma
Lemma 5 tells us that
w:O Ci + (C2 + r3
(1)
X/.
where d1 1 - || and de 1 + - ||Al|2.
27

---

Under the conditions of Theorem and pick 17 S 00( by Lemma 7 and 8 Vt s Tí, A1 sis d.
we know that w.h.p.
S Ô
1C2d 2 á
dīa ) |C2sl
1
Ô
1 do
re
S δ
1
Si
dza-1
1Casl
dia-i
(7)
We first prove that see this, first note that
reaches
for some coordinate i before d
|W/
for Vhi E [d]. To
w = 11(-11 1)
W W 1yT
A
+
17
(t-1)
2
1 1- B
1-1
= -W10
+
1
(T)T
A +
11
Σ 1
1 1- B
1-3 7=0
T=0
= 1
C1
(-1
Xj + C22
Σ
(-1
N/ +
Σ
1-1
(+)
T
A+
7
Σ
L-1
1- B
1- 3 T=0
-=0
T=0
1=0
1-1 Σ
1-1
T
-W
1
(T)
T
C
+
A+
"
Σ
M A
1 3
1- 3
1- B
1-1
t-1
1-1
== (((0)
7
Σ
1
(7)
+
C Σ
+
Σ
(7)
2(1)T + 2000
1- 3
0=
1=0
where (())
1 A and
H-1 Σ
1 =
NOC...
(8)
Moreover, we have that
1/
Who = -GM + (C2+11)) (
- 10 (T))
)d) =
Σ1 15"
Fort S T1, by eq.
we get that w.h.p.,
EN= N5rl 000
x-- No
LI S Sijjd: d:
d
S Ō
(dta-i).
12
1t-1
\o
d
555
1
c(022))))
122
d'a-i
Fort t < Tı, by Lemma 6 A1 < <i,j <d:
d: |110 |3.01
- 3
Then we have that w.h.p.
2 D r [t,5]
Σ (7) |i.00
Σ (T) |
I-B
u 00)
(1)
E A3
Σ
Σ
di
d
Here we used Vi E [4] A, ((1) by Assumption |
Since \ = 1 - - ||A|| |2. we have that |CT:A1| < [Cial S SO
1 7e
and that
t-1
74
Σ
1
A,CCui
C1+
Xf (0
= 1 3
C1.
<
(1 - B1 - A1)
5
||
s
dutt/
7=0
28

---

Using the Gaussian tail bound and union bound, we have w.h.p. L1 S j - d:
[i.
1 Combining the above bounds together yields that for t S Ti and Vi,j E [d].
w [. Rlf (.
+
u 1
(1)
+ ol (t.j).
(9)
22 Rig. + uu
(1)
where for j E [d]. [ & o 1
R
S ( and |
Further we notice that for & S Tı, we have Vj E [d].
nA,
Σ
N
MA,
\/ -
A, (22 -1)
Aj O
1
-
cl
1- -B
X2
1 B 12 1)
X111A112
|| vd
which
yields that
|u
S o
(122) (c(0m.))) (0)
Together with eq. (9) gives us that
22:
reaches
Zu77 for some
E
[d] before
| [K: for vk.j E [d]. i.e. Ti =
inf 2 0: Bi
[d]
(c)
d
Further, we know that at time Ti, means w.h.p.
c(Ti) (Ii)
|C2io 12 = = (o73 for some in E [d]. which
θ
+
T s δ (dto))
17 ((7)) S X1 = 1 1- 3 311A112)
S
+
0
Ti = θ ( log d
|1||A|12
(10)
This is the length of the first phase. As for (i) and u. oj
(T))
(T)) (Ti) for other coordinates, we have
that w.h.p. VI S j.j S d,
T
X
(1) X3 1
u 15
(3)) (2i)
= 1- 3 r=0
17
Σ
C2.A1l
=
1- B
17
A2 - 1
C2
||A112 |C2,A,1
i (ddm1)
1
1
5
-
θ (va)
11a+t
( c ui
(TT) = |C2i 1291 2 ② (g=0/)
dàa
=
d
Here in (i) we used X2= = 1 + 1-3 ||A||
Then we have at time Ti, Vi,j E
[d].
5
δ d and that
< Ô ) Together with eq. g we have the following weight
structure:
w(zi)) = n2(2ii + w (T)T
wi(Ti) = RIT..TT + (())) (T)IT
where w.h.p..
(1.1
visi,j d:
< 0
| ,
1
S Ô
1
m (
((Ti)) (T))
d'a-I
Finally, we consider the loss. Since Vj E [d] (w. L (W(T))) (dd).
-A, -((), we know that
29

---

D.3
Proof of Lemma !
Eq. tells us that (T)) ΣT-J XOOT Lemma g tells us that C2 can be written as C2
(Css + Ca) where Cst, i E [d] are i.d Gaussian random variables and that w.h.p. Vi E [d] :
Cur
I
Cad
Ô ) Combining these two facts together finishes the proof.
D.4 Proof of Lemma
Replacing / by ( 1 in eq.
we get
W (t) W (t-1)
"
AW((-1)
"
B 22
((-1)
(11)
1
Eq. 6 and substituting eq. 5 yield
n72
112
W - w((1) = -W - w-((-1) 1)
((1)
(1 - 3)2 Art (((()
+
+
2
(1 5)2
(t)
(t-1)
+ 1- B
2
r2
w:(++1) = -2W2 (t)
1-
(1 - B)2 |1411/
1/2
W +r/ (1)
-
where ra(t)= = (1-0)2 Ar (t- 15T
2
+ "
1-8
(1)))
22
- 12 (1-11)
For the equation x²
2x +
1/²
|
0, the roots are \1 =1 -1-133 |All2 and A2
(1-3yz
1 + -3 We have that
W - 12 = ~1 (W//0- 12 + +r/
W0 - -22 = 2/ (W/") - 22 + + X x-1-- (T)
== 1 (W!!1) - )o + + (8)
We further have
t-1 Σ X/- - Xi (W/!") -22
1 +> 15-rr/0
w 12
r=0
= =x/ + \/2 X1
(w:(1) >2
+ Σ 7/
A22 X
r=1
{ = =GM + +22l/+ + ΣxTTT 1イ
=G + (C2 + r!!0 Nb.
(t)
= Ei r/ Ci = -,p''_a-:'00 22-
-
-x; w(o)
where r3
and C2 =
22-A1
D.5 Proof of Lemma 6
Write
r(1) = -Bc+1 W (t) A
+
912
+
g/3 + 4i4
(t)
where (
= (1
(3) -'== gr-r (W(r)T - (
A,
413
(t)
= -(1 - () Σt= 3 3-- (-( W w and
4i4 -(1 B) Σt -1= gr 00=
(t)
(7)
And write r2 -3 (0+1AW(())''
(6)
(t)2' + 222
(t)
+ 423
(8)
+ 22
(8)
where
(t) (1 B) Σt 1
Er= gi- A [ (r)T (())
) 923 (1 B) Σt 3 - W W w (T)
(1)
(T)
422
(1) (1 8) Στ =o g 3
(T)
and
924
30

---

Let's first try to bound
912
(t)
[
and
922.2
(t)
For any T S Ti, we have that
d
Vi vi ed]:
(7)
[i,
S
2j
[3.8]
de
do
1
j=1
i=1
and thus 1A E [d] E = -1).. Then we have for all j,j E [d].
W(r+ ( (1
S 11
gr
(k)
E
1
k
0
da/2
k=0
k=0
'70
do/22
221
(T+1)
2221
(T)
S 17
8
W
(8)
[i,5]
Sn Σ
3
k
d0/2-11
1
k=0
==1
k=0
010
dd/2-1
That gives us Vi,j E [d].
(t-((--)
4i2 [i
(1
S (1 3) Σ
g
(w))
w A m1 3) Σo
da/2-1
=0
d=/2-1
d
(1) 922.3
S (1- - 3)
8
Aj
[6,5] (w/))(E.3) - (3
7=0
==1
S (1 B)
Σe
o
3
(1
0 =0( (7072-11)
d0/2-1
0=0
Then we bound
4/ [i
(1)
and
423,i
(t)
We have for Vi.j E [d].
4i3 |
(0)
$ (1 - 33) Σ
g
1226
(r)
(w("
S (1 - 3) Σ
gt-r
1
di d-
-=0
T=0
= =0 ( 1
dda-1
d
(0) 423,4
(1 3)
81-7 Σ
(
(
d
$ (1 3)
g
0
da
d=a-2
iml
Finally use Lemma 32 to bound
(0) 414 [t, .] and 924.1 For & S Ti, the M(O) Mn(0) in Lemma 32
(()
are upper bounded by d In the theorem consider the training period before TSGD,2 so the time in Lemma 32 is set as TsuD In the following sections, we will prove that TSUD,2 S
d² log( /d/e)
Then by Lemma 32 we have with probability at least
1. for Vt S Ti and
31

---

Vi E [d].
= 01 1 [i:
9i
[t,
s 0
d
-3
1
0
ddt1
log -
d
+
d
0
d
log
d
1
da+2
S
0
d²
Dg2
(1)
92
(t)
92i
(t)
S 0
1
0
da+1
log
d
+
0
d 0
ov
da+2
log
d
d²
4
c
c
d
<
d3
1
By picking
<
/0/24I, we have w.h.p. for Vt S Ti and Vi,j E [d]
1//2
(Dg("1111 [i.5]]
S O0
and
1g S n
which yields
4i4 [1,5]
(0)
S (1 - 3)
gi-
Dgi
[1.5] S 3) Σ
=n
d
d
7=0
-=0
924,i
(1)
S (1- 3)
Σ
Dg2
(T) S (1 (1-33 3) Σ
g n
dz
=
n
d
==0
==0
Combining all the above bounds and substituting 7) S gives us for Vt S Ti and Vi,j E [d].
[191 s g (1) A,
A1|+ +
(
1
(0) (2
S g
Σ
A, 1 (9 [1, + dā
dda-1)
(12)
Fort S Ti, we have Vi,j [d].
(1)
W22
Aj
0
and
Σ A A, (1) [i.5]
0
(
which gives us |r.0) [r,]
I 0 (347z)
and
22
(t)
O
La07
Substituting into eq.
and eq.
66)
yields that for 1 < Ti and Vi,j E [d].
[i.
w/ (i, إ|
O
7
((+1)
50 dd//211)
da/2
Hence for t S min|
a loge 1o
Ti }. we have Vi,j E [d].
[i.0 S
w 1w(0 [[..]]
a elog d log(1/3)
O
"1
S
5
do/2
d
(
S
u22
(0)
a log d log (1/3)
dn/2-1
71
S Ō
d -1
Then we know that Ti
>
a log and also get tighter bounds of [1/3
(2) 2226
for t <
log(1/3)
o log
Now we use these new bounds to analyze (110 [E.i]]]]
and
122 again.
(t)
When t <
a log To we have for all i,j E [d]. ge+1
u22 Aj
S
(t) 2
A
S Ô
and
.8+1
E == A, A, [..1]| < E A, (0) (1) [r,5]]
S 3
When
alog d [og(1//
<
1 5 Ti. we
have
g++1
ja suggesting that Vi E [d], g
12/ Ag
(1)
and
ge+1
Ej= A, 1, s 1 Ō 4 s δ
Substituting into 12 completes the
24-1
d
proof.
32

---

D.6 Proof of Lemmal 7
Based on the bound in Lemma
we have
3i
(8)
(1 B)2
11²
d
Agri
(1-1)) [1, + 1 B
(1(()
('2
2i
(2-1)
1=1
r/²
d
Apr (f-11) 1) [1.5]]
[0)
17
(t-1)
S (1 3)2
+ 1 B
'2
+
22
1 - B
j=1
1
1
S 176
+ 2n
n/
dia-2
dža-2
d dja-2
Since \1 = 1 - 1 ||lll2,
A2 = 1+ - || 22 and note that ||A||| = 4 (
(vd)
we have that
t-1 1
(t)
Xx
(7)
31
X -
dāa-2
$
1 - A
d
Ō
di (a-1)
r=1
(5
[0]
4
(7)
11
A
di(a-1)
S
"
d
dš(a-1)
1
= Ö
dia-1
D.7 Proof of Lemma g
For the equation x²² - 2x + 1
(1-03² ||A11/
n/²
0, the roots are A 1 1n9 ||Al|2 and A22
1 ++ - || which gives us
C2 W X1 W
A22 À
+ W
+nAW1
(0)T
10)T
+ 1722
(0)
W(
+ 1-3 2 (0)
"
(13)
2g
||
I-3
0 +
1 - 3
AW (0)T
1 - 3 (0) 2|| 2
2
211A|112
where r/20)
(0) w - -Dg22) Note that this is slightly different from the definition of 6 Now let's bound the i-th coordinate of 7
W1 --
(0)
(0)
r2
in eq.
In Section
D.5 we have shown that w.h.p. for Vt S Tj and Vi.j E [d],
| (t) S n (
x
which also applies to t = O. Using the Gaussian tail bound and union bound, w.p. at
least 1 6, for ever < j.j S d, we have that
1
(0)
S
2
log
2d
W |
S V
2
log
2d 8
d2c
ddn
Then we have that w.p. at least
8, A1
S i, S d :
( (w:50 (0) w/ (
W2 w/ (0) [
<
[0) Waj
w
[
[i, d]
==1
2
2d
2
2d
2
2d
log
log
<
log
dea
da
8
d
8
i=1
(w w1 ) | +
(0)
j=1
(14)
2x
2
2
Ba
log
5 V dí
log
+
d -1
0
d
33

---

Next, we bound the i-th coordinate of W2
1-8
AW O
.ie. (0)
+ 1-3
A (w.(0tt.:)) (2.:) "
||Al|2
1|A|12
By independence under Assumption 2
we have that
Var
(0) W/2
+
1 1-8 3 (wi/
Var
W22
(0)
(1 .32
A
Var Var (w [2,
|A
|
j=1
(1 - -3
1
A3
0
d
|A
d4n
d2a
Using the Gaussian tail bound and union bound, w.p. at least 1 - 6, for ever 1 I i S d, we have that
(0) (w1%
1-
0o
(o log
d
=
W2
|All2
12x
da
Since for ~ (0, )²), we have that P([X|
S 11)
2t V2na
then for a fixed in
P
(0)
1- 3 A (wi
1
0
2/dia
||A|2
dia
V2T
1/daa
di
Then by union bound, have that w.p. at least
for every iiid.. d, d
(0)
1- 3 A (wI(0)
T
1
21 +
2
||A112
dēa
Now define C3 = W/ (0) + 1-8 AW
(012 and C4 = 1-9 1-1
We get that Ca., i E [d] are i.i.d
TAT
Gaussian random variables and that C2 (C3 + C4), where w.h.p. for all E [d].
Cail S Ô
CCal 2 θ
|C4c
(15)
da
dīa
where (i) follows from eq.
(14 and the fact that ||
Vd.
Then we get that w.h.p.
1
Vi
[d]
Casl
dta-t
Ω
Substituting eq. 15 into eq.
13
we get that w.h.p.,
122l
θ
W226
(0)
1 - 3 A (w:10
())
E
1
1
||ll2
dia
da
Similarly, note that
C w Aa 1 (0)
d2- A1
W 2 + m/W
(0)
10)T
(0)T
'//2
(0)
((
-
2 10)
27j
A|l
I-3
(r(0) (0)
1-8 3
Aw-
(o)T
1 1 B (()
|||
2||A||2
T2
we can use the same techniques to get that w.p. at least Vi [d] CCul < Ō (()),, ii) w.p.
least 1 - Vi E [d] CCll
34

---

D.8
Proof of Lemma 3
The proof in Section D.2 tells us that at the end of the first phase (when T1).
w = u o (T))T' + (T)))
W cl2i) (Ti)T ( + R2 R
(16)
where a (21)T
r== X5 ((T))
1 1- 3"
Denote the i-th coordinate of (0, (()) R((a) as u; 1/
(t) (8) R respectively. Denote the (f,
element of 11 [i,],, For 1 2 Ti, prove by induction that,
w = u ( + R
(17)
1h ( R
where
v (1+1)T (1) me E(1)
R((1)) (++1) R Tp RR
(() E(D + ri
(t)
((++1) dl 71 E(t)
1(+ (T) E(E R(Di 12
(6)
9r- Doi F(0) = - w( with 11 = 1t= BC- W (1)T E(() w(T) F(61) 1) -
(1)
A. M4 E 8 and (t) = 1) Er= grー (EWi B (T)
7 o gr Dgz (T)
Note that the (11)) and (t) here are different from those defined in Section D.2, but we abuse the notation and still use and (t)) to represent the error terms.
The base case is already given by eq. (6).
Suppose our lemma holds for then for + 1, using the same techniques as in eq. have that
and eq. (6).
wr(+t)) -w(1) 1) Σ 8r 8
E(r) E(r) 11 - Σ gt-r DD-DDi
==0
T=0
w(()) n W E(0 + +ri (t)
1
W WhO) n
8
Σ 8 DDg2
T=0
w n (O +((0)
(t)
Plugging in the inductive hypothesis yields
(t) (((+1) w((() ni ( E(1) + +((() (t)
(T)) + R(()) - 1]t (n (T)) + R!(1) E(()) +r(())
(Ti) (uitiT (1 E(t)) + R((1) n R E(1) + r!(
(((+1) w/! (,t(())((11)) w,(1) + + (1)
( + R E(t(t) (1 + +R R1(02) + +
(1)
(c(0) - -,,, u (T)T + ( F00 R((1T + +
(1)
It implies that our lemma holds for 1 - 1, which completes the proof.
35

---

Now we analyze the error terms [R("[t.5]]
[ and R/ Eq. (66) tells us that eTTi) and Vi E [d], uf (()
are all positive. We first prove by induction that for all Ti tt s T2, elt) 0, Vi [d] 0.
The above discussion already proves the base case. Suppose at time 4, we have e(t) > 0, Vi E
[d] ( > O. Note that when Ti S t < T2, Vi E [d] : E S 0. then for t + 1,
(t-1) v/ () E! 0, 8
((--1) c(t) n
E (t) 1/3 > 0.
Therefore by induction, we have proved that for all Ti St 1 T2, cl9 3, L.vieed,.0
100 0.
Now we prove that for all Ti
S - Tz
R
[r.
R
L1 S و d:
0
S di
+
0
6 +
(18)
(21)
10)
c(t) m
(Tr)
)
T=n 7
=
where
R1
[i,
(3))
(0
[.
(2)
5, == max
max
c = max
max 3
(Ti) u (211)
e(Ti)
(21)
ui
(T1)
(1)
cl uu
(1))
The left hand sides of the inequalities are trivial since we have proved that c(t)) > 0, Vi E [d], u/ (1) > 0 for all Ti S / < T2. Now we prove the right hand sides by induction.
The base case already verified by the definition of oG. Suppose eq.
[8] holds for T2 S 1 < T2. Then
for + 1, using Vi E [d] E < 0 and (+++)) 2 v(0) clt++)) 2 c(t), we can get that V1 I i, S d
R [i.
(-E!-)
R (Ti) u / ((+)) S
Iu ( + 7/e
19 1 [s. Jl
1/) + m)((()
(-Es(()
( v
(6 + Eint, c1a))
eg/ + 17
(0)
65 + Ei=r e(²))
clt)
(-E
2/ + me(()
(-E!") (1)
6 +
T=
Similarly, we have that v1 < i < d
((++1)
R
1
+ 17c Eド=1 (-E]!) R
(1)
[8.5]
1
(1)
R2g
u (1)
5
c((+1)
[T)]
c(t) + 171 Σラ= -E/(7) 1/0))
21
(Ti)
c(
(6, + Σイ= &i
(T)
( + m/ 05 + -E c:-
(-1 (T)
プラ=
(-15)²)
(1)
(4)
clt) + T/e Σラ= (-E(")
(1)
= 8, +
(7)
Therefore by induction, eq. (18) holds for all in the second phase. So far we have proved the rank structure stated in Lemma 3 The remaining part of the proof is given by the following lemma, whose proof is deferred to Section D.9
36

---

Lemma 9. Under Assumption
2 and suppose S
2012-1. By picking nl S O
x/12
( (
we
have that w.h.p. for Ti S S min
22, Ta).
R2(" [t, (T))
(1)
24
vi d: 5
(((o), 0
S <((e))
(19)
u
c(t)
u;
(T1)
and that when t min (T,, Ts}, we have || 1 O(eod).
D.9 Proof of Lemma
We first have the following lemma which describes the structure of v(() for 2 T1.
Lemma 10. Under Assumption
and for t Tú, can write v(17 w at)) A (!)TT
with (T))
n2 1-3:
R!!T)) [0, 0, ....... 0J. and
(4+1) = (1- ()0 a m))(t)
R(++1) = mo(0
R!!t) mee R
where d(t) : c(0) (000l1/0 |uu + R2 (1) u (T)) R == (0) (TiT (T)T RT + R1 (117 1810
Moreover, we have that
wh w (1)
(000 (t)T + R d (1) A+ (00 /
R(OT + Rs3 R
(20)
We prove Lemma 9 by induction. Denote the i-th coordinate of R and R!. R and R!d respectively. The following lemmas constitute the inductive part.
Lemma 11. Under Assumption
and suppose a I
d0//1-1 and pick n I (
3/2
d Con-
sider any such that Ti < min (T2, Ts}. Suppose for all Tj I T S t, we have Vi,j E
[d] : wol S (d)4), ( [ S OOL (ztt), then we have that Vi,j E [d] r (1)
[i, =
δ (1721144) 22 = Ô (n (3/4) Moreover, we can get that Vi E [d] ( = Ô (1ddiat
where g defined eq.
(88
Lemma 12. Under the conditions of Lemma 11 and pick n S O T
we have that time
t+1. 1,
Re(+1)
[i,
((++1)
LA visijsd 0
I Ō(e)), 0
S Ō(co).
u
(TL) (1+1)
c(t+1)
us
(T1)
where co defined in Definition 2 Lemma 13, Under the conditions of Lemmal 11 and pick v7 0
we have that time
1
L+1.
1 [++1TT (Ti) c((+1) [T))
12(+1)
0
S Ō((o))
[ 0 05 c(++1) |u(Ti)|?? ug9)
S Ō((o).
Moreover,
R
Vj E [d]
< Ō(Eo).
(21)
A
Lemma 14. Under the conditions of Lemma 11 and pick 17 S O
d 1 we we further suppose that
R(2)
R(0)) and ( A are of order Ö(ev), then we have that at time +
Vj E [d] 2 = 0 (+22)
A
37

---

(A) Vi E [d]
(())
(B) Vj E [d]
+/))
cU+
(C) Vi.j E Ed]
2
(1+1)
1 0 (d'/4),
[i.t
S so (
[ Ag
(D) Vj [d].
and
of order Occo).
By combining Lemma 11
12 and 14
can prove by induction that for all Tj <t min Tz}.
eq. (19) holds (which follows from Lemma [2). and
E (
Vi.j E [d]
01),
(22)
E
which follows from the part (A) of Lemma 14 Now the only thing to verify is the hase case, i.e. when
t Ti. More specifically, we want to prove that 1) Vi, E [d]
2
(Th)
S O (d'/1), w [i, S
R and
(1/4) and that 2) Vj E [d] 20 = θ (Ja). and that 3)
a are of order Ō(co).
A,
All of them can be verified by the proof in Section D.2 and the definition of R( R
So far we have proved eq. 19 in Lemma 9 Now let's prove when / min (T2, Ts}, we have that
L00 1 2 (cgd).
min (2,, Ts} = T3. by Definition 3 we have || S e. If min (T2, [3} = T2, by Definition 2
there exists j E [d]] such that F(0) - (Vra). Combining with eq. (22) gives us [d
F.! ー (v mm) Combining these two cases, we get that when / = min [7,, Ts), F3(0 1| S
max(c, θ (cnd) = O (cdd).
D.10 Proof of Lemma 10
We prove this lemma by induction. The base case (1 = Ti) of ( is verified by eq. (6
Suppose time (()T (t)) A + R! then by eq. 17 we have that
w (0 w(t)) = (T)T +R R (T)) + +((Tt)
(t) u (T))
+ R (t)z wT))) 0(0) + (0) (g (T)T R( + R!!T R(
(O + R(OT
(1))
(1) A + (0 R R(O0
where d(t) = clt) (10) + R 1)) u (T))
R ( R1 + R (1) 1
That gives
(++IT (1/T () E(t)
(0) A R(UeT () (2()) (0 R(!eT + R((iT -A)
((1 (11 (0 ((1) (1) + mcll) A+ (1- Rlo00 (0 p(/TT
= a (1+1) A 1(T
Therefore we have proved by induction that for in the second phase, ((0) al1) + [(!)T The above steps also proved eq. (20
38

---

D.11 Proof of Lemma
Write ri [1] (t)
(2)
where we have (1)
911 11 E Ei== Br- (WLer EX(1) - W/ F2(T)
Qi1 + 922
912 = - - Dg Write r -221 + 422 where 922 = -1 - go Dgg
-921
77 0 gr-r( ( E(W(1
Let's first bound gff [i.
(t) and Ø21.
By definition of T2, we know that for Ti S T S t, t, e [d]
| = -1)). Then we have for all j e [d].
[i,5] - -
8 -k
g k 0 (dl/t) = n1 (d!vt).
222
k=
k=0
wbi w// - m0/t -00ff
(T)
S 17
8 k
E
(b) W
(2)
15.00
1 Σ
B k (d811) = no ( (d211))
k=0
k=0
(23)
Note that
F (of+t)=F))
(r+1) w/b') w/ + w6? (w/(+t) - w/
=
+ Σ
((15t-1)
(T-1)
- m (17)) (1.1) - (t.)))
1022
i=1
We can further get that for Vj E [d].
-E/" - ndO (d%/e O (d=2/1) + olh (al/t) 0 (ddi)) + 12d0 (244)(( (dv1)
o (1d0/2 + +zp(²) 0 (m121)
Combining the above inequalities gives us Vi,j [d].
911
[i,]]]
17
Σ
g
W (1) E - Eイ}T)
0=0
8
W22 -w |
(t)
(7) + [1/3
- - -15t0)
7=0
S 772
gi-r (t 7) (d(a))) (r1o (1 010))) (
(d0r2) = 0 (17671).
7=0
(
Σ
3
( ( (1)
4
= 1
T=0
j=1
< Σ
g T Σ
| w( + E - E [1/
T=0
f=1
< s Σ
Tt - 7) (((I)) (d1)) + - (
o (a-1 +
0 (1²2014)
Next let's bound 912 [i. and 922.4
(t)
By the assumption of this lemma and the analysis before
Ti. we know that for all 7 < t, the ((T), (²)") in Lemma 32 are upper bounded by O (a)xx) and O (d-/1), respectively. In the theorem we consider the training period before TSGD, the time T in
39

---

Lemma is set as 2 In the following sections, we will prove that TSuD,2 S 0 d' logl dfc)
Then by Lemma 32 we have with probability at least
2. for VT S t and Vi, E [d].
Dgi 1r 0l [1.51 9-
0
do
da+1
log
d
+ O
dio
d
log
d
da+1
D02
=
92
(T)
92
(T)
S 0
d 4 01
15
da+1
log
d
+
+
d
da+2
log
d
17
d
0
d σ
17
1//2 do/241, we have
Dgi [t.01
1 no (a#) and 1
(T)
S n (dtt). which yields
By picking S
012
(2)
11 j
S
1)
---7
Dgi
(T) [K,5]] 2ō ()pa")
==0
(t)
---
Dgg2 <0 (
(T)
422.8
1
T=0
Combining the above bounds, we get that Vi,j E [d]
1 [i s δ (rp't),
| S ô (1²
By the analysis in Section
D.2 we know that at time Ti. for some in E [d], c(Ti)
vio
(Ti)
= (
and for Vi, j E [d], we have c(T)) u
(Ti)
á
and
(T)) (Ti) 52
1
which gives
d +
us V E [d].
(4) [i,5]
ri
[
Ö (1/dtat
T2
(2)
T2
(0)
Ô (1patat
<
u
(Th)
(2)
u1
(21)
>
(1:)
c(t)
us
(ii)
(TT)
(2r)
Hence we get the bound Vi E [d] : (t) s (((idtaat)))
D.12 Proof of Lemma I2
Let's first try to bound the length of min Ta}. More formally, we prove that under the conditions of Lemma 11 and pick 17 S O
+ we have that min (iinTT,,7 Ts} S d" logí y d/t))
d
Under the conditions of Lemma 11
we know that
Vj E [d:
Σ
22 W (0) [1 = ())at)
10
1=1
(m/50)(9))
<
2 ri [1 = (170M)
(2) (0)
40

---

Combining with eq. (3 we get
E(++1)
=F(() + (wrl w w + w (w((tt)- - -w,") + (wrett)- - w3,') (w/(t3) - m11)
(w/(tt)-w(1))
=F((t (.F. w + 1 1 -
E: + 11 +0 (1(1)))
=b (1- [ w - n | + +00+))))))) (rdd)
Then we have
((-1)
2 2 1 - 1--wii nw w7th - w/ / - (
- $|1000]
Ttt [w) 1001l1 + (1p
< 1 1
When Ti < 1 <
T2.
we have proved that (()) is increasing over time in Section D.8 which implies
that
2
2C C
( since (T)) is the leading term of W Combining with n
2
gives us
1 - C/u (ws | +1)1dt)
-
n
o (1724)
o (17-tt
E(t)
(--nc|
|
n I
ndt
|. (nnc
2 (t (t--i)) 0 (vaa.
exp
|| (1 ?
where (1) uses
F(TT)||2 O(va). By picking 17 S 0
and noticing that | 1| 2
df
Ω (2)) have
ndd"
Ve
Hence when /. Ti 2
log
d/c
have that | S
7) W: (1)
Ve, i.e
|| S..
log V d/e o||w]]
steps from Ti, either t 2 T2. we have || 12 In
That means after at most O
x
log
d/e 7
O d log (V/ee)
other words, min Ta}
< +
"1
Now we are ready to bound eq. 18
Combining min(T2, Ta) < O dd logl loglyd/ec) and Lemma 11 yields that for t + 1 S min Ta
Vi e [d].
so( Indaat log v va) -0 elog log d
s(+11 (172t1))1) so
Lemma 2 tells us that 5, Ô 6
Substituting these bounds into eq. (88) completes the proof.
d=a
41

---

D.13 Proof of Lemma 13
The proof in Section D.8 tells us that for Ti I S T2. clt) > O. E [d] : :0f0) > 0. which gives us
18 F and 0 c01 (
0 S
o/e41) By Lemma 12, we have that
c 1) [
R
R
[i,
LA sisiss:: 0
S Ō(oo,, 0
S Oca),
u
(T:)
Dj
(e+1)
clt+1)
u;
(2x)
which gives us
R u (T))
R Σd 1=
(Ti)
R((+1)
(( (1+)
(Ti) 2
"i
-i=1
W3
Ō(eo).
c(t+1) |1/
c(l+1) D
u,
(7/)
2
cll+1)
Σ
(Ti)
Lemma 10 tells us that
RS (1-1)) c((+1) (T:) R((+1)
R2 11+ R((+!)
And we have that
(((+t) n((+1))
c(t+1) 2i=1 u1 (u R (0
cll+1) uultiil2 (4423)
(((+1) Ei=1
u1 (Ti)
2 ((+22)
(T1) 2 (+1)
((E))(((+1)
Ei=1
d
Ö(co).
d(t1)
(1+1
R 11t+1) B R((+1))
R2
(++1)
R
((+
Ô (cτ) ((++1)
Ei= up
(T])
us
===
oll+1) // ui gett1)
c((+1) Σd
(T)
(1+1)
c(+1) Σ ==
(T1)
2
u/ fett))
"1
"
= =
Therefore
(
(1+2) (T (e((t1)) (1) )
ia) ( (e((t1))
b((+1)
(
ol+1) || 2 seet S Ō(co)
(|
c((+1)
|u(Ti)] 13
(++1)
c(+1) |u(ii| ( (ia
By Lemma 10
w((++) = c(++))
| 1(C + nl 1(t+ u (T) + Rg ([+1IT
Then we have that Vj E [d].
(1 (o((tt)))(()t) = ((+1)
(ubbill' 1/ ((t11) (1 ((t))) where cj (t-1)
< O((o)) (24)
(w(++))//(+)))
Since & < T2, we have Vjj E [d]
= 01), which yields
Az
((++1)
2 (1))
11-1) <0()).
(25)
A
(t+1)
which proves eq. 21 since Vj E [d]
0 < o(++1
()) 104l S Ō(Eo).
42

---

D.14 Proof of Lemma 14
(A) Under the conditions of Lemma 11 and pick 17 S O
we can apply the technique when
| A0 S S((oo)
proving eq. (24) to show that eq. 24 also holds at time t. Since
we get that
o -00= a A 2 at))
(1 + e!9):
where
[3)) S ((oo)
Substituting into the time version of 2q.(44) yields
Vj E [d] (w]" =d(/(l)) (T)
/ |2 Aj (1+e(").
where ( I Ö(EO),
That gives us
Vj E [d] : 53 = A, (a())((0) 4 (T) 2 - 1 + a (
2 (T))
u
{
Since t < T2, we have 1.) <E. Combining with A, = 0(1) gives us a ((1) | 1
-
- (Vo). Then we can rewrite E as Vj E [d].
JelTTil ( -1) (1 1
al)) |
E As
(1())(1)
c(t) | 1| -1 (
Ag (11))(1)) | (
(1
(t)
where
((Voo). Hence Vi,j E [d] [E)
(t)
1 @1))
(B) Note that we assume Vj E [d] :
θ
then we have for j E [d].
vi
((0
ΣS
EE!
(0)
U2
(0)
1
- (va).
θ
cl0) (EE(1)
c(t)
(-rxg")
Vdl
clt) (-E(T)
- Via)
((()
Then for f + 1, we have that for j E [d].
(++)
Dj + me(()) (-Ea!")
(t)
cf++j c(l) + T)6 (-E()) v(() θ
(C) Combining eq. (25) and Vj [d] A₁ O(1), we know that
((+l) | of getl) s 0(1),
which yields Vj E [d].
(u(0)) lucos)/² (15 )
2
c(r+1) ((1 0 (2)) (d(+1)² 2 | < f
(++1)
cl(+1)
O (va).
0(1)
(26)
Hence Vi,j E [d].
(++1) | = 0 (1111)
(Ti)
((-1)
- ((++1)
(u (T)) + (e+1)
(
o (a1/4)
7 uu2
(Ti) (1+1) = 0
$ (n( of (1+))
(T.)
+ ( ((((t2)(2 [1,
(8)
ui 1
d1/4
dn/4
43

---

where (1) and (it) use Lemma
12
1)(++))
(D) The fact that Vj E [d]
I (Cv) was already proved in Lemma 13 in eq 21 To analyze
A
R
A we first prove that (0c00dd0)) 0.
It is not hard to prove that eq 26 also holds for time t. Recall that d(t) c(*) 1| + 12foT (TT)))
i (Ts))
and Lemma 13 tells us that 0 v cleo
I O(eo), then we have
(do))
c0 R (T) <o(vi)
and pick n s o (Z we have that 1 - mell)) d(t) 2
Under the conditions of Lemma 1- mc d(0) > O.
The assumption Vj E |d]:
S (Eo) together with c(t) > 0 gives us
m,l(1)
[
S (Ee).
))(()
|
Combining with the assumption
S (cco) yields
A
R((+1)
(1 me010000) R!(0) + mel()) R
e [d]
S Ö(co).
((++1) Ai
[1 d() al1) A, + ml)( Ag
E
Analysis of Adam
Note that A = L L AZx == 1 XXT. Denote 9h = Vw, L(W((( = We have that
91 (t) W (W." w( 1) (t) 1((0 - -A) w(Orr
Let A A and 9 (1)
k = 1.2 be the corresponding batch versions at time t. Let E() =
(t+1) wh (((() A, and denote F2(1) as the j-th component of E(0), We also denote Aw!y (c) == w2zi
(e)
AW!" (0) ( = w((+1) [i (( [i.5]. By eq. 2 the update equations of Adam are given
by
17: = 17 v 35+1
9 [i. w (4) E
9 = (E N00,,W1'1, 12.))...))
1 gt+1
( [i.
(1- 31) Ed 39 96 [b.5]
[i, (((++)
w/ [i.]] = 77
-h V1 V1- 32) E 8 - (ol)) (1.1))
Volt (2.0
+ S
(1-31 Et= 87-8 " (T) [1,5] + ri 1
= n
(1 32)
E'=o -5- (g.")
2
+ + (. +
(1+1)
(1)
m
(1)
(1- B1) Et
Di 81 (
17t
"I v/1 B2) Σt= 82-7
(8)
021
+
(1 (1-11) Ei= 31- (T) (t)
92s +
T
=-
V(1 32) Ei= 3/ - (r)
92
+ (t)
2d.
+ +
(27)
44

---

where 10 = -0("- -gi
(0) and Dy2 (9)
:- 92 and
;1( (1)
(0) ci = (1- -11 Σ 8 8-- Dgi [1.5],
(T)
==0
(0) 1i.i: = -1- 82)
8/2
20i (7) [i, [r,)) + (Dof")(.3))
7=0
(28)
(9) 2nn, = (1- B1)
Σ
81
D92
(T)
T=0
(1) 2d. -1 -1- 32)
Σ 82-
20 Do (Do!f))
(T)
7=0
Denote the i-th coordinate of W,W1 and A as (W2W1); and A, respectively. By Assumption 2 and the assumption that Vi [d] : A, A, (111,, the beginning, w.h.p., Vi E [d] (W2W)); A; < O. Based on this, we divide the training procedure into two phases (note that these two phases are different from those of SGD+M).
First phase: when the error (W2W1); ); A, is negative and its absolute value is big for all E [d].
2. Second phase: when (W2W1): - A; is close to zero for some coordinate i [d].
More formally, we define the boundary between the two phases below.
Definition 4 (End of the first phase). The end of the first phase (denoted as T2) is defined as T1 = inf { > 0 : 3i [d] E 2 - /nd
In the second phase, we define some time points.
Definition 5. Define Tg = inf { Ti 3i [d] :|00 S <ddn.
For < Ti, we have Vi E [d] : E.(') by Definition For t > Tr some E.") may flip the sign and become positive. For certain coordinate i, we define the following "flip time".
Definition 6. Define T == inf {t= Tı EE") -vna Define T := max; Tj as the largest "slip time" over all i [d], i.e. the "flip time" " of the last Eg which flips the sign. Moreover, denote := min (Tg, Ty}.
We can first show that after few steps in the first phase, W1 will become an approximately rank-1 matrix, as described in the following lemma.
suppose S n2/2g² By picking n S (2 .
Lemma 15. Under Assumption
and
Ba-1. and B2 = 38, there exists time > such that w.h.p. for time St < Tı.
Vi.j E [d]: u2 = sign
U12s
(t line) + R!!!)
w010) [i,1]
sign
W22
(t time) - -2((1) [i,
where [15) 5 (Vn+ n(t-lar)
oit )mn
1
Ö (vT + 1/( lin
Specially, when t Ti, we have that
Vi,j E [d]
(2t) sign
Wi2s
(0)
17 (T; time) + Rizi
(Ti)
w (1)) [i,
(0)
(T lanc) R (T)) [1,
sign
1124
where n) (Ti -line) 0
and
Vil
R
[i-j] {
(Ti)
(vn++ ja-
Riz
- 6Vn+ Aيب
(Ti- tine)
(Ti lime)
45

---

The following lemma tells us that this approximate rank- structure is preserved when 7i ststs.
"//222 ga/s By picking n 0 ( i
Lemma 16. Under Assumption and suppose I 350-I, and 82 = BB, we have w.h.p. for T St 1 T
Vi,j E [d]
2i = sign
((1) + R
[ [ j] = sign
W22
((0)
+ R
i [1,
R2
R [
② ( (T+
1
where
17
c(t)
da-1/2
(()))
d
and that L (W(T))= - Ō (nd+).
Now we are ready to prove the Adam part of Theorem I
E.1 Proof of the Adam part of Theorem 1
Define TAdam,1 =
tine
+
Note that this choice of TAdam, gives (TAdam, - tine) =
1 By
nd
04
picking 77 S 0
E
and 32 BP. we can apply Lemma 15 to get that Vi,j [d]
W/2
(1xim.1)
0
W
T TAdan.1)
[.,]]
d and therefore Vi [d] F.( (Thamm.))
11)
d
and L (W(Tnam,))
(d)) Define T Adam, = T. By Lemma 16
we have L (W(TAmmm))
(ndd*)
For any p 0, by picking p+1 we have (W(Tatm.))
Ô (1b11) Ô (1).
Moreover, combining Lemma 15 and 16
we get that when t E [T [Taum..1 T 2dlam,2), the conditions in
Lemma 31 are satisfied with 6 =
The i-th component of the u vector (denoted as max; max;(M,)² 1. Then we can apply Lemma 31 median (
u4) is sign
That means Vi E [d] of and Wizi
(0)
and get that
(1 - 8)
2 (ux;(u;²
( + 8) 2 max, (u;)²
Reddam (t), RAddu (t) E
1+
median (4))2'
1-8)
median (cd)an(u:)?
1 1 0) 1+
(1=3
1 0 27
(1-8)
Note that by definition, Reddam (L) and Reddam (1) are always larger than or equal to 1, then we have
1
Rdddum (1), R Adiam (1) 110 -10
d
E.2 Proof of Lemma 15
For some time t, we introduce two conditions. Condition 1.
Vr E H sign (1,01)) = 5 [t,)].
(1- - B1)
(7) ( (1-1)
91
[i,
2 2(E))
==0
Condition 2.
H
[H] sign [12 (((-T) (t)
(1- - 31)
Σ
3( (t-T)
42s
2 (EE))
46

---

Next prove that, under Assumption 1 and by picking I O
den . S سعر
and B2 5..
there exists fine > 0 such that for in I. the update of Adam can be approximated as that of signed descent.
[i.]] = " [i. 77 sign (9)* (t.)l + + (1.5))
(29)
W22 w1/s (1)
(1+1)
11 (sign 92s
(1)
(1)
+ 222
where (1)) [i.5] = Ō
| Ô (v)).
Before we dive into the proof, let's introduce useful lemmas.
The following lemma reflects our key idea: converting the exponential average in Adam to a finite-step average, and trying to bound the stochastic error terms in eq. (28).
Lemma 17.
Under Assumption
and 3 and pick B2 B7. Let M =
W [6.i)] - == maxx,jel4
(T)
G == mexx 91 [d.] and
(7)
0he = me( (T)
We have that w.h.p.. for all Ö
vān and Vi,j [d].
(1-32) 31) EM
= 87
[1. 51 + (1
A [i,]] = T)
(1- 32) E+= 35 01 (
[K, (1(--
+ + (1) (i,,) +
(1- B1) Σ 9 92i
(l-T)
+ (1)
Aw2 (1) = -1t
C2n.
Va B2) ΣH 35
$2i (
(1)
+
2d,i
+
max c aalo (elaaa((c0)) (c19)² 2 (c)))²}
where 11 2 1-3,
1 log
and
75
(1) [i s O(E)3) o (Df*).
S 00(k3) + +0 (1 c10 + (D)))² ))
[mmm s0()3 + +o (D!").
(t)
1 d + (D!)) 2) <0()) +
2d,
with
t 5 (M (MS0)² 2
d1/2
d3/2
- Ô
Ma Nv
o
1
17
D0 < Ō di (M/()) Mo
/1/2
dM((")
d3/2
n
1//2 d-a/s-
Corollary 2. Under the conditions of Lemma
17 and suppose
S
Consider any t S
δ ( (van)) M M, <0(t))
δ ch s δ
(Vdd)
(Vd)
then H in Lemma 17 can
be picked as 1-8, log d and we can get that Vi [d]. Ö(ng²).
nga
100 ii. 5
2 f 5
2n,
(4)
2d.1
(
S
The following lemma analyzes the magnitude of weights during short period at the beginning.
Lemma 18. Under Assumption 1 and 3 suppose a S n2/222 Pick E S - then there exists
diz/5.
some time point tim E (H, T1), such that w.h.p., for t S time. for every i.j E [d].
AW (0) [1,5] - 5((),
(t)
< <(())
[[')( < d#a+1
1
1
(0) W
<
<
dia
d
47

---

Specifically, when tines we have sign u22
(fin))
sign (w/ttr))i [..5) sign W'26
(0)
2 S(E). Moreover, Condition and 2 are satisfied for t = time. The $) [2, 5l and 8/28 in the conditions are both -sign (1
Wh [1 = θ (1tieiz) and (ut (imil(is)) 2 S2(E);
(
(922
(tae)
(t)
The following lemma gives us lower bounds of 9i ( [t.5]]]]]]]]]] and 92i
Lemma 19. Under Assumption
/ 2 Land
3 suppose a S "215t. Pick & S V - nf S 00(
130
Consider tine in Lemma [8]
We have w.h.p. for any t E [time: Tı). and for Vi.j E [d].
sign
AW
(1) J1. 5]
sign Aw
(t)
sign (u22)) and that Vi,j E [d] 91 [t,
77
92
(1)
2
① vdd. Moreover, we have VT I t. Vi.j E ld]
W [t.si -
(T)
I Ō and gi [i.c]] - Ō +
S δ Vd).
222
The following lemma shows that when tim I < Ti, we have Vi E [d] :
925
(t)
>
92i
(t)
92i
(t-1)
and that 91 (00
[1.5]]
>>
91
(0)
[i.j]
9i
(1-1))
[
Lemma 20. Under Assumption 7
2 and suppose o S g2/2c²
/is/t Pick & S /
1)
S
0
For lime in Lemmal
18 we have that w. w.p. for tin S 1 < Ti and S t, vi,j E [d].
(t) 9i
[i,
9i
92
(t)
92i
((-T)
= (VVTT)
=Ō (((T)))
(30)
(5¹h [1,5]]
(t)
92i
(of% - -0sit (01- 2
+ 10
91 (0 (i.ii)
(31)
(10)
2
925
92
① (ViiT) + Ō (1T²).
(1)
92c
Equipped with these lemmas, now let's prove eq. (29
For any E [mm Ti), by Lemma we know that M(() M² s ò (1)) and that G(O) <
δ G < Ō [Vd] At the end of the proof for this lemma, we will show that Ti = θ
1
Van
Then we can pick H == I-3, log d and apply Lemma 17 and Corollary 2 to get that, w.h.p., for all
nez
t E [mme T2) and Vi.j j [d]. eq. can be written
(1 B1) E" 8.9 (1
[
(
[t = 1)1 Vi B2) HH
35 91
(t-T) [i,
+ + [ +
(32)
- (1- Σ#= o 82i + 62n,i
(t-T)
(t)
Awll) (1) -1)
V(1 - B2) DH 33
02 (6(-r)
)
2d,i +
(t)
where Vi.j E [c]. [101,,1]]]
(1) id
[ 5]
(
(t)
5 $0(152).
2d,
48

---

Let's first look at the update of W(0 [i. For L in the first phase, we write the RHS of eq. 332 as
(1 31) Σ H
= 9ig
[8. .1 + 1%
(1 (1.11) 2
32) Et= 35 91
+ + +
(1 (1-33 11)91 (4)
[ Σ== 3 + (1 B1) Σ#=o Bi (9i- 1::01-901000 (..l) + (
V - B2) (vf2 Σ 31 + (1 - (1- Σー 8r (sf (1.1) (0f (10"(..)) ) + (0) [3, +
of 11, - 3/ + 11 [i - [
(of 2 (1- + +1 1.
where
H
[ (1 31)
Σ
Bi
(9f
[7.5] - - 1..1..)).
==
H
eid = (1- B2)
Σ
32
(3.11) -
==0
We have already shown that
(i
- ), By Lemma 20 we have that Vi.j E [d]
H
12 S (1-11
Σ
37
91
(t-T)
[i,5]] [i.]] 9و [i.5]]
==U
$ |01 1 [.,]
(vV)) (1- - B1) Σ B [i ŌVV).
Similarly, we have Vi,j E [d].
H
H
Idd [ = (0
2 ((vT)) (1- B2) > B + (o1" 2 ((VT)
(1 B2) >
r=0
r=0
2 =011 Ō(V)).
(6) By Lemma 19 know that 91 [i, 22 (vT). Then we have that
Vi.j [d] :
[ S Ō Ô (VT)) of [3.0]]
Therefore by Lemma 34 in Appendix
we have
9i ( [i, (1 - 31++1) + 11 [7 + (%
V (0 (..i)))) (1 - (1 e + 10.00 +
1 - B/411
(sign (019 + 2 F..i)
1 1 1
where 24 [1.]] -
Since 3 E (0, 1). we know that log 3 < 3 - 1 <
0.
Then our choice of H gives us HI
log 22?
log ne²
T log ng=
2
+ and H > 1-
log
which implies that 3,.3 I ng² /d.
Tog or
log H+
1-8
Hence for fort2 2 time
H, M
n + O())).
V1 1-
g
49

---

Combining all of the above yields that
(1-31) Er= gig (i.
[i.
[t,51 T
V(1 - B2) Er= 85 91 (c 111-7((1,1) +
3/1+1
-w[[(/,.] 1/-
(sign (1 (6.)) + [t.5l)
V 3
[1,] 1/ sign (g) [1.5))) e (3.5))
where ci [i.]]
Ô
(Vi). The proof for w2z is similar.
(t)
So far we have successfully proved eq.
229.. By sign (w("1(2.31) sign (Aw
(1)
sign
(0) 1122
in Lemma 19 we know that sign
(0) (1.51) (-91
sign (-9 (1)
925
sign (12g
(0)
which gives
Vi.j E [d]
22 sign (wn 17 (t time) + R,(4)
W [i,
sign (w2, (0)
17 (1 line) - 12 [1,
(ine Was
where
Ô
v +
and
oet ((t Now it suffices to
show that Vi,j [d]
w124
(mne)
da
w
(mm)
[ē,
O (a) which is implied by Lemma 18
Finally to complete the proof, we show that Ti = θ
(an When = Ti. we have Vj E [d] (21) [i, ((1). Combining with the above results, we know that dn7 time )
ΣL 22 (2i) (T)
@(1), i.e. (Ti tine) = θ ( In Section E.5 we will prove tine -
Then we have
T = θ ( (van)
E.3 Proof of Lemma 17
For certain t and H, we write eq. (27
(1- B1) H
TT9 (1-T) [1 + [1
( = -1)
Va- B2) ΣHH 35 0 ( [b,
+ + [i +
(1 - B1) I 3 921
(t-T)
+ (2)
Aw (1) -1)-
12n
V (1 B2) ΣH o 83
(t-r)
(t)
021
+
2d,i
+
50

---

where
19 [ = (1 Σ
11
((-T)
[1,
[1,
1. 2n. (1 B1)
(1) 2m,6 H+1
1192
(t-T)
T=#+1
1244 = (1- 82)
Σ
2 [(,))
[1,51]
91
+=
(P.s)
(0)
(1 -32 Σ
())
2
(0)
'd,
Ø2i
2d,
9ad
and (00 J1. (0, (1,). 2n,3 2d, are defined in eq. (28)
(t)
(0)
max cta) c.a (ctt)) (cto)²)
Since 32 = 87 < 31, then if we pick H
22,
1
log
we can get that
1-
n
(c
(
H 2 log T
log "e"
H
2 1-8
log
H 2
log
Hence we
DE
can apply Lemma 33
Appendix H get that
9in
(+)
[
91d
[ j.
92n.
(+)
92d
S n
Pick T in Lemma 32 of order Ô Vdn By Lemma
32 have with probability least 1
2.
for all t < T. Vr < t and Vi,j [d].
(00000000 (aeM (M)2)
da/2
d3/2
Do1 (1) [ = [..]] [r,]]
,4 ( [1 5
+
M ov 11
+
1
:= D("
92 S Ō dt (2 Ma oV dl/2
G3/2
Dg2
(T)
/2
(r)
+0
ov 17
=
=D/
Plugging into eq.
28 gives us
in [
S (1 B1) Σ 8
(T) |..1 s so(D!"))
7=0
<11 - (2) Σ 8 2 (t. 1 + [t.1" 2
Lia [
0=0
0 D 1 a + (D( 2)
S (1- - B1) Σ 8: D (T) s {0())))
2n,i
7=0
(() 2d
< (1 - B2) Σ 8 -
292h
Doli
(7)
+
Dgli
(T) 2
00
D l00 c (0(D)
+ +(DS)))
7=0
E.4 Proof of Corollary
2
Since
(0 501 Vi
(t) 0 Vil
log nET then H 1-5.
is bigger than
max ((+)) (()) (c!t)) ( (o
I-3 log
n
51

---

By M/ MM Sō <Õ(²z)
5 (+(+a) c.a <ō and the assumption
22//222
va)
S dd3/s we
get that D and D are upper bounded by D o od 1/2) and D - O od which yields Vi,j [d] | [id
(t) (t) (1)
2d, s ((nE))
E.5 Proof of Lemma 18
The proof is based on the following two lemmas.
Lemma 21. Under Assumption
7 and 2
we have that wap. at least 1
every 1 e.rrr
sidd
d
(0)
2 log
2d and that wap. at least 8 for any given 8 0, 8
d=n
22
2
log
20 g.
suppose a S np/2g Pick B2 = B8. EE (0. 10.1),1 < 4-4
Lemma 22. Under Assumption /
and
Consider any time point t S Ö
Vf S t, E [d]
[t,]] S Ô ( (7 [2z (7)
-
van
Ô va and 9i (T)
[1,5]
20
Ō ( (va). will have
92i
vd
Aww1 (t)
[
S (())
Aw!
2 s Ö()),
where the Ô notation depends on H = 1-81
1 log n)22
Furthermore, for certain j E [d]. Condition 7 (resp. Condition is satisfied, we will have
sign AW("(E.)) =s1 [ú Aw
( [i.5]] θ(n)
resp.. sign A Awil 12) --(0 Amoll
(0)
(c) =6())
18 Define ta := inf { ( Ji ( (2, 51 > 1/2 or w22 > }/. Now we
Now we prove Lemma
want to find a time point Line before for the lemma hold. During the period L < ld have
vj E [d]. E, = (11 (which means ta < T1) and therefore for all i,j E [d], 91 (8) [i. 5] S d and
(0)
1. Then we can use Lemma 22 to get that for t S min { tta Van }. have Aw (1) [i, s
92i
Ö())
Aw!y
(t) < Ön)) Hence ta 2 Ω2 nd
Define Lign
= inf {1 min (as Van]
} 3i E [d]
< 27
}. By Lemma 21 w.h.p. Vi E
[d]
(0)
VF d3
combining with A Aw!f
(
Ö(n) gives us that w.h.p., tsign
1Õ(n)
d3"
Now let's analyze the behavior of W1 during the period t < biignn Consider any i.j E [d]. By
definition, sign W24 sign
(0)
Note that F.(() (1), then we have sign 19 (1.
-sign
(0) and that 91 (9) [5. = Ω
2(E) by our choice of g. Then we know that
Condition
is satisfied with (0) [t
sign W1is
(0) (for all H t fsiga), which by Lemma 22
yields sign A ( (t.i]) = sign
(0)
and
Aw
(4)
[i, 5il O())
Lemma
21
tells us that w.h.p., Vi,j E [d]
W 00) [+-]]]
Ô (z For any i,j, if ini-
(0) ) then for the following steps before bsign we will have tially sign (t.0)) sign 122
sign (w.' = sign (w))))
(0) then after at most (0) If initially sign (w1") [1, (1) # sign
to = Ö TT steps, W1 [t. will flip the sign. Note that 4o = Ô ne
is smaller than 4gn
52

---

Hence we have shown that at some time point to. we have Vi,j e [d] sign [v.5l =
sign (12))) - sign u22
(0)
Now we analyze the period & 2 t0.
When to < t S Exigns we still have sign
(Aw ( 11.51) sign
W2i
(0)
and AW 10 [t. || = (()).
Combining these two with the fact sign
(w("0) (..il) sign
(0)
we know that for all & E
[to, Aigul sign [w/"(t.3]] = sign U22 and that Vi.j E [d]
(0)
W
|i = -
(()). Then at certain step line which satisfies line = l0 + 6
E (H, (sign), we will
nd
have Vtine - 1 S T S tine; V E [d] [
= O
and therefore [
dda+z
E W [1, [b = Eラ= / (1 [1 (
θ
(2()) For & S lines we have
d
j [d] : 1 1 [2.51 = (et2-r)
Since tine < Esign. we have
fire
Ω
(1+)) For t S tine, note that A I Ö(1). tine
(t)
d
to + 2 nd
1
combining with the upper bound in Lemma 21 yields
nd
+ line Ö(1) S Ö (0-11)1 ( s 0
(1)
(0)
W2s
S
d
Moreover, Vliee - H S r S line, Vi E [d] sign 92i
sign
(0)
Then Condition is
satisfied with $2 sign (w)0) for t = tm In the analysis of of [t. j.. we have already shown
that for all & S figgn (and thus for line), Condition T iss satisfied, which completes the proof.
E.6 Proof of Lemma 21
Since for X ~ N (0, d²), we have that P(|X)
St) -
21
then for a fixed i,
V2na
2/7
P
W22
(0)
V2T v 2/d2a
d3
Then by union bound, we have that w.p. at least
d -T- for every S i S d.
W2:
(0)
v
d
As for the upper bounds, using the Gaussian tail bound and union bound, we have w.p. at least 1 - 8,
Vi,j E [d]:
(0)
2
log
2d
[1,)] S V
2
log
2d
12a
8'
día
E.7 Proof of Lemma 22
Now we analyze the magnitude order of AW(()) [.. The analysis of Aw (t)
is similar.
For t S Ō
By assumption, (1 _M (t)
S Ô a < 2 c S Ô (Va). and
vdn
3 3/2g2 2:5/T
Hence we can pick H := 1-3,
1 log
d and apply Lemma 17 and Corollary to get that,
w.h.p., for all t s Ö Von and Vi,j [d]. eq.
be written
(1- 31) r= 87 g
[1, fl
(0) [
Aw("D(. = n
Va- B2) IH 35 gg ( [t,1])
+ Edd 10 [i,
(4)
+
(33)
Awll (t) -
(1- 31) ΣnE 35g2
(L-T)
+ E2n,i
(t)
/
(1 - B2) DHO
o
35
82i
(t-T)
+
E 2d,i
(t)
+
53

---

where j E [d]. / i |]
27 &. .|
2d.e S Ö(nE)).
On one hand, using fin b .
sid
(9
[.,
S OnE2) and B2 -p. and v V
| when
~2, O. we get from eq.
(33) that
- (1)
Σ=
Pioi 7) [1,1]] + +O(15)
Aw s n (1 82) Σー= (
-Ö(VTS) +
(1-B3)VE H + IV Σ/= (9000))1..)))
+ +O(E))
50 <0((TT) Ö(1),
2 1/-
(1- 32) Em=o (8T (1 (t-T)
[2, [1,51)
+
where (i) uses Cauchy- Sauuhy-Schwarz inequality for the numerator.
On the other hand, when sign (ol (((-H)1p (2, sign 0 ((-H+D)(..5)
= sign of (1.51.)))
[r, J], we have
sign Eor 11 1 |4.51 81
(0) [i.
Σ
Big (t (t-T)
[1,5]
2
Σ T9 (1 ( (1.51)
If we further have (1- - B1) EHH 3lt
91
((-T)
[1,5]
2 (2(E), then combining with
(00) [2
ÖneE)) & we will ger
11 (1.
sign (.w[(((..]] = sign Σ B9 (1-T))
[i,5]
12 cM +
= sign Σ 3
51 (0) [t,
Using v .リ v |0| + / |ll. we obtain that
(1 - 31) Em 1(10 91 (0--)) [1,51 - -
( [1,
11
(1 B2) Σm=
(8 (1 (1.
+((vTE) +
1-A 2
DH g(47 01 (1,)]
[(n))
7
max
(1 B2) EHo
(8ig) (t 7)
2 3 [t.
Together with the upper bound completes the proof.
E.8 Proof of Lemma 19
The proof is based on the following lemma, which gives coarse analysis on the magnitude of weights and their increments per step during the first phase.
Lemma 23. Under Assumption
1 and
suppose S n2/20" Pick E S min
{
for
T 2z
time in Lemma 18 have that w.h.p. for all time St Tú Vi [d].
sign (w (1) (1.01) sign (Aullt
(t)
sign (ung0)
(0)
AW (1) [t.5] =O(n), Aw
sign (w(") sign (1022) sign
(0)
| |.. ō
(c)
1022
Vd
/dl
Specially, at the end of the first phase (t Tı), we have Vi.j
E [d].
(Ti)
and
Va
[( = - (ta).
54

---

Now we go back to the proof of Lemma 19
For tine St< Ti, since E -A,
E 22 1 [i, - Aj. we have,
A = E(+11) - E
(1+1) W [t. 22 w [1, + u2 ( [2,5] w (f0 [2.51] (34)
(++1)
(t+1)
AW
[. + Awyz (i.))
W2s
Combining Lemma 23 and eq.
34
gives us Vj E [d].
AF > 0,
AF
(++1)
222
(35)
0 (nvi)
"1 Vd
Let's first analyze 91 ( [t.j]. Note that
Agi (1) [i. 5]
w E
((+1)
(((+2)
w22
11+1)
(8)
w2i
(1+1) (1)
w2 E
(1)
(36)
22
(1+1)
(t)
Aw
(1)
D
where sign
22 A AF(T))
(1+1)
= sign
W12s
(0)
while sign Awwy 10 E
= - sign
Wizi
Now we analyze the sign of gi
[ when time S < Tí. Using
W2i
+1m
= Ō
and eq. (35).
we get that
u22 (fie
+1)
AE (fine)
I Ō da vdn)
While on the other hand,
Awgg
(Ma)) E
= (()).
That means sign sign
Aoi
= - sign
222
(0) Note that sign
91 (fme) [i,
sign
Wi2e
(226)) we know that
91
(1) [i, will increase when t = time.
In the following steps.
91 [ will keep increasing as long as
Aw E1
12:
(e+1)
A
E
Since
( [1,5]]
222
keep increasing while (w [6.1]]
(t)
remain ( by eq. 35
we know that the trend of
An( is to increase. On the other hand,
keeps decreasing since
E/e < O while AES" 0. Then after some time point will have
(t) Aw!
E
(t+1) A
21
w2
(1) [ will have the trend to decrease. Specially, when = Ti we
and in the following steps gi
have E/
nd) and
= Va
2
(++
||
è (+t) (aa
by Lemma 23 which gives us
d 1 ê (IvV).
AE
Σ
(++1)
AW
(2) [1,5]] + Awoy (0)
( (1)
[.
S Σ
11 "Va)
Hence
W22
((+1)
A
(t)
O(n)=
(t) 20500 è (nvnd))
Therefore have proved that when line
1 < Tı, the trend of o [i.J] is to first increase and then
decrease. In order to prove (6 ( [R. á (VT). it suffices show that gi (tom ( [ú. ji|
á (vT) and
91 (TT) j ā (vT).
When t = tine.
9
(tm))
[
22
(fne)
1 (tne)
= Ω
dia
e(1) = (Vn).
55

---

When t = Ti, we have
(T)) [1. -
(Ti) =|
|
(Ja vnd) = 6 (VT).
91
As for 92i since for Vi E [d]. w/ [i for different j have the same sign. Combining with
(t)
vj E [d] E < 0 gives us
d
92
Σ
D3
10 [i,
D3 w (1)
Then it suffices to show that for lime S < Tı, w [6.01] 52 (Vn), which can be proven
using the same technique as above.
Finally, for Vr I t, Vi.j E [d]. note that the upper bounds of |w/ [c.]]] and W/2 are already given
9i (7) [i 5l and 92 ( we have (of of [c.3]] = w E = Ö (.22) 92 (r)
in Lemma 23
As for
Σ j=l
[i Ô (vd)
E.9 Proof of Lemma 23
For any i,j E [d], and any in the interval i Ti], we prove by induction that
(A) w (1) [i.5]] Ō (tt) 22 = =0(ta) (Ja).
(B) VT E [t HHt sign (w[ (r.1)) sign (w))")
sign W2is
(0)
(C) 91 [i 2 S2(E),
92r
(3)
2 S(())
The base case time was already proven by Lemma 18
For & E [line, Ti). suppose (B) and (C) hold for time / and (A) holds for all T E [lie, 4]. From (A),
we get that Vr E [line f] [9
(T) [i..]]
w/ E Ö ( 92
S Σラ E W (T) [i,5]] =
δ (va) Since when f < Ti, Vj E [d]
E
< O, from (B) we know that Vr E [t - H, t]
sign (oi (v.5)) sign
92i = -iign
(W) (2's''
Combining with (C) tells us that Condition 1 and
2 are satisfied.
Van )
In Section E.2 we have shown that Ti A
Then for t E [iiee Ti), can Lemma to
get that V Inc S S 1, vi, E [d].
(1.T1) = sign (AwlT) =sign (w!2
Aw!
Aw! (7) θ(1).
sign
Since when i lines sign (w((am)) (i..i) sign
(Cam))
sign
We get that for line S T S
t,
Vi,j E [d]
|w(++1) [i.5]] (w li, Jl S(n).
(T+1))
( + (())
Now fort + 1. we have Vi,j [d].
sign (w((+t)) (t.s)
sign
(w00))
|w<(+1) [i,
( [i,]]
+ (t + 1 time) (n))
sign 122
((+1))
sign
(0)
wi2s
(++1)
(tm))
(t 1- 1mmme) time) @(())
That means VT E [t 1 H, + 1): (B) for time t + 1.
sign
(
= sign (w/)))
(0) sign (wi25
This proves
56

---

On the other hand, we get that | [i
z W (/(t [i O
and
(++1)
d3 12
22:
Since + S T1 which means Vij E [d] d=a
| E((+t))
z rdd. Then
91
((+1)
[t,5]]
11+1)
E
E((+1))
2 Ω
d
vıdd= 2{}),
d
92i
(++1)|
E
W
[i,
=
Enj
E ((tt)) [1, 2 d0
(datt Vnd==((())
j=1
This proves (C) time + 1.
Since t + 1 & Ti which means vj [d] (w/ (((t))))))
S <0(1), we obtain that
Σ
(14
(((") [i, = Σ uzi
(etD)
[i
22
"(( (fue)
+ (t + 1 - time ) (1w [c,1] + (t 11 )(()) s 01))
1=1
Note that (/((t)) (ton j]
(tuc)
2/2 (since time < ta), we get that (t + 1 - time ( O
vi
which gives
W224
(++1)
ô (t2) and [[(") Ô (1z) and hence (A) holds time L 1.
[1.1]]
Therefore by induction, we can prove that (A), (B). (C) hold for all time St ST.
Ti. Then applying
Lemma 22 we get that for all line Tr, Vi,j E [d] Aw(") [. θ(n),
Au (t)
(()).
Specially, at the end of the first phase, we have Vj [d] : (3(+))(((t))))
=(1). Repeating
the above proof techniques gives W12
(11) 6 va) and w ((("((t,))) = θ val
for Vi,j E [d]
E.10 Proof of Lemma 20
Let's first prove eq. 30
By Lemma 19 for line < l < Ti, we have Vi, E [d],
(1) gi [i, = 52 (vT))
(0)
= â (Vid).
921
Then it suffices to show that for tine S t < Th. 91 (1) [ú,.]] gi (t- ( [i.i]]
= -Ô(n) and
(
02i = (nd)). It suffices to show that when & < Tı, 96 (++1) ((
(1-7)
9i [3, Ō(n)
921
and
/2i
(t+1)
92:
(8)
(nd)..
By Lemma
18 and
23
we know that when t < Tı. Vi,j E [d], (w
S 5(m,
Aw 1
(())
and that w
(0
[.
(
100
v Then the bound |E
(0)
I Ô nvVd d in
eq.
(35 hold for all 1 < Ti (not only Line S 1
< T:). Substituting these bounds into eq.
36 gives us
vt
Ti.
(+1)
[1 - 9 19 [r,
(t+1)
A (0)
- Au |
9i
- u0z
1
= 5
θ (V + (1) =
Vd
Similarly, we have that 921
(++1)
921
(t)
= (nd), which proves eq. 30
Note that for a, b E R:
2
12
a2
(a- b a)
2a(a b) (a - b)2
la bl
la b) 2
< 2
+
a
a²
a²
|a|
[a]
57

---

Then eq.
31 immediately follows from eq. (30).
E.11 Proof of Lemma 16
We divide Lemma 16 into the following three lemmas. Combining them together immediately gives us the whole proof.
The first lemma below gives us the structure of W2 the second phase and that of W under some conditions.
Lemma 24. Under Assumption
and 38 and 32 = 18, we have w.h.p. for T <t
suppose
S
By picking ry
O . S
Vi E [d]
22
(t+1)
w/2
sign
028
where
(0) C2
(v
and moreover
R27
E
(t) was
sign
[222
(0)
cl + R(0 where
=
[v++ =O(v7+
c(t)
da-1/2
As for W1, for certain i.j [d] and certain t. E (Ti,T) we have
( [i, = (vn) then
w (( [1,] = ( [i,1] - sign si [i.i)) + ef (i.j)). where
e [i = Ô (Vn).
The second lemma below also analyzes the structure of W1 but removes the conditions in Lemma Lemma 25. Under Assumption and
suppose S "0 By picking n S 0
a50 E S
V 1__11 and 32 = .3.. we have w.h.p. for Ti St < T, Vi E [4]. w [t, = 5t Va
and for
any j E [d].
1 [i,
w sign (w v/00 + 22
1
where
17
v2(0
d
The third lemma proves the convergence of Adam at time T.
Lemma 26. Under Assumption 2
2 and suppose S ///2a²
By picking n S (
and B2 = 18, at time T, we have that whh.p. Vj [d] E s Ô (dvryd). which implies
B(T)
Õ (nd*).
E.12 Proof of Lemma 24
The proof is based on the following lemma, which gives coarse analysis on the magnitude of weights and their increments per step during the second phase.
n g/
Lemma 27. Under Assumption
2 and
suppose S
dist..
By picking 1j S 500
. s
and 32 = 13, we have w.h.p. for all Ti 5t<
Vi.j E [d]
(t+1)]
A Aw 2
(1)
= O()).
A s Ō(n)
Moreover, we have that V E [d]
(t)
w [i, = Ô (1a)-
Equipped with Lemma 27 we are ready to prove Lemma The proof for w (i. j] uses the same techniques.
We will only prove the results of w.
58

---

Lemma 27 gives us upper bounds of
(0) 22 W (
[...].
as well as
Aw (t)
and AW (1) [i, for
all i,j E [d]. Then we know that eq. 35
still holds, which gives us Vj [d]
B [((+1) E =
δ v dn) Then we can use the same strategy in Lemma 20 to prove that
(++1) - 92. = (
92i
By definition, for Ti S < T, we know that
92s
(1)
2 (dv
Combining with the bound
(1+1)
92 = (nd)). we know that the 925 parts in eq. 330p and eq.31) still hold. Then we can
(4)
(f)
92
use the same strategy in Section
E.2 to prove that the uzz part of eq.
(t)
29 still holds, which gives us
Vi E [d]
W2:
(++1)
W2i
(2)
sign
02i
(e)
E22
(e)
where
e2i
(t)
Ô (VT)).
By Lemma 15 have that the end of the first phase (t Ti).
R2
Ei [d]
W2i
(TT)
sign
W2
(0)
(T+
+
(T)))
where
c(T1)
Ô (vn+ da-1/2 -
Combining with Vi E [d], Vt < T;
sign
92i
(t)
= sign
W125
yields that during the second phase,
fort $ T, we have
R
clt) = 0 =0( do-1/22
Vi E [d]
w
= sign wu2 cll) + R!)
where
E.13 Proof of Lemma 27
By definition of Ty, there exists jo E [d] such that E(T)
- -id for Ti 1 f. We prove by
<
induction that during this period, Vi E [d] sign
W2
sign (wr(01e, [2. sign W22
(0)
and
that j E [d] u2
(0)
W
The base case (t = TD was already proven by Lemma
23 Now suppose for some such that W22 (r) sign (w/²
Ti < T, for all such that Ti- S r S L, we have Vi E
[d]
sign
sign
w22 and that V E [d] (r) é
W
[i..]] = Ô
va Using these bounds,
va
we get that Vj e [d] :
(())(())) [i,5]]
+ |4, = O(). which then yields two upper
bounds 9 ( [i.5] = 222 EJ" = δ
(r)
and
92s
$ d 2j=1 W |..J|
= 5 (va).
By definition of T.g, we know that for all Ti S S t, Vi
[d]
92.
N d = 2(E) and that
sign 92: = sign
W2,
(0)
which implies that Condition is satisfied for Vi E [d]. At the end Together with the upper bound of
of the proof of this lemma, we will show that π ô Van
92 (r)
we can apply Lemma 22 to get that w.h.p. for Ti S S t, sign (Aw = sign
W2z
and
Aw (T)
@g)). Combining with the inductive hypothesis sign (w077 (T) sign
(0)
gives us
that W2
(r+1)
(r)
+ @(7). Specially, when = t, get the lower bound
(e+1)
(0)
Ω and that sign (1/2 ( sign
(0)
Since E/7) - -ndd, we have that Vi [d]
[t.
(G)
E = á (VT) 22(E)
and that sign (0f (T) (6.50) sign
(0) Wh
That means Condition is satisfied for Vi E [d]
and jo. Using the same technique as when we deal with w we get that for Ti < T S t.
59

---

Vi [d]
( (e. jo] jo] w/
1 [ol
+ @(n). sign (w (6.50))
sign u2
(0)
and that
Vi,j e [d]
( = -(())
Now we analyze the magnitude order of u22 //
w (i. [1.51
Let's first analyze
(e+1)
By Lemma
15
when t = T2, Vi,j E [d].
[i,
W2:
1 δ (V1+ -
=1 Ō (0vv (v+ d
(T))
2j
Combining with the facts that for Ti S S t,
[i.5o]]
|w/ jol + ²(7)) and
| [5,00l
(r+1) = wxx (T)
+ @(τ) yields
@11.
Since we just proved Vi E [d]
222
sign W
(++1t) sign (w(Ctt)( [1.5ol)
sign
11:22
(0)
get that
(W2Wi) (c+1))
Wii
(++1) wr((+1)
[o,
=
W2s
(1+1)
W
(1((+)))
[i,
= =O(1,
i=1
which gives us that
Wii
(++1)
Ö Va
Recall that we have shown
(e+1)
2 52
then
W22 (1+1) = +
Now we prove (w [2
= θ
We have proved that Ti S t S t. Vi,j E [d]
Val
(() and 2u22
(1)
= (n)) then Vi,j E [d].
21
[i,j]
W
(T)
[i
+ E
w (7(4
[i, - W (T) [i,
Ti
(1+1)
(T
(r+1)
w/y
W1
W [i,
+ +1+1-11 Ti)
Ō()),
(TT)
+ +1+11
Whi
where the last equality uses
= + θ (vi + -/7) Since we already proved that
(t+1)
1
we get W ((+1)
= v1a
222
(aa)
Therefore by induction, for all t in the interval [Ti, T), we have Vi,j E [d]
w2z
(e)
vd
w: [ 5]
ō
Jaa From the proof we also get Vi e [d]
22
(++1)
2
(t)
and that
Aw (c) @()). AW
(1) [6.]] < Ö()).
Now we verify that T θ ) Combining Vi,j [d]
(T) è (12) and Vt
W22
[Zi, T), W2i ((+1)
(1) ((n). we immediately get that T Ti à
Vdn ) In Section E.2 we
1
have shown that Ti e Van) then we get / -= (Van)
E.14 Proof of Lemma 25
We prove this lemma by induction. The base case (t = Ti) can be verified by Lemma 15 Now suppose for l the interval [Ti, T), we have Vi,j E [d], (( [1.1]] à (1aa)
60

---

For t E (T, by the proof of Lemma 27 (Section E. 13 we know that for VT S t, Vi,j
E [d]
W2s (r)
= θ va
W bi.i] - ô va and that [. <
S Ô vv and
921
that Ti = ④
Then we can pick H == T-M log o0 and apply Lemma 17 and Corollary to
vdn
get that, w.h.p., for all E (Ti,T) and j E [d], the update of W can be written as
(1 - B1) Et T Bi9i 1" (1-7) 7) [1,5] + (1
w [i,
w [1, 1/ VI (1-B3) B2) Et=u 82 (ot (1.1) + 1dd [ +
(4)
(t) [ |/. = <O(ns)) By Lemma 24 we have that for 1 1 i S d.
where
9l [1.1 = (1) E3²) e sign
(w26
(0)
E
(1)
+ R J..
vī (((V++
cl0) E
d/-1/2
H
7 Fig (t- -
[i.
sign
W'2
(0)
Σ
By
cll
Es
7)
Bi
Rl +)
n.1
[1,
(37)
Using the fact that for a, b E R.
ba²-b²|
a²
2
|a|
+
|aabb]
2
we get that
(o1 (3/0) 2 = ((()) sign
sign
(0)
F + R (..1)
+ (0(0) )
R00))
gug, 16.1),
=
where
Ö (VT) + - ). That yields
H
H
H Σ
Σ 7=0
B5
9i
(i.0]
(1 (
Σ BB (c(t)))
gg R s [1. (b.
(38)
==0
Since (c(
2
>
0. in eq.
38 we have that
SHO o 35 ne [5,
[5,
= (VT+ da-1/2
-0(v7
(39)
ΣH= $3 (()))))))))))
However we cannot similarly prove that En S7 R 1) [i, << ΣH cel D/ (t-) 7)
in
eq. 37 because ((f-)){(())) may not have the same sign for 7 = 0, 1, l.... H. To deal with e3.
we need to consider the two cases where I#H 87 R r) [.,]] < E Hill p/ rt-
or
Σ H! 3f R [i.0]] & EH= ((-T) p((-r)
Case 1.
(1 - 31)
11 i 1
( [i.0] + 10) 10
s (1-1))) 3.) (el(t-)) ((((-r)) where
31)
8
(n²²
+
1
d
-
Note that from eq.
39 we have
H
H
(1 - B1)
83
R ((--))
[i
s0
=0(vit 1--1
(1 - 31)
Σ T=U
35 (((
61

---

Combining with (4) </d - =(
d7 x) <0
<², we can apply Lemma 34 to get that
w [i.5] - 1 [i,j]
(1 - B1)sign wiz 3 101 -1 + (1 B1) Σラ= FT T>
3,1 [t, [1, + (4) [.
=-
V (1 - -2) Σ== 32 (()-))))1()7)) 2 + (1 -32) Σサー 33 R / ( . + (i +
1 - B1
sign 1/21
Σ 35 clc F.)
Fry (t-
(1+ 1 (1.1))
7) v1 - 32
V -===0BB B2 (ci-r)) )
+
(0)
(1 + ef"((.5))
= sign
W22
1/3
[c] [1016,5]
ō Ö + ) Since
[i.5] - [/.1 Ö()). we get that
where
=(1).
Case 2. |11 B1- EI 37 R
(i.5] + |M > 01 B1) En o ()
where
5
17 + d xz-T) ) )
R [[ai
Since
Ö (vT) + 2o-1 ). we have that
olt)|
H (1 (1- 31) Σ
H
89
(1- 9.1
[i
s 20 ( (Vn + da-1//
(1 -1-30 (1) Σ 3
(t- ( 1
7=0
r=0
H
(i+ + da-1/2
1
(H + +1)(1 - Bi)
Σ
35 (d
==0
δ
( (
1
(1 (1- B1) Σ 31 (ol (3.5))
11
da-112
T=0
where (1) uses Cauchy-Schwarz inequality and B2 = 88, (ii) uses eq.
(38) and (3)
G [t [| = Ō(n5)) = Ō (vT++ + x ) ( (5-v/ 1 [1.1 gives us
Combining with
"
(1- 31) E 3f R F) [1 + 1 1
(1 (-:3))
Σ 0=0
T (
<
n + 2 7
0 (vV) +
ja 172
Σ 83 (1 (1.)7)))))) 2
n1
-1 [2.01] +
17 1
+
1
d² -
r=
H
<ō
Σ 33 (gi (1.2) + +c
n
+
d !
r=0
which implies
wr((t)) [i.]] - [1
n +
1
d2 d2-1
62

---

Consider certain i.j E [d] and the period from Ti to t. Denote T as the set of time points when Case 1 is satisfied. By Lemma 27, we know that 17(1 -Ti) - Ö which gives us
Σ
17 + d8
= Ô
d8
d2
di+
TET
By the first phase analysis, we have that
( = sign
(0)
V
+ R (1)
[1,
1::2
1
where v((r)) = O (72): R 21) [ = Ô (V + )) Combining with the analysis of Case 1, we have that
W
[1,]+ + Σ 1.. 1 sign (1))0) (r
Σ 505") Rr(ii.). TET
FET
where |i,1, S (1)² + x (|vsei/| + -TET 1..
Since for T E T, | = -O(n). viTh) = o (+zz) and 17(t - -i)) = Ö (+t) we can bound
[RT[i, by
|T|i..il s Ö n) 8(p
1
(o
1 ( + (1 (1- ((-77))mm) 50 di
n
1
d
d=+1
Combining the above results together yields
= (T) + Σ AW (T)
[ = -
+
Σ
[1-5+
+ Σ AW (t,
TET
TET
sign
(w
(0)
(r) Ô
1/2
1
dà
d
TET
(0) (w)g))
p((1) 3
+ +R(2') 1e1'(1,1), [i, where | S 5
11a
== sign
d2
d +
By the inductive hypothesis
| = 0 =(())) we get that |/,00=
= á
which gives us
R
[J]]
17
+
Therefore, we have that for any j E [d] and any i,, i2 E [d],
( [
sign
[ ① 17 + -
1 + δ
n1² +
w [i2. 5l
sign
W2in
(0)
(1)
1 + δ
n 2
d8
(t)
By Lemma 24. we know that W22
(t)
with different i are also roughly equal, i...
"21
1 +
22
δ Ó + do-1/z).
Then we have for any j E [d].
(W2W1)
(1
[i.s]
(t) w
( θ (d wilh |w (k.0))
[1,]]]
= 3 (va w/ ((:3))
where & can be any index in 21,,,.., d} and the last equality uses Vi E [d]
1 = - ((t))
63

---

Now we analyze the lower bound of [
[k: j]
Although may decrease during some period.
we observe that once
w ] decreases to some value of order é
va such that (W2W1),) () <
A Vnd, i... 1
vnd. we can apply the technique in Section
E.13 when analyzing
w Jo] to get that w [ will increase in the next step. This mechanism ensures a Q va
lower bound of W ( [. Since j are arbitrary, we have proved that at time + 1, Vi, E
[d] : [ [..01| 2 s
Therefore by induction, we conclude that when Ti S t < T, for Vi,j E [d], [ [t,5]] = (+2)
The remaining part of this lemma has also been proved by the analysis above.
E.15 Proof of Lemma 26
Lemma 27 tells us that for any E [d]. 2 keeps increasing when / < T. However, the behavior
(t)
of W (1) j is more complicated. The following lemma tells us that [t will increase until
Tyj. After that [w(O [i. and E/) may zigzag. but En will not fluctuate dramatically and will be trapped in a small interval around zero.
Lemma 28. Under Assumption
2 and suppose o S dis/s Pick vj S O ( E -
and B2 = =B. Consider certain coordinate j. For Tı S t < min { T 1T, j. we have vi E [d]
[1 keeps increasing. if < T, then for Ts S t < F, we will have - Ô (Vnd) S F s Ô (Vnd)..
Now we start proving Lemma 26 At time T, denote S := [i: Tys <T). fi.. the set of coordinates
whose E, have passed its "flip time". By Lemma 28 we know that Vj E S, |mft < ô (Vnd). If
S Φ which means Vj [4] E S Ô (Vnd). then our lemma will immediately follow. If
S * の. we have T min
(Ty m Tg and that Vj e S²: E3Sh O. By the definition of Tg. we
know that Big E [d] : 22i0
(T)
S 50 (dvT))
Then
Σ
B
W
(T)
[io.
Σ B W (1) [io..] Σ
E
W
[
JESc
JES
<
|250
(P)
+
E
7 w
s < + dō ( (Vnd)
Va
1
= =Ō (ddv))
JES
By Lemma 25
we know that when Ti St <<< T, for Vi E [d]. W [1,1]] = á (Ja) Since the
update per step
A [.:]] < Ö(n), we know that sign (w/10 [1,]]
remains unchanged during independent of j. Combining
this period and sign
[
(8) [1.51) sign WiTTitE
sign w22
(0)
with Vjj E 50 F.
< gives us that E w (1) [io..]] for different j have the same sign. Therefore
for any jo E S
(dd/)) >
Σ SES0
E
W (1) [io.]]
=
D
W
(f)
(io.
1 ( (T) [io. Fo]
JESc
E (T)
E
s ô ( (a (dv/nd)
ju
=
64

---

| s sÓ(dn)))
Note that the above inequality holds for any jo E S which means Vj E s
Combining with the fact that vj E S: E f Ô (dv/nd) completes the proof.
E.16 Proof of Lemma 28
Consider certain j [d]. when < min
we have that E/ < Vnd.. Therefore we can use the same argument as in Section E.13 to prove that [ [i.j] keeps increasing, and sign (w. [i.o) = sign w22 for all E [d].
T, T. j.
At time the "flip time" t Ty. r. by definition, (1)
- -vdd. After that E may oscillate. Now prove that once E 2 /nd (or 5 I -nd,,, after short period E will decrease (or increase) until E s Ô (Vnd) (or Etd - (Vnd). Moreover, during this period, E(" won't change too much.
We first recall that when Tj t T, Lemma 27 gives us for all [d].
and [ [. = Ō (tt)) Then eq. (35) we obtained in the first phase analysis still holds, which tells
1 (aa
us that the change of E per step satisfies
Ef(+1) = = Ô (vvV) for all Ti < T.
We divide the analysis into cases, based whether these E." 2 Ind E nd. By
Lemma 25 we know that when Tj S 1 < < V
[d]. |w [r, = Since the update per
step
A AW
O [R,s]| & Ō()). we know that sign [t,3] remains unchanged during this period
and sign [2.0l = sign (w/T"
sign 2222
By the analysis of w2 in Lemma 24 we have for all E [d]. was = wi() + sign
(t+1)
(0) uu
A
25
where 1/0 = 1 (1
Case 1. Consider some time point t such that E S -vdd. Note that for all
E [d].
01 00 [t
11 (1) E(() á (vn) and that sign 9i [1.5)) sign
(1)
(t) sign
(0)
W22
By Lemma 24 for all i [d] we have ( [i,]] w (('(i.0] + sign
Wiz
(0)
A (0) [i.
with
A ( = 7 1 + δ (vm) That gives us
E((t1)
[++1) (0(1))(11)))) (++1) [i,j] Aj
W2s
iml
Σ
W2
(t)
sign
W22
(0)
A (1))
(w(')(i [i, sign (w A [G...) A
Σ
W2 W/(h (1)
(1)
[i.]]] +
sign
(0)
(t) 1 [ W/ 11 + Af - A₁
0 m + Σ
1 19 [.. + 1
11244
+
1=1
E 1(1+1) > E/0
65

---

where (i) is because sign
W22 (t)
sign
[i,)) sign (w2.
(0)
Therefore we have proved
that F will increase in the next step. After that for 2 + 1, long F I -vnd, the above
analysis will hold and Ef will keep increasing until E > -nd or we reach T.
Consider some time point such that =000) 2 Vnjd. We will prove that E will decrease after a short period, and during this period, the change of is at most Ô (Vnd)).
Case 2.
By similar arguments as in Case I, we can get that (f = [i]] sign
W2: (0)
A 1 [.
where A [1. = 1 (1 + ① (vm)
Then
E (1+1) Σ
(++1) W (++1) [i, A,
222
i=1
W2 + sign
(t)
(0)
A 21
(1)
(w/ [i. sign (w!²) 10) (i.)) -A,
d
21
W10) (1) [i.}]]
sign
W22
(0)
(22 1 [1. - 1 - -444411.000
Aj
(t)
[. [.
10
1
" [
+ A 1 [.,1))
W2i
where (i) is because sign
(1) sign (w(²) (1.ii) = sign
(0)
E E
may not be smaller
than E0 but we will show that after at most t, steps for some ta we will have E. toftte+t)) 1 Ee(tee)
To see this, first note that by the bounds of A (1) [ and (t) 2 we get ( [i, 2 A - n (VT).
(t) Since 222
increases by O()) per step, and [(]](
[ keeps decreasing, then we have either i)
after te steps for some t.. Vi E [d] 222
(++++>)
N W (t+
[ + v ii) we reach T.
For i), if E((tt)) < Vnd, then it's already what we want. Otherwise we will have A ((ttta =
(1 (1 + Ö (vT) Hence
El(th)) - ES (+tte++1)
Σ ( u22
(e+e)) 1 (1tt.)
A/4 ((++a) (w(ttt>)(( [i.5] + +
(1++t) A ((+t))) (i.
Σ (w/(ttr>) [1:01] (4 ((ttt))) [1,1] - - + vTA (1+te)) [1.5]] + 24
(t+te) A ((:tt2) (1,
2
(-nō (VT) ( (1 + + + 2i
((+e) ((tt)))))))) > 0,
where the last inequality uses Vi,j [d]]
| |6. ō (1z). Therefore E E((tte--))
E E((--))) After that for 2 21+ L,4 + 1, as long as FT(T) > Vidd, the above analysis will hold and En(T)
will keep decreasing until EX(T) < /nd or we reach T.
Now we prove that during these t, steps, the change of E, is Ô (VTrd). Since at each step the
difference |22x| - [W](..]] will be enlarged by Ω(n), then we know that 4 = vn/Ū(n) = ô ( (#)
66

---

Combining with the fact that for all Ti S S T, E / = (πvá) gives us
Elttta) - E 5 S Ô
n m1.Va) = Ō ((vnd)
For ii), we reach T before Vi E [d] 1''
(e+ta)
2 W [1,]]]
v7. Then we have P 1
vn/9(()) Ö (..)) which yields E,") E s (n(T-t)) -tlvd s (Vnd).
Combining the above two cases, we find that for some 1, F..") 2 vnd, then after at most is steps Ej will decrease and keeps decreasing until Eg vijd increase at most Ô (Vnd). . If for some t, E - -vndd then after one step it will increase and keeps increasing until Ej > vnd or we reach T., That means once for some coordinate j. Eú overshoots, will zigzag small region around zero, which - (Vnd) Ö (vnd))
reach T. During these steps, Ej can
Hessian tends to become more and more diagonal during training
In this section, we empirically demonstrate that the trend of loss Hessian in practice is to become more and more diagonal during training. We also give rigorous theoretical analysis on two-layer network under Assumption aand]
F.1 Empirical Results
Let's first define the diagonal domination of the i-th coordinate time f.
EEf (H,))))
OPT diag.
(t)
To measure the diagonal domination of the whole Hessian, we need to consider the distribution of .PPT (t) for different Figure 15 shows the mean and median of ,.GDM/ diug.. (t) and y-
diag, (t) on the sentence classification task (see Section [.T]. Here we chose 4 layers (Layer #6, 12, 17 and 22) and computed the Hessians across these 4 layers. Since the number of parameters is very large, we did the computation by random sampling. As we can see, for both SSGDM (t) and raag. (t), the trend of their mean or median is to decrease over time, although there might be some oscillation.
T 'diug,
diag.
SGD+M
SGD Adam
Adam
602
70009
6000
30cen of 100
5000
median of
4800
403
503
1000
1502
2000
500
1020
1500
2000
lleration
Iteration
Mcan
(b) Median
Figure 15: Mean and median of diag. 12,,,2,),,22)
SGDM
(t) and "dlggn (t) for the full hessian across the four layers
F.2 Theoretical Analysis
To simplify the theoretical analysis, consider the mean of digg,()) over all coordinate and define
OPT ding
(t)
mean (rding, (t))
OPT
(40)
67

---

We consider a 2-layer network under Assumption ! and 2 and have two goals in our proof:
1. To show that ding after training is smaller than that before training (t =0).
2. Note that in our setting (see in Assumption I the Hessian is (dd-d) x (d +d) matrix. For
a completely "uniform" matrix with the same size, we have that dingg() ding (t) e (vd+
(d). Hence our second goal is to show that the Ragg (1) after training is lower order than (d).
Theorem 2. Consider the ratio Rigg(t) (t) defined in eq. 40 that before training (t = 0), with high probability,
Under Assumption / and 2 we have
diag (0) $2 (daa-3).
(41)
For SGD+MM defined in eq. (3) For any > O, by picking the same hyperparameters as in The- orem 7 for 1 TsGD;2 mentioned in Theorem |
we have with constant probability, for any
t EE [T SGD,1: TSGD.2).
RSDD (t) (t) 10 + g00
(42)
where the trend of (()) is to decrease over time and (Team <O ( o(dd)
For Adam defined in eq. (3) For any p > by picking the same hyperparameters in Theorem for
T Adam, 1, T Adam,2 mentioned in Theorem] 7
we have with high probability, for any t E [T [Tadmm, T Adam. 2)
R (t) S 2
(va) rl())
(43)
where the trend r()) is to decrease over time and y (T 5O
= (0())) 1'' (
(va).
F.3 Proof of Theorem 2
Lemma 4.3 of (aa 16 gives the following forms of Hessian.
For any & E 21,2, ...; H + 1}, we know that V Veec(We) (V Vee((W))L(W) equals
((W W ... Wk+1) W (WW++1.... ......... Wx+1) @ W.W W.),,, and for & E ( H + 1},
V Vonc(W)) (V bee (e)
(WH+1.... Wx+1) & WV4-1...W.y..
+((W+-1... "". 3 / 8 (r(Wn+1--. Wk+1).. ... Ta @ (r(WH+ Wk
where r = (WH+1 W1 A)", C = WH+1 WH- ... W2.
For the 2-layer linear network, write the Hessian as
H22 H
H
H21
H11
then we have that
H1 (W?W;) @ la E o xxd
Hy WaWF E Rddd
H21 = W @ WT + la & (W - AA e xdd
Intuitively, before training the elements of W1 and W₂ are very close zero, and W2W, A -A Since the elements of A are (1), we know that the magnitudes of elements of Hz are much bigger than those of H11 and Hzz.
After training, for both SGD+M and Adam, W,W, R 0. Then H21 R (W2)² @ (W1)² and the magnitudes of its elements are no longer much larger than those of H1 and H22. From the formula of H1, we know that all the diagonal entries are nonzero, and among the d d² off-diagonal entries, there are only d - dl nonzero entries, which helps us bound ROPT (t).
68

---

F33.1 Proof of eq. (41
Let's first analyze the weights and Hessian before training (t = 0). For ease of notation, we omit the superscript (t).
For the i-th row where 1 S i < d, i.e. the i-th row of the submatrix [H22
H2i we have
da
Σ j+1
Σ
H2g ol
H i
H = Σ (w2: (w.:W + (W2W, =)))²=
(=(+-)dd On the other hand, for the diagonal elements, we have w.h.p.
1i Σ W 20 ( (
[M E]1 - |1/22 all
Then we have that for 1 sisdd d,
V Eijes (1,6
v [d(
52 (atn-3)
[H i]]
O (
For the (id- + k)-th row where 1 i [121 H11]. we have
k S d, i.c the ((i - L)d + k)-th row of the submatrix
CHi-l Σ 1iil(i 1)d + k.
+ + Σ H2 [(i 1)d + k.j]
j
s((i-1)d+k
2 H2 (i - 1)d + k = (wr Wi 14, k] + (W2W, - A)x) = 2 O(1)
On the other hand, for the diagonal elements, we have w.h.p.
Hid + k id + R:] = [ - 1)d + k (i 1)d + ki) = w s Ô d20
S d. Then we have that for 1 S d,
v Enpiat Hi..
v (1) = 52 (d2a)
H Hid + k, id + k]]l
Taking the average, we obtain that before training, i.c. when t = 0,
dit (dda-) $
+ + (dfaa = á (gta 8).
ROPT (0)= N ding
d2 +d
F.3.2 Proof of eq.
The proof is based on the lemma below.
Lemma 29, Suppose the weight matrices have the following structure:
W1 = uv T + R1,
W₂ = cu² T + RT
where V1 si S d:
[i-l S 5,
< 8,
8 E (0. 1).
|0,0,1
Jow,
Then we have for 1 Siss d,
v 5+4 H2 (6,50
$ 1+6 8 (1
|cl
|
Hii :]
8
/0/12/
u²
(1 - 8)
for S < d, 1 1 k d.
/ ESfsdt+ H2 (1
+ 8
[ux]
Σ 1= u
|Egl
H id + k, id + ki] S 5
+
|c|
(1
69

---

Now we are ready to prove eq.
42
By the analyses in Section D.1 with momentum satisfy
we know that for / E [TsGD,1, TsGD,2], the weights obtained by GI)
w = u (T)) (1)T ( + _R
w (0) (21)2' R
where TSGD,1 Ti and
1 [t,
R
11 S Misijs d:
I O(co),
I Ö(co).
(2i) (e) ui 03
c(uu (T1)
Here Eo is defined in Definition 2 Since u (Ti) doesn't 't depend on time in the period (TsGD, T).
we write uTTi) for ease of notation.
Hence by Lemma 29
when E 1T T2GD.2]. we have for Usisc d.
V Eiye (11 (1,1)
1 O(co)
ΣS j= up
|| | l12
1
+
H1)) [t,
1 O(co)
up
11 Ö(en)
uz
|e(0)|
Σラ=
uz
( uz? |010
FX()) 12
=0
1 | 1000)11)))
V
0
u²
(44)
and for 1 si si S
V Esgiatr (1.0)1)))))
1 + Ö(Eo)
)
Σ f=1 u²
+
H idd + ki id + k:]
1 Ō(co)
(c())
u²
(1-O(Ee)) 2 (cll)) uz
(t)
E
uk
Σ j=1
u²
= 0 1+
|e(t)
un
(clo)) u²
(45)
By Lemma 4. we have u = X + Y where X,, i E [d]
are i.i.d Gaussian random variables and w.h.p..
Yi X
Vi
[d]
< θ
= 8zy-
(46)
dia-
which yields that
d 1
u²
1 day day
v
Eラ X²
1
2 1
Vi E
[d]
5
(47)
Xil
Jay
X?
By the proof in Section The induction in Section
D.8
we know that for t E [TsGD, TsGD,2l Vi [d] : 2
c(()) are positive.
(())
D.9
further gives us that for & [TsaD,1; Ts. w.h.p. Vk: E [d] 2oo
(t2)))
which yields tes
( Combining with eq. we obtain
|c(t)
プド X?
1
[000)
u
X
(48)
uk
(1)
v
V
X?
1+
ct0
u
Xx,
70

---

By the proof in Section D.8 we know that for L [TSGD. T 2.. Vi E [d]
(1) c0(0 are positíve
and monotonically increasing. On the other hand, the proof in Section D.2 and D.9teell us that w.h.p.
|B0012 (resp. Ak E [d], E decreases from ((v)) (resp. (())) when 1 TsaD, ((vend)
| and
(resp. ((((o)) when t. = T- Therefore, the trend of time, and when 2, we have w.h.p.
is to decrease over (e(1) u
u 0000 l2
4k E [d] 1 O (VEo),
| = =0 -(Vod)
(49)
cl e @(v) and
Moreover, when & = 2, 2, the inequality in eq. (26) becomes equality, i.e.
Vj E [d] || = θ
Using a X + Y and eq. (46), have
1 = (vd).
kj
Va
1
=
| = θ (va).
which together with the second inequality in eq.
(47
yields
Ej= X3
e
12
1
6ay
X
2
12
x2Vd
1
2
Σ X?
u
1
Owy
c²x?
X3Vi
Combining with eq. (49)
we get that
X2
Σ1 j=
X²
j-
6D
S
O
VeD
0
(50)
u² 10
X2
(c(0)
uf
X2
d
Substituting eq.
48 and (50) into eq.
44 and
(45)
gives us
Σje
(1,)
v
V X?
A1 sisd::
5
9if
(1)
Hi
X2,
(T where the trend of g1x is to decrease over time and Ti
(1)
0
E X²
X?
50
We also have
V
V
X2
Hi k, id + kc]
50
Xii
92
where the trend of 92. (1)
is to decrease over time and 42
(Tann,
<0
E x² x?
v
Hence
X²
RSDDM (t) diag
0
5=
4is
(t)
+
d
42b
X11
d + d
d
+
d
1=1
1=1
d
X²
1
0
+ t1)
Xi
1=1
71

---

where the trend of (1)) is to decrease over time and
(Tsm.aa) =
d2 +
j=1
X?
X?
Vco
+ dc
d
d 1=1
Σ
0
2j=1
d
X?
X3
V d
c
i=1
j=1 X
X2
0
Verd cod
O
Σ
dc
+ d
x?
d
X?
V
i=1
Denote o² as the variance of X; for i [d] By concentration of chi-squared distribution, we know that with probability at least 1 - 8 for 8 > 0.
x od 9'0 (
By Lemma 36 in Appendix H we know that with constant probability
à E [x,] = = log d).
1
Then with constant probability, Ei
X7 S |
O ( log d). Hence
d
1=
Ō(va),
1
d
2i=1~
d
X3
= 0 (d.)
X,1
X?
i=1
i=1
Therefore with constant probability,
ding ReCE (1) = 0 Ö + t')")
where the trend of (()) is to decrease over time and (.rrm,))) S <O =ddveg)). For any > 0, by picking the same hyperparametens in Theorem [ have egd < Ö (t)) and hence g (Too..2) s
o(d).
F.3.3 Proof of eq.
(3
By the analyses in Section E.1 Adam satisfy
we know that for t E [TAdam, TAdam,2), the weights obtained by
w = (() + +R(
w nn + R((0T
where Vi E [d] : u, = sign(w
(w24
(0)
E {+1} and
1 [i,
A1 Uiiijs S d:
< 5 Ö
< 8.
1 (t)
d2
c(())) m1
Hence by Lemma 29 when t E [TAdam.1 TAdam,2], we have for siis d.
Σ j+i (H (1,))
++0 8
|clo|
j=
uz
|50000
1100 [
1-8 8
1+
(|0001112
u?
(1 5)2 up 1000
(51)
=o (
|cl0|
||
1+
Vd+o
2
||0(001122
|
and for 1 5i d, < k < d,
v Eitiaak (i,s))
1 + 5
%
)
Σラ= uz
Fr
1+
Hillid + k, id + li|
cit)
uf
(1
8)² (c(t)) ut
(52)
=001 1+
(
(t)
ok
vd+o
c(t)
(clc))2
Recall the following facts of Adam.
72

---

(A) By Lemma
15
we know that for t E [Tadam,1, Ti] (where Ti is defined in Definition
w.h.p. vk
E
[d]
(4)
c(t)
(tt
line). Specially, when & TAdam, Vk
[d]
" = c
(9)
Lemma
25 and 27
tell us that for t E (i,mm,,,,) w.h.p. Vi, [d]
d
[i[i,,]]
=
which gives us Ak E [d]] 1K
+ Va
and
va
c0 θ
1
That means when & E [T Adum,, 1: TAdum.2], vk [d]
| and le(t)
increase from 1 to and
(1).
1o
1
ct13
|000 1
(B) Lemma and 26 tell us that w.h.p. || (resp. Ak E [d],
E decreases from ((d)
(resp. (1)) when t = 1 to Ô (cf VT) (resp. Ö (dvnd) when t = TAdam,2-
1 100
Combining (A) and (B), we get that the trend of when & = T we have w.h.p.
and
is to decrease over time, and (cte)
E(1)
L
(d
dl Vndd
(53)
|
2
(clt))
Substituting (A) and eq.
53 into eq.
(51) and 52
gives us w.h.p..
Ej (HO
0o <v(vi) +r
dsiss::
H1 il
where the trend ris is to decrease over time and ri We also have
(Trtmm 2)
0 VT)
Essate (i.)) HHW + ki id + ki]
S 5Ō(vd) +
(t)
(t) where the trend of r2 is to decrease over time and r2
(Taam.))
(d /nd).
Hence Rangg (t) = Ô (vd) + Ei ri + 4
E 1= 12 = Ô (vi) +r(() where the trend
(1)
of (() is to decrease over time and
(Tsamm))
1 (o(em))
+ de + d Eo d (d
d
nd <ō (d
dd
For any p > O, by picking the same hyperparameters as Theorem
we have mju* < Ô (1) and
hence y (Thaamn)) <
(Vd)
F.4 Proof of Lemma 29
By the assumed weight structure, we get that
E [d] :(1- (0))())))) < (wx) s (1 + 0))
(1- (1, I.. < (1+ (1,
73

---

For the i-th row where 1 R i V d, i.e. the i-th row of the submatrix (H22 H2g by triangle inequality, we have
dst [
Σ J002
H2 +
Σ
H
j=1
d
(Wi :2. Wi (..12 +
Wf[t, k] + ||
2j
k=1
- ||Wi[t,;]lb
Σ jイ
Iiib..:113
Σ
W2z
+ -ll22
j=1
Then we have that for 1 siidd d,
Eitt H2
[W Σj
+ V E w3y
||
H (1
[ 11,2112
[W(2;:]
EAEr W...31111
- waj
d
|51|2
+
+
W1., :1112
|| :1112
1| [t. :112
(1 + 6)2
ΣjS
u 11
(1 8)2
c² E n2
| 12
u 1
12
(1 8ر2
(1- 8)
1 ++0
8
lel
1-1 u3
| l
+
|00|2
u
(1 (1- (1-8) u) 1011
For the (id + k)tth row where s d, 1 S k S d, i.c. the ((i - 1)d + k)-th row of the submatrix
[Hn H11. by triangle inequality again, we have
Σ jfid+k
H
Σ
UPiiki ()dl--1)d+ + k, +
H2 1)d + k.
d
W 2
+
W
WPLi k] + Ex
j=1
|
Σ
w +
WPb. k]
+ +EEx|-
Then we have that for 1 S i 1 <d,
Ejfid+h H2
|uzil / Ejft way
v Eラnm WP[iAA] k]
Ex
H lid + kc id + xi]l
W2
2
Σ jzt wby
j=1 W k]
EE
W22
(1 + 8)2 (1 -882
Ejtt Puz Zuf
(1 + 8)2
0 E3= u2
|R
1=
V
(1 8)22
d2u
+
(1
6)
1 ++6 8
(
|uA|
Σ E uz
|Eg
<
1+
+
8
[c|
(1 - Pu
74

---

G
Connection between diagonal of loss Hessian and weights
The partial derivative at W; of the cost function for each given by:
Ww; I(W) = W (W WH W1 AJWT W
(54)
In our experiments, we were interested in the diagonal elements of the hessian. These are given by:
( (Vw a(W)),, V(w.),, (W .... ( WH Wi A)W? W (-1)a, 1),,b
for each possible i, b. For ease in notation, define for each the quantities M; == W #+11 and N; WT ...W..-- Then we have the following lemma.
Lemma 30. The diagonal elements of the hessian of the cost function are given by:
((W.). (VW))..WW. [,W)) (M,M) (a,a(N
(N? N:)
for each possible i, a,
Proof. We have:
VwW)(W) WH ...
W1 AJWT W
-M WWW .W.W -A),;
= MWW++1 M
This implies that:
V (V (W))) V(W.),, (M,W WW WIN; - M,Y MYNN)..
VWWDa.. W+ WiN.)a.
where the last step follows since Mj and N; are functions of W.,
Since Mj := WT #+1 WT N₁ == WT WT,, by defining C; == M;WH+1WW ... W41 = M,MT and D; == W,-1- W,WNN, N? N; we have that:
V a a (Vw. L(W) )a,b ((w)). (CWWDD))abb
where C2 and D; are not functions of W, Now, Equation 74 in the Matrix Cookbook for any matrices A and X have:
shows us that
V (XA)is Anj.
Note that W2 E Rddxd.-1, then we can apply this to obtain that:
((W),, (V. (WW)a(W)) V(W.).. (C.WDD)a..
-(W.aas
Σ (W,
Σ (Ci)a. VaWae (W;D))
Σ (CCaaakba k= (Ci)a,
(M:M)) )a N:)o.
This completes the proof.
Shttp:://// http://wwww.aaooooooooooooooeeea
75

---

For case of notation, let's now drop the superscript OPT and (t) and write F (f) as Rmad,1 and m (t) as 2. For 2-layer linear network, H 1. Consider the Hessian w.r. W1, have MrM MM" = W W2 and Nj N1 is an identity matrix. Under Assumption 1 we know that W2 is a row vector, which can be denoted as W2 [221, W22, .... W24, Then we have
max, (W2s)²
(Mi Ja,a 2 (NT 1
Rmed,
median (W22)) )2
Similarly, consider the Hessian w.r.t. W2, we have that M, Mf an identity matrix and NT N₁ W,W7" Therefore,
max; [ [t, :113
(M, )a,a (NV (1) b1). [| :112
Rmad,2 median [W [t.::1]11
Hence we have related the uniformity of diagonal Hessian to that of weight matrices. In the detailed analysis, for both GD and Adam, we can prove that W1 converges to an approximately rank matrix. The following lemma allows us to use this rank I structure to compute R and Rmat,2-
Lemma 31. Suppose W1 Radd and W2 R have the following structure:
Wi uv R1,
W2 cu + R2
where u E Rd, U E Rd, R4 E Bddd
R₂ E IR
lxd and that
[R]
Rad
AI vsiiidd:
S 6,
S 6,
8 (0, 1).
ccu
Then we have
(1-5)²
max; u?
8)2
max, up
Remed, Bmod,2 E
(1+ 8)2
median u. (1 - 8)= median ui
Proof. Let's first consider R we have
Vi [d] : (1 - 8) S w2z S (1 + (cu,)
(1 - 8) max(cu;)² & max w S (1 + 8)2 max (cu;)
(1 8) median (cu)² < median w < (1 8) median (cu;)² )
which yields
(1 8)2
max; uf
max; w
(1+ 01) max; ui median u2?
< Rmed, (1 + 8)2 median uz
< (1 8)2
median w
Similarly, for Rmed,2. We have that
Vi.j E [d] (1- (11-8) (), (1
= (1- (1-5)32?10013 < [W :112 s <(1 (1+0 61+6)2u2
(1 - 8) max uf || max [W :11B < (1+ 6) max uf 1|
(1 - 8) 2 median u²² |1/ < median [[ [1: S (1 + 8) median uf m
which yields
(1 6)2 (1+ 6)2
max; uZ |
max, | [1.
(1 + 8)2 max; u 101 (1 8)2
Remed
median u
median [1.::]
median u² 10
That means
(1 8)2
max; u median uz
(1 + 8)2
max; uz median uf
Rmed
(1 + 8)2
(1 - 8)2
76

---

H Auxiliary lemmas
Lemma 32. Let A = m
LYX", Am - _0 = Vv,,LWW)) 15==1, 1,2. Denote Ao,, All
and 5% k 1, 2 as the corresponding batch versions at time I. Let M max | [6.i]] and
(t) _ mmx;; W'2
Under Assumption we have with probability at least 1 1 for Vt S T and
Vi, E [d].
(0) 91
[i.j]] - 06 [i,i]]] S d (M (M²")² 2 oVVT++++ + M
(0)
(0) dTT
92: S dt (M M. ov aT Vo + d 1 o d..
(0) (1)
Proof. By Assumption 3 and Chebyshev's inequality, we have for fixed i E [d] and t S T,
o²
P ((
A.
2)
P [ Aza [i. 2
x2
x2
Applying the union bound gives us
T
P e 31 ST:
| A > >>) 5
X²
Td
ji.j e [d] TT : | - Ae >>) 5
P
X2
which gives us with probability at least 1 - 1.2 for Vt S Ji.
E [d].
A
A. S SoodeT
[ | S SodvaT.
Now we are ready to bound ((()
g (9)
for k = 1,2 and t T.
Note that for all t S T and Vi [d].
d S
(
Σ
(
6.
Σ
22
w [i S (
j=1
Then we have with probability at least 1 - 4. for all t < T and Vi E [d].
[ w110 (AM) Aaw)) b)=m))
(w!" (w/"w
AW i] A bi
S de (((
(2) Ma!0 oV
Combining with g(²) (0) = whor w((") (i - Azx) - (A(0) - -))) we get that with
probability at least 1 - 1. for all 1 < T and Vi,j E [d]
o [i 10 | s [0221 (w (t) (A - -A2z) )ex)), ) (1)
+|000 A
A,
S M() (M oveT + -1 oVdeT.
928 = (w w/(1) (A00 - AAx) - -(Ā1)) a)) w(1') we then have
Similarly, note that 02
that with probability at least 1 - J.2 for all t < T and Vi.j E [d].
d
d
02
924
S Σ
W
W (AL - Amx)) | [1:33]
+ Σ (g(0)) A, 1 1/
j=1
3=1
s dt (2/
2 " ovdT + LaMl"ood'
77

---

Lemma 33. Consider two sequences {a /120, 1b(0) 1200 which satisfy
a (1 31 Σ bd m, E (0,1).
Suppose Vr 1:: |b:00 < B, then for any >> 0, the following truncated version
H
(()) = (1 3)
B ((-
7=0
with H 2 B log = 52 (1-B) satisfies
( - 2(1)) Sc. c.
Proof. We have that
10 200 S (1 1 3) T=#+1
Σ bit-)-)
S (1
Σ
B & B3
T=H+1
To make it less than e, it suffices to choose H 2 log( / log B.
Since 3 E (0, 1), we know that log S 3 - < We also have log B choose
< 0. Then it suffices to
log(e/B)
log(c/B (
H 1 8
B
1
H
log
3 3-1
log
1 1-3)
Lemma 34. Suppose a, b, c, Eas cb. ea E R. > 0, > 0 satisfy + eb + e, > 0, leal S Sola), S
5b, leal S 822 with 0 < 8 & 1, then we have
c
(1 (1+ R), where R
Vb + c tc
vb+c
Proof. We have
Co
Ea vb+c eb+ebb +0+e+tett tc vb+c v
Vb+ c + ee t
vb+c - 1+ Eo
vb+c
1
cb+c
vb+et + + +b+eb + Cc +
m
42
Define R = Q1 + 92. The term |9ıl can be bounded by
Vb- vb c c
|m1l
vb+ eb tec +c
les + edl
(vb c +
Vbb+bbbe /b +bb
16b
1ea
+
(Vb cb e + +c)
Vbb+ vb vb+eb eb tee)
(Vb+eo+ee c) Vb+ vb+eo+es
leal
+ Vleel
Vleal
(vb + eo + Eo + c)vb
Vb+ +b+
leal
Vleal
+8 Vb+ vb+eb +
Eb + Eo + c)vb
43
96
78

---

where gss can be bounded by
(Vb c vleel + c)vb < Vb(1 8) + c(1 6) S 8vb
5b
5vb
Ugl
vb(1 - 0()).
Here the denominator of (i) uses b + c 2 b(1 6) > 0 and vx t 2 v2 Vlul when x 0, y 2
Now let's bound q4). If e > O, we have e = leel and loal
vcc vee
1 since c 2 b(1 8) >
If eg S 0, note that b+ e + ee > 0, we have leal < bb+ be v b( + 5). which yields 19a1 S v %cl
(1). Combining the above bounds give us 1 S [gal + 8|44 O6))
On the other hand, |/2| can bounded by
vb +
vb bb+c
122l S Vb+eb- veel+ +
58
0(6).
Vb(1 - 8) + c(1 -))
Then |R| < 1/ + 1221 O(0)
Lemma 35. Suppose X2, Xa are i.i.d Gaussian with mean 0 and variance o², then for < we have with probability at least -, 5,
(Gloogd -1lood-C22lop C2 log 1
max X? 2 0² Isid
for some C1 C2 > O.
Proof. It suffices to assume that 0² = 1 and prove that w.p. at least 1 - 6, maxigiid X? 2 C₁ log - C2 log log }.
First, by the lower bound of Gaussian tail, there exists a, B > 0 such that P(|X,|| > x) 21P > r) 2 ae- 9nn for x 2 0. Then by ii.d., we have
IPmmax Xi| S = ( sa)) ((xx
d
P(X xx) = (1- > x))"
[=1
$ (1 ae - Bx yd
< exp(_ddaa -5x²)
where the last inequality uses - < - for [0, 1}. Let exp( -dpe -8x2) 5, we get that w.p.
at least 1 - 8,
log(ad) log log 1 !)
max
X 2
Ist5
B
Then we have w.p. least 1 5,
2
[oggad)) - log log 1
max X? 1sid
=
max
XXi
>
1sdd
Lemma 36. Suppose X1, X2, NA are i.i.d Gaussian with mean 0 and variance 0², then we have with constant probability,
d
5 logd)
X
79

---

Proof. It suffices to assume that x² and prove that with constant probability, E 1 O (log d).
S
X
Then we know that with probability at least - θ (d 1), [ 2 for some > Then by union bound, with constant probability, Vi E [4] : [X,] 2 -
Consider X; for some fixed i. Since X, ~ [0 1), we have P([X,) =) Σ
2r V2r"
Now we split the interval 12 1] into several subintervals I = {i : X E [2-k-1,2-4) for Let ps = | E
k = 0, 1, ...., flog2 g
we know that [Tg] ~ Binomial (d, Pk) and pk - C1 2 k- Then by the concentration of binomial variables, we have w.p. least d
for > 0, [Ixl = O (dpk + /dpa log d log d) = O (d-2-k-1- + d 2 log d + log d)
Then we have
Σ
|2 O (d Vd. 2 log d + 2 logd) k = = [og 7
(d
REZ
Therefore, with constant probability,
d
[log, 1
X
Xi
X
k=
TE
|x,>1
[oga -1-1 Σ
d + dd22" +1 log d + 2 logd)
k=0
d dloge d ))
+ o (V2) (1+c2) + 2 (ogs 1ltl log d + d
o
= O (dlog d),
which means with constant probability, 1 [X = =0 (logd).
80


Reviews:
[{"review_id": "cDCNz5Guzz", "reviewer": "Reviewer_PbqS", "summary": "The paper provides a new explanation of why adaptive gradient methods perform better than SGD with momentum though extensive experiments and theoretical analysis. The key insight is that adaptive gradient methods especially Adam bias toward solutions with uniform Hessian diagonal values, and this property may contribute to faster convergence. Theoretical analysis is conducted on a simple setting with 2 layer neural network, it is proven that Adam guarantees to converge to solution with more uniform Hessian diagonal values than SGD.", "questions": "Could the authors provide more intuition on how uniformity in Hessian diagonals can induce faster optimization?", "limitations": "Limitations are adequately discussed in the paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The insight that Adam bias toward solutions with uniform Hessian diagonal values is an interesting and novel observation and this measure is easy to compute compared with Hessian singular value-based metrics. The observation is supplemented with comprehensive experiment results and an analysis in simplified setting of neural network, which make the work pretty solid in both theory and empirical sides. ", "weaknesses": "To my understanding, the weaknesses have two folds:\n\nThe experiments in the paper mainly focus on settings where Adam outperforms SGD, in which more uniformity in Hessian diagonals is observed with Adam. However, it lacks a comparison in settings where SGD outperforms Adam. It remains a question whether uniformity is still a good measure of performance in this case. In other words, would SGD iterates observe more uniformity in Hessian diagonals instead?\n\nIn the explanation of the paper, Adam biases towards solutions with more uniformity in Hessian diagonals, and this contributes to faster convergence. Thus, Adam converges faster than SGD. The bias in Adam has been well experimented with and studied. However, it is not very clear why uniformity in Hessian diagonals implies faster convergence. The paper does not provide a theory on it, and the experimental support of this claim is not strong. Additionally, intuitively, it is hard to connect uniformity with faster convergence since regions with more uniformity are allowed to have either large curvature or small curvature. The folklore intuition on a flat region encourages better optimization and generalization does not seem to hold for this measure.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and answer their questions below.\n\n1. **Comparison in the opposite setting.** In Appendix A.8, we have empirical results on image tasks, which provide the comparison in the opposite setting. As is shown in A.8, on image tasks, Adam does not converge faster than SGD+M and in the meantime, $R^{\\text{Adam}}\\_{\\text{med}}$ values are no longer smaller than $R^{\\text{SGDM}}\\_{\\text{med}}$ during training. This reveals the connection between the local diagonal geometry and the convergence speed from another perspective. That is, when the diagonal of Hessian of Adam is not more uniform than SGD+M, its convergence speed is not better, either.\n\n2. **Contribution of uniformity in Hessian diagonals to fast optimization.** We actually discuss and demonstrate the contribution of small $R^{\\text{OPT}}\\_{\\text{med}}$ (more uniform diagonal of Hessian) to fast optimization. Please see the second and third paragraphs of Section 4.3 and Appendix B for more details.\n\nRoughly speaking, to rule out the possibility that small $R^{\\text{OPT}}\\_{\\text{med}}$ is just a byproduct of adaptive methods and is unrelated to fast optimization, in Appendix B.1, we add a supplementary experiment similar to that in Figure 1. We select two iterates $x_1$ and $x_2$ from two trajectories that both come from SGD+M (instead of one from Adam and one from SGD+M in Figure 1), such that the loss $f(x_1)=f(x_2)$ but $x_2$ has a smaller $R^{\\text{OPT}}\\_{\\text{med}}$ than $x_1$. We then run SGD+M with the same configuration twice, once from $x_1$ and once from $x_2$. Under this setting, we get a similar observation: running SGD+M from $x_2$ (with smaller $R^{\\text{OPT}}\\_{\\text{med}}$) achieves faster convergence than from $x_1$. This suggests that the uniformity of the diagonal of loss Hessian (measured by $R^{\\text{OPT}}\\_{\\text{med}}$) reveals some intrinsic trajectory property beyond the algorithm choice and is indeed a contributing factor to fast optimization. In Appendix B.2, we theoretically prove the contribution of small $R^{\\text{OPT}}_{\\text{med}}$ to fast optimization in a simplified setting.\n\n3. **More about uniformity and curvature.** The uniformity measured by our statistic $R^{\\text{OPT}}_{\\text{med}}$ can be viewed as some variant of the diagonal condition number. We want to emphasize that the condition number is also a relative measure instead of an absolute measure of large or small curvature and is well-known to have a direct connection to the speed of optimization. As an analogy of the condition number, our statistic $R^{\\text{OPT}}\\_{\\text{med}}$ also has an intuitive connection with the optimization speed. In the theoretical analysis in Appendix B.2 on a simplified setting, we can see this connection more clearly."}, {"review_id": "Sh1VJBVC2q", "reviewer": "Reviewer_pmE6", "summary": "This paper aims to study the interaction between optimizers and the local loss landscape. The authors introduce the notion of $R_{\\mathrm{med}}^{\\mathrm{OPT}}$ to characterize the uniformity of the Hessian diagonal. Through extensive experiments, they show that, compared with SGD,  Adam biases the trajectory towards regions with higher Hessian diagonal uniformity, which, they argue, contributes to faster optimization. As a theoretical support, the authors demonstrate that Adam indeed possesses this property in a two-layer linear network setting.", "questions": "Regarding the comparison of different training runs, specifically run 1 (using Adam throughout), run 2 (using Adam initially and then switching to SGDM halfway), and run 3 (using SGDM throughout), Figure 1 illustrates that run 2 achieves a considerably lower loss level compared to run 3. However, I am interested in the gap in the final loss between run 1 and run 2.\n\n If this gap is acceptable, it raises a possibility: Could we design an algorithm that employs Adam during the initial phase and then seamlessly switches to SGDM halfway through the training process? Such an approach may offer the potential benefit of memory savings.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well-written, with extensive experiments and solid theoretical analysis.\n\n2. The interaction between optimizers and local loss landscape is an important and interesting topic. The findings of this paper may inspire future studies on developing a deeper understanding of the success of adaptive methods.", "weaknesses": "I was wondering whether $R_{\\mathrm{med}}^{\\mathrm{OPT}}$ can characterize the degree of the loss landscape being ill-conditioned. While it is true that $R_{\\textrm{med}}^{\\textrm{OPT}}$ conveys the similar message as the condition number when the Hessian is dominated by the diagonal ($\\nabla^2 \\mathcal{L}(\\theta)\\approx\\mathrm{diag}(\\nabla^2 \\mathcal{L}(\\theta))$,  it remains unclear whether $\\nabla^2 \\mathcal{L}(\\theta)\\approx\\mathrm{diag}(\\nabla^2 \\mathcal{L}(\\theta))$ holds throughout the training trajectory.  Can the authors provide some empircial evidence?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and answer their questions below.\n\n1. **Diagonal approximation.** We actually add empirical evidence on the diagonal approximation in Appendix F. In Appendix F.1, we conduct experiments on a language transformer model and demonstrate that the trend of loss Hessian is to become more and more diagonal during training. In Appendix F.2, we give a rigorous theoretical analysis on a two-layer linear network for this trend.\n\n2. **Gap between run1 and run 2**. We add a figure in the global rebuttal to demonstrate that the gap between run 1 and run 2 is not acceptable. The convergence speed of run 2 (using Adam initially and then switching to SGDM halfway) is slower than that of run 1 (using Adam throughout). A possible reason is that after switching to SGDM halfway, the trajectory geometry becomes worse and worse, which harms fast optimization."}, {"review_id": "A3Epei06gQ", "reviewer": "Reviewer_cQDf", "summary": "This study makes the comparison between the adaptive optimization method (Adam) and the non-adaptive one (stochastic gradient method with momentum) through a lens of ``uniformity'' of diagonal components of the Hessian denoted by $R_{\\rm med}^{\\rm OPT}(t)$. The comparison is conducted experimentally for various deep learning models such as BERT-small and theoretically for the two-layer linear neural networks under several assumptions. Both comparisons conclude that the uniformity will be smaller when using the Adam, which results in faster convergence.", "questions": "- Although the initialization order required in Theorem 1 seems to be very small, can the authors import the theoretical result or its insights into the practical settings? If not, what is the difficulty?\n- The setting in Section 5 requires $d_1=d_0=d$, but can the overparameterized settings (i.e., $d_1\\ge d_0$ holds) be treated in the same way?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The experimental comparison is conducted carefully.\n- The theoretical results have a novelty, which gives new insight into analyzing the implicit bias of various optimization methods.", "weaknesses": "- I couldn't fully understand the motivation for employing $R_{\\rm med}^{\\rm OPT}(t)=\\frac{{\\rm max}\\{|H_{ii}^{(t)}|\\}}{{\\rm median}\\{|H_{ii}^{(t)}}|\\}$ as for the notion of uniformity. I realize this value is more stable than the condition number; as explained in footnote 1, many types of variants could be considered. Indeed, in the experiment, there are some cases where $\\frac{R_{\\rm med}^{\\rm SGD}(t)}{R_{\\rm med}^{\\rm Adam}(t)}$ is smaller than 1.\n- In the theoretical comparison (Theorem 1), there are several unclear points to me. First, the training dynamics will change by the selection of hyperparameter and affects to $R_{\\rm med}^{\\rm OPT}(t)$, but some parameters (such as $\\alpha$, $\\sigma$) can be taken in a different order. Will it be a fair comparison? Moreover, I think it will be more helpful for readers if the authors clarify the dependence of $T_{\\rm OPT,(1,2)}$ on other parameters.\n- (minor) The reference to Figure 4. seems wrong. The authors write ``see Section 4.1'' in the caption, but it is explained in Section 6.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and answer their questions below.\n\n1. **About our metric $R^{\\text{OPT}}_{\\text{med}}(t)$.** First, we did consider another type of variant, which is a singular value-based metric. More discussions and experiments can be found in the paragraph starting from Line 79 and Appendix A.9. The short conclusion is that when measured by that singular value-based metric, our observation still holds: the local geometry obtained by Adam is more uniform than that obtained by SGD+M.\n\nSecond, we want to emphasize that although in some cases, $\\frac{R^{\\text{SGDM}}\\_{\\text{med}}(t)}{R^{\\text{Adam}}\\_{\\text{med}}(t)}$ are smaller than 1, this only happens on a small number of layers. For most layers, $\\frac{R^{\\text{SGDM}}\\_{\\text{med}}(t)}{R^{\\text{Adam}}\\_{\\text{med}}(t)}$ are larger than 1.\n\n2. **Selection of hyperparameters.** Theorem 1 holds for hyperparameters (such as $\\alpha,\\sigma$) in certain ranges instead of just particular values. The ranges of SGD+M and Adam overlap with each other. That means we can choose the same hyperparameters for SGD+M and Adam in the overlapped region to make a fair comparison, for example, the same $\\alpha,\\sigma$ such that $\\alpha\\ge 4(p+2)$ and $\\sigma\\le\\min\\left(\\frac{\\eta^{3/2}}{d^{\\alpha/2+1}},\\frac{\\eta^{3/2}\\xi^2}{d^{13/4}}\\right)$.\n\nWe provide the dependence of $T_{\\text{OPT},(1,2)}$ on other parameters in the detailed proofs in Appendix D and E. For example, we prove that $T_{\\text{Adam,2}}-T_{\\text{Adam,1}}=\\Theta(\\frac{1}{\\eta\\sqrt{d}})$ and provide an upper bound for SGD+M: $T_{\\text{SGDM,2}}-T_{\\text{SGDM,1}}\\le\\mathcal{O}(\\frac{d^\\alpha\\log(d/\\epsilon)}{\\eta})$. However, due to the page limit, we didn't write these details in the main text.\n\n3. **About initialization.** Empirically, our observation still holds under standard initialization. That means the small initialization assumption in our theoretical analysis is not essential. Below are the empirical results on a 2-layer linear network under standard initialization. We trained the models for 20 epochs to convergence and tuned and chose the best learning rates of SGD+M and Adam.  As we can see, our empirical observation still holds, i.e. $R^{\\text{Adam}}\\_{\\text{med}}(t)<R^{\\text{SGDM}}_{\\text{med}}(t)$.\n|   |  Epoch0|Epoch0|Epoch5|Epoch5|Epoch10|Epoch10| Epoch20|Epoch20|\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| Layer # | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ |$R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$|\n|1|57.18 |57.18|59.85|35.37 |61.20 |38.16|75.10|59.13|\n|2|1.54|1.54|3.68|1.58|3.33|1.60|3.97|1.86|\n\nActually, the small initialization assumption is only for technical reasons to make the proof of low-rank structures easier. Besides, it is a common choice for theoretical work to make the initialization close to 0.\n\n4. **Overparameterized case.** The answer is yes. The assumption $d_0=d_1=d$ is to make the notation simple and is not essential. For the overparameterized case $d_1\\ge d_0$, the bounds in our theorem will have a more complicated dependence on $d_1,d_0$ (for example, $\\frac{1}{poly(d)}$ becomes $\\frac{1}{poly(d_1,d_0)}$) but the main message will not change, i.e. $R^{\\text{Adam}}\\_{\\text{med}}(t)<R^{\\text{SGDM}}_{\\text{med}}(t)$.\n\n5. **Reference to Figure 4.** Sorry for the confusion. In this caption, we write \"Singular values and $R_u$ of the weight matrix in the 27-th layer on the translation task (see Section 4.1)\". Although Figure 4 is explained in Section 6, actually the note \"see Section 4.1\" is in terms of the translation task because the setup of the translation task is in Section 4.1. Thanks for pointing out this confusion. We will edit the caption to clarify this."}, {"review_id": "IQMUW89HUf", "reviewer": "Reviewer_8Quz", "summary": "This paper aims to study the connections between an optimization algorithm and the geometric properties observed during the training of a neural network with that algorithm. The authors argue that when analyzing these connections, it is important to consider the local iterates. To address this, they propose the utilization of a metric called $R^{OPT}_{med}$, which involves the local condition number of the Hessian calculated at the iterates. The authors conduct experiments to illustrate their findings that fast convergence is often associated with low statistics, and adaptive methods such as Adam often has low statistics. They proves that for some small theoretical setting, Adam has better statistic than SGD with some probability. ", "questions": "Please see above.\n\n---\nI thank the authors for your rebuttal. My recommendation remains the same. ", "limitations": "NA", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Studying the influence of algorithms on the loss geometry is interesting. This paper provides some insights on the uniformity of diagonal geometry by the $R^{OPT}_{med}$ metric where the values found by Adam are smaller than those found by SGDM in the empirical results. \n\nThe theoretical setting shows that the diagonal of loss Hessian for Adam has good uniformity while the diagonal of loss Hessian for SGD+M is less uniform.\n\n", "weaknesses": "While there may exist a correlation between the proposed statistic and the optimization algorithm, the reasons behind the favorable behavior of these algorithms remain uncertain. Is it better due to the uniformity or is the diagonal uniformity a byproduct of Adam's inherent algorithmic properties that are unrelated to its overall effectiveness? As a result, it is more important to demonstrate the contribution of the proposed statistic to achieving successful convergence rather than focusing on the relationship between Adam and the the diagonal uniformity.\n\nWhy the large batch setting are needed to prove the correlation of Adam and diagonal uniformity effectiveness? It is not clear how the algorithms behave in other settings and how high the probability can be in the theoretical analysis. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and answer their questions below.\n\n1. **Contribution of our statistic to fast optimization.** We actually discuss and demonstrate the contribution of small $R^{\\text{OPT}}_{\\text{med}}$ to fast optimization. Please see the first paragraph on Page 7 and Appendix B for more details.\n\nRoughly speaking, to rule out the possibility that small $R^{\\text{OPT}}\\_{\\text{med}}$ is just a byproduct of adaptive methods and is unrelated to fast optimization, in Appendix B.1, we add a supplementary experiment similar to that in Figure 1. We select two iterates $x_1$ and $x_2$ from two trajectories that both come from SGD+M (instead of one from Adam and one from SGD+M in Figure 1), such that the loss $f(x_1)=f(x_2)$ but $x_2$ has a smaller $R^{\\text{OPT}}\\_{\\text{med}}$ than $x_1$. We then run SGD+M with the same configuration twice, once from $x_1$ and once from $x_2$. Under this setting, we get a similar observation: running SGD+M from $x_2$ (with smaller $R^{\\text{OPT}}\\_{\\text{med}}$) achieves faster convergence than from $x_1$. This suggests that the uniformity of the diagonal of loss Hessian (measured by $R^{\\text{OPT}}\\_{\\text{med}}$) reveals some intrinsic trajectory property beyond the algorithm choice and is indeed a contributing factor to fast optimization. In Appendix B.2, we theoretically prove the contribution of small $R^{\\text{OPT}}_{\\text{med}}$ to fast optimization in a simplified setting.\n\n2. **The large batch assumption.** Our proof relies on a very elegant analysis and tracks detailed dynamical properties of weight matrices during training, which is a big technical challenge. To overcome the difficulty and simplify the proof, we require large batches to reduce the stochastic variances.\nMoreover, we want to emphasize that people in practice use very large batches when training language models, especially when using multiple GPUs."}, {"review_id": "mSgkmi7dm2", "reviewer": "Reviewer_qhhF", "summary": "This paper aims to understand why Adam works better than SGD for language model training from the perspective of local geometry of the training loss around the algorithm iterates.", "questions": "See above", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The problem that this paper considers is indeed a very important question, and this work sets out to answer the question from first principles.\nThis paper has many good intuitions. This paper also has some partial theory that backs up their empirical observation. Given their empirical observation is quite hard to theoretically analyze, I think their theory is quite valuable although it only shows a weak result for a very specific model.", "weaknesses": "- Many plots have too small fonts. Please adjust the font sizes for better readability.\n- In general, the plots have legends and axes that are quite hard to read, please adjust them.\n\n- In Figure 8, how can the spectrum on the right plot approximately low-rank? All singular values are greater than 0.5 and compared to the left plot, I wouldn't consider this approximate low rank. Please explain\n\n- Regarding Section A.6, the plots seem a bit inconclusive. Given how noisy the plots are, I don't see this as alignment. Do you have any intuitions behind why the layer gradients align with diagonal of Hessian?\n\n- Are the results established in Theorem 1 empirically tight? I think the experiments on this very simple setting to see if your empirical observations hold true for this very simple setting is required. In particular, for this two layer MLP, do you see lower R-value for Adam than SGD-M?\n\n- In Figure 9, I see that the observation doesn't hold for the right singular vectors.. This makes the observation quite brittle in my opinion. Could you explain why it doesn't hold for the right singular vectors?\n\n- In general, this paper lacks empirical investigation on smaller/toy models. I think to what extent this observation holds for simpler models would help readers comprehend how universal this observed phenomenon is.\n\n- Given that the language model training is done mostly with the cross entropy loss, I'm curious whether the conclusion you made in Theorem 1 similarly holds true for the logistic loss?\n\n- In the Gaussian initialization, I see the variances of the different layers are chosen quite carefully. Also the variance looks quite small. Could you justify this assumption as to why making these assumptions is practical. Also, could you explain intuition behind why these assumptions are required for the proof?\n\n- In the Theorem statement, how big are $T_1^{sgd}, T_2^{sgd}$ compared to $T_1^{adam}, T_2^{adam}$? Are they comparable? Or as you observed in the empirical results, are $T_1^{sgd}, T_2^{sgd}$ provably larger than $T_1^{adam}, T_2^{adam}$?\n\n- Could you share some intuitions as to why training behavior of vision models (CIFAR 10, IMAGENET) is so different than that of language models? I don't find the explanation that it has to do with transformer architecture quite conclusive...\n\n- Minor comment: It would be helpful if the proof sketch of Theorem 1 appears in the main text. I think it has many good intuitions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and answer their questions below.\n\n1. **Fonts of plots**. Thanks for your advice. We will increase the fonts and adjust the legends and axes of the plots.\n\n2. **Low-rank structure.** Figure 8 is to depict a low-rank trend of the spectrum instead of a mathematically rigorous low-rank property. Actually for practical deep models, due to various noises, we don't expect it to have a very nice low-rank property.\n\n3. **Alignment in Section A.6.** Though not very significant, if we focus on *large* gradients, we can find that they align with *large* diagonals of Hessian. Intuitively, when the diagonal of Hessian is large (sharp region), we should use a small update. However, for SGD, the corresponding updates are usually large due to large gradients. This gives us some intuitive evidence of why a non-uniform diagonal Hessian will harm fast optimization for SGD. \n\nHowever, things are different for adaptive methods such as Adam. As is shown in the right columns of Figure 10-12, their *adaptive updates* do not align with the diagonal of Hessian and are more uniform. Actually, we don't want to demonstrate the alignment between *true gradients* and diagonals of Hessian. Instead, we want to emphasize the uniformity of Adam's *adaptive updates* and the misalignment between them and the diagonals of Hessian, as it provides intuition on the optimization advantage of adaptive methods. Finally, the figures in Section A.6 are just intuitive and we don't want to draw a conclusion. That's why we put them into an appendix.\n\n4. **Empirical results on the setting of Theorem 1 and other smaller/toy models.** Below are the empirical results on a 2-layer linear network. We trained the models for 20 epochs to convergence and chose the best learning rates of SGD+M and Adam.  As we can see, our empirical observations hold true for this very simple setting, i.e. $R^{\\text{Adam}}\\_{\\text{med}}(t)<R^{\\text{SGDM}}\\_{\\text{med}}(t)$.\n|   |  Epoch0|Epoch0|Epoch5|Epoch5|Epoch10|Epoch10| Epoch20|Epoch20|\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| Layer # | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ |$R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$|\n|1|57.18 |57.18|59.85|35.37 |61.20 |38.16|75.10|59.13|\n|2|1.54|1.54|3.68|1.58|3.33|1.60|3.97|1.86|\n\nMoreover, Appendix A.1 presents empirical results on a shallow transformer (8 layers), which is a simpler model than those in the main text. Our observation still holds for this shallow transformer. The above results reveal the universality of our observation on simple models.\n\n5. **Right singular vectors.** We planned to use a paragraph to present the behavior of right singular vectors and discuss possible explanations. However, due to an editing error, we forgot to put that paragraph into the appendix. Roughly speaking, we noticed that the right singular vectors do not have uniformity for Adam. We are not sure of the reason and one possible explanation is that for a weight matrix, its right singular vectors are closer to the input data than left singular vectors and more easily influenced by the data, therefore may not show uniformity.\n\n6. **Logistic loss.** The answer is yes. Below are the empirical results on the setting of Theorem 1 but using the logistic loss. We trained the models for 25 epochs to convergence and chose the best learning rates of SGD+M and Adam. As we can see, the conclusion similarly holds, i.e. $R^{\\text{Adam}}\\_{\\text{med}}(t)<R^{\\text{SGDM}}\\_{\\text{med}}(t)$.\n| |Epoch0|Epoch0 | Epoch5 | Epoch5 |Epoch15| Epoch15|Epoch25|Epoch25|\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| Layer # | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$ | $R^{\\text{SGDM}}_{\\text{med}}(t)$ | $R^{\\text{Adam}}_{\\text{med}}(t)$|\n|1|47.75| 47.75| 18.5|15.55|32.08|14.54|56.42|12.58|\n|2|2.06 |2.06| 3.57|1.74| 2.28|1.81|4.34|1.60|\n\n7. **About initialization.** Empirically, our observation still holds under standard initialization. The empirical results in the fourth point are conducted under standard initialization. That means the small initialization assumption in our theoretical analysis is not essential. Actually it is only for technical reasons to make the proof of low-rank structures easier. Besides, it is a common choice for theoretical work to make the initialization close to 0.\n\n8. **Comparison between $T\\_{\\text{Adam,2}}, T_{\\text{Adam,1}}$ and $T_{\\text{SGDM,2}}, T_{\\text{SGDM,1}}$.** In the detailed proofs in Appendix D and E, we prove that $T\\_{\\text{Adam,2}}-T_{\\text{Adam,1}}=\\Theta(\\frac{1}{\\eta\\sqrt{d}})$ and provide an upper bound for SGD+M: $T_{\\text{SGDM,2}}-T_{\\text{SGDM,1}}\\le\\tilde{\\mathcal{O}}(\\frac{d^\\alpha}{\\eta})$. We can see that $T_{\\text{SGDM,2}}-T_{\\text{SGDM,1}}$ is larger than $T\\_{\\text{Adam,2}}-T_{\\text{Adam,1}}$, but actually, they are not comparable because the bound for SGD+M is just an upper bound.\n\n9. **Results on image tasks.** As is discussed in A.8, the behavior on image tasks seems different, but the underlying correlation between $R^{\\text{OPT}}\\_{\\text{med}}$ and optimization speed is actually consistent with that on language tasks. As is shown in A.8, on image tasks, Adam does not converge faster than SGD+M and in the meantime, $R^{\\text{Adam}}_{\\text{med}}$ values are no longer smaller than $R^{\\text{SGDM}}\\_{\\text{med}}$ during training. This reveals the connection between the local diagonal geometry and the convergence speed from another perspective. That is, when the diagonal of Hessian of Adam is not more uniform than SGD+M, its convergence speed is not better, either."}]

Author rebuttal:
The attached figure is to address a question raised by Reviewer pmE6.

Meta-review: This paper analyzes how adaptive gradient algorithms affect the local geometry of the neural network learning landscape. The key contribution of this paper is to introduce a statistic that characterizes the "condition number" of the Hessian (which is defined as the ratio of the maximum diagonal value and the median diagonal value of the Hessian matrix) and shows empirically that this new statistic is related to the optimization speed. The authors further prove that in a two-layer linear network, Adam and SGD training trajectories have different hessian uniformity behaviors during the training, which may explain why Adam is faster than SGD in neural network training on NLP tasks. 

One of the reviewers criticized that it is unclear whether the proposed statistic is really related to fast optimization. Although the authors do not provide a theoretical answer for this, they have provided experiments in the Appendix to show that they seem to be related. The large batch assumption is indeed a bit unrealistic in practice, but I believe the new trajectory-based analysis will be appreciated by the NeurIPS community and it might inspire future work studying this.

Therefore I recommend accepting this paper.
Decision: Accept (poster)