How Does Adaptive Optimization Impact Local Neural Network Geometry?
Authors: Kaiqi Jiang
Dhruv Malik
Yuanzhi Li
Venue: Submitted to ICLR 2023 (2023)
Keywords: optimization
adaptive algorithms
neural networks

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, 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 convergence. By contrast, vanilla gradient methods like SGD bias 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.

Under review as a conference paper at ICLR 2023
How DOES ADAPTIVE OPTIMIZATION IMPACT LOCAL NEURAL NETWORK GEOMETRY?
Anonymous authors Paper under double-blind review
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, trajectories produced by running a generic optimization algorithm OPT, we intro- duce R a statistic that is analogous to the condition number of the loss Hessian evaluated at the iterates. Through extensive experiments, we show that adaptive methods such Adam bias the trajectories towards regions where Radam is small, where one might expect faster convergence. By contrast, vanilla gradient methods like SGD bias the trajectories towards regions where ASCD is comparatively large. We complement these empirical observations with a theoretical result that provably demonstrates this phenomenon the simplified setting of a two-layer linear net- work. We view our findings evidence for the need of new explanation of the success of adaptive methods, one that is different than the conventional wisdom.
advocate for local trajectory analysis. For iterate
INTRODUCTION
The efficient minimization of parameterized loss function is a core primitive in statistics, optimiza- tion and machine learning. Gradient descent (GD), which iteratively updates parameter vector with a step along the gradient of the loss function evaluated at that vector, a simple yet canonical algorithm which has been applied to 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 non-smooth. The optimization landscape of such problems is thus extremely challenging, and in these settings gradient descent often suffers from prohibitively high iteration complexity.
To deal with these difficulties and improve optimization efficiency, practitioners in recent years have developed many variants of GD. One prominent class of these GI) variants is the family of adaptive algorithms (Duchi et al. 2011 Tieleman et al.) 2012) Kingma & Ba, 2015). 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 optimize deep neural networks, especially on Natural Language Processing (NLP) tasks with transformers (Vaswani et al. 2017 Devlin et al. 2019).
From a theoretical perspective, adaptive methods are well understood in the traditional context of convex optimization. For instance, Duchi et al. (201I) show that when the loss function is convex, then the Adagrad algorithm yields regret guarantees that are provably 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 sparsity or a coordinate wise bounded Lipschitz constant), and the algorithm adapts to this global geometry by adaptively selecting learning rates for features that are more informative.

---

Under review as a conference paper at ICLR 2023
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_d.. Hence it only makes sense to study convergence by looking at the local geometry of the loss along the trajectory of the optimization algorithm (Arora al. 2018).
82-
Figure 1: (left) Training losses of SGD+M starting from TSGD and Txdan- (right) The 10th largest value over median in the diagonal of loss Hessian (which can be viewed as a variant of med (t) defined in eq. (I) for Adam and SGD+MM. Since the full Hessian is too big, here we selected several layers and randomly sampled 200 coordinates per layer compute.
300
2005
150
20R
1540
Figure 2: Training losses of Adam and SGD+M on the sentence classifi- cation task descrihed in Section 4.1
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 hyperparam- eters (Lewkowyez al., 2020 Cohen et al., 2021). For instance, [ohhen et al. (2021) 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 2/(step size). The viewpoint from convex optimization, where a loss function has some (potentially) non-uniform but fixed underlying 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 ZAdam from the Adam trajectory and iterate ZSGD from the SGD trajectory, such that (1rAdam) = (ISGD) We then run SGD+M twice, once from & Adam and once from ZSGD If the underlying geometry of the loss function was truly fixed, then would not expect significant difference in the performance of running SGD+M from either of the two iterates. However, as shown in Figure Figure[[left), running SGD+M from AAdam achieves lower loss than that from TSGD, suggesting that Adam may bias the 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 a trajectory produced by running an optimization algorithm OPT, where the Hessian of the tth iterate is given by H+) Rdxd, we define the second order statistic ROPT (t) in the following fashion. For the fth iterate in the trajectory, let ROPT (t) be the ratio of maximum of the absolute entries of the diagonal of (1),, to the median of the absolute entries of the diagonal of ( Concretely, we define
max max(111)40 (11
ROPT
median
is 1
This statistic thus measures the uniformity of the diagonal of Hessian, where a smaller value of ROPT (t) implies that the Hessian has a more uniform diagonal. It can also be viewed as a stable variant of
'Consider the case where one parameter has little impact on the loss, then the second derivative w.r.t. this parameter almost zero, making max[1H] (1dd)
infinity. So we consider median which more stable.
(in( 1
2

---

Under review as a conference paper at ICLR 2023
the condition number. Instead of eigenvalues, choose diagonal entries because adaptive methods used in practice are coordinate-wise, which can be viewed as the diagonal scaling approaches In Appendix A.9 we discuss this intuition in detail and compare Rmed(t) (t) with singular value- based metrics. As a supplementary result, in Appendix [ 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 (Jacot et al., 2018) etc. could yield faster convergence rate for (S)GD (Liu et al. 2022). As for (diagonal) adaptive methods (e.g. Adagrad), they were original designed to adapt to the nonuniform diagonal geometry. Intuitively, a smaller Rome (t), which implies more uniform diagonal geometry, could lead to faster convergence.
Armed with this statistic, we make the following contributions:
On a wide variety of neural network transformer architectures and language modeling datasets, we conduct experiments to compare how Amed (t) and Rngg (t) 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 Radaa (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 [right). This suggests that relative SGD+M, Adam biases the optimization trajectory a region where the Hessian diagonal is more uniform. We call this phenomenon the uniformity of diagonal geometry for adaptive methods. As an aside, we observe that larger improvements in optimization performance of Adam over SGI++M are correlated with larger gaps between Remea (t) and ReGPPM)) (t). This suggests that region where the Hessian diagonal is more uniform is also a region that is more amenable to rapid optimization.
We complement our empirical results with a theoretical analysis of this phenomenon in the simplified setting of large batch Adam and SGD+M, on two-layer linear network with d-dimensional input and hidden layer, and one dimensional output. We show that for a wide range of t, RAdim (t) + o(1) but R (t) = 2llog d). Our proof reveals that Adam induces the weight matrices to have low rank whose leading singular vectors have certain type of uniformity (see Section 6 for discussion), fact that we also observe empirically in large scale neural networks, suggesting that this may be a mechanism by which adaptive methods bias trajectories to have uniformity of diagonal geometry.
RELATED WORK
Existing analyses of adaptive methods. The vast majority of prior theoretical work on adaptive methods has focused on the blackbox setting (Duchi et al. 2011 Kingma & Ba 2015 Chen et al. 2020) Reddi et al., 2018 Ward et al., 2020 Défossez et 2020 Ene et al. 2021). )). 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 (governed by parameters such as the smoothness parameter) analysis is 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.
assumed to hold over the entire domain. Hence this style of
Existing analyses of (S)GD on neural networks. There an extensive literature on the analysis of GD/SGD in the non-blackbox setting, e.g, overparameterized neural networks, (Du et al.) 2018 Ji & Telgarsky 2020 Allen-Zhu et al., 2019a b) Arora et al., [211a]] Liu et al. 2022). However, is unclear how to translate these analyses of GD/SGD, 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, e.g. linear or logistic regression and the Neural Tangent Kernel (Jacot et af., 2018), the loss geometry usually fixed and not influenced by learning algorithms. However, in neural networks the interaction between
Recall that the main theoretical bound in the original Adagrad paper Duchi et al. 2011 is in terms of the diagonal scaling.
3

---

Under review as a conference paper at ICLR 2023
algorithms and loss landscapes is more complicated. Lewkowycz et al. (2020) find a so-called catapult effect of initial learning rate on the training trajectory of SGD and related loss curvature. [Cohen et al. (021)) 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 et al. (202T) leave open the problem of whether similar interactive phenomena occur in algorithms that are not GD, including adaptive methods.
3
OVERVIEW OF RESULTS AND SETUP
3.1 ISSUES OF PRIOR ANALYSES ON ADAPTIVE METHODSS
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 | To see this, we initialized transformer mode with default initialization in Pytorch but chose a large gainf and computed the smoothness parameter (denoted as 1) and the condition number (denoted as k) of loss Hessian on one layer. We observed that setting the gain as a large constant (e.g. 800) results in extremely large and x (1 2 10" and A 2 10"H) 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, is unclear that when two optimization algorithms are used, they will have the same geometry in local trajectory. In particular, although in theory, adaptive algorithms can yield to a 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 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 that of SGD (with momentum)?
3.2 OVERVIEW OF THE EXPERIMENTS
As is discussed in Section 1 we consider (() (t) defined in eq. ①I as measurement of the uniformity of the diagonal of the loss Hessian. We conduct experiments on different NLP tasks to examine pOrT (t), as in language models, adaptive methods have shown significantly faster convergence than SGD (with momentum). The details of these experiments will be shown in Section 4 To explore potential different patterns of different layers, do the computation layer by layer. On wide variety of transformer architectures and language modeling datasets from the same initialization, we observe that:
When we train the neural network using Adam, the uniformity of diagonal geometry, mea- sured by POPT (t) smaller than that when we train using SGD+M from the same initializa- tion, except for first several layers.
Table shows a typical example of R (t) compared to ASGDM (t) on sentence classification
task using BERT-small (Turc et al. 2019; Bhargava et al. 2021) (See Section 4.1 [.1 for details). We
repeated the experiments for 12 times starting from the same initialization. Table D shows the
averaged RAdam (t) and RSGDM (t) in some randomly selected layers (except for the first several). We
RSGDx (1)
also report the averaged
and their standard deviations in the brackets
Figure 2
shows the
corresponding training losses of one in these 12 experiments.
s httpre//errraanoooeeettttorrrrrr htmm
"This refers the gain parameter in some commonly used initialization functions of Pytorch, e.g- torch.nn.ing www.tiiavieerunifii.i >PSCDM (t) values in Table i for most layers roughly 1.4 to 2 times RAAm"()) in corresponding layers. In practice, it can be considered sigmificant because might imply 1.4 to 2 times faster convergence.
4

---

Under review as a conference paper at ICLR 2023
Table
Renean (t) and Ro
(t) some layers, on the sentence classification task (see Section 4.1
Layon
bemiion
beratirr 750
Betation 1250
(())
"vel
(1
R0me1 (1)
15.7
1.45 (0.3)
14:34
(0.40)
22.63
15
9.35
9.35
80.57
53.52
145 (1,6)
10005
63.80
2401 (1.00)
83.32
405.0
223.56
181 00.33
197.3
(0.8)
31.32
15.88
22
47.13
47.13
233.72
72.07
158.38
93.13
2,5 (1.18)
24
31.12
17.93
17.34
1:130000
16.33
(0.50)
To understand this phenomenon in a more principled point of view, we also provide a formal proof of the statement in a simplified setting: large batch Adam and SGD+M on a 2-layer linear network. Although simple, the choice of 2-layer linear network to understand learning dynamics is common in below describes the theoretical setup.
prior works (e.g. (Tian al. 2021)) Section 3.3
3.3 SETUP OF THE THEORETICAL ANALYSIS
Notation Let [d] {1,2, d}. We use || || 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 (p, o²) be the one-dimensional Gaussian distribution with mean μ and variance o². For scalar (vector, matrix) A which evolves over time, we use ( denote its value time I.
Let there be m data points. The data matrix is E assume that the input dataset is whitened, i.e. Aga= m2
IRd2 xm and the label matrix Y E Rd, xm We xdz is identity matrix.
_XX E Rd,xd=
The parameters of 2-Layer linear network are given by W == (W2, W1). Assume W3 E Rd for i == 1, 2. We have d2 = dy. do = dz. We consider the square loss L(W) = 2m 2 W, I
Denote A = 1 x E Rddxxd xdx Arora et al. (019b) show that with whitened dataset,
L(W) == 2m [WW YIl = L(W) +c, I(W): = 2 WW2, AA1
(2)
where e does not depend on W. We consider the following model with small Gaussian initialization. Assumption (Setup). The input covariance Agz := LXXT IRdaxdd an identity matrix. The input and hidden layers are both of dimension d, i.e. d1 du d. Without loss of generality, can assume that A is a row vector (i.e. d2 1) whose coordinates are positive andd ((1) in terms of d.
Assumption (Gaussian Initialization). Vi,j wiz NO ( NO0 daa) are independently initialized with sufficiently large 0.
Denote À and Âgæ the batch versions of A and Agz. We make the following large-batch batch assumption. We emphasize that large batches are commonly used in NLP tasks (e.g. (Brown et al. 2020)).
Assumption 3 (Large Batch). For the randomly selected batches, assume E[Å] A, E [AAz Agz-
Vi,j E [d] : E (A. A.) o² E [(Ann [i, - Aza [i. < o2, and の² = O poty
Denote g(t) the batch gradient at time t. The update rules of SGD+M and Adam are given by SGD+M: (++1) ()) =(()), w(++1) w(t) m))!)
Adam: 1h1 1|. / 1- 3/511 /(111)+1)) ml++1) (1m + (1 (1))((1)
(3)
(((+)) (t) + (1- (t) c g()) w(++)) = - =(()) T/- Ve( +
m
where v7 is the learning rate, 8,,3,, B2 momentum parameters, and & is for numerical stability. All operations on vectors are element-wise. Here and throughout, the notation f(x) Og(())
II Assumption 2 we assume Gaussian initialization. Due the rotational invariance of Gaussian distribution, we assume that all coordinates of are positive without loss of generality.
5

---

Under review as a conference paper at ICLR 2023
(resp. f(x) = ((g(=))) =(g(x)) means that there exist constants C1. 0 such that f(z) & (2g(z) (resp. f(z) Cıg(z), Cig(z) & f(z) S C2g(z)). We will also use the notation with ~. i.e. ( ( to hide factors that are logarithmic in d. In our theoretical analysis, "with high probability", "w.h.p." for short, means that with probability at least poly(d)
THE UNIFORMITY OF DIAGONAL GEOMETRY
As is mentioned in Section 3.2 we computed I (t) defined in eq. ①1 on different language models. In this section, we present the results of SGD+M and Adam on different architectures and datasets. In Appendix 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 (Ture et al. 2019 Bhargava al. 202T) on the IMDB dataset (Mass al., 201T): the task classify whether movie reviews are positive or negative." The momentum parameter 33 in SGD was set as 0.9. The two momentum parameters (31, (2) 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 Section 3.2 Figure and Table show the training
2
losses and the comparison between R2yddam (t) and (SCD (t). respectively.
Translation task We trained Seq2Seq network that uses Transformer to solve a machine translation task on Multi30k (Elliott et al.) (CC BY-NC-SA 4.0): this task is to train German to English translation model The momentum parameter in SGD was set as 0.9. The two momentum parameters (3,, 32) of Adam were set (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 lefft) shows the training losses for one among them. Table 2a shows the averaged Reddam (4), RSGDM (1) and
(t)
(with standard
Reed
deviation in the brackets) in some randomly selected layers.
Figure 3: Training losses of Adam and SGD+M for the translation task (left) Multi30k (see Section 4.1 and (right) data with randomly generated targets (see Scction
Figure 4: Singular values and R. of the weight matrix in the 27-th layer on the trans- lation task (sec Scction 4.1
4.2 EXPERIMENTS ON 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 on real targets, the model was trained

---

Under review as a conference paper at ICLR 2023
Table 2: Rend (t) and Rmmdl (t) in some layers for the translation task. (a) on Multi30k and (b) on data with randomly generated targets.
(b)
-
(00%))
Lapet
(45%))
38-07
tpxck
127
4.25
F:14
2:41
216/3.5
3:14
4.52
2.55
135
235903
3799
1.81
2.2933.5
129
206
3611
18:33
20012-40)
5233
16.99
546-3.96
1.64
8.64
4016
45.96
1.8 0111:
41:01
36.19
1.4121.00
576
$29
19
142
149
1788
1:5 90
3.99
5.00
46.7
352
1658 (2.0)
46.5
3.4
12:11.90
IK11
18.1
3500
2834
1231220
18
5/8
1818
1128
1128
317.44
35A1
531 8276
301.85
Tal 1160
12
111
111
135
098
(230)
927
506
158 (3.90
1929
1929
629
7540
103
136
1929
3600
30W
7631
6600
787.0
3699
2100 Hh
K00
1653
6.38
600
13812.30
583
34
36.22
36.2
465.5
633
60040
351
1843
19:07
199:1:00
19
1.0
120
34
1501
1500
188
133
169
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 3 right) shows the training losses for one among them. Table 2bjsowss the averaged Ra (t), RSGDNT and
med
Reaa (t) mad
(with standard deviation in the brackets) of the same 10 layers as in Table 2a
4.3 SUMMARIZATION OF THE EMPIRICAL RESULTS AND DISCUSSION
Overall, through extensive experiments on language models, we demonstrate that starting from the same initialization, the ( (L) values found by Adam are smaller than those found by SGD+M, except for the first several layers. This suggests that Adam is biased towards region with more uniform diagonal Hessian than SGD+M. In Appendix A. we also validate this observation on the in-distribution test data.
Positive correlation between uniformity of diagonal Hessian and fast convergence. We observe that on random dataset, SGD+M plateaus after about 400 steps and thus converges much slower when compared to Adam than on real dataset (see Figure 3}. On the other hand, the gaps of RSSDD(()) and Redaamm (t) are more significant on random data than real data (see Table 2bb as well. In Appendix A..4] we conduct another experiment where we switch from SGD to Adam 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 ROPT (t) are gradually closed after switching (see Figure 7 and Table S). Hence we find positive correlation between fast convergence and the uniformity of diagonal of loss Hessian, suggesting that a region with more uniform diagonal of Hessian is also a region that is more amenable to fast optimization. In Appendix [A] we study other adaptive algorithms (Adagrad, RMSprop and A 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 R (4), suggesting that the uniformity of diagonal Hessian might be a universal mechanism (partially) explaining the faster optimization of adaptive algorithms than SGD (with momentum).
More discussions on the trajectory difference. Considering the fact that our comparison between R (t) and RSGDM (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 than SGD+M. To rule out this possibility, in Appendix A.3 we add comparison of our measurements Reddam (1) and RSGDM(L') (t'), where t, t' are picked such that ith Adam iterate and t'th SGD+M iterate have the same training loss. The results (in Table 7 show that Radam (t) < RSGDM (t') for most layers, thus demonstrating that the trajectories of Adam and SGD+M are truly different and that the difference because Adam biases the local geometry (as opposed to faster convergence).
Adding regularization. People in practice usually add weight decay (equivalent to 12 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 9 suggest that in this case, the positive correlation between SROPT (t) and convergence speed still holds: Adam converges faster than SGD+M and in most of the layers except for the first several, RAddum (t) values
TTo prevent ROPT (t) from getting too large due to tiny median, we added an additional term 0.001 ( Dr the denominator of eq. (T) when computing.
7

---

Under review as a conference paper at ICLR 2023
smaller than A This reveals the robustness of our observation under weak regularization. However, under large weight decay parameters, we observed cases where Adam still converged faster but (e (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.
Image tasks. Although in this paper focus on language models where Adam shows significant fast convergence, we also add supplementary results in Appendix A.8 on image tasks where SGD+M performs better. On residual network trained on CIFAR-10, we observed that Adam did converge faster than SGD+M (see Figure 13b) and in the meantime, / (1) values were no longer smaller than R (t) during training (see Table 10). 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.
5 THEORETICAL ANALYSIS
In Section 4 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-dimensional output.
Since the weights and Hessians in different layers may have different magnitudes, compute the p (t) layer by layer. We denote RSGDM (t) (resp. A (t)) as the (OPT(() (t) found by SGD+M (resp. Adam) w.r.t W5 at time where k
2,
Theorem 1. Under Assumption 1 and 3 consider the weights { (with) } SGD 120
(resp. {Walom)) }
120
obtuined by SGD+M (resp. Adam) defined in 33.
1. For any p > O, pick 0 < < TF, n S O (Trafss)) and a 2 4(p + 2). Suppose S
1///2
then
/u/24I
there exists TsGD,1, TsaD,2 such that w.h.p., L (W(Too..)) (dd) L (w(Tsoo2))
< Ô and
SGD
11 E [TsGD, TsGD2)
Resen (1) 220lg d), 1,2.
For any p > 0, pick 7 S O (2 ) E
p+1 The a a 3 and Bg = .R. Suppose o S
n 3/2g2
Then IT'aaom 1. I Adomm such that w.h.p., L
( (Tam,))) =
) (d)) L ( (Tham.z) < ① 26 and
Adam
Aduu
Beden (t) 1t)=1= 00O+ 11 d
Vt [Tadam, 1 I [ddam,2]
1
1
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 L (w/lt) E Adam
[T TSGD,2] and t' e we have
[TAdom, 1, I AAdom,a) 2) such that L (W (Wit)) (GD)
( (1b) (2)) O(d)]
nn n di
1
RSCOM (t) = 12(log d),
Reldum (t') = (t)==+00
k = 1,2.
Theorem D and Corollary tell us that during long training period when the loss decreases from (d) to O (1)) the diagonal of loss Hessian for Adam keeps nice uniformity in the sense that for
each layer, its diagonal elements have roughly the same value, i.e. Rahhaam (4) = + o(1). k 1, 2. On the other hand, the diagonal of loss Hessian for SGD+M is less uniform. Appendix B giies
proof sketch of Theorem [] The detailed proof can be found in Appendix Cnd
THE LOW RANK STRUCTURE OF WEIGHT MATRICES AND UNIFORMITY OF LEADING SINGULAR VECTORS
The proof sketch in Appendix B highlights one crucial intuition of Theorem T After TSGD.1 (resp. T Addmm,1) steps, W1 of SGD+M (resp. Adam) becomes approximately rank-1 matrix. Consider
8

---

Under review as a conference paper at ICLR 2023
the left singular vector == [u1, U2. ..., Ud]² which corresponds to the leading singular value d We can show that the distribution of u². ui.....uf 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 [ for more details.
Similar low rank bias after training has been studied in prior works (e.g. (Gunasekar et al. 2017 Li et al. 2018 Chou et al.) 2020)). For more complicated models, we want to 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 E &Rmx, we want to check
(A) Whether W E Rm is approximately a rank k matrix with k < min m, n). and
(B) If (A) is true, then consider the top singular values 01, F% and corresponding left singular vectors 11, tug. Define a new vector ū = Σ of O u, == [ā], ū2,... ūd]² and compute
R,, = max, ū,,, which is a generalized version of max; "i in the rank 1 case. We want see whether median is,"
median uf
R2 obtained by Adam is smaller than that of SGD+M.
After reviewing the weight matrices we got different settings, we observed that (A) and (B) hold for many layers in those models. For example, on the translation task mentioned in Section 1.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 shows the result on one typical layer. Results of more layers can be found in Appendix A.5
Remarks 1.
The definition of R, is based on the connection between diagonal of loss Hes- sian and weight matrices. Appendix shows that for a 2-layer linear network, ØOPT (t) = max; [w("),311 When W1 E mx is approximately rank k. i.e. W1 ~ Ci=,, de-
median|| [ ( [6.3]112"
note [4:1, U42,..., timl 1T and = Pu, Us2,..., viml" we have that for the j-th row,
2
o By defining Ex ofu, Uu we have that |ibb, = ūy. Although in multi-layer nonlinear neural net- works, the connection between diagonal of loss Hessian and the weight matrices is more complicated and ROPT. (t) may depend on the product of many weight matrices rather than one single matrix, still believe that this definition of R., is reasonable ratio to consider.
:113
EL ofy
[ū1 u2 ..., úd
2
2. We may also want to consider the right singular vectors V1 U2. and corresponding = [ö1, U2, ..., TaJ"" = Σ o² @ V; and compute R2 == max, for Adam and SGD+M. However, on this translation task, among the 12 layers which were approximately low rank, for only 6 of them, R2 of Adam were smaller, i.e. we did not observe uniformity of the leading right singular vector for Adam. Results of Rg can be found in Appendix A One possible reason 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.
median 5.
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 in a simplified setting of two-layer linear networks. Although our findings may not directly lead 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 a method that performs better in the bad loss geometry, our findings suggest that we can design algorithms which implicitly avoid regions with bad geometry. There are a lot of future directions along this line. For example, our theoretical results on the two-layer linear networks may be able to generalize to multi-layer networks. In fact, people conjecture that the key-value-query structure in language models can be approximated by a three-layer linear network. Hence the generalization to 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 to relax our large- batch assumption (Assumption 3 and prove similar results in the general stochastic setting.

---

Under review as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
The training details (e.g. hyperparameters) of experiments are specified in Section The source code is provided in the supplemental material. For the theoretical results, Section 3.3 states the full set of assumptions. Section C and Section D in the appendix provide complete proofs.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems, volume 32, Curran Associates, Inc., 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp- 242-252. PMLR, 2019b.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameteriiaaton. In 35th International Conference on Machine Learning, ICML 2018, pp. 372-389. International Machine Learning Society (IMLS), 2018.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine- sine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332. PMLR, 2019a
Sanjeev Arora, Noah Golowich, Nadav Cohen, and Wei Hu. A convergence analysis of gradient de- scent for deep linear neural networks. In 7th International Conference on Learning Representations, ICLR 2019, 2019b.
Prajjwal Bhargava, Aleksandr Drozd, and Anna Rogers. Generalization in nli: Ways (not) to go beyond simple heuristics, 2021.
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, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.
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 (ed.). Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI 2020, pp. 2267-3273 ijeai. 2020. doi: (0.24963/ijcai.2000/44 URL. httpss://oi.. org 24963 Scai 2020/452
Hung- Hung-Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep ma- trix factorization: Dynamics and implicit bias towards low rank. arXiv preprint (XXi:2001.11372) 2020.
Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
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.
Jacob Devlin, Ming-Wes 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 Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume I (Long and Short Papers), pp. 4171-4186 Association for Computational Linguistics, 2019. doi: 18618/ v1/n19-1425 1423. URL https:: //doi http:/ooooorrrg///oooo wtt.:///...111111111//1111111110
10

---

Under review as a conference paper at ICLR 2023
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes pwer-parameterized neural networks. In International Conference on Learning Representations. 2018.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Desmond Elliott, Stella Frank, Khalil Sima' an, and Lucia Specia. Multi30k: Multilingual english- german image descriptions. In Proceedings of the 5th Workshop on Vision and Language, pp. 70-74. Association for Computational Linguistics, 2016. doi: 1..18663//1/W16-3222 URL http://www. ac web. org/
logo/ W16 6 3210
Alina Ene, Huy L. Nguyen, and Adrian Vladu. Adaptive gradient methods for constrained convex optimization and variational inequalities. Thirty-Fifth AAAI Conference on Artificial Intelligence,
AAAI 2021, Thirty-Third Conference Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, pp. 7314-7321. AAAI Press, 2021. URL https://o/) /ojs
htt./////aatttteeaaaaaaaaaaaeeae http://////aa//taaateee/////eeee aaai. baaii.og//ae index...ex... http:///nexxxpppp.e.eeeee
16898 16898
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
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.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 30,30, 2020. OpenReview.net, 2020. URL htt.tps:///ppeeeeeeeeeeeee
net /ene/fooruee forum? id=Hygeqyr
Kenji Kawaguchi. Deep learning without poor local minima. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, vol- ume 29. Curran Associates, Inc., 2016. URL httpss://rroceeeingg..eee neurips.........p....... 211199922266222222222222922 1eeeee222444555e http::eeccccccaa1aaa3332000000-0
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. htRLhttp::// /ttp///axxiii httpaaxxvvv.. /aarr///b////// 1112. 6980
Matthicu Lerasle. Lecture notes: Selected topics on robust statistical learning theory. arXiv preprint arXiv:1908. 10761, 2019.
Aitor Lewkowycz, 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.022 18, 2020.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over over--aramette parameterized matrix sensing and neural networks with quadratic activations. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet (eds.), Proceedings of the 31 Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pp. 2-47. PMLR, 06-09 Jul 2018.
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.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142- 150. Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL http:
7/ w/www. ac web. www..rrggantholoogo/oeeeeee
11

---

Under review as a conference paper at ICLR 2023
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th Interationall Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. | //openrevien ew net For.mmidudpss 72udxe)
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. URL https /openrevi ew net/ Ftroo//rmmmmrruuuf------
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 Proceedings of Machine Learning Research, pp. 10268-10278. PMLR, 18-24JJ Jul 2021.
Tijmen Tieleman, Geoffrey Hinton, al. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning. 4(2):26-31, -1,1 2012.
Tulia Ture, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv: /908. 008.68292 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017
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.
12

---

Under review as a conference paper at ICLR 2023
A MORE EXPERIMENTS OF THE UNIFORMITY OF DIAGONAL GEOMETRY
A.1 VANILLA SGD VS. ADAGRAD
In this section, we present the med (t) values defined in eq. T obtained by vanilla SGD and Adagrad on a language modeling task 10| The task is to assign a probability for the likelihood of given word (or sequence of words) to follow sequence of words. We trained transformer model to solve this problem on both Wikitext-2 (Merity et al., 2017 BY-SA 3.0) and random dataset (generating random integers as targets). This model has roughly layers (not counting normalization and dropout layers)
The setup is the same in Section 3.2
We used the same learning rate schedule (constant or 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, 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.
We tried different initialization (normal and uniform) by using different gains of the Pytorch initial- ization schedule.
A.1.1 EXPERIMENTS ON REAL DATASET
Figure 5a shows the training losses on real dataset (wikitext-2). Table 3 (resp. Table 4 shows the ROPT (t) for Adagrad and SGD under uniform (resp. normal) initialization with different gains.
30.7
300
100
303
4.0
1990
3003
(a) Wikitext-2
(b) Random dataset
Figure 5: Training losses of Adagrad and SGD on wikitext-2 (left) and random data (right)
Table 3: ROPT (1) of Adagrad and SGD under uniform initialization with different gains
(a) Gain =
(b) Gain = 0.5
Epoch
Epoch 20
Epoch 40
Fpoch
Epoch 20
Epoch 40
Layer#
SGD
Adugrad
SGD
Adugrad
SGD
Adagrad
Layera
SGD
Adagrad
SGD
Alagrad
SGD
Adagrad
6.07
6.77
5.91
9.77
5.16
10.37
69.36
78.60
15.26
7.74
18.22
7.23
4.60
6.26
3.43
1.66
3.44
1.88
2
24.12
24.36
4.05
2.30
3.70
2.04
5.15
4.35
4.84
4.98
9.47
10.78
9.76
3,54
8.67
3.14
4
5.25
4.74
3.83
5.68
3.11
4.81
12.54
13.96
10.31
9.79
66.49
67.83
88.75
19.31
63.01
15.64
4.92
5.25
7.21
2.33
7.94
2.28
6
6.54
6.91
3.57
2.08
3.50
1.97
5.73
5.45
40.56
4.57
21.24
4.76
7
3.22
3.73
13.03
3.97
9.55
4.07
9.39
8.87
37.95
4.50
46.03
3.19
8
6.12
5.99
6.73
7.82
5.43
6.98
13

---

Under review as a conference paper at ICLR 2023
Table Rmmd (f) of Adagrad and SGD under normal initialization with different gains
(a) Gain =
(b) Gain 0.5
Layerh
SGD
Epoch
Adagrad
SGD
Epoch
Adagrad
SGD
Epoch 40
Adagrad
Layert
SGD
Epoch
Adagrud
SGD
Epoch
Adagrad
SGD
Epoch
Adagrad
8.27
12.28
9.12
10.90
10.19
9.51
6.61
3.19
1.87
3.21
1.73
2
10.70
15.42
8.52
2.12
8.88
2.04
3.38
5.73
10.16
6.05
4
18.02
15.63
6,45
4,86
7.49
4.44
4
16.62
12.94
8,90
3.91
8.12
4.14
12:70
11.69
11.23
16.98
42.57
18.45
6
12.76
11.86
3.84
2.32
3.20
2.09
4.84
6.46
7,92
2.66
5.30
2.46
11.79
8,58
17.95
4.32
14.99
107.51
136.38
17.09
12.73
26.70
5.16
26.91
6.73
8
8.79
8.20
337,34
5.18
315.21
4.48
1.1.2 EXPERIMENTS ON RANDOM DATASET
Figure 5b layers.
shows the training losses on random dataset and Table 5 shows the Me (t) in different
Table H (t) of Adagrad and SGD for random data
Epoch
Epoch
Epoch 40 Adagrad
Layer#
SGD
Adagrad
SGD
Adagrad
SGD
10.88
10.98
9.99
18.66
9.67
22.37
9.47
12.15
14.98
4.43
13.01
3.99
7.45
8.52
459.71
6.09
451.16
5.11
9.84
10.42
135.37
7.22
126.91
6.04
5
7.09
7.88
103.60
353.89
184.61
190.17
6
7.68
8.58
18.38
4.08
18.69
2.73
7
7.81
5.40
294.68
62.72
229.25
29.76
8
13.51
9.16
329.12
20.59
203.70
9.57
A.2 RMSPROP AND AMSGRAD
In this section, present the results of RMSprop and AMSGrad and compare them with SGD+M.
The experiments were conducted on the translation task described in Section 4.1 The learning rates used were 2.5e-5 for RMSprop, 0.0005 for AMSGrad and 0.03 for SGD+M. Both RMSprop and
SGD+M used momentum parameter 0.9. The two momentum parameters (31, 82) of AMSGrad were (0.9, 0.98). Figure shows the training losses and Table 6 shows the corresponding Hoed (t).
50C-M
Traniin losk
200
850
1000
1202
1502
1802
Teration
Figure 6: Training losses of RMMprop, AMSGrad SGD+M the translation task described in Section 4.1
14

---

Under review as a conference paper at ICLR 2023
Table PoMMpppp (t),
RAMSGad
(t) and A med
(t) in some layers, on the translation task described Section 4.1
Epoch 40 (
Layort
(()
(()
(1)
(
(1) Fexd
(
Foad
2.69
1.89
5
26.17
21.19
11.36
37.11
17.83
10.85
$1.94
10.22
12.31
7
4.10
6.98
6.12
3.94
4.95
2.92
7.58
2.29
2,58
35.72
25.86
19.89
16,90
30.68
16.24
12
4.93
12.67
6.61
6.44
15
85.06
33.63
19.51
140.99
12.22
6.72
44.07
6.9%
5,37
18
8.71
2.99
9.48
3,86
244
4.16
3,51
2.10
2.35
24
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 RSCDM (t) and R (t) conditioned on the same training loss. More precisely, we make comparison of R (t) and RSGDM (t'), where f, are picked such that tth 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 R (L) and RSSDDM (t') some layers.
Table 7: Ronda (t) and (eaaa(M) (t") in layers. Dataset and task: (a) sentence classification task BERT- small, (b) translation task on Multi30k.
(b)
(a)
Layend
Loss 0.251
RAlaa(1)
Los 0,179
Loss 0.133
Layerk
4.01
Laxo 3.22
Rele (()
4.45
5.40
Loss 2.78
3,00
2.44
Leos 1.90
Rile(()
2.28
16.77
13,69
14.14
12.71
15.17
9.96
5
31.19
27.50
44.29
21.46
$7,83
19.52
12
16.68
8.29
9.98
8.31
9.50
5.42
7
5.80
4.38
7.51
3.91
5.25
2.87
15
18.64
7.79
$1.39
46.43
8082
40.97
9
21.23
53.65
28.99
20.92
44.36
28.13
208.39
464.30
115.58
498.36
13
17.27
35.50
14.43
19.17
1926
15
186.41
33.87
257.32
188.5
11087
21
100.43
23.66
23,45
12.96
14.32
1643
1808
3.48
9.54
Table 8: R (L) of SGD. Adam after SGD and Adam some layers after roughly 2160 iterations
150
Layer#t
SGD
Adam
Adam SGD
Adam
13
294.76
150.02
332.96
150.02
14
14.34
5.84
5.33
5.84
15
36.38
16.66
11.86
16.66
6.47
7.05
3.76
7.05
17.17
6.05
4.76
6.05
26
5.68
3.53
2.30
3,53
27
14.33
15.93
21.76
15.93
28
9.10
1.71
1.71
1.71
Figure 7: Training losses of SGD, Adam after SGD and Adam for the translation task
29
8.22
3.04
2.82
3.04
11.39
5.12
5.29
A.4 EXPERIMENTS OF SWITCHING FROM SGD TO ADAM
In this section deseribe another learning schedule: the "Adam after SGD" schedule, where switch from SGD to Adam in the middle to see whether the loss and HOPT (t) can catch up with the model trained by Adam from the very beginning. Again, used the same model as in the translation
15

---

Under review as a conference paper at ICLR 2023
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.
task in Section 4.1
As is shown in Figure 7 and Table 8 both the loss gap and the gap of m (t) 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.1 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 8 shows the examples of layers with or without the low rank structure.
Laper
Taper
SGD
5G3
Adem
Adam
1.5
sigaarvsss
63
332
400
033
0328/0ote
Figure 8: Examples of layers with approximately low rank structure (right) and without low rank structure (left)
We then studied the uniformity of leading singular vectors of these 12 layers, i.e. computed R. and RG defined in (B) and the second remark of Section [6] The observation is that for 10 out of these 12 layers, Ru, values of Adam were smaller those of SGD, which implies the uniformity of leading left singular vectors of Adam. However, we did not observe significant uniformity for Adam in terms of leading right singular vectors (R,). The second remark of Section 6 discusses possible reasons.
Figure shows how R., and R, changed over time in some layers.
#11
Layer #13
2.5
2.5
SGD Adam
5mpubr labies
SGD
snoultr alues alues
SGD
smpulorro solves
SGD
Shiuler
1.5
0.5
0.5
500
503
coodinate
coodirate
cocdinana
coodinate
20-
25
10
560
SGD
32
SGD
SGD
50
50
10
50
10
50
epoch
12
12
o7 10
5GD Adom
S.DD
SCD
SGD
Adam
8
6
50
50
50
50
epoch
epuch
Figure 9: R., and Rs for Adam and SGD with momentum some layers
16

---

Under review as a conference paper at ICLR 2023
A.6 How THE (ADAPTIVE) GRADIENT ALIGNS WITH DIAGONAL OF LOSS HESSIAN
In this section we present the uniformity of diagonal geometry of adaptive methods from another perspective. Denote Hü as the (i. i)-th element of the loss Hessian and 9i as the i-th element the gradient. It is conjectured that when Hal is large, the corresponding [gs] is usually large well. For adaptive methods, we can regard the update per step as the learning rate times the "adaptive gradient". Let's use Jadape, to represent the i-th component of the adaptive gradient. Through experiments on language models, we found that 19ndapt, for different i are quite uniform and do not align with [ as the true gradient [g:| does.
In the experiments, we first sorted |Ha] in the ascent order: |H,, S |Ha.,bal -... S HHaad (suppose H E jddd),, and then plotted the corresponding 19iu and [andapt for A E [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.1 We observed that the figures of all layers are quite similar so we select one layer as an example, as is shown in Figure 10
True gradient of SGD
True gradient of Adagrad
Adaptive gradient Adagrad
10
epoch 22
0.5
100 coordinate
150
200
50
100
150
200
100 coorcinate
150
200
coordinate
-10
apoodh 20
10 atee
0.5
50
100 coordinate
150
200
50
100
150
200
50
100 coorcinate
150
200
coordinate
-010
-10
0.05
epach 40
0.5
50
1000
150
200
50
100
150
200
50
100
150
200
coordinete
coorcinste
Figure 10: How the true gradient (119,, 1) ) and "adaptive gradient" (Agaatpt,s 1B align with diagonal of Hessian (i 1) Here coordinates are sorted such that (H $ |Hyy S ... S [Hag (suppose H E R Rdxa), Experiments were conducted on the model described in Section A.1 This figure shows the results the 12-th layer.
A.6.2 SGD WITH MOMENTUM VS. ADAM
Here we compare Adam and SGD+M on the tasks described in Section 4.1
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.
A.7 ADDING REGULARIIZATION AND OTHER TRICKS
In this section, we add weight decay to both Adam and SGD+M on the translation task described in Section 4 The momentum parameter in SGD was set as 0.9. The two momentum parameters (31, (B)) of Adam were set as (0.9, 0.98). For both algorithms, we set the weight decay parameter as 0.001. We trained the model using constant learning rates for 60 epochs (1800 iterations). We tuned and chose the best learning rate 0.03 for SGD+M. The learning rate of Adam was set as 0.0001, under which Adam converged faster than SGD+M with its best learning rate 0.03. Figure shows the
17

---

Under review as a conference paper at ICLR 2023
True gradient of
True gradient
Adappive gradient of Adam
0.5
0.05
5000
50
100 coordinate
150
200
50
100
150
200
0
50
100 coordinate
150
200
coortinate
0.5
0.02
1.5
0.01
opec
100
150
200
50
1.00
150
200
100 coordirate
150
0c0 coordinate
coordinate
x10
apoo =5
20
0.5
0.5
100
coordinate
coordinate
Figure 11: How the true gradient (lg,, 1) ) and "adaptive gradient" ((lgaapt,,, IB align with diagonal of Hessian ( H 1) Here coordinates sorted such that H S |Hi ... S H it (suppose
H E Id Experiments were conducted on the sentence classification task described in Section 4.1 This
figure shows the results on the 12-th layer.
True gradient of IM <10
True gradient Adam
Adaptive gradient of Adam
5-
ep0cs m
0.5
0
150
200
coordinate
coortinate
coordinate
x10
x100
10000 12
2
0
50
100
150
200
50
100
150
200
50
100
150
200
cocrdinate
coordinate
coondinste
<10
x20
ApOiL 55
0.5
50
100
150
200
30
100
150
200
50
100
150
200
coordinate
coordinate
coordinate
Figure 12: How the true gradient (19a, 1) and "adaptive gradient" (llhattpt D2=1) 1) align with diagonal of
Hessian (1. 1) Here coordinates are sorted such that [H S |Hig. < |Hag.al (suppose
H e R Rdxa), Experiments were conducted on the translation task deserihed in Section 4.1 results on the 5-th layer.
This figure shows the
training losses and Table shows the values of Reddam (1), RSGDDM (1) and selected layers.
PSSGEM RATEE
in some randomly
A.8 RESULTS ON IMAGE TASKS
II We trained a ResNe on CIFAR-10 dataset and compared the convergence speed and ROPT (1) of SGD+M and Adam. The momentum parameter / in SGD was set as 0.9. The two momentum parameters (81, 32) of Adam were set as (0.9, 0.98). The model was trained using constant learning
11 We borrowed the implementation here ht 1 Tatest/tuu wettutrrialcccaptteet ier___itteeeeddee the "layers" aray [2,2,] with [1,1,]]
iorroccs..o//rre io/en/ wtt.mmddaeee//eeee ate/3 at and replace
18

---

Under review as a conference paper at ICLR 2023
25
Trannoo
05
1300
2000
Itration
(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 ResNet trained CIFAR-I0.
Table 9: Re
(1) and R ResssM (t) (both using weight decay) in some layers for the translation task.
Epoch
Epoch 30
Epoch 55 aas
Layent
pOODM (1)
Roda (A)
NE0DDM (1)
NSSCMM
(0000 2M
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
494.27
494.27
150.14
123.79
1.21
301.89
119.53
2.53
1.5
632. 10
632.10
277.18
175.34
1.58
334.48
282.88
1.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
28
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
rates for 41 epochs (2050 iterations). We tuned and chose the best learning rates for both algorithms: 0.5 for SGI++M and 0.005 for Adam. Figure shows the training losses and Table T0 shows the values of RAdam (1), RSGDM(U) and
RSSEp (*)
Resa (c)
Table 10: Raddam (f) and RSCDM()) for ResNet on CIFAR-10.
Epoch 10
Epoch 20
Epoch 40 RAla"(1)
Layert
R800M(())
(e))"(()
WAT 7052-47
p000 700
R550M (1)
Jo (c)
705
6,88
25.34
0.27
3.74
39.35
0.09
4.39
15.80
0.28
10.19
3
40.89
16.92
2.42
13.98
15.92
0.88
11.70
37.01
0.32
4
28.56
13.47
23.78
0.57
8.64
12.07
0.72
6.52
14.23
0.46
6
18.72
12:49
1.50
12.19
8.80
1.38
8.96
21.69
0.41
18.85
9.00
12.81
13.87
11,42
8 13.79
19.91
0.69
8.87
11.72
0.76
7.48
9.34
0.80
10
14,89
14.53
1.02
8.15
5.80
1.41
6.21
8.89
0.70
A.9 COMPARISON BETWEEN ROPT (t) AND SINGULAR VALUE-BASED METRICS
In Section through extensive experiments on language models, we demonstrate that when we train the neural network using Adam, the uniformity of diagonal geometry, measured by ROPT (1) is
19

---

Under review as a conference paper at ICLR 2023
smaller than that when we train using SGD+M from the same initialization, for most of the 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 make comparison between diagonal-basee metric and singular value-based metrics.
First, we believe that metric has natural connection the mechanism that underlies adaptive methods. Adaptive methods practice choose coordinate- wise adaptive learning rates. From a high-level perspective, this procedure can be viewed adapting 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 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 F (t) is a reasonable proxy of singular value- max(a,(t)))
based metrics. Define a singular value-based metric S (t) == median (a.(1)) as an analogy of our diagonal-based metric F (f), where (t) [i== denotes the singular values of loss Hessian (1) E IRd at the /th iterate. We compare 5 (4) along the trajectories of Adam and SGD+M in translation task described 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
R (t). This is expected because in Appendix E demonstrate that the loss Hessian approaches diagonal during training. The fact that our diagonal-basee -based 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.
Finally, there is strong reason why our metric is often casier to compute empirically and analyze theoretically than singular value-based metrics such SOPT (t).
1. From the empirical computation perspective, suppose the loss Hessian is d x Then computing its singular values, in general, requires computing the whole matrix with d? elements. However, our metric only requires computing the diagonal entries.
2. From the theoretical analysis perspective, in Appendix F we show that the diagonal of loss Hessian linear networks can be connected to weight matrices by simple formulas. These straightforward formulas simplify the analysis and allow us to 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 it almost impossible to theoretically analyze any singular value-based metrics.
Table Schdam (t) and 5SDM (t) in some layers for translation task.
Epoch
Epoch 30
Epoch 55
Layertt
SSGDD (t)
(1)am(1))
SSGDM (1)
5ydamm (1)
SSODM SPSPT())
SSGDM (t)
S 5tidam (1)
med
3
4.53
4.53
6.08
2.87
2.12
5.92
2.58
2.30
5
14.64
14.64
40.01
15.38
2.60
52.28
15.05
3.47
7
6.91
6.91
5.06
1.94
12.25
4.22
2.90
9
24.12
24.12
42.02
30.89
1.36
33,20
21.54
1.54
12
19.07
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
28
55.38
55.38
6.18
2.77
2.23
4.84
2.01
2.41
20

---

Under review as a conference paper at ICLR 2023
A.10 THE UNIFORMITY OF DIAGONAL GEOMETRY ON IN- IN-DISTRIBUTION TEST DATA
In this section we compare Remea (L) and M (t) on the in-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 region with smaller med (t) than SGD+M, similar to what happens on the training data shown in Table 2a This expected because of the same distribution. One other thing we 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: Rised (t) and RSGDM (t) in some layers for the translation task.
Epoch O
Epoch 30
Epoch 55
Laycr#
medd (t)
(t)
R ( mee
(t)
P (t)
M (t)
R (f)
Rog (E)
3
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
4.61
1.30
5.53
3.20
1.73
9
25.09
25.09
28.81
17.17
1.68
16.67
14.85
1.12
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
14.78
14.78
7.15
0.55
5,42
6.04
83.25
83.25
26.04
5.44
4.79
13.11
5.57
2.36
24
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
B PROOF SKETCH OF THEOREM
Now we give proof sketch of Theorem which contains three major steps. The detailed proof can be found in Appendix
and
First we relate the diagonal of Hessian to weight matrices W1, W2. Under Assumption 1 denote Wi : as the i-th row of Wi and W2 == [w2s, u122 W2d). Since the input dataset whitened, we can show that
max;
(t) 1222
max;
ww(0
[
OPT (1) 1
OPT
med med.2(() (1)
median
(t) Wuiss
21
median
|w [1,
2
Next, due to the one-dimensional output, prove that W1 converges to approximately rank-1 matrix. More precisely, we have
w ( + R(()
w (T) (()
where c(t) is a scalar, (t), (lt) v()) (t), (t), R as (1) (0)
12 E Rd and 12(2) E ddd dddxdDDnntte the i-th coordinate of v R(0) respectively. Denote the (i. element of R(()) as R [i,
We have that Vi.j [d]
R
(1)) u and R (1) [i,
U
(1)
Using the rank 1 structure, we can further simplify ROPT, (t) and ROOPT (t) by
(t)
max;
ROPT (t)
2 k: = 1,2.
(4)
median
21

---

Under review as a conference paper at ICLR 2023
The final step is the detailed analysis of u (*)
For SGD+M, we can prove that (t) = C X2,
Xa)" where C(t) E IR and X,, E [d] are
Li.d. Gaussian variables. Then we have with high probability,
modian ux Oloog d). For Adam,
(
we can prove that Vi E [d] ( E {+1}, which gives us completes the proof.
max,
= 1. Substituting into eq.
(u)))²
median
C ANALYSIS OF SGD+M
Note that == LYX", AZz == 1
( m XXX. Denote Ih = Vw. I(W()), 1,2. We have that (1) = W (w: - -A) A). 92 (0) (w/" N - -A) w
91
Let A As2 and 9h (1) k 1,2 be the corresponding batch versions at time t. Let E(9) w w A. and use E,", A, and (w:") (" represent the i-th coordinates of E(t) A and
w/ w respectively. By eq. 2). the update rules of W1 and W2 for SGD+M are given by:
w((-1) = -W1 1/
Σ g gh (w -A) Σ
0=0
T=0
w((t-) = w-(1) 11 2
ge- (w<() - A) w(T) Σ gt-r 30- (T)
-=0
r=0
where
1911 (1) (11) (1)
91 91 w
(w('0(() (A - -Ax)) - - -)))
Dg2 02 02
(8) (1) (t)
(w w((") (AO2 - ARx) -xx) (100-1) W(AT
Based on the magnitude of W2 and W1, we can intuitively divide the training procedure into 2 phases.
1. First phase: the first several iterations when W1 and W2 are "small" so that W2W1 A R
A.
2. Second phase: later iterations when W2W, 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 Tr) is defined as Ti inf { t2 0: ji E [d]
(t)
L or ( (1) [. j]
W2
d
By Assumption 2 and the assumption that Vj [d] A, 0, O(1), the beginning, w.h.p.,
Vj E [d] (W:W)) - A, < O. During the training, each (W2W1); increases and approaches Aj- We hope that by choosing small learning rate, when (W2Wi)s overshoots for some coordinate j. i.e. (W2 W1); > A,. it will be close to convergence. To analyze this overshooting issue more carefully, let's first define the following "almost overshooting time".
Definition 2 (Almost overshooting time). For > denote co म + e clog V Define
T3 : inf { { 0: 3j [d] (w!"W((")
Definition (Convergence time). For > we define the "convergence time"
T3 inf { 22 0: | <e).
We can first show that after the first phase, i.c. when t = T,. W, will become an approximately rank-1 I matrix, as described in the following lemma.
22

---

Under review as a conference paper at ICLR 2023
Lemma 1. Under Assumption and suppose o S du1zTI- that when = Ti, (W(T))= (d), and that
1//2
By picking 1j S (+)), we have
w - 127T1) + u (11) (
w R + (
where e(Ti) E R. (TT)), w(T), Rg Rd and R(Th Raxd, Denote the i-th coordinate of u((t), w(T), Ri.i. Then w.h.p.
(TL) (T)) n respectively, and the (i, element of RT as ((T)), n as us 1
R T1) [i
(T))
Rizi
S iijdd:
dia )
S0
us (
(Ti)
e(Ti)((T))
(Ti)
d
d=a-1
The following lemma tells us that this approximate rank-1 structure is preserved when Ti S tt min(T2, T3}.
Lemma 2. Under Assumption have that w.h.p. for T1 S < min [2, T3).
2 and 3
suppose S 70/20T By picking 1l S O
w u (Ti) (1)T + R2
w ca TT/T +R
where
R
L1 S i,j
Ö(EE),
Ō(Eo),
(Ti) (2)
c(t)u: (T1)
and is defined in Definition Moreover, when t = min (2,n (3). L (w(t))
O
The following lemma gives us more detailed description of u (T1).
Lemma 3. The (T)) in Lemma 7 and can be written as (T)) X + Y where Xi, i [c] are
i.i. Gaussian random variables and that w.h.p. Vi
[d]
M X
Ô
Now we are ready to prove the SGD+M part of Theorem
C.1 PROOF OF THE SGD+M PART OF THEOREM 1
Define TSGD, TsGD,2 min Ta}. By picking S O 4
we can apply Lemma 1
and 2 to conclude that L (WTTsso.1) ((d) and L (W(Txo0))) = ((cyd). For any p > 0, by picking 0 < <<<< 7 and a 2 4(p + 2), we have L (W(Tkm.))) = ((odd) S Ö (+).
Moreover, when & [Tsa Tap,2) the conditions in Lemma 30 are satisfied with Ō(co). Then we can apply Lemma 30 and get that
2 max; (u(25)) (25)
Ō(Eo)
RSDDM (t), R SSDM (t)
+ Ō(eo) median u (Ti) 2
By Lemma 3 (T) X + where w.h.p. ViE [d]
Y X
θ
)
This fact yields
(Ti))
max; u1
max; X? median X?
Vi E [d]
2
median
(Ti)) u1
2
1 + θ
23

---

Under review as a conference paper at ICLR 2023
Here X, [d] are i.d Gaussian random variables by Lemma To prove the concentration of
median X2, we borrow the Proposition 12 in Chapter 2.3 of (Lerasle in this proposition, we have
2019). By setting K = N =
((median x?- - -[ > 2 Var(X?) S se #
Denote o² as the variance of X;, [d]. Then B[X?] o² and Var Var(X?) 20 Hence
P (Imedian x? - o²| > 2V/20² S e
That means with high probability, median X² S Co for some C > 0. By Lemma 34 in Appendix G we know that w.h.p.
max x? = o (2 d), Isidd
which gives us w.h.p.
max Siid X? = Si(oo d). median X?
Hence we have proved that ( (t), m (1) 2 22 d).
C.2 PROOF OF LEMMA 1
In the first phase, W2 is "small", and we write the update equations in the following way
w(((+1) = -W110
W. (r)
w/1 - A) -))- -n B --Dg"""" (T)
r=0
T=0
=W(" +1
g B--TW6I
(FT A -
Σ g -- (T)
W (T) w/(T) (7)
Σ
8 "Dg (7)
17
+=0
T=0
7=0
= -W(O +n
(t)
A
g 1)
Σ g
( (T) - : A
(5)
2
r=0
T=0
g
W:(
(T
W ((())(
g D2i (T)
11
7=0
r=0
( (1)
(1)T
A +
"
(t)
1
1
where
1 = -3+ w (1) (1)
A + (1 - B) Σ ge_ (w/+tr -W :((()) A
7=0
- (1 33- 3)
g - g- (T) w( (1 B) 0=0
3 Dg (T)
0=0
Similarly, we have
wh = w 7 Σ 8' (//)))((TT)) - A) w(+)T = w/( +
7 (WT
+ 7
+1
1 B
1- B 2
(6)
1=0
where
(1) = -p++1 - + (1 - B) Σ ge- A - -
0
- -1- - 3) Σ 8 W W 1 -
(1- - B) Σ 3 DDg2 (T)
T=0
==0
The following lemma gives us explicit formula of w(())
24

---

Under review as a conference paper at ICLR 2023
Lemma 4. Let X1 < X2 be the two roots of the quadratic equation x² 2x (1-カ || Lhm
1-8 then we have that Pick 11 [|]]22'
W = Gff + (C2+ +
X/2
((1) 2 C2 ( -A, 2
(0) will be specified the proof.
where C1
d2-
is "mmall" More specifically, denote its i-th coordinate as Go and the i-th coordinate of C22 as Cz.. Then the following lemmas tell us that vi E [d], 7 S
We can prove that in the first phase, r (1)
where w.h.p.
min
|C2il-
AplaT
We first have the following bounds of
1
(0)
(t) (+) and r5s for i [d].
Lemma 5, Under Assumption
2
and
3 suppose o S
/ and pick S O (z)) We have w.h.p.
2/2
for all t < Ti, vi E [d]
r [ <
(8)
dda
26
Lemma 6. Under conditions of Lemma
we have that w.h.p. for all t S Tú, E [d]
T51 S
δ d3 )
Next we prove upper and lower bounds of CCiil and |Czil for E [d].
Lemma 7. Under Assumption
and 3 suppose S 2/22
/6/24T Pick
1-B TAA
we have that
i) w.h.p., Vi E [d] |Cuil ó
C2al
(#);
fi) C2 can be written C2
(Cs + C4) where Csi i [d] are i.i.d Gaussian random variables
and that w.h.p. Vi [d]
[Crd
< Ö d=a 5
[Cal
iii) w.h.p., Vi E [d] |Cail
|C2sl 52
Now we are ready to prove Lemma
Lemma tells us that
Wh = - + (C2 + r(0)) X/.
where A₁ = A' and d2 = 1 + - ||A1|2.
Under the conditions of Theorem and pick v S O (1),, by Lemma and Vt d.
we know that w.h.p.
(t)
(0 s θ
)
CCr 2 Ω d3a
(11 1C2i s 7
()) da
C2a
rst
<
dia-1
1
25

---

Under review as a conference paper at ICLR 2023
We first prove that w22: see this, first note that
(0) reaches
1 d=/s
for some coordinate i before |W/ [k,5]] for jk. E [d]. To
w (((11)) 1)
17 Who-1)T
1)T
A
+
"
(1-1)
3
1 1-3
1-1
"1 1-1
= (
7
W
(
A +
(7)
1
1-3
T
- (0)
7
C1
+ C2
Σ
1/+ + Σ
15r (7)
A
71
(r)
1-3
B
0=0
r=0
H
-w
1
Σ
(
π
C1
Σ
x
+
Σ
X5r( (7)
4+
M A
1 1-3
1 1-3
1-8 ==
T=U
7=0
[-1
== -w
71
(7) +
Ci
Σ
x7
+
Σ
\5r(
(+)
u(1)T + (O!TT
1- 3 T=0
T=0
T=U
where (()T
1-3 A and
t-1
(0 = Σ
N.C...
(8)
7=0
Moreover, we have that
W CGA + (C2+r5] (0)
2/= C1 C= +
1 600 00
+
(T,
e0)) =
==
For L S Tí, by eq.
we get that w.h.p..
A5r
7=0 X1
0
55
1
LA S; d:
S
U
)
(t)
Ω
-==
N5
dia-i
12r (0)
(1)
5i
1
cl0m( (0)
1024
d
For / < Ti, by Lemma
A1 < ز 2 d: |11 177100))
< Ö
Then we have that w.h.p.
-3 1
" Er= ri (b.]]
Σ Σí= ri (T) [i..]]
Σ Σm (T) [i,
1-8
(1) (e)
Σ A2
u 0/
= Ö 1
Here we used Vi E [d] : A, = O()) by Assumption] !
Since \1 = =1 1 ||A11 we have that |C1sA1| S |Ctl S 50 (1) and that
1-1 C1
Σ ATo (0) n
C+ 1-1
A > AT S (1 - -1))
A,CCii ||Al|/
Sō (diit) ).
1
1 3
datb
Using the Gaussian tail bound and union bound, we have w.h.p. VI S i S d: | [i.0]
=
Combining the above bounds together yields that for < Ti and Vi,j E [d].
1(O [i. = [ [i + un (1) (0)
oj (1 + e (0 (t.)).
(9)
w = R2i. + (0)) (1 (1+00 (1+ (
26

---

Under review as a conference paper at ICLR 2023
where for j E [d]. |i - Ô
(
R
S ( and
di
Further we notice that for & S T1, we have Vj E [d].
n7
-
1-11 X3
n
\ -
A, (22 -1)
Aj 0
1
[(0)]
1 B
N/2
1
1₂ - 1)
1111A112
||A||
vd
which yields that 10/00/11
(t)
0
|cft) uu!
Together with eq. 9 gives us that
(t) W2s
reaches 2u77 for some
E
[d] before
| [R: for vk,j [d]. i.e. T
inf (22 0: Di
[d]
22
(t)
Further, we know that at time Ti, means w.h.p.
c(Ti)
(Ti)
1Czio 1171 = = (2217)) for some in E [d]. which
θ
θ
Ti S ò (12a).
d
" = (dd3) S 2m = 1 1- 3 (111111))) 11+
0
Ω
Ti = θ ( logg d
11|lA|12
(10)
This is the length of the first phase. As for ely that whh.p. V1 S í.j S d,
(T) and (T)) (Ti) for other coordinates, we have of
(T)) (TT)
1-8 3 7=0
11
T
Σ
X5
C2:
=
1- B
"
X3: 1
x2
C2i A₁I
( AZ. -1
||A||/
1C2:4,1
i (2
1
2
1
θ (dd
dia
dia+z
(TT)a u
(Th)
1C2i 12/1 > á (d//2)
52
1
5t
1
dāa
dia
Here in (i) we used \2 = 1 1-8 114
Then we have at time Tı, Vi,j E [d],
m (T))
R
δ d
and that
S 6 -=
Together with eq. 9 we have the following weight
structure:
wT(hi) (T(TT) u (T)) w((r) = (TT + (T)T,
where w.h.p.,
R Ti) [ j]
[T))
IA iijidd::
5 20
( ,
1
50
1
(i)) (
(Ti) vj
((Ti)) (i))
ddt--
u
Finally, we consider the loss. Since Vj E [d] : (w.(t)
A, = -0(1,, we know that
T. (W(T))- (d))
27

---

Under review as a conference paper at ICLR 2023
C.3 PROOF OF LEMMA 3
Eq. 8 tells us that (T) = - X502. Lemma tells us that C2 can be written as C2 ==
1 (C3s + C4) where C3i, i [d] are i.i Gaussian random variables and that w.h.p. Vi [d]
C3il -
Cu.l
d1 Combining these two facts together finishes the proof.
C.4 PROOF OF LEMMA 4
Replacing t by t 1 in eq. we get
w = 2
(tt-2)
+
AW
+
1- B 2
(1-1)
(II)
1- B
Eq. 6 and substituting eq. 5 yield
W. (((++)) w -w (=(t-1 1)
n/²
|1 (1-1)
(1 3)2 Art (r((-))
172
(1 3)2
1
(1)
((-1)
1 1-
22
r2
11(+1] =2W2 (t)
(1-
1
n
l1al1]
(1)
2
(1- Byz
where rz(() = (I- Ari (1-1)1
2²
+
"
(1) 12 2-11
I-1
For the equation x² 2x + 1
n²²
||a1 = 0, the roots are \1 = =1 - -1-1" ||All2 and X2 =
(I By=
1 + T-3 12 We have that
w((+1) - 22 = >1 (W/²) )211(-111)
(1)
+ +r²
1-1
W - \2W/1-11) = 21-1 (W/" 22 12))
+ Σ X 1 (T)
= - (W/') )24210)
+ +r
We further have
t-1
wh6 13
+ Σ - Xi (W//") 22
Σ 1 r4
(T)
T=0
= =>$W! (0)
+ X/2 1
-
-
+ Σ (1
A22 X
r=1
1
= CMM + +22/++ + - Σx Nイ
7=1
=ci + (C2 + r55
(0) N/.
where (0) = = No T4 C = A=-
W (())
and C2= =
A2-
C.5 PROOF OF LEMMA 5
Write
(1) = -3 W (1)T)
A +
012
+ 913 +
44
(
where 912
(1
9) E g6- (W/T)T - W A,
(t)
= -(1 - 3) Σイ= 3 3 w (T) W (T)
W and
4
41
(t) -(1 B) Dt gr-r DD1 (T) And write (1) gt+ AW (t) + 422
(t)
+ 923 + 22 where
(t)
(t)
(0) (1 8) Σt gi-rA A (W(r)? w((l)T 923
(1)
(1 B) Σイ - - W
W w (r)
42
(1 (1 () Ei=o 3- Do2 (T)
and
124
28

---

Under review as a conference paper at ICLR 2023
Let's first try to bound
912
(9)
[.
and
922.2|
(t)
For any r S Ti, we have that
d
d S
d
di E [d]
[5,
W
[i.)
[7.
>
24
22
de
do
j=l
im1
and thus 1A E [d] E = (1). Then we have for all i, E [d].
W W [8.5] [t.
S 17
g
W22
E
I Σ
gr
0
1
do//2
k=0
k=0
110
d///2
2/
(T+1) (T)
0222
S 17
8 &
E
W
(6)
[2,5]
S Σ
8r
0
da/2-1
1
k=0
k=0
= 010
d0/2-1
That gives us Vi,j E [d].
((-(((-))
412
S (1-
-3)) 3) Σ 81
(w/TT w A S 5 3)
da/2-1
0=0
17
da/2-
(1) 922.3
51 3)
8
Aj (1/1->(0,3) - -W
7=0
g g (1-T) (t-))))))))))))))))))))) = =0 (f0/2-1)
$ n1 - 3) Σ
7
T=0
Then we bound
4is [i
(t)
and 423,i
We have for Vi,j E [d].
Via [1
(0)
(1 3)
Σ
(w/5t
511 3) Σ
g
1
di
d
1=0
7=0
1
0
dia-1
d
(0) 923
[1 3)
3
( w(())
j=1
d
(1 3)
8
-- 1+ = 0
d
T=U
Finally we use Lemma 31 to bound 4ia [i and 921.1 For S Ti, the M("), , in Lemma 31
(t)
are upper bounded by In the theorem we consider the training period before TSGD,2 so
time T in Lemma 31 is set as TsOD22 In the following sections, we will prove that TSaD,2 <
d" log( dde)) Then by Lemma 31 we have with probability at least 4. for Vt S T₃ and
O
7
29

---

Under review as a conference paper at ICLR 2023
vi,j E [d].
D91 1 [1.5]]
=
g [t:
[t,
= 0
(
d
V
datl
log
d
+ +
(
d3
oV
dd+z
17
log
d
1
da+2
di
1
D92
(0)
=
921
((t)
92
(0)
S
d
1
0
d
log
d
+
d3 -1
V
d
log
d
da+2
S
d
By picking 0 S
du/27I, we have w.h.p. for Vt S Ti and Vi,j E [d]
1//2
Dg [t.
s 20 (1/
and
Dy22 (9)
S n
which yields
914 [.,c]
S (1 3)
g-
Dgi [ 5 - 3)
Σ
8
1
=n
d
d3
T=0
7=0
(1) 921.1
< (1 3)
3
Dg2
(T) S - B) g 0
Ou
d3 -1
d3 -1
Combining all the above bounds and substituting 1) S 0 (2a) gives us for 42 I Ti and Vi,j E [d].
[1 S g/ /w/2 A, + +
1
(t) 12:
< ge+l
Σ 1, [i,]]] +
1
dia-1
dāa-2
(12)
Fort fort= S Ti, we have Vi.j [d].
222
(t)
Aj
0
and
D A, A, (1)
[i.0]
0
20.7
which gives us |r.5) 1r [i..i]] 5 O (TozZ) and
(1)
O
Substituting into eq.
5 and eq.
(6)
yields that for Tı and
Vi,j E [d].
1 [i.i]]
((('((,)))
[i, ||
O
7
((+1)
wiz
(1)
50
d
da/2-1
Hence for 1 S min a loge Tog(1/35.
Ti }. we have Vi,j E [d].
[i.0 s
[[.c]] |
alog d log(1/3)
0
(1122) S
d
(
(0)
a log d (1/3)
71
<
S
do/2-1
d
2
Then we know that T;
a log d and also get tighter bounds of ( log(1/8
[
(t)
2226
for
t
<
log(1/3)
a logd
Now we use these new bounds to analyze r1 1 [
and
12 again.
(!)
When t
log[1/3], we have for all i,j E [d]. g++1
a log
122 A
(1)
I
W2
(9)
A
S θ
and
g8+1 Σj=1 A, ( [..1]| < Σ A, W|(d (1) [i,
When
a log
1 5 Ti,
[og(1/
have
ge+1 < d suggesting that Vi,j E [d], gr+1
w/2
(4)
A
1
and
da
31+1
Σ A, A, 10) [ < _ (1 S δ
Substituting into 12 completes the
proof.
30

---

Under review as a conference paper at ICLR 2023
C.6 PROOF OF LEMMA 6
Based on the bound in Lemma 5 we have
(8)
172
Agr
1
(1-1) 1) [1,5] +
10 (1)
(1-1))
31
3)2
1 1-
2'
2i
j=1
n/²
(1 8)²
Σ
Agr
(1-1) 1
[1:1]]
1 3
71
2'
(0)
+
1 8
77
2'
(t-1)
1=1
S 1
+ 2
dza-2 = n/ō
1
1
dia-2
dja-2
Since \1
= 1 1-3 |||122 A2 = 1+ - ||A
1,, and note that ||A||2 = - (
(Vd
we have that
44
11 1 (7)
r3
<
1 1
(-
d#a-2
1 A1
dda-2
d3
(a-1)
X3 0
1
7 θ
1
0
1
34
d3(a-1)
A2
dā(a-1)
d2a-1
T=1
C.7 PROOF OF LEMMA 7
For the equation x²
2x +
(1 (1-8ya ||A1|2 = 0. the roots are A1 =
|1-A112 and A22
1 + -"3 |2 which gives us
C2 W X1 1
d2 A1
W + +nA
(O)T
+ +
W
+
" 1-9
A
2 (0)
(13)
2n
|||||
1-3
1 1 (0)
+
1 3
AW
(0)T
1 1- (0)
2||A||12
2||A|||2
where 1/2
Note that this is slightly different from the definition of Now let's bound the i-th coordinate of r22
W/ w/ (0) w/ON
Dyy2" (0)
(0)
in eq.
6
In Section C.5 we have shown that w.h.p. for Vt S Tj and Vi,j E [d],
Dgz (t)
S O (1² 1 - =
δ which also applies to f O. Using the Gaussian tail bound and union bound, w.p. at
least 1 8, for ever و.j S d, we have that
2d
2
2d
Wiz
S
log 8
W1 [.
S $V
log
ca
día
Then we have that w.p. at least 1
8, V1 < i,j < di
( (
( w/
(0)
(()
(0) w/ ( [5
(0)
w [')
22
<
W2j
2
2d
2
2d
2
2d
log
log
log
d2a
d 1
8
100) 24
<<
(w:' 1r)0))
(w:")
| [i:5]|
+ 1gg0
(0)
j=1
(14)
d
log
2
log
2
+ δ
(0(
=
-
1
da-
8
d4a
1=1
31

---

Under review as a conference paper at ICLR 2023
Next, we bound the i-th coordinate of W 1-8 AW1 (0)
i... w (0)
A (w:
+ |||la
By independence under Assumption 2
we have that
Var
(0) W222
+
1-3 (w/1
= Var 1222
(1
3)2 d
1 12
Var
(W710 (2,
A
(1 3)2
1
A3
d2a
A
día
d2a
Using the Gaussian tail bound and union bound, w.p. at least 1 - 6, for ever sis S d, we have that
1-3 A (1
=0(v log
-
d
W222 +
[All2
da
Since for X ~ ~N(0,,7) we have that P([XI
1()
<
21
then for a fixed i,
V25a
(0)
+ |A|2
1- 1-B
(w/")..3)
(..
T
S
1
0o
2/dia
+
di
2
1/d2a
di
Then by union bound, we have that w.p. at least
for every sisd.. d. d
212 ||A|12
(0)
1- 3 A (c::))
7
θ
1
az
Now define C3 = W +
1-8
AW and C4 = 1-9 (1))
2T We get that Cas, i E [d] are i.id
TTAT
Gaussian random variables and that C2 = (C3 + C4), where w.h.p. for all E [d].
1 di
(11 (
CCil S Ö
CCsi 2 e
dia dia)
|C4
S
d
(15)
where (i) follows from eq.
14 and the fact that ||Al|2 /d. Then we get that w.h.p.
Casl
Vi E [d]
S
Cssl
dia-
Ω
Substituting eq. (15) into eq.
[3 we get that w.h.p.,
122l +
((ve
(0)
- 8 A
(w:
")
1
1
||ll2
dēa
da
Similarly, note that
w((1) (1) 1 (0)
G
A2 - A
1 + nAW
(O)T
W:(0)
1-9 1- 1A1l//22
"
2n T-3
|Al|2
(
1- 3
AW (oyT
1 B ((0))
12
A||
2||A1||2
we can use the same techniques to get that i) w.p. at least Vi [d] |Cal s Ō _(+))) ii) w.p.
at least 1 - a7 1. Vi E [d] 2 Ω di
32

---

Under review as a conference paper at ICLR 2023
C.8
PROOF OF LEMMA 2
The proof in Section C.2 tells us that at the end of the first phase (when = Ti),
w u ( + +
W cl
+ R2 R
(16)
nA
((r)) Ag
where a (211)
1 1- 3
T== X
Denote the i-th coordinate of u0 v(0), -1(5)) as u; U R respectively. Denote the (f,
(t)
element of R [t,"i, For t Ti, prove by induction that,
w1oh (6) u ( (0)T + R
(17)
W = (1) u (2:)2' + +R R
where
r (1+1)7 = (t) niC E(0,
1(((1) R The R E(t) + ri
(t)
((++1) = d To E (t))
11( R E() ((I)T + 12
(8)
with 11 = -r ( (t)T (xt) - ()T (rT) - - Dz (T) E(t) = W (1) ((t))
A, "
- 3" and r2
(t) Ei g'- (FOW F(O
Et= go- Dg22" (r)
Note that the (1) and (t) here are different from those defined Section notation and still use ri
C.2 but we abuse the
(6) and v!!0) represent the error terms.
The base case is already given by eq. [16).
Suppose our lemma holds for t, then for t + 1, using the same techniques as in eq. have that
and eq. (6),
w/(t+)) -w(()) 7)
3 B--Twlo W (T)T E(T) E(T))
Σ
g
Dgf (T)
T=0
-w(1) W ( E(() + (t)
t
w((+2) -w(1) n 3 F
(x)
Σ
8 3--DD92 (T)
0=0
wh E O((O + +r2
Plugging in the inductive hypothesis yields
(t) (((+1) w:(!) n7 ( E(() + ri
(t)) R((1) 17t (0 (T,) + +R((') E(()) + +r(()
= (Ti) (u - c Elt)) + n² - RR E(4) + +
w((+1) w Th w((r) + +r (1)
(T)T + R(fii E(t) (1 + -1((0)) +
(t)
= =(d()) m, E()))))))))) u TT]T + R
+ -1) (()T (t)
+
33

---

Under review as a conference paper at ICLR 2023
It implies that our lemma holds for t - 1, which completes the proof.
Now we analyze the error terms |2 2 [i. and |// Eq. (16) tells us that cl and Vi E [d], v/
(T)
are all positive. We first prove by induction that for all Ti S 1 S T2. elt) > 0, Vi [d]. > 0.
The above discussion already proves the base case. Suppose at time 4, we have c(u) > 0, Vi [d]. og > O. Note that when Ti
S tt= < T2, Vi E [d] : E - 0, then for t + 1,
1/ V/ mic (1) F. > 0,
(t-1) (1)
((-1) = e T)t
10 (0) E
0 > 0.
Therefore by induction, we have proved that for all Ti St t < T2. c00 > 0, Di E (d), (0) > 0.
Now we prove that for all Ti
S 565
T2,
R (0) [
R(O)
1-1
Riz
V1 S iiijisd
05
S
6i
+
0
6: +
(18)
(T))
c(()
(T1)
r=
r=
where
R
((T))
[i.j]
R
[(0
[6.
(0)
6 := max
max
(1) = max
max ]
"
(Ti) (Th)
c(T2)
(Tt)
(Ty)
c(t)
"1
(T))
The left hand sides of the inequalities are trivial since we have proved that c(t) > 0, Vi E [d], u (1)
for all Ti < t S Tz. Now we prove the right hand sides by induction.
The base case already verified by the definition of 6G. Suppose eq. [8 holds for Ti S t < Tz. Then for L + 1, using Vi E [d]
(E(A) < 0 and (1++)) 2 v(0), c((+1)) 2 c((), we can get that V1 S i < d
R (1) [
R(
1. (TT) us 1 (01)t))
R
m
21
5/) + m/(lt)
(-E((1)
(Ti)
8z Σ c
(7)
(1)
+ 7/1
05 Et= ((7)
d(0)
-0(())
oi + m,c(()
(2)
(-D)"))
6 +
Similarly, we have that V1 S / < d
(-E=(") [ [i.5]
R2g
(((+1)
{1)
+ 17c شØ=1 5=
1
22
(2)
S
c((+1)
ui
(27)
(( + m
Σイ= (-E(() 1
21
(TT)
c(t)
8; + (-1
(T)
(10) + !/ (5 + Σt=r, (T)
(-P2))) (1)
-
(1)
(
c + m Σラ= (-E(()
(t) 0/51)
= 6, +
Therefore by induction, eq.
(18) holds for all in the second phase.
So far we have proved the rank I structure stated in Lemma 2 The remaining part of the proof is given by the following lemma, whose proof is deferred to Section C.9
34

---

Under review as a conference paper at ICLR 2023
1p/2
Lemma 8. Under Assumption
and
suppose a S do/zet.
By picking
ll S 0
we
d
have that w.h.p. for Ti 215 S min
T2, T3).
R [i
R(9) (T1)
LA Msi d: 0
Ö(co); 05
S s((e))
(19)
(2i)
dlt) ui
"
ojj
and that when t [T,, Ta}. we have | 1| O(ead).
C.9 PROOF OF LEMMA
We first have the following lemma which describes the structure of w(t) for t 2 T1.
Lemma 9, Under Assumption
and for t 2 Ti, we can write v(I7 as v(t)T a*)) R!!TT
with (Ti) I-3. Re p!(T)) [0, 0. 00. and
(1++1) mol (²)) mpelt)
R((+1) = m()))
R!!!) m0 ng!
where d(t) = c( ||/ R u (1)) RO ( + R R(
Moreover, we have that
Wh W ((0) +
R d (*) (1) A (0)) ROOT T + R
(20)
We prove Lemma 8 by induction. Denote the i-th coordinate of R and R!. as R and R(( respectively. The following lemmas constitute the inductive part.
Lemma 10. Under Assumption sider any t such that Ti S min (T2, Ta}. Suppose for all Ti S l, have Vi,j
1 and suppose S J.ा.. and pick n I ( Con-
3/22
[d] | S (d [ [ I O (atr)) then we have that Vi,j E [d] ri [2,
Ö (1Pdu//) 22 = - (1333111
Moreover, we can get that Vi E [d]]
(9) = Ô (
where e, is defined in eq.
18
Lemma 11. Under the conditions of Lemma 10 and pick S O
we have that at time
d
t t+1,
R
((++1)
[i
R
R(f+1)
V1 misjjsd 05
I O(co), 0
< Ō(c)).
u;
(Ti)
(1+1)
c(t+1)
u
(Ti)
where co is defined in Definition
2
Lemma 12. Under the conditions of Lemma 10 and pick v S O t+1.
( ( (
we have that time
1 (++IT (T))
1((+1)
0 0 c((+1) ||]
S O((o)) Vj Vjeld [d]
00 c(e+1) [[T]]] u/)
S Ō((o))
Moreover,
(((+1)
Ri
Vj [d]
Ō(Eo).
(21)
A
Lemma 13. Under the conditions of Lemma 10 and pick 17 S O
( if we further suppose that
(9) रσ = 0
R())
Vj E [d]
and
A are of order Ö(eo), then we have that at time I. + 1,
Vd
A
35

---

Under review as a conference paper at ICLR 2023
B (A) Vi.j E [d]
O1))
E.
(B) Vj [d]
(va).
20+1)
(C) Vi,j E [[]]
(++1)
S O (dt/1))
( [.,5] 0 (2tt))
R(f+ty| A are of order (co))
(D) Aj E [d].
and
A₃
By combining Lemma 10
and 13 we can prove by induction that for all T1 I <t> min((T: Ta). and
eq. (19) holds (which follows from Lemma
E E
Vi.j E [d]
= @(1).
(22)
which follows from the part (A) of Lemma 13 Now the only thing to verify is the hase case, i.e. when
t = Ti. More specifically, we want to prove that 1) Vi, j E [d]
(71)
S O (d'/4),
W (T) [i, s
22
(T))
| and
O (2//4) and that 2) Vj [d]
20
θ
vi and that 3)
aLTAA of order (cc)).
A,
All of them can be verified by the proof in Section C.2 and the definition of R(*) R(t)
So far we have proved eq. 19 Lemma
Now let's prove when min
min (T2, Ts}, we have that
| 1/ ((ed))
If (2n, = T3. by Definition 3 we have | S e. If min [T,, T3} T2, by Definition
there exists [d] such that E - (VEo). Combining with eq. 22 gives Vi [d]
En! - (Von)- Combining these cases, get that when & min Ts), 1001
m O (End) O (end).
C.10 PROOF OF LEMMA
We prove this lemma by induction. The base case (1 = TT) of ((!) is verified by eq. (6
Suppose at time I, (C)T
(1) A + n!
(0) then by eq. 17 we have that
Wh w(())
(0) (T.TT RT n
(T)
+ +R(()) R(())
(1) (Ti) 2 R ulli)) ((()) + (10) (1:2 R(()) + R R²(")
(ITT R
= d00 (1) A + d R(T + p(i)
where d(t) = (1)) 1 (1) + n (t) u (11),
R = ( + R R R
That gives us
(1+1)T = ((1 ml E(1)
= (1) + R(CiT - m (1 [(0)p(t) + R(O0T -A)
(11-meh me ((())) all) + tndl)) A (() RSOT (0) (
(++1) A (((+))T
Therefore we have proved by induction that for t in the second phase, (( a(1) + R(!) The above steps also proved eq. (20)
36

---

Under review as a conference paper at ICLR 2023
C.11 PROOF OF LEMMA 10
Write r1 4i1 + Q12 where we have 911 1) - (W/O/ F.(1) -W F2(T)
(1)
(t)
(0)
(t) - 3 Dof"' Write r 921 + 422
(t) (1)
where (t)
922 1) E g Dg22 921
(T) (1)
012
1 0 3 (E() W/ (0) - -E
Let's first bound of [2.3]] and
(t) (21.2
By definition of T2, we know that for for s [d]
| = 0(1) Then we have for all i E [d].
w [i..]] - (/( [8:5]]
S 17
Br
W22
(1)
R
S 11
gr x 0 (d = 110 (al/t).
k==
k=0
d
(((+1)) -
In -2ff
8 -
E
ki)
W
(4)
[1,0]
5n >
B k (8811) = n/0 (1811)
22
km0
jml
km
(23)
Note that
d ((
122 (wu))
(r+1) - w w/ [2, + uug (1r(=t) [1,c] - w//
1=1
Σ
(T-1) -mb?)) (wr (3.10 - ww("
+
1122
i=1
We can further get that for Vj [d].
b - -E5² s ndO (1/+) 0 (a- (d-1/=) + mdo (a)14) 0 (2/1t) + a00 (D/4) 0 (d/14)
= xZ) = 0 (nd912)
+
Combining the above inequalities gives us Vi E [d]
M11
(4)
[ 5]
7)
Σ B
(t) E. w22 mT
(r)
0=0
8--7 [w/2 -wbس |,: + (+)
(1)
(7)
- | |
- -155))
-
< 1)
7=0
BL-T(t (t- 7) (0 (1811) 0(1) + (1bit) o (1812)
0 (17771).
1)
(
Σ
3
1 - 11.1
112
-=0
j=1
< Σ
g
|w/ - -1 + E - En5T| (1/-((,2]]
0=
< Σ
gr (t - 7) (o(1)0 (a(1))((1))) + o
= -0 (((")))))))
0
()²611))
7=0
By the assumption of this lemma and the analysis before Ti. we know that for all 7 < t, the M(TT), Mn5T) in Lemma 31 are upper bounded by O (z++) and O (dl/4), respectively. In the theorem we consider the training period before TsGD,2 the time T in
Next let's bound
v12 [i, ii and (c) -
37

---

Under review as a conference paper at ICLR 2023
Lemma 31 is set as 2. In the following sections, we will prove that TsGD,2 S d logl vafc)
Then by Lemma 31 we have with probability at least
3. for VT S t and Vi, E [d].
Dgi (T) [t.3]]
[r.]]
-9
S
0
d o
da+1
log
d
o ddog
da+2
d log
1/
dat1
I δ
d
13
Do2 =
92i
(T))
(T)
- 0
d 0
d
log
d
+ +0
dàa.
d
log
d
17
dat1
= Ö
d d
17
By picking az S
d0/241, we have Dgi 1i.51
1/22
S n (a#) and
Dggi (T)
S n (d+).
which yields
012 012 (t, 5
S 1)
8
Dgi (T) [K, 56 =0(Pttt
==0
422
(t)
S 1)
---
Dgg2 S 6 (
(T)
Combining the above bounds, we get that Vi,j j E [d],
1" [1 < ô Ō (n"
(t) 22
S Ô (1²aa²))
By the analysis in Section C.2 we know that at time Ti, for some
in E [d]. (TT
(T))
θ
1
uio
and for vi, j E [d]. we have diT)) u
(Ti)
á
and
u; (T;) (Ti) 52
1
which gives
d
us Vi,j E [d].
(1 [E
ri
[i
Ö (1/d
18
22
(0)
2'
(t)
Ô (1/dfat4)
5
<
(Ti) u
(t)
u1
(7i)
(T:)
c(e)
(Ti)
e(Ti)
11
(25)
Hence get the bound Vi E [d] : :e") 2ō (
C.12 PROOF OF LEMMA
Let's first try to bound the length of min Tz}. More formally, we prove that under the conditions
of Lemma and pick 7) S O ( we have that min Ta} s d" log( /d/+)
Under the conditions of Lemma 10
we know that
Vj E [d:
(1!0w(0))
5
2i w [3.1]] = (17(t),
(M/50/99)
22 ri [3 (17())
(1) (0)
38

---

Under review as a conference paper at ICLR 2023
Combining with eq. 23 we get
b((+1)
=E(h + (w w w + W -1((t + (w - -w:")
w7,0)
-E(t) - -m,(() W (0 w + (1) w( m. W2
(1) E + w, 11 +0
=F(()) (1 T-- ( w( - n | (3:"((?)) + +
Then we have
E(1-1)
1)
- [E(0]
|E10|]
1 - -m w 7 //
11 - (120))
12
2
< 1 1 Ttt [w/]1()) 100011 + (rp
When T1 St t < T2, we have proved that e( is increasing over time in Section C.8 which implies
that (/ 2 C (T) since el0) (Z)) is the leading term of (( Combining with т)
1
gives us
2 1 1///
|101 + +
0
(1721))
|
o (17-44)
=
z
+
C
W
n |
ndd
(-nc ( 2 (t (T-) 0 (va).
0
exp
where (i) uses
| ||2 = O( (vvd). By picking 1 S 0
( and noticing that [ (11) 2
++
log Va/e)
Ω ).),, we have
ndd 1
L Hence when / /- T- 2
we have that |E(0||22 S
1 w ?
| 11 Sc
i.e
log V d/e 1
That means after at most O
steps from Ti, either t > T2. we have || In
x
log Vale
< 5 dd log (Vafe)
other words, min (T,, 73} < T +
"1
Now we are ready to bound eq. 18
d" log( (og(y/o dd/e)
Combining min (T,, TT} O
and Lemma 1i that for t + 1 S min(T2, Ts).
Vi E [d].
1+1
(1700a-111)))) so nndaat log va) = =o elog log V d
Σ st+11
T=T1
Lemma I tells us that 5, =6 -
Substituting these bounds into eq. (18) completes the proof.
39

---

Under review as a conference paper at ICLR 2023
C.13 PROOF OF LEMMA 12
The proof in Section C.8 tells us that for Ti S S T2, clT)
> 0, Vj E [d] : : > O. which gives us
0 S
d(++1) [ and
of sarry. By Lemma
we have that
o(++1)
F R((+2) [i,
R
2:
LA S di
0
O(e)),
0 x
(ce)),
(Ti)
(1+1)
c(++1) (Ti)
u;
u
which gives us
R2 u (T))
d
(2i)
R
((E))((++3)
D 1=
(Ti) 2
-==1
ui
O(Eo).
c(++1) |1
c(t++) Σd
Ei=1
"
(Ti)
c(t+1) Σ
41
(T))
Lemma
9 tells us that
R((-1)1 1)
(T))T R((11)
11 R(((t1)
And we have that
(((++)) (TV)T R )
c((+1) Σi=1 11
(Ti)
R (1(+))) [i,5]
clt+)) Tu(Tii]] (1t1))
clt+1)
Ei= 4
(21)
2 (1+2)
((E))((++1)
Ei (Ti) (11))
Ō(co)
(T1) 2 ((+1)
c(++1)
Di=1
( (1111)
(2
((1t1))
(1+1)
R
R((t))
(c)) (1++))
i=
R2
-i=1 "
ole+1) (1+1)) [11 get1)
c((+1) Σ d
(21)
of ((11)
c((+) d
(Ti)
Σi=1
1=1
= -
Therefore
(est))ai (T)) R( (
R((+1)
R [[(+1]] R
R(((T1)
3j
c(++1) ((+t)||u(tii) 1| sgett) S Ō(c))
elt+1)
[++1]
((+1) ( ( ((41))
By Lemma
W
1(i+1) |
+ R (1+ a (t)
R (1+1TT
Then we have that Vij E [d].
(w((tD))((+t))) ((+1)
1e (TTii)]' 13 (141) (1 1 () where (ep |o0-1// zO((o). (24)
(w(+D))((t)))
Since 1 < T2, we have Vj E [d]
((), which yields
((++))
(c-1)
< 01).
(25)
A;
((+1 11
which proves eq. 21 since Vj E [d]
0 S 00+1)
S Ō(co).
12 (etry
40

---

Under review as a conference paper at ICLR 2023
C.14 PROOF OF LEMMA 13
(A) Under the conditions of Lemma 10 and pick 1) S O
we can apply the technique when
proving eq. (24) to show that eq. 24 also holds at time t. Since
a S Ō(eo).
we get that
0g = (1) A,
(0) = a
A + (.))
where
c - Ō(o)).
Substituting into the time version of eq. (4
24 yields
Vj [d] (w)0w(")) a())(()) (T) || Aj (1+ where ( I Ö(E)),
That gives us
Vj E [d] : E/ Aj (a()2(1)
(T) ( 2 1 + a ((t)
| (
)
Since t < T2, we have 1350) < VEb Combining with A, 01), gives us ((1) || 2i) - -1
-St (v). Then we can rewrite EZST) as Vjj E [d].
(1 d(t) |uu
E = As
=A
1e taiii (a))(())))
-1)
1 c(t) |u(ti>||=
(00))
= Ag (0)) 11 |2
- (1
(1)
where
(1) CE)
Ö(V Hence vi,j [d]
(1).
(B) Note that we assume Vj E [d]
θ
then we have for E [d].
c
Σ
(-E:!))
u
(t)
d
1
= 0 (Vi)
= θ
clt)
(-1)00)
clt)
(-E((0)
va
-=1
1=1
c(t)
(12)
(1()(())
Vi
Then for t + 1, we have that for j [d].
1/ + mcll)) (-E(-)
(++1)
c((+1] c(t) + m (-E(t) v(() θ 1
Vd
(C) Combining eq. (25) and Vj E [d] : A₁ O1), we know that
1++1) 4 m))))/ 2 2j < O1),
(++1)
which yields Vj E [d].
uu m// 11 ( 2
"/ c(t+1) 0(1) = o (22) (1tth) 2 |u/
(++1)
llt+1)
I (+))
0(1)
(26)
Hence Vi E [d].
=
c+++)) |u [Ti] 0 (d!')
+ u22
(1-1)
- ((++))
(uu [T +
o d1/4)
u
(T)) 1++1) 0 di/4
=|
1 m (T)))
of
+
| [ [c,]]
d1/4
41

---

Under review as a conference paper at ICLR 2023
where (1) and (it) use Lemma 11
|
(D) The fact that Vij E [d]
(cu) was already proved in Lemma 12in eq 21
To analyze
A,
R A we first prove that 7000dd d(0) O.
It is not hard to prove that eq 26
26
also holds for time f. Recall that d = c( || +
(Th) (T))
1 (T)))
and Lemma tells us that 0 v c0e3
(Eo)), then we have
(1) = (c())
1 + n R
(1T
u (T)) 00
(vd)
Under the conditions of Lemma 10 and pick n S O 1 we have that 1 - ml(())())) 2
1 mell)) 1- d(t) > O.
R (EE)
The assumption Vj E [d]
together with c(t) > 0 gives us
Ag
m
Raj
(1))
S Ō(Eo).
Combining with the assumption
A( (EE)) yields
I((+1)
(1 (c()) (((0) + me(()
p(0)
Ps
Rigi
Vi [d]
< Ō(e)).
a(++1) A
(1 d(( all) A; + mcl() A;
D ANALYSIS OF ADAM
Note that A = m _Y AZx = m L Denote g = Vmn 1,2. We have that
o W (w!"'/(" -A) 92 (t) (w(" w/( - -A) W/O0
Let A(0, A and (1) k = 1.2 be the corresponding batch versions at time t. Let E())
w (1)) A. and denote E//t) as the j-th component of g((), We also denote Aw := u2
(t+1)
AW [i,]] = w(t+ 1) [
[i, By eq.
2b. the update equations of Adam are given
by
35
91 [1 = (0) E00
92 = (F(() w/(1) w
(
1/t = 1)
gtt1
1 = 11 mi (1.
(1 - 31) Et 34 9i (T)
[i.
(((+2)(1,3)
= =-/
1 V 32) Et= B5- (617)1.11)
Vo [i,5]
+
(1- B1) Σt= 87 9l (7) [1,5] + (.
V(I (1 32) Σt 3/ (91 (1.11)
(r)
+ + [..]) +
((+1)
(1)
m2z
(1)
(1 B1) Σt=
8(-Tg(()
925
(7)
1)t
7/1
V 0/2
(t)
(11 32) Σt 8/
(x(7)
921
+
(1 - 31) Ei= 3
- (T) (t)
92s + +zzmi
=-t
B2) Σt= 3/
92
(T)
+ +r +
(t)
(27)
42

---

Under review as a conference paper at ICLR 2023
where Doj = 010 - and Dye (9)
(1)
(8)
: 92
-22 and
t
10 [ibii:] = (1- 31) Σ
81 - "D9 [r,j],
T=0
[ibi = - -2)
82 7
22 2t, D9 [t,j] (Dol (Doi))
7=0
(28)
22n, = (1- B1) Σ
8 81-1 DDg2
(T)
T=0
(1- 32) (1) 22,,
Σ
ot
292i Dg2i + (Do:)))
P2
7=0
Denote the i-th coordinate of W2W1 and A as (W2W1); ); and A,, respectively. By Assumption 2 and the assumption that Vi E [d] A, > 0, A; (2(1), at the beginning, w.h.p., Vi E [d] (W2W1); A < Based on this, we divide the training procedure into two phases (note that these two phases are different from those of SGD+M).
1. 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 E [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 T₃) is defined as
T inf { > 0 Bi [d] E 2 -Vnd
In the second phase, we define some time points.
Definition 5. Define Tg == inf { Tj Ei e [d]
92i S d
For < Ti, we have E [d] : E < by Definition For t > Ti. some E may flip the sign and become positive. For certain coordinate
define the following "flip time".
Definition 6. Define Ty. = inf { Ti E.( 2 --vdd Define T₁ == max; Tj as the largest "flip time" all i [d], i.e. the "flip time" of the last E, which flips the sign. Moreover, denote T := min (Tg. Ty}.
We first show that after few steps in the first phase, W1 will become approximately rank-1 matrix, as described in the following lemma.
Lemma 14. Under Assumption
2 and
suppose S
n g222
d/3/4 By picking 17 S O (55 . S
and 32 = B7, there exists time such that w:h.p. p. for tim St Tí.
Vi.j [d]
Wai
(e)
sign
112s
(0)
17 (t
time)
(!)
[(9 [i.
sign W22
(0)
17 (t time) R [1,
where (F.)) mlt-(ar)
5 (Vm olt far
R
Ö (vn ml (m0
7(-hma)
Specially, when t. = Tí, we have that
Vi,j E
W2 sign
(2r)
(0) (0)
W22 17 (Ti time) + R2/3 (T1)
W Ti
[
sign
W (T< time) R (T( [i.
(0)
43

---

Under review as a conference paper at ICLR 2023
where 1i (Ti - timc) = vil
and
(Ti) ( [2,
(Ti)
(
Ö
77 (Ti time)
da-
(Ti
time)
The following lemma tells that this approximate rank-1 structure is preserved when Tj st T. gist.. By picking 1J S 0 س3ر i S
Lemma 15. Under Assumption
2 and
suppose <
and 32 = 18. we have w.h.p. for Ti St <<< T.
Vi,j E [d] 2
(t)
sign
22
(0) + RR2
W [ = sign
1226
(0)
vy(0)
+ R
[2, [i,
1 [i
where
=
Ō (V7+
1
I <
17 +
|d0
da-1/2
d d
and that L (w(t)) S Ö (nd*))
Now we are ready to prove the Adam part of Theorem 1
D.1 PROOF OF THE ADAM PART OF THEOREM 1
Define TAdam,1
line
Note that this choice of AAdam, 1 gives 1/ (T'dam, -ine)
By
nd
picking 1) S 0
S
and 32 = BP. we can apply Lemma 14 to get that Vi, E [d]
22: 1)
- d7
W
Tadmn
[ = θ d7
and therefore vi [d] E (TAam,)))
Q1)
and L (Tatm...))
(d). Define T A'damm T. By Lemma 15
we have L (W
Ô (yd')).
For any p > 0, by picking a 2 Moreover, combining Lemma 14
D+4
we have (W (W(Tadmma))
Ô (nj") S 0 (#).
3
we get that when E [Txdam, T 2), the conditions in
Lemma 30 are satisfied with 8 = =
17 + d
The i-th component of the u vector (denoted as and fax)
u&) is sign W22
(0)
That means Vi E [d] uf
1. Then we can apply Lemma 30
median(w,])
and get that
1 8" 1
2 max; max)(u,)) (u,)
+ 8 2 max, max,(u))
Radam (t), Adddam (f)
median(us, (u
median (1;))
1 1 6)
2
1 8) i
27
1
RAddu (4), Radda (1) = 1 =()) ①
1
d
D.2 PROOF OF LEMMA 14
For some time I, we introduce two conditions. Condition 1.
Vr E [ sign (9) (1.0)) = (0)) (t.)),
(1 - B1)
Σ
(7) (1 (1-7) [i,
0i
2 (({))
7=0
Condition 2.
1
VT E [H] sign
92s
((-T)
=
826
(t)
(1- - 31)
Σ
(T) (t-r)
42
2 S(E)
r=0
44

---

Under review as a conference paper at ICLR 2023
Next prove that, under Assumption | and by picking 0 paa i م and 3== B. there exists time > such that for tine- S Ti. the weights can be approximated in the following way.
w 10 [i,5] = ( [i
sign (6) [t.3]) + ((.
(29)
(1+1) = W/2 (1) (sign 922
(1)
+ (1) 22
where (1) (1.5] = (VT),
(1) Ô (vī).
Before we dive into the proof, let's introduce some useful lemmas.
The following lemma reflects our key idea: converting the exponential average in Adam to finite-step average, and trying to bound the stochastic error terms in eq. (28).
Lemma 16.
Under Assumption w
2 and and pick B2 = BR. Let M
[i,5]
M == max,,j
G( == max, je(d, "l [..i] and
Wi
d == mex
92i
(7)
We have that w.h.p.. all S Ô Van
and and E [d].
(1 31) Σ 8ig
7)
[t. 5l + (i
( = -
Va- B2) E 35 9i
2 (1-T)) [i:51)
+ Eia ii +
(1 (1- EtH 35 3792i
((-T)
(1)
Aw (t) -tt
92i
+ E
Va- B2) E/H 85 (2i (1))))
(t)
2d,t +
max 61" Nal,
(alp((((')) (cf9)² 2 (cd')')
where 11 2 1-8 log
and
75"
[1/11 s O(nE3) 0 (D(*),
12 20(n)² +0 (D00D/)))) |/ + (D())
(0) (2m. s0((5))) + + (D!").
(8)
50( + 0
D _ + (D!9) 2)
2d,
with
1 so d A (M)))²
0
d1/2
M2
(t)
d3/2
1)
1
1 < Ô dz (M)))² Ma oy
dL//2
dM(() (t)
d3/2
1
n
2222g2 d.3/4.
Corollary 2. Under the conditions of Lemma
16 and suppose
<
Consider any t S
Ö (
Van M M, <ō =0(t)) G <
G Ô 0((v) (Vd) then H in Lemma can
be picked I- log d and we can get that Vi,j [d]. (
ngz
c ( 5]
(4)
(
2n.
(t)
2d,5 S
The following lemma analyzes the magnitude of weights during short period at the beginning.
n3/2za Pick & S _ then there exists
Lemma 17. Under Assumption 1 and 3 suppose a < dis/4.
some time point time E (H, such that w.h.p., for t S tines for every i,j E [d]. Aw [ < <O(n)
Awge
(t)
< <((n),
( 50 d3a+t
Ω
<
(t) W
0
d
45

---

Under review as a conference paper at ICLR 2023
Specifically, when t = tines we have sign W
W22
(lax) = sign (wifhel = sign
(0)
[1,5] = θ (3-=1)
(1.mm) and 0f (ad [1.5] z S2(E), (2 2 Q(E). Moreover, Condition
and 2 are satisfied for t = time. The s [i and 82i in the conditions are both -sign
(0) 222
The following lemma gives us lower bounds of [2. and 92i (0)
/ 2 andd 3 suppose a S ga/zz2 Pick & S v VT n1 S 0 (
Lemma 18. Under Assumption
Consider tine in Lemma 17
We have w.h.p. for any t E [time Tı). and for Vi,j E [d].
sign
Aw1 1w [..j] = sign
Aw
(t) = sign (1 and that Vi,j E [d] 9 [t, 2
(
7
92
(t)
2 Q (vd... Moreover, we have VT S t. Vi.j E [d] [ S
(T)
I Ō and
91 [5.00 < 6 (
(r) 92
S Ō Vil).
The following lemma shows that when line S < Ti, we have Vi E [d]
92
(0)
>
92i
(0)
92:
(t-1)
and that gi (1)
[i
9i
(1)
[i.J]
9i (t- (t-1)
Lemma 19. Under Assumption 1 and suppose <
213/4 Pick { S v
17 S 0 (7te)
For tim in Lemmal
17 we have that w.h.p. for time S t < Ti and Vi E [d].
9i
(t)
[1
(
[1,5]]
92i
(0)
92i
(1-T)
= Ö((TT)
Ô (v/TT)
(30)
91
(4)
[6,5]
92
(4)
(of (1.33))
- (c.5))) 2
01
2
(31)
2
92i
82i
( + Ö 0
(1)
92.
Equipped with these lemmas, now let's prove eq.
For any t E [ime Ti), by Lemma 18. we know that M _a S Ô ((t)) and that G <
δ G < 5 Vd) At the end of the proof for this lemma, we will show that Ti θ
Vdn
Then pick H := T-9, log d and apply Lemma 16 and Corollary 2 to get that, w.h.p., for all
nez t [ime, R3) and V E [d]. eq. can be written as
(1- B1) Σ" F.. (1 [i + (t)
(W = m
Va- B2) Σt1 35 0i (t-r) [i,
+ + (i +
(32)
Aw (1) = -t)
(1- ΣHH
3 F2i
(t-T)
+ C2n.
(2)
(a B2) ΣH 83 12
(t-r)
+ (t)
)
2d,i +
where Vi.je E [d]. 1. 19 [1, 3/
< &((g2)).
2d,
46

---

Under review as a conference paper at ICLR 2023
Let's first look at the update of W/0 [E, For L in the first phase, we write the RHS of eq. (32) as
(1 - 31) 11
[1, j] + (M
(1 B2) E+= 35 [1 (3.1))
+ 44b,,++
(1- - B1)gg (1) [ 5] E H1
3. + (1 B1) Σn= Bi (9f (i.5) - 96 (..1) + [(hli.c]
(1 B2) (110 [t. Σ 03 +(1 - 32) -HHO. B3 (0f- 2 (0l ( (1.11) ) + (e +
2)
(1-
g 11 11 - / +
(1)( + [. - (1) (.
V 13.01) (1- + + 4. [i +
where
H
[. : (1) Σ
1 (1.0)) - -i (..i..
r=0
11
(B.)= - (of%
(i := (1- - (2) Σ
B2
0=0
We have already shown that [00
</
1id |
- Ō( By Lemma 19 we have that Vi, j E [d],
f b. s (1-31)
Σ
37
91
a1
[1.]] - [r
r=0
S o
ò (v (1 31) Σ 8 9i (0)
[i.j] [
Ō Öv
Similarly, we have Vi,j E [d],
H1
Id [1 < (9i (3.1)) " (vVi) (1 B2) Σ B + (0i (b)) ( (1 (22 > 8
T=0
r=0
(of 10 (0/0).))) 2 Ō(VT).
By Lemma 18 we know that (0)
[i (1015.3]]
12 (2(((()) (VT). Then we have that
Vi.j [d] :
(1 || S Ō < Ô (VT) of [3.1]..
Therefore by Lemma 33 in Appendix G we have
sith (0) [i (1- (= + e
[. [. + + 1
0 19 (..i))
2 (1- ( + 1 [. + [0 +
1 0/141
sign (..))))) + ह (..i))
where 12(0) [v.]] :0 (VT).
Since B E (0, we know that log 8 - 8
O.
Then choice of H gives H
log 28²
log 25²
1 log d 2 T-3
and H > 1- log d
which implies that 8 , S ng
nga
nga
log
Hence for 2 tine
H, Mit VI- का 1) 11 g.
1-
0
TT
g ++
n + O(r)).
1- 3
47

---

Under review as a conference paper at ICLR 2023
Combining all of the above yields that
(1 (1-31) 1-- Tig1 (9-(-~1 (1--)) [i,
(1(+"(1s))) = [t,) - 7/2
(11- B2) Ex= 35 [g) (1. [1.51) +
= =w(4"(t,3l - 1/-
(sign (2.5) + ही (t.5l)
V 354
-w("(e,,1) 17 sigg (01"((.,3) t (i.1),
where (6) [1:
Ô (V7). The proof for 2221 is similar.
(t)
So far we have successfully proved eq. (29.. By sign (wr(' (1.1]) sign (Aw? (1)
sign
(0) W22
we know that sign (-9)² [..]) sign -22 (1) sign (226 which gives
in Lemma 18
Vi.j E [d]
22 sign w1 1j (t line) + R!4
sign 1 [.
122x (0)
17 (t time) + 11( [1,
(ine (/(t-ba)) Now it suffices to
where
(T+ +
and
Ô (T
(ic)
show that Vi,j E
[d]
W2i
(fnx) 0
do
w (1m)
[2,5]
O (2a), which is implied by Lemma 17
Finally to complete the proof, we show that Ti θ (.nn)) When Ti, have Vj E [d] (T)) (T2)
EL 12 ((T)) [ [
(1)) Combining with the above results, we know that d line )2
@(1). i.e. (Ti time) e (aa) In Section D.5, we will prove time = O 2
Then we
have Ti = 0 (1)
Vin
D.3 PROOF OF LEMMA 16
For certain and H, we write eq. (27 as
(1- B1) DM
8T9 (1-7) [i, + (
A [1,J] = Th
V (1- B2) ΣH 35 (0i² ( [1,
+ + +5
(1- B1) Σ/r 3
(t-T)
+ (1)
92
C2n
Aw = =1
(11 - B2) HH 35 12
(1-T))
(t)
+
2d,5 +
48

---

Under review as a conference paper at ICLR 2023
where
[i == (1) Σ
"9 [1,5] + (1-B1) B1) Σ
(t-T)
++++
H+1
94 1. = (1- B2)
Σ
B2 91 ((.57)
2
T= H+1
1=9ld [0,5]
(0) 2,, = (1 - (1)
Σ
((
2 (0)
02i
T= H
92d
and r (1) (2)
22,33 2d, are defined in eq. (28
(0)
c( _ (C (c99)
max
Since 82 = =9 < B1, then if we pick H 1-5
( 2 log
we can get that
7/5
G 1-8 n
((
H 1 log
H 32 log
H
log
H 2
log
Hence
ng²
1-31
ng²
1-
can apply Lemma
32 in Appendix G to get that
Vin
[.,]]]
Hid
[ .
92n.
92d
S ng
Pick T in Lemma 31 as of order Ô
By Lemma
31
we have with probability at least 1
Vdn
for all $ T, Vr= < f and Vi, j E [d].
D91 (r) [1.5]] = oof [ .ं x [i. 5 ( (M
d1/2
+
M
d3/2
v "1
:= D
D922 (T)
(TT) ( S ô d' (2 "1300 oo 17
d1/2
+
ov
d3/2
:
D
17
Plugging into eq. (28 gives us
10 in
< (1 - B1) Σ
Dgl (0)""()))] S (D(t))
r=0
S (1 - B2) Σ 8 |0| " ( |1.11 + Dol [t.1]
s 0 1 c(S) (D!-)) ")
(t)
S (1 - B1)
Σ
8 DoLE
(r)
s o (D(").
2n,
(t)
< (1 B2)
Σ 8 - 29 (
(7)
(7)
Dg (r)
so (D + +(D5)))
2d,5
D.4 PROOF OF COROLLARY
Since 000 (0( va ch 2o (V)) <ō(va)
then H 1-0,
log ng²
is bigger than
max (c9. (.))) (c(1) (c!!)) 21
I- log
n
49

---

Under review as a conference paper at ICLR 2023
By M _M I =ō (tt) .t).... of s (tt) _..² =o <vva) and the assumption o= S
12/2z²
//3/t we
get that D and 1/") are upper bounded by 1 I Ö (d 1/2) and D - Ô (011/41-112)))
which yields Vi,j E [d] 1i 1 al cia [i s]
(0)
2d, - O
(1)
22n
D.5 PROOF OF LEMMA
The proof is based on the following two lemmas.
Lemma 20. Under Assumption and 2 we have that w.p. at least 1 -.. for every 1 Usis
vT d3n
24 2i
S
V
log 8' 2
2
2d and that w.p. at least 1 - 5 for any given > 0,
s
V
2 7
log
Lemma 21. Under Assumption /
and
suppose S 17 2200 Pick B2 = 88. g E (0.
1).
< 1
4.
Consider any time point t Ô
Vs t, E [d]
(t.5) S Ô (t)
(7)
dan
Ô V and g (7) [t.5] S
s Ō ( (va) we will have
92
|ww1 (t)
[
< (())
AW
(9) s O((),
where the Ô notation depends on = - log T
Furthermore, for certain i E [d], Condition 7 (resp. Condition is satisfied, we will have
sign (1.1)) = -s [. Aw [1,5]] = θ(n)
resp.. sign Aw'2) (t)
= -822 A (t) = Ō(())
(1)
Now we prove Lemma
17 Define La == inf {1: Ji,j ( [i, 51 >
w22
(t)
d z
11
Now we
want to find a time point Time before tu for the lemma to hold. During the period t < tdd we have
vj E [d], E, = (1) (which means ta < Ti) and therefore for all i,J E [d], g [ S à
and
92 < 1. Then we can use Lemma 21 to get that for 1 S min { (d: J we have AW [i, S
Ö()).
Awlz - O()). Hence ta 2 Ω vd
Define Exign
= inf {t < min (td, vān
1
3i E [d] 122
S
1 }. By Lemma 20 w.h.p. Vi E xx-1
d
[d]
(0)
v d
combining with A Aw!y < Ó()) gives us that w.h.p., taign
(t)
OO(n)
23a
Ω
gd
Now let's analyze the behavior of W, during the period L < biign Consider any i.j E [d]. By definition, sign
W224 Note that F.(() (1),, then we have sign 19 (1.))))
W24 sign
-sign
W22
(0) and that 9i (1) [5, Ω
((E) by our choice of g. Then we know that
Condition is satisfied with $ [t,5]
sign Whiz (0) (for all H <t fiign), which by Lemma 21
yields sign Aw (t.))
sign
(0)
and
Aw
[ 5| O(n)
Lemma
20 tells us that w.h.p., Vi.j E [d]
Ô (2tz). For any i.j, if ini- (0) then for the following steps before Esigns we will have
W 00) |.-
tially sign W [t.2)) = sign Wiys
sign (i.))) = sign (w2)) (0) If initially sign (w [i.33) * sign W then after at most
la = δ ाक steps, Wi [t..]] will flip the sign. Note that 40 = Ô na
is smaller than bsign
50

---

Under review as a conference paper at ICLR 2023
Hence we have shown that at some time point to, we have Vi, E
[d] : sign [w1'1.,31) =
sign w2 = sign (ubb))) Now we analyze the period t 2 to.
When to < t S fsign we still have sign (AW (..5)) = sign
(0)
and A |1.31| = ((n).
25
Combining these two with the fact sign (w(t) [i,51) sign
22
(0)
we know that for all t E
(0) (lo, Lign), sign (t.il sign W2y and that Vi,j E [d]
w
| +
(()). Then at certain step tim which satisfies time to +
(H, taign). we will and therefore |o5
ydba+1
have Vine 11 S T S mi Vi,j E [d] | [6,5] = Q
d
E w [1.)) 18 Σ (..0)))) o
(2(E)- For L S 4in we have
d
V E [d] : w [i,s]] 0 dialt)
Since time < Esign= we have lubur
For S tines note that Aw I Ö()). line
la + θ nddat1 θ
combining with the upper bound in Lemma 20 yields
S
22
+ Iine Ö()) S dda+1
d
Moreover, Vtiie - H S r I ti Vi E [d] sign
sign
W22 Then Condition is
satisfied with 8/ sign W22 for 1 Line In the analysis of gi (2)
(0)
[t,j], we have already shown
that for all & S fiign (and thus for t tine), Condition is satisfied, which completes the proof.
D.6 PROOF OF LEMMA 20
Since for X N (0, o²), we have that P([ St) S
2t
then for fixed
/2na²
2 2/d
P
(0)
π
din
V2T
2/
d
Then by union bound, we have that w.p. at least
d - for every S i S d,
[22
(0)
_===
As for the upper bounds, using the Gaussian tail bound and union bound, we have w.p. at least - 5,
Vi.j E [d]:
(0) W224
2
log
2d
(w()(. [i-]]
S V
día log
2
2d
2a
8
D.7 PROOF OF LEMMA 21
Now analyze the magnitude order of Aw (1)
(t) [i.j]. The analysis of Aw is similar.
For t < 3 Van
By assumption, _ MM (t)
s Ô (1))) c s {
S 2 {va). and
< 2/2222 /13/4.
1 log d and apply Lemma and Corollary 2 to get that, Hence we can pick H == 1-3,
w.h.p., for all t < δ Van) and vi, j E [d]. eq. (7) can be written as
(1- B1) HH Fig [i
(w [1, = -1 Va- B2) IH 35 9i
2 (t-T) [i,
+ t (.. +
(33)
Aw (t)
(1-11) B1) E#H 9 95 922
(1-T)
+ 62n.
(1)
-1)
/ (1 B2) DH1 35
(t-r)
(0)
921
+ Ed, +
51

---

Under review as a conference paper at ICLR 2023
where V E [d]. <in 1 51]
(0)
(1)
- Ö(n5)).
Lid
(2m.)/
:
(dd.)]
On one hand, using in [t 0 + セリ
1 [
- Ō(n5? and B2 = Bp, and v + y 2 vI
Vlvl when
we get from eq. (33) that
(1 -1- 31)
Σ = 87
(t- 7) [1,3]] + +O(EE))
S n
(1 (2) En= (8i0) (1.0))
-(((T)) +
( 11 (1 (--B))V H I Eh= (BT91-1))1)) (1.51)
+ +((nE2)
50 (V(W)) Ö(1),
32) Σ (8 1
[2. [2.51) +
where (i) uses Cauchy-Schwarz inequality for the numerator.
On the other hand, when sign (oi-- (1.5)) sign [6
(1(tH++((( ... sign (0f [(.5))
(9) [[. we have
sign E.T 1990""11.3 [1,5] = 81 (i,1,
(1)
Σ 3
[t
[1,5]
2
Σ
(i) 2
[i.
If we further have (1 - B1) E 01
H
1 [3.5]
2 (2(E). then combining with [olli-31 s
(1n52) < we will get
H
H [1. + cf2(e..)) = -ign Σ
sign (A = -iign
Tig (c-)) Σ
TT
(1,51
=- (1..5.
Using V + S Vll + /lll. we obtain that
(1 31) 31) EH 1/10 9/" 1/00-) 11,51 - -Ō(n53
B2) Σ/H (5T (1,51)
+((TTE) +
1-9- El= )/()) 91
(1-7) (i,5)
a(n))
7
(1 B2) EHO (819) (t-m)) -) [1. 2 36
max
Together with the upper bound completes the proof.
D.8 PROOF OF LEMMA 18
The proof is based on the following lemma, which gives a coarse analysis on the magnitude of weights and their increments per step during the first phase.
Lemma 22. Under Assumption and suppose a I ". Pick
E S min
{ {v 75u 220 for
1 }
Lin in Lemmal 17 we have that w.h.p. for all line S1 ST,T. Vi,j E [d].
sign (AW(") (1 sign (Ault)
(t) sign (0)
AW [i.5]]
=()),
Awi
(t)
= =
sign (w/( = sign
(w:2)
(t) =
sign
(0)
|W/
[.. ō
(c)
Vd
Vd
Specially, at the end of the first phase (t = Ti). we have Vi.j E [d] W
(Ti)
and
(T [i-1]] = (tt))
52

---

Under review as a conference paper at ICLR 2023
Now we go back to the proof of Lemma 18 For tine St< T1, since E." = =(w -A, =
Ei W2i W/ [i, Aj. we have,
(1)
AE j = 14 b((11 - Ei
10
W2i W
(++1)
(++ (1. [i,5]] w2i
(2+1) ( [i,j] + u2 w "et""/"" [i, wf? (34)
W2i (w [. + Au ( (i.sl)
(++1)
Combining Lemma 22 and eq.
gives us bj E [d].
AE > 0,
AE
(++1) " [1 + A("
222
(35)
1 3 (nva)
11 Vdl
Let's first analyze gf [b.j]. Note that
Agi (t) [i. 5] = w12s E
11+1)
(t+1)
W/2
(++1)
E
+
Wiz
(++1)
wigi E)
(36)
222
(t+1)
(E
+
Aw
(1)
(t) ;
(++1)
AR = sign
(0)
while sign (Aw (1)
r2(T))
(0)
where sign
W2s
sign
Now we analyze the sign of gi
(t)
[i, J] when Line S Ti. Using
W22
(fmc+1)
Ö c and eq.
(35)
we get that
W22
(1i
A
(fftta)) S Ō de
vdn) While on the other hand, Awga (tree Lr((tw))
(n).
That means sign
Agg (fun)
[1:5]]
= sign
112 Note that sign
gg (fua) (1 [1,5]] = sign
Wi22
(fux)
-sign
we know that (1/2
[i,
will increase when = line-
In the following steps,
91 [
will keep increasing as long as
!
Aue2 21
E
(1)
(e+1)
A
AEf"
Since
w [1,
(
keep increasing while [w Aui
(1)
remain θ(r). by eq. 355
we know that the trend of
A is to increase. On the other hand,
(4)
E
keeps decreasing since
E 0 while AE(") 0. Then after some time point will have
Aw
(t) E(0)
((+1)
AE (t)
21
and in the following steps [61" [ will have the trend to decrease. Specially, when Ti we
have E E = (vndt) and [
= Ô + 22
(e+1)
= â ( by Lemma 22
which gives us
AE
(++1)
AW
1
(2)
[2,51
+ Amm w (1) [2.01|
d
6 17 V2
1
e (Iva)
Hence
((+1)
A (
Amyl
(1) F. é (nvnd).
Therefore we have proved that when time S t < Ti, the trend of 9 [i.5] is to first increase and then
decrease. order to prove 9 [i.
á (vT). it suffices show that g ((=m) ji
á (vT) and
9 (T)) [i = (vT).
When = tinc.
91
(1me)
[1,5]] =
112
Fs (nc)
= Ω
dīa
0(1) = (v
53

---

Under review as a conference paper at ICLR 2023
When = Ti, we have
(T) [2. = -112
(Ti)] | ê
(vā Vg..vdl = θ (vr).
As for 92i since for Vi [d]. w [i, for different j have the same sign. Combining with
(t)
Vj E [d] F < 0 gives us
!1) 021
Σ F w/ [i, [t.]]
Σ E3 W [1,3|
Then it suffices to show that for line S l < Tı, Fr) W
(1. |1.01|
á (Vi).
which can be proven
using the same technique as above.
Finally, for Vr S v E [d]. note that the upper bounds of [w/"'(t.]]
[i.5] and w22 are already given
in Lemma 22 As for
(T) [ j] and 92 (T)
we have [t.
[t. .i = w ET) = Ō (Ja) 9 S
[ = δ (va)
D.9 PROOF OF LEMMA 22
For any i,j E [d], and any in the interval Mime, Ti], we prove by induction that
(A) w/ ('') [i. Ō va)
(t) = (taz)
(B) VT E - [t, sign (w(") (1.51) = sign (w()) = -sign 2222
10 [.-] > 2(()) 925 > 2(E).
(C)
9i
The base case t = line was already proven by Lemma 17
For E [liee, Ti). suppose (B) and (C) hold for time 1 and (A) holds for all r E
[Lines 4.. From (A),
get that V/ E [tine, t) |91
(T) [
(T) (T)
W2: E (T)
- vaa
(T)
d s 1=
E (
W (T) [
923
Vda Since when t < Tı. Vj [d]
E O, from (B) we know that Vr It H, t)
sign (g(*) [t.5l) sign
12i
(T) -igg
(w22
(U)
Combining with (C) tells us that Condition
and
2 are satisfied.
In Section D.2 we have shown that Ti = 2 ).n)
Then for t E [tine: Tı), we can use Lemma 21 to
get that Vine S T I t, vi E [d].
sign Aw = sign Awwz)) = sign
\w [i..]
O()) |W = θ(1).
Since when f = lines sign (W/(tm) (i..]) = sign
(1:2)
(tan.)
=sign
(m)(0)
We get that for line S r S
Vi,j E [d] |w/(+1)) [i.5]] = ( L .
+O())
(T+1)
w + +O(r)
Now for + 1, we have Vi,j E [d].
sign (w(ttt)(i,) (1.1))
(0)
|w<(++1) [i.
(r(taw) ( [1,1]] + (t+ Line θ())
sign
sign uu (((+1))) sign
(0)
(e+1) =
wu (tme)
+ (t + 1 tine) (())
That means VT E [t+ 1 H, + 1] (B) for time 1++1. 1.
sign (w(7) (1/1))1.3))) = sign (wb7)) = sign
(0) W2
This proves
54

---

Under review as a conference paper at ICLR 2023
On the other hand, we get that | ( [. z | [6,5]
(1+1)
= θ
and
222
2
W2s
(Pisc)
d3
Since 1 + S Ti which means Vij E [d] [B
2 rjd. Then
91
[z.5]
((+1)
= 1u
(1+1) E((+1)
E
2 Ω
dda )
1
vıdd= 2({);
d
(++1)|
E
W
(
[i,5]
1 W (1 [i, 2 d0
Vnd=
(EE))
92:
=
=
dda+t
j=1
This proves (C) at time L + 1.
Since t + 1 $ Ti which means Vj E [d] (w/(:1/((11)))
1 01), we obtain that
Σ
(++
w( [i,
(c+1)]
( [t.i||
((wss (fac)
+ (1+ 1-+1-lin line )
(W + + 1- (ine) (())))(m))))) S0()).
1=1
Note that ( [
(ta))
(since line < La), we get that (1 + 1 - line (m)) = O (t)).
(++1)
θ and | [1, ò (+a) and hence (A) holds at time ++] 1.
which gives us
Therefore by induction, we can prove that (A), (B), (C) hold for all time St S Ty. Then applying
Lemma 21 we get that for all time s Tú, Vi,j E [d] Aww j 5| 6(n), ((n).
Aw
Specially, at the end of the first phase, we have Vjj [d] :
( = -I(1.. Repeating
the above proof techniques gives 222 = 6 Va and w [
(T))
= - ( for Vi,j E [d].
D.10 PROOF OF LEMMA 19
Let's first prove eq. 30
By Lemma 18 for time S t < Tı, we have Vi,j E [d],
91 (i. 51 (vT)
(1)
921
(0) 0 (Viid)
Then it suffices to show that for time I t Ti 91 [1.]] 9i -
10
[i.5] =Ō(n) and
(t-7)
(((d). It suffices to show that when < Ti, 91 [i
[i..] g [1 Ō
92
92i
and
921
(++1)
922
(t)
Ö(dd).
By Lemma
17 and 22 we know that when
< Ti, Vi,j E [d], |w S Ō()),
Aw S
(t)
(()) and that w [i
22 < Ô
va Then the bound |
< Ô 1 in
eq.
(35) hold for all t Ti only line St t
< T.). Substituting these bounds into eq. 36 gives us
Vt
Ti
9 [
(++1)
[i,]]
9 (1) [.,il
<
W22
((+1)
A 0
- Aw |
1
(()) = -())
Va
Similarly, we have that 92i
((+1) (8)
921
((jd),, which proves eq. 30
Note that for a, b
E
R:
2
b2
a²
(a- b 2)
[a b)
(a
-)2
la
bl
aa
bl
2
<
+
a
a²
a²
|
55

---

Under review as a conference paper at ICLR 2023
Then eq.
331
immediately follows from eq. (30).
D.11 PROOF OF LEMMA 15
We divide Lemma 15 into the following three lemmas. Combining them together immediately gives us the whole proot.
The first lemma below gives us the structure of W2 in the second phase and that of W1 under some conditions.
Lemma 23. Under Assumption
2 and
suppose a
S
13yt
By picking rj S
.5 S
and 32 = Bi, we have w.h.p. for T <t<
Vi E [d]
W22
(e+1)
w/
(1) sign
828
e22
(1)
where
(0)
(VT).
and moreover
R
(0) 2
= sign (w/) ( + Rll where
v d=-1/2 1/2
cl
As for W1, for certain i.j [d] and certain t [Ti, T) we have
[t. 52 (VT) then
w( [ [i 17 (sign (i.)..))) + 1d (1.5)).
[1. Ô (VT)
where
The second lemma below also analyzes the structure of W1 but removes the conditions in Lemmal 23
Lemma 24. Under Assumption and suppose S
" n dis/t.
By picking 17 S 0
2 . S
and 32 B3. we have w.h.p. for Ti St < T, Vi,j E [d]. | [..1 52
1 and for
any E [d].
R [1.
[i.j] [i.)
(0) sign (w v,(0) + 2
where
0
17
1
17(0
d3
The third lemma proves the convergence of Adam at time T.
Lemma 25. Under Assumption and suppose S n ///2²2
By picking n S {
i S
/1344.
and 82 = B7, at time T, we have that w.h.p. Vj [d] | s Ô (dv/nd) which implies
B(T) Ô (nd*).
D.12 PROOF OF LEMMA
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.
Lemma 26. Under Assumption
and
suppose S
a
By picking 1j S
. S
and B2 = B4. we have w.h.p. for all Ti S 1 <
Vi.j E [d]
(t+1)]
A Au
(1)
Ō()).
[w S Ö())
Moreover, we have that Vi E [d]
| [i. = ō (
2
=
Equipped with Lemma 26 we are ready to prove Lemma
23
We will only prove the results w/
The proof for W (i. [i uses the same techniques.
56

---

Under review as a conference paper at ICLR 2023
Lemma 26 gives us upper bounds of w2i W [1.21. as well as
Awiz
(t)
and AW [1,5] for
all i.j E [d]. Then we know that eq. 35 still holds, which gives us Vj E [d]
E (((+)) - E
δ v dn) Then we can use the same strategy in Lemma 19 to prove that
(++1) (t) = ((jd).
92i
92i
By definition, for Ti S tt> < T, we know that
92i
(t)
9 (dv
Combining with the bound
92: (++1)
924 (nd), we know that the 924 parts in eq. 30)
(t)
and eq. eq.31 still hold. Then we can
use the same strategy in Section D.2 to prove that the Wizi
part of eq.
(29 still holds, which gives us
Vi E [d]
22
(++1)
22
(e)
sign
(6) 926
(t)
where
c22
(t) Ö (VT).
By Lemma 14 we have that the end of the first phase (t Ti),
RT
Ö (vi do-1/2)
Vi E [d]
2
(TT)
sign
W2x
(0)
(T))
+ (25)
where
c(T)
Combining with Vi e [d], Vt s Ti : sign
J2i
(1)
=
sign (w25
yields that during the second phase,
for S T, we have
(()
Vi E [d]
(t) W2
sign
(0) W2i
(t)) +
where
Ö @Vi
1 da-1/2
c(l)
D.13 PROOF OF LEMMA 26
By definition of Ty, there exists jo E [d] such that E(7)
Vid for Ti S15 f. We prove by (0)
induction that during this period, Vi E [d] sign
W2
sign (2.10 = sign (22
and
that Vi,j E [d] 22
(t) = ē (122) w 1. -
The base case Ti) was already proven by Lemma
22
Now suppose for some such that
Ti <t T, for all such that Ti I r S t, we have Vi E
[d] sign
W22 (7) sign (w jol)
sign
W22
(U) and that Vi E [d] W
=
va
W [i
Using these bounds,
we get that Vjj [d] : 5 S < d
(r)
(((T)
[i.]]
+ |4, = 01). which then yields two upper
bounds
91 (7) [i.i]]
wii E = δ
(r)
and
92
$ j=1 1 ( [i.i = (vv).
By definition of Tg. we know that for all Ti S S t, Vi E
2 d = 22E) and that which implies that Condition is satisfied for Vi E [d]. At the end
[d]
92.
(0)
sign
U21 = sign
Was
of the proof of this lemma, we will show that T 6 Vdn
Together with the upper bound of
92 we can apply Lemma 21 to get that w.h.p. for Ti S < 1, sign
(r)
(mm
sign
W2:
and
Aw!
2i
(T)
(1)). Combining with the inductive hypothesis sign
sign
W
(0)
gives
us
that w
(r+1)
(r)
@(n). Specially, when t, get the lower bound
10
(c+1)
(0)
and that sign
(112 ((+11) sign
(0)
Since ElTT -ndd we have that Vi [d]
gi
w22
(T)
á (VT) 22(E)
and that sign (9) (7) (1 bol) sign
(0)
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 S r 1 1
57

---

Under review as a conference paper at ICLR 2023
Vi E [d]
[ [2, jol = | Jol + @()). sign (w." (w( [1. sign
(0) u22
and that
ji,j E [d]
|w]"" [ = O(n)
Now we analyze the magnitude order of u2
(1(+ [t,51|
Let's first analyze (c+1)
By Lemma
when t Tí. Vi,j E [d].
(Tr)
| W 17.
Wii
=1 (vT+ d
=1 Ô ( do-112)
(Ti)
(TT)
W2j
W22
Combining with the facts that for Ti S r S t, W1 ( [i.5o] = |wf + ( and
W [5.00]
22:
(T+1)
w2:
(T)
+ @(1)) yields
[1)(111
@11. Since we just proved Vi E [d]
sign
((+1)) sign
(w((+t)
sign
(0)
we get that
(W2W1) (e+1)
W2i
(1+1) (1+1)
w
jo] =
Σ
22
(++1)
W
(1((+
[o.
= O(),
which gives us that M22
1e+1)
Recall that we have shown
(e+1)
2 Ω v7
then
Vil
(++1)
22:
Now we prove W
(1+3)
= δ
We have proved that Tj S T S t. Vi E [d]
va
[1,
(1) and 1122
(++1)
= Ō()), then Vi,j E [d].
W (1t1) [i,5]
W1 ( (Ti)
[i
+ Σ
W (
[i,
- W (T( [i.j]
11+1)
(2:)
Σ
(T+1)
(T)
W/2
W2i
W12
10:2
W ( [
(++1- Ō(1)
W22
(Tt)
+ (t+1
where the last equality uses
1 ②
zt 1/2) Since we already proved that (VT 1
(11)
(1+1)
= 3 1
we get W (1(++]]
W22
va
=
Therefore by induction, for all f. in the interval [Ti, T), we have Vi,j E [d] w
pt = Ō (2.). From the proof we also get Vi E [d]
w
(++)
22 and that
Aw (2) (()), Aw (0) [i s Ö(n).
Now we verify that T θ ) Combining Vi.j [d]
(T)
W22
and Vt
(Ti, T), W22
(e+1)]
(t)
(()). we immediately get that T - Ti Ô
In Section D.2 we
Vdn
have shown that Ti θ
Van then we get - Van
D.14 PROOF LEMMA 24
We prove this lemma by induction. The base case (1 = Ti) can be verified by Lemma 14 Now
suppose for L in the interval [Ti, f). we have Vi,j E [d], [ á
58

---

Under review as a conference paper at ICLR 2023
For E [Ti, T)
by the proof of Lemma 26
(Section D.13 we know that for Vr S Vi,j
[d]
22s (7)
= θ
1
W [1,11| = O vaa and that 91 (r) [i
o
921 Ō (d))
and
va
that Ti θ
Then we can pick H == I-B log d and apply Lemma 16 and Corollary
to
Vdn
get that, w.h.p., for all t E [Ti. T) and vi,j E [d], the update of W1 can be written as
(1 (1- B1) D T 1i91 7) [1,1] + 1i
w(tt [t.5]
w/' [r. T1e
(2) (1- B2) -== 3y (5) (1.06 + cid [..ii++
where
(in [
(4)
ci S $((15)) By Lemma 23 we have that for 1 1si S d.
R [
Ö 8
91 [i.j]]
10)
wyyD sign
W22
(0)
1 +
2o-1/3
c( E
H
3
Σ gig ( [i.5]
sign
(0)
Σ
BT
clt
E((-T)
Σ
87
R
9.l
(37)
Using the fact that for a, b E R.
ba²-b"
< 2
la-bll
+
la-bl
we get that
a²
|a|
(49B..5))
2 ((()
sign
(0)
E + R (..51)
= (000) (c0EE E 2
+ esqr, [i,
R(e) [al1l (()gg()
where
3 (VT + アo-17z) ). That yields
H1
H1 Σ
11
B2
91
18 (t-7) [1,5]
BB
((-)))((-)))
Σ
35
R
gsqr, [i,j].
(t ((-r)
r)
(38)
0=0
T=o
0
(c(t-r)
2
Since
> 0, in eq.
38 we have that
Σ 33 Rog.. [1,5]
(e-T)
- -OVntt
1
(39)
Σ/' 35 (c(-))))(T)
da-1/2
However we cannot similarly prove that Σn 87 (2(lT))
9.1 [i,
< E 84 cll- p
in
eq. 37 because clt-T) ((--)) E((-") may not have the same sign for T =
0, 1l.... H. To deal with eq.
we need to consider the two cases where Σ/I 37 (((-T) [b,j]] << Σm oft -)( or
ΣH
H o 97 gg [t.5]] * E of clt p(i-r)
Case 1.
(1 - B1) Σ"
35 R
(1.0 +
s 6- 3r) IM (ellt [((-)))
where
8
d -
Note that from eq.
39 we have
(1 - B1)
Σ
35
[i,j]]
5
1 )
(1- B1) Σ 3r (((
"
g
T=0
59

---

Under review as a conference paper at ICLR 2023
Combining with (2 = Ō(153) s 'ō )
)
2 "²,
we can apply Lemma 33 to get that
Wr( - 1 [
(1 - B1 )iign
222
(0)
0 35 c F4 - + (1- - 31) E. H
87 R
9.l
1) [t.5] Li
=- 7 V1- B2)
(clt-r) ))) 2
-+= By
H
+ (1 - (22) ΣHH 85 R
(t. 5 + 17 [1,] +
goqr.
(0)
Σ 35 clc ( E
B1
sign W2i1
(1+
1/ v1 - 32
V -#= 83 (c ) )² +
(0)
(t)
1 + e [(]]).]).
== sign
Wiz
ō (nt + 25-5) Since |w([+()), [t,5] - B/.l Ö()). we get that
where
0 Ö(T).
Case 2. |11 (1-31) (1) HH
34 - (1,5] + 1/11.01 > >8(11 - (1) 01(1-11)27 _. B9 (((=))/((()))
where
9.1
5 = 17a +
1
d d )
1
Since
Ö (vT) + Ta-17z) ), we have that
olt)
H Σ
H
(1 - 31)
89 (t 7) [i,
S 2 +
1
(1 - B1)
Σ o (1--ggt
da-1/2
7=0
==
1
(H + 1)(1 (I]+ 10l--30 BB) Σ 35 (()-)) 2
Ö
+ d=-1/2
r=0
/ + d 1/2 1/2 )
1
(1 B1) Σ 83 ( [i. i1 2
\
where (i) uses hnwaus and 32 = 87. (ii) uses eq.
38 and 39)
Combining with
(2) c |1
S 0(18)) = Ô (VTi + ) ( - v/
sid
(t)
[i.
gives us
(1- B1) E Sf R R(TT - [i,5]
(1 3,) Σ
3 ct b( f)
<
n + d8 -
0=0
② 7
da 172
Σ r=0
11
91
(1 (1-7)
2
V sid [i.i]] [id.
+E
11
+
d3
-
<
n
1
Σ 33 (0) 2
+ +1001.11
+ {
d +
T=0
which implies
wr((t2) [1.5] - [2,0)
$ n (n) +
1
d7
60

---

Under review as a conference paper at ICLR 2023
Consider certain i.j [d] and the period from T1 to t. Denote T as the set of time points when Case
1 is satisfied. By Lemma 26 we know that ry(t - Ti) = Ö 1 which gives us
Σ
(t-
n² + d
)
Ô
n
1
d
dd++
d
TØT
By the first phase analysis, we have that
w [i =
sign
(0) v(Ti) Ui2i
+ R (2)
[i,
where V = O (12) (
R T1) [2.5i| = 0 (v3 + )) Combining with the analysis of Case 1,
we have that
W1 [, + Σ A [i, ] = sign w ) (v
-
-
Σ TET
RT(i,,),
TET
where R 50 (n) + a4-7) (|veeri|=Erss
+ ETE (
Since for T E T, |, = =O(n). v,(h) = o [/a and nt( - Ti) =6 (1²). we can bound
|RT[i, by
[R|i. s + d
0 Val
+ (1 ((-TT))(m)) 50
n1
1
d2
dd+1
Combining the above results together yields
w [i.]] = w(th) (1(h) + Σ
Awij- = [1:5=
[1,11
+
Σ
Aw
[b,5+ +
Σ
(t.1l
TET
TET
sign (1/0)
(()))
Ô
n
1
di
d
TET
= sign
W22
(7(1) + +R((' Ji where (t, S 3
dh
n
dit!
By the inductive hypothesis w [1.11 = (v²)) we get that |/,00 = ÿ (1/2)
(1a).
which gives us
v2
Therefore, have that for any j E [d] and any i,, is [d].
w (4) [i.ji sign
(0)
(1)
0 17 - -
0
17²
w [
(0)
v
[ [ ō n -
d d
sign
W2in
By Lemma 23, we know that w22 with different i are also roughly equal, i.c.
(t)
1 +
(1)
Ô (VT++ Je-1/= Then we have for any j E [d].
d
(W2W1)
(0
[1,3]
Σ
(t) Wiz
[i, 31 = θ (a w | (k,5l))
= 6 (va (
where k can be any index in 21, ..., d} and the last equality uses Vi E [d]
| = 6 (t2a)
61

---

Under review as a conference paper at ICLR 2023
Now we analyze the lower bound of W [ [k, ol|. Although may decrease during some period,
we observe that once
W1 k decreases to some value of order
such that (W2W1))) <
A vid. i.e.
v nd. we can apply the technique in Section D.13 when analyzing
w1" [ol to get that
W1 [k will increase in the next step. This mechanism ensures a 2
lower bound of [ 1) [ Since j.j are arbitrary, we have proved that at time + 1, Vi, E
[d] W 1) [. 5| 2 S2 (t))
(++1)
Therefore by induction, we conclude that when Ti <t< T, for Vi,j E [d], ( [i.i]] = á (vz2
The remaining part of this lemma has also been proved by the analysis above.
D.15 PROOF OF LEMMA 25
(!) Lemma 26 tells us that for any E [d]. 212
keeps increasing when t < T. However, the behavior
of W [..j] is more complicated. The following lemma tells us that [ will increase until
Ty.- After that [ and E may zigzag, but E will not fluctuate dramatically and will be trapped in a small interval around zero.
Lemma 27. Under Assumption / and
2nd suppose < /13/4. Pick v S ( S v ///-.. n2/2g2
and B2 = -.. Consider certain coordinate j. For Ti S 1 min { (7,T35 j. we have Vi E [d]
( [..]] keeps increasing. Tf <T, then for Ts3 S < F. we will have Ô (Vnd) s E s
Ô nd).
Now we start proving Lemma 25 At time F, denote S =
(3 {s: T < <T}, L.e. the set of coordinates
whose E, have passed its "flip time". By Lemma 27 we know that vj E S, | - Ô (Vnd). If
S の. which means Vj E [d] D S Ö (Vnd), then our lemma will immediately follow. If
S # é we have i min (Tg. g: Ty} Tg and that Vj E S": E < 0. By the definition of Tg,
(
know that Big E [d] : $2i0
00 (dvm). Then
d
Σ
E
W
[i0.5]
= Σ E W (T) P/oosll - E
JES=
JES
<
$2i0
(f)
+
Σ
W
(T)
(T)
S 50 (d + dō (Vra)
Vd
1
= = (ddT)).
jES
we know that when Ti S < T, for Vi E [d]. /w/ (i, = Since the
By Lemma 24
( [i < Ö()). we know that sign (w(") [1,
update per step
remains unchanged during independent of j. Combining < 0 gives us that F w (T) (T [to. for different j have the same sign. Therefore
( (+) (6,3)) sign ( sign 12:
this period and sign
with Vjj E S E(T) for any jo E Sc,"
ô (ddT)) 2
Σ
E
W
W1
[i0.1]
Σ
E
(io. Jl
D STDwi () [:o.
jo]
jES
jESc
F (?)
(d pd
+
Va
62

---

Under review as a conference paper at ICLR 2023
| s (dvidd).
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 (T)
< Ô (dvijd) completes the proof.
D.16 PROOF OF LEMMA 27
Consider certain j [d]. when L < min T, TTss
}. have that F.j0) < -vdd. Therefore
can use the same argument as in Section D.13 to prove that ( [i, keeps increasing, and
sign (w." = sign w for all e [d].
At time the "flip time" T.. by definition, F.j vdd. After that F. may oscillate. Now prove that once E2(() 2 /nd (or F/(0) S Vidd), after short period F750 will decrease (or increase)
until F < 2 (vndd) (or F7 - (Vnd)) Moreover, during this period, E.5)) won't change much.
We first recall that when Ti S < T, Lemma 26 gives us for all i E [d].
(1) = ē (+) and
12xx
W [r.0] Ö ( (122) Then eq. (35 we obtained in the first phase analysis still holds, which tells
us that the change of E per step satisfies | F Ô (nvaa for all Ti SL<
We divide the analysis into two cases, based on whether these E 2 Vnd or E S vnd. By
Lemma 24 we know that when Ti S t < I, Vi [d], |w/" [r,]] á Since the update per
AW 1) [t,0]| S On)), we know that sign (t.3] remains unchanged during this period
step
and (w(0) (v/01 sign (w("" (i.5))
(0)
sign
sign 222
By the analysis of (1) W22 in Lemma 23 we have for all [d], 222 w/2
((+1) (0)
(0)
(t)
+ sign [u24
A
where A!l = 17 (1 1 O
Case 1. Consider some time point t such that E S -vdd. Note that for all [d]
91 (1) [ = w E = Ω (vr) and that sign (oi [1,55) = - sign w (1) = sign 1
By Lemma
23 for all i E [d] we have w [i,s] = [('
[i,j] + sign W22 A (0) [ with
[1,] = n 1 + δ (vm) That gives us
E
E((+1) Σ
w22
(++
(++1) (++1) [t,5] Aj
W
(t)
sign
w
(0)
A (0))
(w (i.
sign (w!0) - A
(1) W2
W
[ + sign
(0)
(1)
(0 [. -
A ( - -A,
E 1222
1 [i + -
+
+ (
1=1
b( > E/0
63

---

Under review as a conference paper at ICLR 2023
where (i) is because sign W'2 (1)
= sign (w ( (1.31) = sign (122
(0)
Therefore we have proved
that Di will increase in the next step. After that for T 2 t + 1, as long as E S <-nnd, the above
analysis will hold and E will keep increasing until E > Vnd or we reach T.
Case 2.
Consider some time point such that F.00 2 Vnd. We will prove that F will decrease
after a short period, and during this period, the change of it is at most Ô (Vnd).
By similar arguments as in Case 1, we can get that (t, = 1 [t, sign
(0)
4 [1
w1ns
where 1 [1,1] 17 1 = δ (v))) Then
E E
(
[1+ (++1) [ Aj
W2i
d
(t)
(0)
> | (w/ sign (w))²) 1 11.3)))) Aj
W2
+ sign
W2
W2:
[
sign
(w (0) (W 19 1/9 (3.01) - -1.1)) - Aj
d
(1)
(t)
1 [i..]
(0)
/ [ |
1 A (t.1))
where (i) is because sign
(8)
sign
(w(") [1.5i) = sign (122
(0)
E may not be smaller
than E but we will show that after at most t, steps for some t,, we will have E Ef(tee+t) E Esetee)
( [t,5] 2 A - S (VT). To see this, first note that by the bounds of A (0) [1. j] and 2r we get A
(!)
Since 22
(t)
w (t) [ keeps decreasing, then have either increases by θ(n) per step, and
after 1g steps for some L,, Vi E [d]
[+++))
W ((++)) [i, + vi ii) we reach 7.
E (1+))) vnd, then it's already what we want. Otherwise we will have A (((t2)) [i,5]
For i), if
7 (1 (1 (vT)) Hence
F/(t() +1(+++++1)
((+c)) u22i
1 ((tea) [i.5] - - (1++) ( [i.]] + + (((H)) A ((+ ((tt))))))))
(1++)
(Iw(ttee) [2,11 (A (1t22)) [i,5] - -l(tte))) + /T (1+ee)) [2. + A 24
(t+ea) ^ (1.t5))
(-nō (
(v (1
[2,1] + +ル + A
(atte) A (a+tt)) ( 0,
where the last inequality uses Vi,j [d]
| [3. Ô (7)) Therefore pf(+t(-11)
E((-t)) After that for 2 1++ L, + 1, as long as E > v Iıd, the above analysis will hold and E(T)
will keep decreasing until E < Vnd or we reach T.
Now we prove that during these 1, steps, the change of E, is Ô (Vrd). Since at each step the difference |wmi| - [ii]i, will be enlarged by ②(1)), then we know that 4, = vn/Ω(n) = ô (+)
64

---

Under review as a conference paper at ICLR 2023
Combining with the fact that for all Ti S s E) E(r+1)_____10 - -E}" | = (vvaa
gives us
D - E S Ô (nlovdd Ö (Vnd).
For ii), reach T before vi [d] W2ys
(1tea)
2
[i,j] vT Then we have T t S
vn/((()) O (T)) which yields F.g)) F s (n(T- - 1[] S SO (Vnd).
Combining the above two cases, we find that if for some t, E z v /nd, then after at most is steps Ej will decrease and keeps decreasing until E < v or we reach T. During these steps, Ej can increase at most O (Vrd). If for some t, E / then after step it will increase and keeps increasing until Ej > Viid or we reach T. That means once for some coordinate j. Ey overshoots, will zigzag in small region around zero, which is (Vnd) (Vndl)
E 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 two-layer network under Assumption T [and]
E. EMPIRICAL RESULTS
Let's first define the diagonal domination of the i-th coordinate time t.
VEAA (H)1)))
OPT ding.
To measure the diagonal domination of the whole Hessian, we need to consider the distribution of TPPT "ding, (t) for different i. Figure 14 shows the mean and median of ding." (t) and ding, (dlle(()) on the sentence classification task (see Section 4.1 Here we chose 4 layers (Layer #6, 12, 17 and 22) and computed the Hessians across these layers. Since the number of parameters is very large, we did the computation by random sampling. As
(SGDM(t)
see, for both rSGDM (t) and rdd (t), the trend of their mean median is to decrease over time, although there might be some oscillation.
300
2000
300 280
30
1.00
(a) Mean
(b) Median
Figure 14: Mean and median of ding,s (t) and raiggg (t) for the full hessian across the four layers (#6,12,17,22)
E.2 THEORETICAL. ANALYSIS
To simplify the theoretical analysis, we consider the mean of OPT
diag. (t) over all coordinate and define
OPT Ro (t)
mean (r (rding, (t)
(40)
We consider 2-layer network under Assumption 1 and 2 and have goals in our proof:
65

---

Under review as a conference paper at ICLR 2023
1. To show that diag (t) after training is smaller than that before training (t 0).
2. Note that our setting (see in Assumption D the Hessian is a (d² - d) x (d2 + matrix. For
a completely "uniform" matrix with the same size, we have that ding (t) = θ (v +d)
(d)) Hence our second goal is to show that the Rudu (t) after training is lower order than (d).
Theorem 2. Consider the ratio ( defined in eq. that before training (1 = =0, with high probability,
40
Under Assumption and 2 have
R (0) 2 S2 (daa
(41)
For SGD+M defined in eq. (3). For any > O, by picking the same hyperparameters as in The- orem 7 for T TTGD22 mentioned in Theorem t t [T T.GD.2).
[ we have with constant probability, for any
RSGDM (1) S Ô (v((v))))))))))) + g000 diog
(42)
where the trend of (t) is to decrease over time and q (TToo.))) SO (0)tt) o(d))
For Adam defined in eq. 33. For any
0, by picking the same hyperparameters as in Theorem 1 for we have with high probability, for any [Tadom, 1. T Adum, 2).
Tadam T Adam, mentioned in Theorem
Tediag (1) S Ô (va)
+ +1)))
(43)
where the trend (()) is decrease over time and y ((aassee
<0 1
<0(-))=
(va).
E.3 PROOF OF THEOREM 2
Lemma 4.3 of (Kawaguchi 2016) gives us the following forms of Hessian.
For any h E {1, .... H + 1}, we know that V Veec(Ws) (V ((W)) equals
((W H+1 ... Wx+1) (W+1))((HHL) (WW+11... ... Wx+1) @ wi)--....::)WW..c... W.)², and for & E (2,3, 3 .... H + 1}.
Veec(Wa)) (V we) I(W))
Wx+1) @ (Wa-1... Wi)
+[(W+-I... ..) 4/ (r(WH+1- ... Wk+1))., Ia. (r(WH+1 ... We+1))
... where r = (WH+1 W1 A)", C = WH+1WH W.......... W2.
For the 2-layer linear network, write the Hessian as
Hza H H2
H ==
H1
then we have that
Hn (W. W₂) La E &dd
H2 = WWW E dxd
H2 WT WT la (W,W, A)T &dd
Intuitively, before training the elements of W, and W2 are very close and W A Re AA Since the elements A are ((1), we know that the magnitudes of elements H21 are much bigger than those of H11 and Hzz.
After training, for both SGD+M and Adam, W.W, - A zs 0. Then H21 = (W₂)7 & (W1) and the magnitudes of its elements no longer much larger than those of H1 and Hyz. From the formula of Hn, we know that all the diagonal entries are nonzero, and among the dt - dl off-diagonal entries, there only d d2 nonzero entries, which helps us bound ROOT(1).
66

---

Under review as a conference paper at ICLR 2023
E.3.1
PROOF OF EQ. 41
Let's first analyze the weights and Hessian before training 0). For ease of notation, we omit the superscript (t).
For the i-th row where 1 S d, i.e. the i-th of the submatrix [H22 Hgg. we have
d
Σ
Σ
High
+
H
j+b
j=1
xd
H2
Σ
(wzi W. Wi (W2W1 A)2) (dd).
j=(i-1d
On the other hand, for the diagonal elements, we have w.h.pp
nTo-T (
H 81 (/ a]I
[[ 11. = Σ WWiiss
Then we have that for 1 isisd. d,
V Eige (1,
v (ddd
52 (dda-
|H[i.
24n
For the (id + k)-th row where i d, - k S d, i.c. the ((i - 1)d + k)-th row of the submatrix (H2 H11, we have
Σ Hb j
=
Σ 1i - h; + H2g - 1)d h.
+((-1)
11 (i 1)d +
k, k, (w> 16 k] + (W2W, A)a) = O())
On the other hand, for the diagonal elements, we have w.h.p.
1
Hid k id ] H [( 1)d + k (i d + k]] = w σ
22a
Then we have that for 1 < i S d, < S d.
Σ Ejta4t H2 Ji
21) = 52 (d2a)
H [id + ks id + k]]
Taking the average, we obtain that before training, i.c. when t = 0,
dit (dia-b) 1
+ d (d2a)
Rain (0) 2
= ÿi (2
d2 + d
E.3.2 PROOF OF EQ. 42
The proof is based on the lemma below.
Lemma 28. Suppose the weight matrices have the following structure:
W1 = uv² T + Rı,
W2 cu² T' + Rh-
where V1 Si S d:
[R [uynyl
S 5,
Rai
<
8 (0, 1).
Then we have for i d,
V H H2
1 8 (
lc| )
j u
||
< - 8
Hii i||
u
(1 2) 1
67

---

Under review as a conference paper at ICLR 2023
and for 1 isis k S d.
V /ijuu+ H jil I + k, id k]]
1 + 5 6
2kl
|&
1
+
|c|
uf
(1 11=
Now we are ready to prove eq.
By the analyses in Section C.1 with momentum satisty
we know that for t E [TsaD,1, T the weights obtained by GD
w1 = u (00) + +R(",
w = c00 (T)T (T + RR
where TSGD.1 = Ti and
R2(2) [i
LA S Misijjd sd:
I O(Eo),
S O(Eo).
(Tt) (2) u
c(0) m (T))
Here en is defined in Definition 2 Since (Ti) doesn't depend on time in the period (TsoD,1, 2s)
we write (T)) as u for ease of notation.
Hence by Lemma 28 when t E [TsaD, TsGD.2), we have for
I i S d.
Σ i+i (H (2, 2
1 + Ö(eo)
(1
c(0)
>
j-= u²
| 1/2
1 (1+ |010
+
H
1 Ō(5o)
u²
(1 Ō(co)) u²
1000
= =0 (1+
|c/0||
Σラ=
uz
( ||
[ 11
||0(0/1/2
V
u²
0
u² ||0000112
(44)
and for 1 sis d, S k S d,
Ejfiat&& (H.0().)))) 2
1 + Ō(Eo)
)
اラ 1 u3
E((2)
+ cct))
1110 iid + +k, id + k]]
1 1 O(Eo)
up?
(1 1 (1-O(co) (clo)) uz
2
(p()) |o(t)
E((t)
Σラ=
u3
= =0 1
+0
uf
(cl1) 2 uz
(45)
By Lemma 3. we have 2 = X + Y where X,, i E [d]
are i.i.d Gaussian random variables and w.h.p..
Y
Vi
[d]
I Ô
== 82gs
(46)
Xi
dia-
which yields that
V Ej= x²
Vi E [d]
1=
1 1 day
1
<
1
1
(47)
day
X
8ay
X7
By the proof in Section The induction in Section e
C.8
we know that for t e [TGGD, 2.. Vi E [d]
: (.
.((t) are positive.
further gives us that for t [TSGD, TSGD. 2/. w.h.p. Vk E [d] _/0
which yields lett
(L)). Combining with eq. 447) we obtain
va
|c(t)|
u
v Σラ= X?
1
+
0
|0(0)
u?
Xx,]
2
(48)
(2)
v / x?
uz
1+
1=
0
u
Xx,
68

---

Under review as a conference paper at ICLR 2023
By the proof in Section C.8 we know that for / [Tsan, 2s..,2], Vi [d]
(8)
clt) are positive
and monotonically increasing. On the other hand, the proof in Section C.2 and C.9
tells us that w.h.p.
|B0 (resp. Vk E [d], E decreases from (vv)) (resp. (1)) when TSGD, to ((vend)
| and
xx (ete))
(resp. ((v(a)) when & = T. Therefore, the trend of u 0 12 time, and when t Ts we have w.h.p.
is to decrease over u?
wk E [d] 1 O (VEo),
|
2 = =0 (Veod)
(49)
(26) becomes equality, i.e. c/u = θ @(va) and
Moreover, when f = 2, 2, the inequality in eq.
Vj E [d] || = θ va
Using u = X + and eq. (46), we have
2 = θ (Vi),
1x120000
1
1X - (Va).
=
which together with the second inequality in eq.
47
yields
1
2
Σj X?
1 1
1
dzy
X?
1
1
X? Vd
1
2
1
Ej=r X?
c
1
8zy
X²X?
X7Vd
Combining with eq.
44
we get that
E
Σ ==
X?
6u
u
|
11
O
X?
E0
(c(0)
u
0
X2
d
(50)
Substituting eq.
48 and (50) into eq.
44
and
(45)
gives us
sisid: V Ejts
(HO (1))2
V V x?
5o
+ 9if
(2)
H
X2
where the trend of g1s is to decrease over time and
(1)
4is
(2
< <0
2is x² x²
VED
We also have
VEn (H00).))) (1.
V Ei=z X2
LA sissdIssss
[id + k, id + k]]
0
X
924
where the trend of 9/5 is to decrease over time and 42
(1)
(TaaD,2)
Ed x3 x?
/ /20
Hence
d
x2
d Σ
E==1
SGDM (1)
diag
0
X,1
dc d
T +
p +
d
421
1=1
1=1
カ X²
+ (1)
X1
69

---

Under review as a conference paper at ICLR 2023
where the trend of (1) is to decrease over time and
q (san, 1
d +
1
60
+
dc
d
0
Σ
Σラ= X3
X
ED V 7
1=1
X?
0
End
o
d d
X
X?
Denote 6² as the variance of X; for i [d]. By concentration of chi-squared distribution, we know that with probability at least 1 - 8 for d 0,
(
( (Vllog 3
x² S so'd+
By Lemma 35 in Appendix G we know that with constant probability
1 = i=l X = ( log d).
Then with constant probability, Ei= X 2 1
O (/t2 - log" d). Hence
O(Va),
Σ ā
1
d
d
x?
= Ö(P)
X.1
Therefore with constant probability,
RSSDM (1) = 7 (vd)) + (t))
dig
where the trend of (0)) is to decrease over time and (aaD,)) S Ö (dVEgd)) For any p> 0, by picking the same hyperparametens in Theorem I have egd < Ô (F and hence y (TaoD.2) <
d o(d)
E.3.3 PROOF OF EQ.
By the analyses in Section D.1 we know that for [Tadam,1, 2), the weights obtained by Adam satisfy
w ut/ R(())
w = d02 + (()T
where Vi E [d] : u = sign(w!) (0)
E {+1} and
R [i,
n(() |e(U)ss 8..
A1 vsiisd d:
28 = Ô n7
ueu! (t)
d
d
Hence by Lemma 28 when E (Tdam, 2. we have for siid.. d.
Ejsi (H (1,1))
1 + 5
|c62/
d
j=
u2
|5
119 [
<
1 - 8
1
||o00
u?
+
(1- - 5) ut ||100)1
(51)
|cl0|
1|
1+
||l(0)||2
Vd+o
2
and for 1 5is d, 5 S
Estidoh (i.1))
+ 8
==1
uz
F
1+
Hid + k + R]]
((0)
uZ
(1-5)² (c(t))² un
(52)
[e)
1+
Vd+o
c(t)
(c(t)) 2
Recall the following facts of Adam.
70

---

Under review as a conference paper at ICLR 2023
(A) By Lemma 14
we know that for t E [ Ti] (where T1 is defined in Definition
w.h.p.
vk E [d]
u
clt)
(t
fime) Specially, when t = TAdam, Vk E [d]
o/ (t) c(t)
Lemma
24 and 26
tell us that for E [Ti, TAdam. w.h.p. Vi,j [d]
d
Wi Wili,jll
=6
which gives Vk [d]] [0]]] = va
and
va
va
c θ
vtt
That means when [T'dam, 1. I'ddam.2) 2). Vk [d]
|0 and |c(t)
|00/ e (ta))
increase from to and
(1),
(B) Lemma and 25 tell us that w.h.p. || (resp. Ak E [d],, / decreases from (dd
(resp. (1)) when t = T TAdam, to O (d /T) (resp. Ô (dvrnd) when = TAdm,2
|
Combining (A) and (B), we get that the trend of when & = 2, we have w.h.p.
and
is to decrease over time, and (cto))
E (cto))
(dl
vT),
5O (8
(53)
|000
2
2
Substituting (A) and eq.
53 into eq. (11) and 52 gives w.h.p..
vsiid:: V jzt (HO
<o(va)+ +r
H1 i]
where the trend of ris is decrease over time and (Trim,.2 rii
(1)
5 vn)
We also have
Espatt (HW),,)))
L1 ssisdiskss
S 2 (Vd) (vV +
(t)
Ho k. id + k]]
where the trend of r2 is to decrease over time and (Taam,
(d Vnd).
2
Hence Radg (f) =
Vd
उ Σ1
d im1
d
(1)
= Ô ( (vd) + +r(() where the trend
1i
J4ta
22 iml
of (.) is to decrease over time and
(Tatm,)) 2)
Σ
(
d d
Σ 0 (d s (
L
For any p > O, by picking the same hyperparameters as Theorem we have nd < Ö (1)) and
hence y (Tnam,,)) <
(dd)
E.4 PROOF OF LEMMA 28
By the assumed weight structure, we get that
Vi [d] (1 - (1-)) (cu < (wx))² S (1+
- (u. < :1 = =1 (1)))
71

---

Under review as a conference paper at ICLR 2023
For the i-th row where 1 S i V d, i.. the i-th row of the submatrix (Hz inequality, we have
Hiil by triangle
ds
11 S V
H32 +
d + +||B||
Wi i2, :), :10
W2z
WPLi, k]
1m
d
I1 W.1b.-112 +
Σ j=1
w
+
-1112.
Then we have that for 1 isisd. d.
E 5# H2
|[W][1,)]l2 Σラ+ 11 11:
+ V w3
|
H 2]
[ :112
[[,
Eift [..
Σj=1
+ |F||2
+
11112,1111
[[ :113
[( -112
(1 + 5)2
E Σj 4
1012
(1 8)2
c² "
| 1 (1
(1 5)2
n
2
(1 --8
+ 8 $
(1+
lcl
u²
|
|
u
(1 - 1)
For the (id- + k)-th row where 1 < d, 1 < k < d, i.c. the ((i - 1)d + k) row of the submatrix
(121 H11. by triangle inequality again, we have
Σ jfsd+k
Σ
-i k, +
Hi - 1)d + k.
i
j=1
d
Σ
w wzy +
Σ
wg. WFG. A:] + Eg
j=1
= |w22ll
Σ
w2y +
Σ 1=1
W?Li k]]
+ L.
Then we have that for 1 S i 5 d, 1 dd
Ejtnt+ H]
|w2s| // ΣΣ EEit: w
+ V プØ WP
Ex
H lid + k,id + kill
w
Ej Way
d
WP[sk]] k]
|F
W3
w
(1 + 6)2
Ej41 dug
(1 + 8)2
2/ E3= u
EE
V
(1 - -)²
2uF
(1 - 8)/2
dau2
(1
2)
1 8 8 (1
(++ |0%]
20 j= uz
EE
< 1-8 5
+
[cl
un?
(1 - 2)
72

---

Under review as a conference paper at ICLR 2023
F CONNECTION BETWEEN DIAGONAL OF LOSS HESSIAN AND WEIGHTS
The partial derivative at W; of the cost function for each i is given by:
Vw, (WW) W 1 [WH+1Wa ... W1 A)WT W
(54)
In experiments, were interested in the diagonal elements of the hessian. These given by:
V (i (V a( = V(W.)., (W - H+ WH... Wi A)W? W 1)a,5
for each possible i, a, For ease in notation, define for each the quantities M, == W ... W and N; == WT... W Then we have the following lemma. Lemma 29. The diagonal elements of the hessian of the cost function are given by:
(WWa,, (V (Vw:L(W))... = (M,M) )a.a (N, N.)o.b.
for each possible i, a, b.
Proof. We have:
Vw,L(W) W ... WT
#+1 WHHWHHH....... ... AJWT ... W
M,WW++ WH...WW W1 -A)N;
= M.W++ WT.... .,,N - -MYN,
This implies that:
V ( (WW)L(W))abb V(W).., (M.W H+ W.... WIN, - M.) VWWJ.. (MW H+> WH W.NNN)..
where the last step follows since M; and N; are functions of W.,
Since M; == W #+1 WT #+ N == WT W by defining C₁ == M,W #+ WH W41 M,M and D2 = W. WoW, N, N? N₂ we have that:
V & (V L(W))a,b
where C; and D; are not functions of W, Now, Equation 74 in the Matrix Cookbook for any matrices and X have:
shows us that
X (XA); Anj-
Note that W, E Rddxdd11, then we can apply this to obtain that:
V V(W)a, (VW L(W))a)) ((W) (C,W.Di)a.
Σ (C)
Σ (Co)a. ((Wa... (W;D.)a.
Σ (aaa
(Ci)a.(Di)b.. (M,M)") (N N:)o.
This completes the proof.
whttps://ww// hwtp://wwwooooaoaaaaaa/aaaaaaaaa
73

---

Under review as a conference paper at ICLR 2023
For of notation, let's now drop the superscript OPT and (t) and write Femed (t) as Rmed,1 and med, (t) as R- For 2-layer linear network, H 1. Consider the Hessian w.r.t W1, we have MrM' = W2 W2 and NÝ N1 is an identity matrix. Under Assumption | we know that W2 is a row vector, which can be denoted as W2- [221, W22, .... W24,] Then we have
max, (W2s)*
(M,M)) a.. = 2 (NTN =1,
= Rmed,
median (1223) (ediaan(w2:) )
Similarly, consider the Hessian w.r.t. 22, we have that M, M, is an identity matrix and N? N1 WTW?". Therefore,
max; [W [t, :113
(M, 1, ()) b1). I [1.
Rmad,
median Wi =1,
Hence 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 to use this rank I structure to compute Rmed, and Rm
Lemma 30. Suppose W1 Raxd and W2 R have the following structure:
W1 uv R1,
W₂ cu² + R2,
where u E Rd, 2 E RR, Rr E Rdddd
R2 E TR. 1xd and that
RR |
R2d
sisjjsds
S 6,
S 8,
8 (0, 1).
|u:yy|
ccus
Then we have
(1-8)² 8) (1 8)2
(1+ 8)² max, uf median u?? (1 -)2 median uf?
max; u²
Bemed,1 1: Rmed, E
Proof. Let's first consider Rmed, we have
Vi E [d] : (1 (),, $ w S (1 + (cu;)²
(1 - 8)3 max < max w < (1 + 8)2 max (cu;)²
(1 8) 2 median (cu)² < median w (1 8) median (cu:)²
which yields
(1 - 8)2
max; u < Rmed.
max; u
(1 (1+00 8)2 <
max, un median up!
(1 + 8)2 median 21
median W
(1 8)2
Similarly, for Rmod,2. We have that
Vi,j E [d] (1 -8) (u,yy) sw?[i.] S (1
(1 (1- (1 1 < [[i]i. < (1+ 0) |
- 8)2 max u² 1111 max [[ < (1 + 66² max ut 1011
(1 median ui E1 < median | [i. :11 < (1 + 6² median u² 1
which yields
(1 - 8)2
max; uf ||
max; | [i:
(1 + 8)2
max; uf 1000 median u, 1/
Reeed median W1[t;]]
(1+57) 5)2
median u
(1 8)2
That means
(1 - 8)2
max; uf median u?
(1 (1+00 (12
max, u² median u²
Rmed <
(1 + 8)2
(1 -1-60 8)2
74

---

Under review as a conference paper at ICLR 2023
G AUXILIARY LEMMAS
Lemma 31. Let A = m LYX"" Agz := _ oh = -w,,LWWo) = 1,2. Denote A A
and 0% k = 1,2 as the corresponding batch versions at time t. Let M = maxi,; |W| [1, and
_ = max; t
Under Assumption 3 we have with probability at least 1 - 2.2 for Vt - T and
vi,j E [d].
[i.jl - 9i 1 [1,5]] - ( (M/))² 2 oVaT
dT + M.MooooV oveT,
) 2 M VT oVoT dM uMMoVVdT 92. 932 s
a d
Proof. By Assumption 3 and Chebyshev's inequality, we have for fixed i.j [d] and S T,
P ((4 A
2) -
o²
P ( (AA [ - A2al > 2) = X2
o²
A2
Applying the union bound gives us
e [d]. 31 ST:
| Ao > >>) S
Tdo²
1²
P [i.j E 14: [ - Aew 11 >>) 5 Td²
X2
which gives us with probability at least 1 1. for Vt v E |d],
A
A. S OVTT,
[ - Am 5 S odVoT.
Now we are ready to bound g% gk for = 1, and t TT
Note that for all t S T and Vi E [d].
d
(w. -
Σ why W (t) [.
(t)
| 1s a0 s (aM("MM))
S
2
j=1
Then we have with probability at least 1 - 4. for all t < T and Vi E [d].
[(w:( w//0 (A Aux)) )) Σ (w:"
AWli o] (i
5=1
S af (t) Ma oV
Combining with g(D) (1) = W (W: ( (() (A - Aex) - (6(00-a) we get that with
probability at least 1 - 1. for all t < T and Vi E [d]
1 - ( [9 < w 000 (I
(1) (Λ cs - AZz) - 1)) + mige
A,)
d (( (M!")) 2
ovdT + M, ovdeT.
Similarly, note that (4)
92i = (w! (1 (r(t)
(1) (1 - Arr) - (A -4]) w we then have
that with probability at least 1 z. for all /. S T and Vi,j E [d]
d
d
( (t) A A -
((--Ax=)) 1) (0 [i- + >
| 4/
92
92
<
==1
2 dd )2
ovdT + dM oVeT.
75

---

Under review as a conference paper at ICLR 2023
Lemma 32. Consider two sequences {a(*) 120, 16(0
120. which satisfy
a (1 3) Σ
3 e (0.
Suppose VT St: |b600| B, then for any 0, the following truncated version
H
a(t) = (1- - 3)
B ((-))
7=0
with H 2 1-B log =Ω 52 (1-3 satisfies
(t) a(() Se
Proof. We have that
a 2(0) | s (1- B)
Σ BTb(t
S (1-
Σ
BTB & B38+1
+=
-=
To make it less than e, suffices to choose H 2 log( // log B.
Since 3 E (0, 1), we know that log 3 < 3 - 1 < 0. We also have log < 0. Then it suffices to choose
log(c/B)
log(c/B) B)
B
1
H
H 3
log
5
3
log
1-
Lemma 33, Suppose a, b, c, Cas eb. e
E R. e
0, > 0 satisfy + eb + ec 0, 1,leal S 8lo.). s
5b, leel < 822² with 0 < 8 & I, then we have
a t
+ot vo+c (1 + R), where R
Proof. We have
a a+
ea vo+c /b+eo+eete vb+c vb+eo
a
+o+
1 vb+c
vo+c
11 ca
Vb+c
+o +c
a
vb + cb + cc +
42
Define R = Q1 + 22. The term |qıl can be bounded by
vb vb + es + Co
1911 = Vb+ eb e +
leb + tel (vb+es e + Ee +c)
Vb+ vb+
l6b
led
(vb+ e + e +c)
(vb+ (vb+ ep Vbb+vb+e+ee, +
(vb+es t + t) Vb+ t +ec)
leol
+ Vicl
Vedo
/b + e, + ee + c)Vb
Vb+ +p
leol
Vial
bb + Co + ce + c)/b
Vb +b+ vb t
43
94
76

---

Under review as a conference paper at ICLR 2023
where 19sl can be bounded by
(vb E Vleel + oivo S vb - 6) + c( 6) & dvb
8b
8vb
Igal
66(1- (6))
Here the denominator of (i) uses b + cb 2 b(1 6) > 0 and vx ty 2 Vx= - Vull when
x 2 0, + y
Now let's bound qal. If ee > O, we have ee = leel and |.a| S
vcc ve.
= 1 since + c 2 b(1 5) > 0.
If ec S O, note that b + es + E > O, we have |eel < b b. s 6(1 + 6), which yields 1/al s y
mel vb
(1). Combining the above bounds give us 1 S lqal + 8/9 = 0(6).
On the other hand, 1421 can be bounded by
vb+c
5 vb+c c
12l 1 Vb+eb vlel+ +
= 0
v6(1- 8) c(1 -)
Then [R] S 111 + 1-21= 0(6)
Lemma 34. Suppose X2, .., Xa are i.idd Gaussian with mean 0 and variance o², then for 0 <<6 8 < we have with probability at least 1 -6 8,
max X? 2 o² -gl log log 1 Isid
for some C1, C2 O.
Proof. It suffices to assume that o² 1 and prove that w.p. at least 1 1 1- -, maxigisa X? >
C1 logd - C2 log log J.
First, by the lower bound of Gaussian tail, there exists a, B > 0 such that x(|x,|>x)) =
<) 2 we Oxx for x 2 O. Then by i.i.d,, we have
P(mmax X2J] x) P ((IxI12))
P([X; 5 - (1- 2))a
< (1 ac An² yd
S exp(- - doe =maP).
where the last inequality uses - x S e x for x E [0, 1]. Let exp( doe - Bx2)) = ō. we get that w.p.
least 1 d,
log(ad) log log !
max X 2 ISidd
3
Then we have w.p. at least 1 - d,
2
log(ad) log log log 2
max X? Isid
max
[
2
ISisd
Lemma 35. Suppose X1, X2, Xa are i.i.d Gaussian with mean 0 and variance 0², then we have
with constant probability,
d d X <o
log
77

---

Under review as a conference paper at ICLR 2023
Proof. It suffices to assume that o² and prove that with constant probability, EE= |x O (logd).
S
Consider X; for some fixed i. Since X; (00,, 1), we have P([X;) S = S
Then we know that with probability at least 1 - θ (d-"), [X;] 2 C for some C > 0. Then by union bound, with d
2r
V2F
constant probability, Vi [d] [X,,] 2 C-
Now we split the interval [ 1] into several subintervals Is {i : [xal [2-*-1,2 2=1 for 0, 1, ..., [log, č
1. Let pt: = |[| E 12- 21),),, we know that [Za] ~ Binomiial (d. ps) and pv - C1 2 k-1 Then by the concentration of binomial variables, we have w.p. at least d=p
for p 0, [Txl
O (dpk + /dps logd log d) O (d-2 2dk2-k-1- v d. 2 k-1 log d + log d)
Then we have
Σ
= O (d Vd. 2 log d 2 log logd). k k= 0, 1, [oga ō 1.
d
X2
REZ
Therefore, with constant probability,
[lgg 11-
X2
X2
X
0=
kEIx
|x,/>1
[ogg -1-1
O d Vd.
2++ vd.2++
log 2 log d)
k=
dlog2 2) +
d 1 @VVllogd (V2) flog2 & )l+(2) + 2 logg Eltl logd +
O (dlog d),
which means with constant probability, 1 [X]]
O (logd).
78


Reviews:
[{"review_id": "w3Q0RtjNXT", "reviewer": "Reviewer_ZKz6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The manuscript introduces a new statistic metric to analyze the local trajectory of SGD methods like SGDm and various adaptive methods.\nThe paper argues the positive correlation between the uniformity of diagonal Hessian and fast convergence, and that adaptive methods such as Adam bias the trajectories towards regions where the metric is small (corresponding to a faster training/convergence).", "summary_of_the_review": "The main concern of the review lies in (1) the missing experiments to justify the novelty and significance of the proposed metric and observations, and (2) missing justifications on some arguments/statments.", "strength_and_weaknesses": "## Strength\n* The manuscript is generally well-written and well-structured: the main message is clear, with some well-supported empirical evidence.\n* The studied problem and the insights therein are quite crucial to the optimization community. \n\n## Weaknesses\nThe reviewer is generally happy with the manuscript. However, the authors must address the following questions/comments\n1. Some arguments in the manuscript need more justifications. For example, \n    * The authors argue to use the diagonal entries of Hessian, instead of the corresponding eigenvalues, to measure the local curvature; however, no empirical/theoretical justifications can be found. Even though in Appendix E, the authors demonstrate that loss Hessian approaches diagonal during training for Adam and SGDm, it is insufficient to support the arguments like \"we believe the diagonal of Hessian is more relevant than the spectrum\" and diagonal entries are better.\n    * If the reviewer remembers the related work correctly, some existing studies consider using $\\lambda_{max} / \\lambda_{10}$ (where $\\lambda_{10}$ is the 10-th sorted eigenvalue of Hessian, and 10 can be replaced by other values?). In order to justify the novelty of the proposed metric, the manuscript is required to compare these metrics in the same experimental setting, otherwise, it is hard to justify the significance of the submission.\n2. The generality of the proposed metric needs to be evaluated on more optimizers, e.g., SGD, other more recent adaptive methods, or even SAM and look-ahead.\n3. The explanation for the inconsistency behavior on the first several layers and the latter layers is expected.\n4. The manuscript mainly studies the training behaviors of different optimizers. As Hessian is also closely related to the generalization performance, it is necessary to validate the proposed metric on both training and (in-distribution) test data, through some comments/remarks.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "d8MxyKZ9iAH", "reviewer": "Reviewer_m6my", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper tries to explain why the adaptive methods converge faster than (S)GD.\n\n Specifically, this paper first defines a measure called $R_{med}^{OPT}$ which can be viewed as a stable version of the condition number. This paper then empirically observes that over tasks where Adam achieves empirical success, Adam and SGD follow quite different training trajectories -- the $R_{med}^{OPT}$ of Adam is significantly smaller than that of SGD.\n\nThe experiments in this paper show: (1). With the same loss, the parameter provided by Adam has different underlying geometry from that provided by SGD. Switching to SGDM respectively will lead to different convergence speeds. (2). With the same loss, Adam has smaller $R_{med}^{OPT}$ than SGD. (3). With the same iteration number, Adam has smaller $R_{med}^{OPT}$ than SGD. (4). Switching SGD to Adam leads to smaller $R_{med}^{OPT}$ compared to SGD which is not switched. (5). Adding small regularization, the forementioned observation continues to hold. But using large regularization, the observation reverses, while Adam still converges faster than SGD. (6). Over image tasks where Adam does not have an advantage in terms of convergence over SGD, the $R_{med}^{OPT}$ of Adam is no smaller than SGD.\n\nTheoretically, this paper proves that over a two-layer linear network, $R_{med}^{OPT}$ of large-batch SGD is larger than that of Adam when the input dimension is large.\n\n", "summary_of_the_review": "Overall, I think that this paper provides a new perspective on Adam's empirical advantage in terms of convergence. Such a perspective may be of interest to the optimization community. Still, the paper can be further polished. Please refer to the questions in the \"Strength And Weaknesses\" part for details.", "strength_and_weaknesses": "Strengths:\n1. To the best of my knowledge, this is the first work explaining the empirical advantage of adaptive optimizers over SGD from the perspective that the trajectories of Adam and SGD have different geometric properties. I find this perspective interesting.\n\n2. This paper conducts extensive experiments to explore the training trajectory of Adam and SGD. The experiment, which compares the $R_{med}^{OPT}$ between Adam and SGD while controlling the loss, rules out the possibility that Adam and SGD follow the same trajectory but a smaller loss leads to smaller  $R_{med}^{OPT}$.\n\n3. The theoretical result is solid.\n\nWeakness:\n1. (Please also regard this as a question) Theorem 1 cannot explain why the $R_{med}^{OPT}$ of Adam is no smaller than SGD over image tasks. Based on the experiments, I expect a theory showing that the relationship between the $R_{med}^{OPT}$ of Adam and SGD changes with tasks. To me, a two-layer linear neural network looks more like the neural networks used in image tasks than those used in language tasks (e.g., the four-layer network used for MNIST classification).\n\n2. (Please also regard this as a question) Why does smaller $R_{med}^{OPT}$ (or condition number) lead to faster convergence for deep neural networks? To the best of my knowledge, such a conclusion has only been derived for strongly convex objectives. \n\n3. I would expect a brief proof idea after Theorem 1 given the proof in the appendix is quite lengthy.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "M4Bmt2Owe4", "reviewer": "Reviewer_djbz", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the impact of adaptive optimization on local deep network geometry. The main discovery is that, adaptive\nmethods such as Adam bias the trajectories towards regions where the diagonal elements of Hessian are more uniform (compare to SGD type algorithms), and moreover such uniform diagonal Hessian tends to imply faster convergence. The experiments in this paper mostly use transformer architectures and language modeling datasets. In addition, the author also provides a theoretical justification of this phenomenon, by considering linear 2-layer neural networks with Gaussian data.", "summary_of_the_review": "See above", "strength_and_weaknesses": "Strength:\n\n-This paper's main discovery is supported by large-scale experiments on big models with real data.\n\n-This topic is important for understanding the success and convergence speed of deep learning optimization algorithms. \n\nWeakness/Questions/Comments:\n\n-The paper defines the $R$ statistics based on the diagonal elements of the Hessian matrix computed at each iteration. While this connects to the adaptive algorithms and is easy to compute, one may wonder how different is this statistic from that defined via singular values of the Hessian. If they are very different, does that imply that the convergence of adaptive algorithms is not very much correlated with the spectrum of Hessian?\n\n-In most of the language model experiments, the results show that the Adam trajectory tends to have Hessians with more uniform diagonal elements, when compared to SGD with momentum. But this does not necessarily hold for the first several layers. An interesting question is if the comparison fails to hold on shallow models (e.g., shallow MLP on MNIST). \n\n-The paper also presents a result on ResNet with CIFAR-10, which shows a different picture, i.e., SGD-M trajectory is more uniform than Adam's and yields faster convergence. Since most other experiments are focused on language models, I am wondering if the main claim of this paper applies to only language model/tasks. Put in another way, will the claim be consistent in standard image classification tasks? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject