Title: How Does Adaptive Optimization Impact Local Neural Network Geometry?
Venue: NeurIPS 2023
Accepted: yes
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 on language models where adaptive algorithms converge faster than vanilla gradient methods like SGD, we show that adaptive methods such as Adam bias the trajectories towards regions where $R^{\text{Adam}}_{\text{med}}$ is small, where one might expect faster optimization. By contrast, SGD (with momentum) biases the trajectories towards regions where $R^{\text{SGD}}\_{\text{med}}$ is comparatively large. We complement these empirical observations with a theoretical result that provably demonstrates this phenomenon in the simplified setting of a two-layer linear network. We view our findings as evidence for the need of a new explanation of the success of adaptive methods, one that is different than the conventional wisdom.

How Does Adaptive Optimization Impact Local

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 OPT med , a statistic that is analogous to the condition number of the loss Hessian evaluated at the iterates. Through extensive experiments on language models where adaptive algorithms converge faster than vanilla gradient methods like SGD, we show that adaptive methods such as Adam bias the trajectories towards regions where R Adam med is small, where one might expect faster optimization. By contrast, SGD (with momentum) biases the trajectories towards regions where R SGD 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.

1 Introduction

The efficient minimization of a parameterized loss function is a core primitive in statistics, optimization and machine learning. Gradient descent (GD), which iteratively updates a parameter vector with a step along the gradient of the loss function evaluated at that vector, is 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 encounters 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.

Dhruv Malik

To deal with these difficulties and improve optimization efficiency, practitioners in recent years have developed many variants of GD. One prominent class of these GD variants is the family of adaptive algorithms [DHS11, TH + 12, KB15]. At a high level, adaptive methods scale the gradient with an adpatively selected preconditioning matrix, which is constructed via a moving average of past gradients. These methods are reminiscent of second order gradient descent, since they construct approximations to the Hessian of the loss functions, while remaining computationally feasible since they eschew full computation of the Hessian. A vast line of empirical work has demonstrated the superiority of adaptive methods over GD to optimize deep neural networks, especially on Natural Language Processing (NLP) tasks with transformers [VSP + 17, DCLT19].

From a theoretical perspective, adaptive methods are well understood in the traditional context of convex optimization. For instance, Duchi et al. [DHS11] show that when the loss function is convex, then the Adagrad algorithm yields regret guarantees that are provably as good as those obtained by using the best (diagonal) preconditioner in hindsight. The key mechanism that underlies this improved performance, is that the loss function has some global geometric property (such as sparsity or 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.

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 a factor of 1 . 5 , then the global Lipschitz constant of the network is scaled up by a factor of at least e 10 . Hence it only makes sense to study convergence by looking at the local geometry of the loss along the trajectory of the optimization algorithm [ACH18].

0

500

1000

1500

Iteration

0.095

0.1

0.105

0.11

0.115

0.12

0.125

0.13

0.135

Training loss

0

500

1000

1500

Iteration

0

50

100

150

200

250

SGD+M

Adam

0

500

1000

1500

2000

2500

Iteration

0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Training loss

SGD+M

Adam

Moreover, the interaction between an optimization algorithm and neural network geometry is highly complex - recent work has shown that geometric characteristics of iterates encountered during optimization is highly dependent on the choice of optimization algorithm and associated hyperparameters [LBD + 20, CKL + 21]. For instance, Cohen et al. [CKL + 21] demonstrate that while training neural networks with GD, the maximum eigenvalue of the Hessian evaluated at the GD iterates first increases and then plateaus at a level that is inversely proportional to the 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 x Adam from the Adam trajectory and an iterate x SGD from the SGD trajectory, such that f ( x Adam ) = f ( x SGD ) . We then run SGD+M with the same configuration twice, once from x Adam and once from x SGD. If the underlying geometry of the loss function f was truly fixed, then we would not expect a significant difference in the performance of running SGD+M from either of the two iterates. However, as shown in Figure 1, there is a noticeable difference in performance, and running SGD+M from x Adam achieves lower

loss than running SGD+M from x SGD. This suggests that Adam may bias the optimization trajectory towards a region which is more favorable for rapid training. This motivates the following question.

How does adaptive optimization impact the observed geometry of a neural network loss function, relative to SGD (with momentum)?

The remainder of this paper is dedicated to answering the above question. To this end, for each iterate in a trajectory produced by running an optimization algorithm OPT, where the Hessian of the t th iterate is given by H ( t ) ∈ R d × d , we define the second order statistic R OPT med ( t ) in the following fashion. For the t th iterate in the trajectory, let R OPT med ( t ) be the ratio of maximum of the absolute entries of the diagonal of H ( t ) , to the median of the absolute entries of the diagonal of H ( t ) . Concretely, we define

This statistic thus measures the uniformity of the diagonal of Hessian, where a smaller value of R OPT med ( t ) implies that the Hessian has a more uniform diagonal. It can also be viewed as a stable 1 variant of the condition number. Instead of singular values, we choose diagonal entries because adaptive methods used in practice are coordinate-wise , which can be viewed as the diagonal scaling approaches. 2 In Appendix A.9 we discuss this intuition in detail and compare R OPT med ( t ) with singular value-based metrics. As a supplementary result, in Appendix F, we demonstrate that the loss Hessian approaches diagonal during training for Adam and SGD+M. There has been prior theoretical work on overparameterized neural networks showing that a smaller condition number of Hessian, Neural Tangent Kernel [JGH18] etc. could yield to faster convergence rate for (S)GD [LZB22]. As for (diagonal) adaptive methods ( e.g. Adagrad), they were original designed to adapt to the nonuniform diagonal geometry. Intuitively, a smaller R OPT med ( t ) , which implies more uniform diagonal geometry, could lead to faster convergence.

Armed with this statistic, we make the following contributions:

1. We focus on language models as on NLP tasks, adaptive algorithms show significantly faster convergence than SGD (with momentum). On a wide variety of neural network transformer architectures and language modeling datasets, we conduct experiments to compare how R Adam med ( t ) and R SGDM med ( 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 R Adam med ( t ) values that are significantly smaller than the R SGDM med ( t ) values found by SGD+M. We show a simple example of this phenomenon in Figure 1(right). This suggests that relative to SGD+M, Adam biases the optimization trajectory to a region where the Hessian diagonal is more uniform. We call this phenomenon the uniformity of diagonal geometry for adaptive methods. Moreover, we demonstrate that a more uniform Hessian diagonal, characterized by smaller R OPT med ( t ) , is a contributing factor to faster optimization (see Section 4.3 for discussion). This suggests that a region where the Hessian diagonal is more uniform is also a region that is more amenable to rapid optimization.

2. We complement our empirical results with a theoretical analysis of this phenomenon in the simplified setting of large batch Adam and SGD+M, on a two-layer linear network with d -dimensional input and hidden layer, and one dimensional output. We show that for a wide range of t , R Adam med ( t ) = 1 ± o (1) but R SGDM med ( t ) = Ω(log 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), a 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.

2 Related work

Existing analyses of adaptive methods. The vast majority of prior theoretical work on adaptive methods has focused on the blackbox setting [DHS11, KB15, CZT + 20, RKK18, WWB20, DBBU20,

ii

i

=1

ENV21]. 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) are assumed to hold over the entire domain. Hence this style of 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 on neural networks. By contrast, our focus is on analyzing the local trajectory that is induced by running the optimization method.

Existing analyses of (S)GD on neural networks. There is an extensive literature on the analysis of GD/SGD in the non-blackbox setting, e.g. overparameterized neural networks, [DZPS18, JT20, AZLL19, AZLS19, ADH + 19, LZB22]. However, it 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 [JGH18], the loss geometry is usually fixed and not influenced by learning algorithms. However, in neural networks the interaction between algorithms and loss landscapes is more complicated. Lewkowycz et al. [LBD + 20] find a so-called catapult effect of initial learning rate on the training trajectory of SGD and related loss curvature. Cohen et al. [CKL + 21] demonstrate that while training neural networks with GD, the maximum eigenvalue of the Hessian evaluated at the GD iterates first increases and then plateaus at a level that is inversely proportional to the step size. However, Cohen et al. [CKL + 21] leave open the problem of whether similar interactive phenomena occur in algorithms that are not GD, including adaptive methods.

3 Overview of main results

3.1 Issues of prior analyses on adaptive methods

As is mentioned in Section 2, existing work on adaptive algorithms has mainly focused on blackbox 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 1. To see this, we initialized a transformer model 3 with default initialization in Pytorch but chose a large gain 4 , and computed the smoothness parameter (denoted as l ) and the condition number (denoted as κ ) of loss Hessian on one layer. We observed that setting the gain as a large constant ( e.g. 800) results in extremely large l and κ ( l ≥ 10 7 and κ ≥ 10 10 ), which makes the convergence rates in prior black-box analysis vacuous.

The failure of global worst-case analysis implies that we need to focus on the local trajectory of algorithms. However, it is unclear whether when two optimization algorithms are used, they will have the same geometry in local trajectory or not. In particular, although in theory, adaptive algorithms can yield to 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 to SGD (with momentum) in the paper. Motivated by the diagonal scaling of Adagrad and Adam for neural network training, we ask the follow main question in our paper:

How does the local diagonal geometry (diagonal of the loss Hessian) along the local trajectory of adaptive algorithms compare to that of SGD (with momentum)?

3.2 Overview of the experiments

As is discussed in Section 1, we consider R OPT med ( t ) defined in eq. (1) as a measurement of the uniformity of the loss Hessian diagonal and conduct experiments on different NLP tasks where adaptive methods converge faster. The detailed experimental setup will be stated in Section 4. To explore potential different patterns of different layers, we do the computation layer by layer. On a

wide variety of transformer architectures and language modeling datasets from the same initialization, we observe that for the vast majority of layers:

When we train the neural network using Adam, the uniformity of diagonal geometry, measured by R OPT med ( t ) is smaller than that when we train using SGD+M from the same initialization.

Table 1 shows a typical example of R Adam med ( t ) compared to R SGDM med ( t ) on a sentence classification task using BERT-small [TCLT19, BDR21] (See Section 4.1 for details). We repeated the experiments for 12 times starting from the same initialization. Table 1 shows the averaged R Adam med ( t ) and R SGDM med ( t ) in some randomly selected layers. We also report the averaged R SGDM med ( t ) R Adam med ( t ) and their standard deviations in the brackets. 5 Figure 2 shows the corresponding training losses of one in these 12 experiments.

To understand this phenomenon in a more principled point of view, in Section 5 we provide a formal proof of the statement in a simplified setting: large batch Adam and SGD+M on a 2-layer linear network with 1-dimensional output.

4 The uniformity of diagonal geometry

As is mentioned in Section 3.2, we computed R OPT med ( 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 A, we present the results of other adaptive algorithms.

During training we started from the same initial weights and used the same learning rate schedule (constant or decreasing) for SGD+M and Adam. We tuned and chose the best (initial) learning rate of SGD+M. The (initial) learning rate of Adam was set as a value under which Adam converged faster than SGD+M with its best learning rate. The concrete values will be stated in later parts of this section. We used large batch sizes to make the training procedure stable. When computing Hessian, we also used large batch sizes. Due to the extremely large dimension, we did the computation on some uniformly selected coordinates, more precisely, 200 coordinates per layer.

4.1 Experiments on real datasets

Sentence classification task on BERT-small. We fine-tuned BERT-small [TCLT19, BDR21] on the IMDB dataset [MDP + 11]: the task is to classify whether movie reviews are positive or negative. 6 The momentum parameter β in SGD was set as 0.9. The two momentum parameters ( β 1 , β 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 in Section 3.2, Figure 2 and Table 1 show the training losses and the comparison between R Adam med ( t ) and R SGDM med ( t ) , respectively.

Translation task. Wetrained a Seq2Seq network that uses Transformer to solve a machine translation task on Multi30k [EFSS16]: this task is to train a German to English translation model. 7 The

momentum parameter β in SGD was set as 0.9. The two momentum parameters ( β 1 , β 2 ) of Adam were set as (0.9, 0.98). We trained the model using constant learning rates (0.03 for SGD+M and 1e-4 for Adam) for 60 epochs (1800 iterations). The experiments were repeated for 8 times starting from the same initialization. Figure 3(left) shows the training losses for one among them. Table 2a shows the averaged R Adam med ( t ) , R SGDM med ( t ) and R SGDM med ( t ) R Adam med ( t ) (with standard deviation in the brackets) in some randomly selected layers.

0

200

400

600

800

1000

1200

1400

1600

1800

Iteration

1

2

3

4

5

6

7

8

9

10

Training loss

SGD+M

Adam

0

200

400

600

800

1000

1200

1400

1600

1800

Iteration

9.05

9.1

9.15

9.2

9.25

9.3

9.35

9.4

9.45

9.5

Training loss

SGD+M

Adam

0

100

200

300

400

500

600

coodinate

0.5

1

1.5

2

2.5

3

singular values

SGD+M

Adam

0

10

20

30

40

50

60

70

80

7

8

9

10

11

SGD+M

Adam

(a)

(b)

4.2 Experiments on random datasets

We used the same model and momentum parameters as in the translation task described in Section 4.1 but generated random integers as targets. Similar to the setting on real targets, the model was trained using constant learning rates (0.015 for SGD+M and 5e-5 for Adam) for 60 epochs (1800 iterations), and we repeated the experiments for 8 times starting from the same initialization. Figure 3(right) shows the training losses for one among them. Table 2b shows the averaged R Adam med ( t ) , R SGDM med ( t ) and R SGDM med ( t ) R Adam med ( t ) (with standard deviation in the brackets) of the same 10 layers as in Table 2a. 8

4.3 Summarization of the empirical results and discussion

Overall, through extensive experiments on language models, we demonstrate that starting from the same initialization, for the vast majority of layers, the R OPT med ( t ) values found by Adam are smaller than those found by SGD+M. This suggests that Adam biases the trajectory towards a region with more uniform Hessian diagonal than SGD+M. In Appendix A.10 we also validate this observation on the in-distribution test data.

Contribution of uniform Hessian diagonal to fast convergence. We observe that on dataset with random targets, SGD+M plateaus after about 400 steps and thus converges much slower when compared to Adam than on real dataset (see Figure 3). On the other hand, the gaps of R SGDM med ( t ) and

R Adam med ( t ) are more significant on random data than on real data (see Table 2a and Table 2b) as well. In Appendix A.4, we conduct another experiment where we switch from SGD to Adam in the middle and compare it with the model trained by Adam from the beginning. The observation is that both the loss gap and the gap of R OPT med ( t ) are gradually closed after switching (see Figure 7 and Table 8). Hence we find a positive correlation between fast convergence and uniform Hessian diagonal (small R OPT med ( t ) ). In Appendix A we study other adaptive algorithms (Adagrad, RMSprop and AMSGrad) and get similar observation : all these adaptive methods converge faster than SGD or SGD+M and also bias the trajectory to regions with smaller R OPT med ( t ) , suggesting the universality of our observation.

Despite the above positive correlation, it is reasonable to ask whether small R OPT med ( t ) indeed contributes to fast convergence or is just a byproduct of adaptive methods. To address this concern, in Appendix B.1, we add a supplementary experiment similar to that in Figure 1. We select two iterates x 1 and x 2 from two trajectories that both come from SGD+M (instead of one from Adam and one from SGD+M in Figure 1), such that the loss f ( x 1 ) = f ( x 2 ) but x 2 has a smaller R OPT med ( t ) than x 1 . We then run SGD+M with the same configuration twice, once from x 1 and once from x 2 . Under this setting, we get similar observation as before: running SGD+M from x 2 (with smaller R OPT med ( t ) ) achieves faster convergence than from x 1 . This suggests that the uniformity of the diagonal of loss Hessian (measured by R OPT med ( t ) ) reveals some intrinsic trajectory property beyond the algorithm choice and is indeed a contributing factor to fast optimization. In Appendix B.2 we theoretically prove the contribution of small R OPT med ( t ) to fast optimization in a simplified setting.

More discussions on the trajectory difference. Considering the fact that our comparison between R Adam med ( t ) and R SGDM med ( t ) is conditioned on the same iteration when SGD+M has larger training loss than Adam, there is a potential alternative explanation of the Hessian diagonal uniformity. That is, the global minimum has uniform Hessian, and Adam simply converges faster to it than SGD+M. To rule out this possibility, in Appendix A.3 we add a comparison of our measurements R Adam med ( t ) and R SGDM med ( t ′ ) , where t, t ′ are picked such that t th Adam iterate and t ′ th SGD+M iterate have the same training loss . The results (in Table 7) show that R Adam med ( t ) < R SGDM med ( t ′ ) for most layers, thus demonstrating that the trajectories of Adam and SGD+M are truly different and that the difference is because Adam biases the local geometry (as opposed to faster convergence).

Adding regularization. People in practice usually add weight decay (equivalent to l 2 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, our observation still holds: Adam converges faster than SGD+M and in the vast majority of layers, R Adam med ( t ) values are smaller than R SGDM med ( t ) . 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 R Adam med ( t ) values were larger rather than smaller. In the case of strong regularization, the adaptivity of Adam requires further exploration and we hope to find new mechanisms in the future.

5 Theoretical analysis

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 a two-layer linear network with 1-dimensional output. Although simple, the choice of 2-layer linear networks to understand learning dynamics is common in prior works ( e.g. [TCG21]). Moreover, in language transformer models when the weights are small, the softmax in the key-value-query structure is near the linear regime. Then this structure might be approximated by the product of 3 linear operators, similar to a three-layer linear network. Hence the theoretical analysis of this phenomenon on linear networks would be a good starting point for further understanding of more complicated language models.

5.1 Problem setup

Notation Let [ d ] = { 1 , 2 , ..., d } . We use ‖ · ‖ 2 to denote the l 2 norm of a vector, and ‖ · ‖ F to denote the Frobenius norm of a matrix. Let 〈· , ·〉 be the Euclidean inner product between vectors or matrices. Let N ( µ, σ 2 ) be the one-dimensional Gaussian distribution with mean µ and variance σ 2 . For a scalar (vector, matrix) A which evolves over time, we use A ( t ) to denote its value at time t .

Let there be m data points. The data matrix is X ∈ R d x × m and the label matrix is Y ∈ R d y × m . We assume that the input dataset is whitened, i.e. Λ xx := 1 m XX T ∈ R d x × d x is an identity matrix.

The parameters of a 2-layer linear network are given by W := ( W 2 , W 1 ) . Assume W i ∈ R d i × d i -1 for i = 1 , 2 . We have d 2 = d y , d 0 = d x . We consider the square loss L ( W ) := 1 2 m ‖ W 2 W 1 X -Y ‖ 2 F .

Denote A := 1 m Y X T ∈ R d y × d x . [AGCH19] show that with whitened dataset,

where c does not depend on W . We consider the following model with small Gaussian initialization.

Assumption 1 (Setup) . The input covariance Λ xx := 1 m XX T ∈ R d x × d x is an identity matrix. The input and hidden layers are both of dimension d , i.e. d 1 = d 0 = d . Without loss of generality, we can assume that A is a row vector (i.e. d 2 = 1 ) whose coordinates are positive 9 and Θ(1) in terms of d .

Assumption 2 (Gaussian Initialization) . ∀ i, j : w (0) 2 i ∼ N (0 , 1 d 2 α ) , W (0) 1 [ i, j ] ∼ N (0 , 1 d 4 α ) are independently initialized with sufficiently large α > 0 .

Denote ˜ A and ˜ Λ xx as the batch versions of A and Λ xx . Wemake the following large-batch assumption. We emphasize that large batches are commonly used in NLP tasks ( e.g. [BMR + 20]).

Denote ˜ g ( t ) as the batch gradient at time t . The update rules of SGD+M and Adam are given by

where η is the learning rate, β, β 1 , β 2 are momentum parameters, and ξ is for numerical stability. All operations on vectors are element-wise. Here and throughout, the notation f ( x ) = O ( g ( x )) (resp. f ( x ) = Ω( g ( x )) , f ( x ) = Θ( g ( x )) ) means that ∃ C 1 , C 2 > 0 such that f ( x ) ≤ C 2 g ( x ) (resp. f ( x ) ≥ C 1 g ( x ) , C 1 g ( x ) ≤ f ( x ) ≤ C 2 g ( x ) ). Here C 1 , C 2 may depend on β, β 1 , β 2 . We will also use the notation with ∼ , i.e. ˜ O ( · ) , ˜ Ω( · ) , ˜ Θ( · ) to hide factors of order log d . In our theoretical analysis, 'with high probability', or 'w.h.p.' for short, means that with probability at least 1 -1 poly ( d ) .

Since the weights and Hessians in different layers may have different magnitudes, we compute the R OPT med ( t ) layer by layer. Denote R SGDM med ,k ( t ) (resp. R Adam med ,k ( t ) ) as the R OPT med ( t ) found by SGD+M (resp. Adam) w.r.t. W k at time t where k = 1 , 2 .

5.2 Main results

Theorem 1. Under Assumption 1, 2 and 3, consider the weights { W ( t ) SGD } t ≥ 0 (resp. { W ( t ) Adam } t ≥ 0 ) obtained by SGD+M (resp. Adam) defined in (3) .

1. For any p > 0 , pick 0 < glyph[epsilon1] < 1 d p , η ≤ O ( glyph[epsilon1] d 7 α/ 4+4 ) and α ≥ 4( p +2) . Suppose σ ≤ η 3 / 2 d α/ 2+1 , then there exists T SGD , 1 , T SGD , 2 such that w.h.p., ¯ L ( W ( T SGD , 1 ) SGD ) = Θ( d ) , ¯ L ( W ( T SGD , 2 ) SGD ) ≤ ˜ O ( 1 d p ) , and

Remark Theorem 1 holds for all values of hyperparameters (such as α, σ ) in certain ranges instead of just particular values. The ranges of SGD+M and Adam overlap with each other. That means we can choose the same hyperparameters for SGD+M and Adam in the overlapped region to make a fair comparison, for example, same α, σ such that α ≥ 4( p +2) and σ ≤ min { η 3 / 2 d α/ 2+1 , η 3 / 2 ξ 2 d 13 / 4 } .

An immediate corollary of this theorem below gives the difference between iterates of Adam and SGD+M that have the same loss.

′ ∈

Theorem 1 and Corollary 1 tell us that during a long training period when the loss decreases from Θ( d ) to ˜ O ( 1 d p ) , the diagonal of loss Hessian for Adam keeps nice uniformity: R Adam med ,k ( t ) = 1 ± o (1) , k = 1 , 2 . On the other hand, the diagonal of loss Hessian for SGD+M is less uniform. Appendix C gives a proof sketch of Theorem 1. The detailed proof can be found in Appendix D and E.

6 Low-rank weight matrices and uniformity of leading singular vectors

The proof sketch in Appendix C highlights one crucial intuition of Theorem 1: After T SGD , 1 (resp. T Adam , 1 ) steps, W 1 of SGD+M (resp. Adam) becomes an approximately rank-1 matrix. Consider the left singular vector u := [ u 1 , u 2 , ..., u d ] T which corresponds to the leading singular value σ 1 . We can show that the distribution of u 2 1 , u 2 2 , ..., u 2 d for Adam is more uniform than that of SGD+M. This property, we call the uniformity of the leading singular vector , is related to the uniformity of the diagonal of loss Hessian, see Appendix G for more details.

Similar low rank bias after training has been studied in prior works ( e.g. [GWB + 17, LMZ18, CGMR20]). 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 ∈ R m × n , we want to check

(A) Whether W ∈ R m × n is approximately a rank k matrix with k glyph[lessmuch] min { m,n } , and if true,

(B) Consider the top k singular values σ 1 , ..., σ k and left singular vectors u 1 , ..., u k . Define ˜ u := ∑ k i =1 σ 2 i u i glyph[circledot] u i := [˜ u 1 , ..., ˜ u d ] T and compute R u := max i ˜ u i median ˜ u i , a generalized version of max i u 2 i median u 2 i in the rank-1 case. We want to see whether R u obtained by Adam is smaller than that of SGD+M.

After reviewing the weight matrices we got in different settings, we observed that (A) and (B) hold for many layers in those models. For example, on the translation task mentioned in Section 4.1, we found 12 layers which had approximately low rank structures and for 10 of them, R u values (defined in (B)) obtained by Adam were smaller than those found by SGD+M. Figure 4 shows the result on one typical layer. Results of more layers can be found in Appendix A.5.

Remark The definition of R u is based on the connection between diagonal of loss Hessian and weight matrices. Appendix G shows that for a 2-layer linear network, R OPT med , 2 ( t ) =

max i ‖ W ( t ) 1 [ i, :] ‖ 2 2 median ‖ W ( t ) 1 [ i, :] ‖ 2 2 . When W 1 ∈ R m × n is approximately rank k , i.e. W 1 ≈ ∑ k i =1 σ i u i v T i , denote u i = [ u i 1 , u i 2 , ..., u im ] T and v i = [ v i 1 , v i 2 , ..., v in ] T , we have that for the j -th row,

∑ i =1 σ 2 i u i glyph[circledot] u i , we have that ‖ W 1 [ j, :] ‖ 2 2 ≈ ˜ u j . works, the connection between diagonal of loss Hessian and the weight matrices is more complicated and R OPT med , 2 ( t ) may depend on the product of many weight matrices rather than one single matrix, we still believe that this definition of R u is a reasonable ratio to consider.

2 k Although in multi-layer nonlinear neural net-

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 to an improved algorithm for practical use, they provide a new way of thinking when designing new algorithms: in contrast with the traditional view which tries to design 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. As is discussed in Section 5, the key-value-query structure in the transformer models might be approximated by a three-layer linear network and the analysis of multi-layer networks might provide more connection to real deep models and could be an interesting and challenging future direction. Moreover, it is also possible to relax our large-batch assumption (Assumption 3) and prove similar results in the general stochastic setting.

Acknowledgments and Disclosure of Funding

This project was supported by NSF CCF 2145703.

References

[ACH18] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In 35th International Conference on Machine Learning, ICML 2018 , pages 372-389. International Machine Learning Society (IMLS), 2018.

[ADH + 19] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning , pages 322-332. PMLR, 2019.

[AGCH19] Sanjeev Arora, Noah Golowich, Nadav Cohen, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In 7th International Conference on Learning Representations, ICLR 2019 , 2019.

[AZLL19] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

[AZLS19] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

[BDR21] Prajjwal Bhargava, Aleksandr Drozd, and Anna Rogers. Generalization in nli: Ways (not) to go beyond simple heuristics, 2021.

[BMR + 20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner,

Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems , volume 33, pages 1877-1901. Curran Associates, Inc., 2020.

[CGMR20] Hung-Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep matrix factorization: Dynamics and implicit bias towards low rank. arXiv preprint arXiv:2011.13772 , 2020.

[CKL + 21] Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

[CZT + 20] Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu. Closing the generalization gap of adaptive gradient methods in training deep neural networks. In Christian Bessiere, editor, Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI 2020 , pages 3267-3275. ijcai.org, 2020.

[DBBU20] Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. arXiv preprint arXiv:2003.02395 , 2020.

[DCLT19] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pretraining 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 1 (Long and Short Papers) , pages 4171-4186. Association for Computational Linguistics, 2019.

[DHS11] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

[DZPS18] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2018.

[EFSS16] 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 , pages 70-74. Association for Computational Linguistics, 2016.

[ENV21] Alina Ene, Huy L. Nguyen, and Adrian Vladu. Adaptive gradient methods for constrained convex optimization and variational inequalities. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021 , pages 7314-7321. AAAI Press, 2021.

[GWB + 17] 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.

[JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

[JT20] Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020.

[Kaw16] Kenji Kawaguchi. Deep learning without poor local minima. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016.

[KB15] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015.

[LBD + 20] Aitor 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.02218 , 2020.

[Ler19] Matthieu Lerasle. Lecture notes: Selected topics on robust statistical learning theory. arXiv preprint arXiv:1908.10761 , 2019.

[LMZ18] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in overparameterized matrix sensing and neural networks with quadratic activations. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet, editors, Proceedings of the 31st Conference On Learning Theory , volume 75 of Proceedings of Machine Learning Research , pages 2-47. PMLR, 06-09 Jul 2018.

[LZB22] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis , 2022.

[MDP + 11] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies , pages 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.

[MXBS17] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings . OpenReview.net, 2017.

[RKK18] Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018.

[TCG21] Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 10268-10278. PMLR, 18-24 Jul 2021.

[TCLT19] Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv:1908.08962v2 , 2019.

[TH + 12] Tijmen Tieleman, Geoffrey Hinton, et al. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning , 4(2):26-31, 2012.

[VSP + 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

[WWB20] Rachel Ward, Xiaoxia Wu, and Léon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. Journal of Machine Learning Research , 21(219):1-30, 2020.

[ZCS + 22] Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification on update rules. In NeurIPS , 2022.

A More experiments of the uniformity of diagonal geometry

A.1 Vanilla SGD vs. Adagrad

In this section, we present the R OPT med ( t ) values defined in eq. (1) obtained by vanilla SGD and Adagrad on a language modeling task 11 . The task is to assign a probability for the likelihood of a given word (or a sequence of words) to follow a sequence of words. We trained a transformer model to solve this problem on Wikitext-2 [MXBS17](CC BY-SA 3.0) with real targets and randomly generated targets. This model has roughly 8 layers (not counting normalization and dropout layers)

The setup is the same as 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 a 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 initialization schedule.

A.1.1 Experiments on real datasets

Figure 5a shows the training losses on Wikitext-2. Table 3 (resp. Table 4) shows the R OPT med ( t ) for Adagrad and SGD under uniform (resp. normal) initialization with different gains.

0

500

1000

1500

2000

2500

3000

3500

4000

Iteration

4

5

6

7

8

9

10

11

Training loss

SGD

Adagrad

(a)

0

500

1000

1500

2000

2500

3000

3500

4000

Iteration

9.9

10

10.1

10.2

10.3

10.4

10.5

10.6

10.7

Training loss

SGD

Adagrad

(b)

Table 3: R SGD med ( t ) and R Adagrad med ( t ) under uniform initialization with different gains

(a) Gain = 2

(b) Gain = 0.5

Layer#

1

2

3

4

5

6

7

8

(a) Gain = 1

Epoch 20

t

)

R

SGD

med

(

t

)

R

8.27

3.19

8.61

6.45

11.69

3.84

17.95

Adagrad med

(

12.28

1.87

3.38

4.86

11.23

2.32

4.32

SGD

med

(

9.69

3.21

9.25

7.49

15.07

3.20

14.99

t

)

R

Adagrad med

(

11.17

1.73

3.94

4.44

12.18

2.09

4.50

SGD

med

(

6.76

9.51

7.38

18.02

12.70

12.76

11.79

17.09

5.16

26.91

t

)

R

Adagrad med

6.06

6.61

7.35

15.63

9.35

11.86

8.58

12.73

26.70

6.73

A.1.2 Experiments on random datasets

Figure 5b shows the training losses on Wikitext-2 but with randomly generated targets and Table 5 shows the corresponding R SGD med ( t ) and R Adagrad med ( t ) in different layers.

A.2 RMSprop and AMSGrad

In this section, we present the results of RMSprop and AMSGrad and compare them with SGD+M. The experiments were conducted on the translation task described in Section 4.1. We used learning rate 2.5e-5 for RMSprop, 5e-4 for AMSGrad and 0.03 for SGD+M. Both RMSprop and SGD+M used momentum parameter 0.9. The two momentum parameters ( β 1 , β 2 ) of AMSGrad were (0 . 9 , 0 . 98) . Figure 6 shows the training losses and Table 6 shows the corresponding R OPT med ( t ) .

0

200

400

600

800

1000

1200

1400

1600

1800

Iteration

0

1

2

3

4

5

6

7

8

9

10

Training loss

RMSprop

SGD+M

(b) Gain = 0.5

R

(

t

)

R

t

)

A.3 Comparison conditioned on the same loss

In this section, we compare R SGDM med ( t ) and R Adam med ( t ) conditioned on the same training loss. More precisely, we make comparison of R Adam med ( t ) and R SGDM med ( t ′ ) , where t, t ′ are picked such that t th 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 7 shows the results of R Adam med ( t ) and R SGDM med ( t ′ ) in some layers.

(a)

0

500

1000

1500

2000

2500

Iteration

1

2

3

4

5

6

7

8

9

10

Training loss

SGD

Adam

Adam after SGD

A.4 Experiments of switching from SGD to Adam

In this section we describe another learning schedule: the 'Adam after SGD' schedule, where we switch from SGD to Adam in the middle to see whether the loss and R OPT med ( t ) can catch up with the model trained by Adam from the very beginning. Again, we used the same model as in the translation task in Section 4.1. In this section, we did not add momentum term to SGD in order to get a 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.

As is shown in Figure 7 and Table 8, both the loss gap and the gap of R OPT med ( 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.

















/

/



⊕∐glyph[arrowbt]˜√(]

⋃⊔〈〉⊕

̂∐⌉











/

/



/



⊕∐glyph[arrowbt]˜√(]

⋃⊔〈〉⊕

̂∐⌉

̂}}̂]{∐√˜

̂}}̂]{∐√˜

We then studied the uniformity of leading singular vectors of these 12 layers, i.e. computed R u defined in (B) of Section 6. The observation was that for 10 out of these 12 layers, R u values of Adam were smaller those of SGD, which implies the uniformity of leading left singular vectors of Adam. Here we may also want to consider the right singular vectors v 1 , v 2 , ... v k and corresponding ˜ v = [˜ v 1 , ˜ v 2 , ..., ˜ v d ] T := ∑ k i =1 σ 2 i v i glyph[circledot] v i and compute R v := max i ˜ v i median ˜ v i for Adam and SGD+M. However, on this translation task, among the 12 layers which were approximately low rank, for only 6 of them, R v of Adam were smaller, i.e. we did not observe uniformity of the leading right singular

vector for Adam. 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. Figure 9 shows how R u and R v changed over time in some layers.

0

500

coodinate

0.5

1

1.5

2

2.5

singular values

Layer #5

SGD+M

Adam

0

50

10

12

14

SGD+M

Adam

0

50

6

8

10

12

14

SGD

Adam

0

500

coodinate

0

1

2

singular values

Layer #11

SGD+M

Adam

0

50

8

10

12

SGD+M

Adam

0

50

8

9

10

11

SGD

Adam

0

500

coodinate

0.5

1

1.5

2

2.5

singular values

Layer #19

SGD+M

Adam

0

50

10

12

14

16

18

SGD+M

Adam

0

50

8

10

12

14

SGD

Adam

0

500

coodinate

0.5

1

1.5

2

2.5

singular values

Layer #25

SGD+M

Adam

0

50

10

15

20

SGD+M

Adam

0

50

6

8

10

SGD

Adam

A.6 How does the (adaptive) gradient align with diagonal of loss Hessian

In this section we present the uniformity of diagonal geometry of adaptive methods from another perspective. Denote H ii as the ( i, i ) -th element of the loss Hessian H and g i as the i -th element of the gradient. It is conjectured that when | H ii | is large, the corresponding | g i | is usually large as well. For adaptive methods, we can regard the update per step as the learning rate times the 'adaptive gradient'. Let's use g adapt ,i to represent the i -th component of the adaptive gradient. Through experiments on language models, we found that | g adapt ,i | for different i are quite uniform and do not align with | H ii | as the true gradient | g i | does.

In the experiments, we first sorted | H ii | in the ascent order: | H i 1 ,i 1 | ≤ | H i 2 ,i 2 | ≤ ... ≤ | H i d ,i d | (suppose H ∈ R d × d ), and then plotted the corresponding | g i k | and | g adapt ,i k | for k ∈ [ 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.

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 11 shows the results on the sentence classification task and Figure 12 shows the results on the translation task.

0

50

100

150

200

coordinate

0

1

2

|g

i

|

10 -3

True gradient of SGD

epoch = 2

0

50

100

150

200

coordinate

0

0.5

1

1.5

|g

i

|

10 -3

True gradient of Adagrad

0

50

100

150

200

coordinate

0

0.5

|g

adapt, i

|

Adaptive gradient of Adagrad

0

50

100

150

200

coordinate

0

2

4

|g

i

|

10 -4

epoch = 20

0

50

100

150

200

coordinate

0

0.5

1

|g

i

|

10 -3

0

50

100

150

200

coordinate

0

0.05

|g

adapt, i

|

0

50

100

150

200

coordinate

0

5

|g

i

|

10 -4

epoch = 40

0

50

100

150

200

coordinate

0

0.5

1

|g

i

|

10 -3

0

50

100

150

200

coordinate

0

0.05

|g

adapt, i

|

0

50

100

150

200

coordinate

0

0.5

|g

i

|

True gradient of SGD+M

epoch = 1

0

50

100

150

200

coordinate

0

0.05

0.1

|g

i

|

True gradient of Adam

0

50

100

150

200

coordinate

0

2

|g

adapt, i

|

Adaptive gradient of Adam

0

50

100

150

200

coordinate

0

0.5

|g

i

|

epoch = 3

0

50

100

150

200

coordinate

0

0.01

0.02

|g

i

|

0

50

100

150

200

coordinate

0

0.5

1

1.5

|g

adapt, i

|

0

50

100

150

200

coordinate

0

0.5

|g

i

|

epoch = 6

0

50

100

150

200

coordinate

0

0.5

1

|g

i

|

10 -3

0

50

100

150

200

coordinate

0

0.5

|g

adapt, i

|

A.7 Adding regularization 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 ( β 1 , β 2 ) 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 13a shows the training losses and Table 9 shows the values of R Adam med ( t ) , R SGDM med ( t ) and R SGDM med ( t ) R Adam med ( t ) in some randomly selected layers.

0

50

100

150

200

coordinate

0

5

|g

i

|

10 -4

True gradient of SGD+M

epoch = 6

0

50

100

150

200

coordinate

0

5

|g

i

|

10 -4

True gradient of Adam

0

50

100

150

200

coordinate

0

0.5

1

|g

adapt, i

|

Adaptive gradient of Adam

0

50

100

150

200

coordinate

0

1

2

|g

i

|

10 -3

epoch = 31

0

50

100

150

200

coordinate

0

2

|g

i

|

10 -3

0

50

100

150

200

coordinate

0

0.5

|g

adapt, i

|

0

50

100

150

200

coordinate

0

1

2

3

|g

i

|

10 -3

epoch = 56

0

50

100

150

200

coordinate

0

1

2

|g

i

|

10 -3

0

50

100

150

200

coordinate

0

0.5

|g

adapt, i

|

A.8 Results on image tasks

Although in this paper we focus on language models where Adam shows significant fast convergence, we also add results in this section on image tasks where SGD+M performs better. We trained a ResNet 12 on CIFAR-10 dataset and compared the convergence speed and R OPT med ( t ) of SGD+M and Adam. The momentum parameter β in SGD was set as 0.9. The two momentum parameters ( β 1 , β 2 ) of Adam were set as (0.9, 0.98). The model was trained using constant learning rates for 41 epochs (2050 iterations). We tuned and chose the best learning rates for both algorithms: 0.5 for SGD+M and 0.005 for Adam. Figure 13b shows the training losses and Table 10 shows the values of R Adam med ( t ) , R SGDM med ( t ) and R SGDM med ( t ) Adam .

As we can see, on this image task, Adam does not converge faster than SGD+M and in the meantime, R Adam med ( t ) values were no longer smaller than R SGDM med ( t ) during training. This reveals the connection between the local diagonal geometry and the convergence speed from another perspective. That is, when the diagonal of Hessian of Adam is not more uniform than SGD+M, its convergence speed is not better, either.

0

200

400

600

800

1000

1200

1400

1600

1800

Iteration

1

2

3

4

5

6

7

8

9

10

Training loss

SGD+M+wd

Adam+wd

(a)

0

500

1000

1500

2000

Iteration

0

0.5

1

1.5

2

2.5

Training loss

SGD+M

Adam

(b)

A.9 Comparison between R OPT med ( t ) and singular value-based metrics

In Section 4, through extensive experiments on language models, we demonstrate that when we train the neural network using Adam, the uniformity of diagonal geometry, measured by R OPT med ( t ) is smaller than that when we train using SGD+M from the same initialization, for the vast majority of layers. We are aware that people also usually consider Hessian singular values instead of diagonal entries to measure the loss geometry. Hence in this section we make a comparison between our diagonal-based metric and singular value-based metrics.

First, we believe that our metric has a natural connection to the mechanism that underlies adaptive methods. Adaptive methods in practice choose coordinate-wise adaptive learning rates. From a high-level perspective, this procedure can be viewed as adapting to the loss smoothness with respect to each coordinate. The smoothness of certain coordinate is measured by the second derivative with respect to this coordinate and therefore corresponds to the diagonal entries of loss Hessian. Our metric, which measures these diagonal entries, is thus fundamentally intertwined with the mechanism that underlies adaptive methods.

Next, we empirically demonstrate that our metric R OPT med ( t ) is a reasonable proxy of singular valuebased metrics. Define a singular value-based metric S OPT med ( t ) := max { σ i ( t ) } d i =1 median { σ i ( t ) } d i =1 as an analogy of our diagonal-based metric R OPT med ( t ) , where { σ i ( t ) } d i =1 denotes the singular values of loss Hessian H ( t ) ∈ R d × d at the t th iterate. We compare S OPT med ( t ) along the trajectories of Adam and SGD+M in the translation task described in Section 4.1. Table 11 suggests that if measured by singular values, Adam is also biased to a region with smaller S OPT med ( t ) than SGD+M, similar to the observation for R OPT med ( t ) . This is expected because in Appendix F, we demonstrate that the loss Hessian approaches diagonal during training. The fact that our diagonal-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 easier to compute empirically and analyze theoretically than singular value-based metrics such as S OPT med ( t ) .

From the empirical computation perspective, suppose the loss Hessian is d × d . Then computing its singular values, in general, requires computing the whole matrix with d 2 elements. However, our metric only requires computing the d diagonal entries.

From the theoretical analysis perspective, in Appendix G, we show that the diagonal of loss Hessian in linear networks can be connected to weight matrices by simple formulas. These straightforward formulas simplify the analysis and allow us 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.

A.10 The uniformity of diagonal geometry on in-distribution test data

In this section we compare R Adam med ( t ) and R SGDM med ( 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 a region with smaller R OPT med ( t ) than SGD+M, similar to what happens on the training data shown in Table 2a. This is expected because of the same distribution. One other thing 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.

B Supplementary result on the contribution of small R OPT med ( t ) to fast optimization

B.1 Supplementary empirical result

As explained in Section 1, Figure 1 demonstrates that Adam biases the trajectory towards a region which benefits fast optimization, and in the meantime, it has more uniform diagonal Hessian (characterized by smaller R OPT med ( t ) ). However, since the two trajectories in Figure 1 come from two different algorithms (SGD+M and Adam), it is reasonable to ask whether the benefit of Adam's trajectory is indeed due to smaller R OPT med ( t ) or it is due to other properties of Adam and small R OPT med ( t ) is just a byproduct. To address this concern, we add the following experiment where we select two trajectories that both come from SGD+M and have different R OPT med ( t ) .

We ran SGD+M twice on the translation task described in Section 4.1 with learning rates 0.01 and 0.03 and got two trajectories. Then we selected an iterate x 1 on the trajectory with learning rate 0.01 and x 2 on the trajectory with learning rate 0.03 such that x 1 and x 2 have the same training loss and different R OPT med ( t ) . We then ran SGD+M with the same configuration twice, once starting from x 1 and once from x 2 . The observation is that x 2 has a more uniform diagonal Hessian (characterized by smaller R OPT med ( t ) , see Table 13) than x 1 , and in the meantime, SGD+M starting from x 2 also converges faster (see Figure 14). Hence we demonstrate that a small value of R OPT med ( t ) is not just a byproduct of adaptive algorithms (since SGD+M with different learning rates can also yield trajectories with different R OPT med ( t ) ), and it is indeed a contributing factor of fast optimization.

0

10

20

30

40

50

60

70

80

90

Iteration

2.8

2.85

2.9

2.95

3

3.05

3.1

3.15

Training loss

B.2 Theoretical analysis on a simple setting

In this section, we add the following lemma on a simplified setting where smaller R OPT med ( t ) provably contributes to faster optimization.

Lemma 1. Consider a diagonal quadratic loss f ( x ) = ∑ d i =1 λ i 2 x 2 i where λ 1 ≥ λ 2 ≥ ... ≥ λ d ≥ 0 . Run GD with learning rate η = 1 λ 1 and produce { x ( t ) i } t ≥ 0 , 1 ≤ i ≤ d . We initialize the parameters as x (0) i = ± 1 , ∀ i ∈ [ d ] . Without loss of generality, suppose d is even. Then we have that for any p > 0 , after T = p 2 · R OPT med · log d = p 2 · λ 1 median { λ i } d i =1 log d steps, the loss f ( x ) will shrink by a half. More precisely,

Lemma 1 demonstrates that under this simplified setting, even when we just use the vanilla GD with no adaptive methods, a more uniform diagonal geometry (characterized by smaller R OPT med ) can still lead to faster optimization.

Proof. By the update of GD with learning rate η = 1 λ 1 , we have that

Then after T = p 2 · λ 1 median { λ i } d i =1 log d steps, we have that

On the other hand,

where ( i ) is by definition of median and the initialization ∀ i ∈ [ d ] : x (0) i = ± 1 . Combining with the fact ∑ d i = d/ 2+1 λ i ( x (0) i ) 2 + ∑ d/ 2 i =1 λ i ( x (0) i ) 2 = f ( x (0) ) yields that ∑ d i = d/ 2+1 λ i ( x ( T ) i ) 2 ≤ 1 2 f ( x (0) ) . Hence we have

C Proof sketch of Theorem 1

Now we give a proof sketch of Theorem 1, which contains three major steps. The detailed proof can be found in Appendix G, D and E.

First we relate the diagonal of Hessian to weight matrices W 1 , W 2 . Under Assumption 1, denote W 1 [ i, :] as the i -th row of W 1 and W 2 := [ w 2 i , w 22 , ..., w 2 d ] . Since the input dataset is whitened, we can show that

Next, due to the one-dimensional output, we can prove that W 1 converges to an approximately rank-1 matrix. More precisely, we have

where c ( t ) is a scalar, u ( t ) , v ( t ) , R ( t ) 2 ∈ R d and R ( t ) 1 ∈ R d × d .Denote the i -th coordinate of u ( t ) , v ( t ) , R ( t ) 2 as u ( t ) i , v ( t ) i , R ( t ) 2 i , respectively. Denote the ( i, j ) -th element of R ( t ) 1 as R ( t ) 1 [ i, j ] . We have that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ R ( t ) 2 i ∣ ∣ ∣ glyph[lessmuch] c ( t ) ∣ ∣ ∣ u ( t ) i ∣ ∣ ∣ and ∣ ∣ ∣ R ( t ) 1 [ i, j ] ∣ ∣ ∣ glyph[lessmuch] ∣ ∣ ∣ u ( t ) i v ( t ) i ∣ ∣ ∣ .

Using the rank 1 structure, we can further simplify R OPT med , 1 ( t ) and R OPT med , 2 ( t ) by

The final step is the detailed analysis of u ( t ) .

For SGD+M, we can prove that u ( t ) ≈ C ( t )[ X 1 , X 2 , ..., X d ] T where C ( t ) ∈ R and X i , i ∈ [ d ] are i.i.d. Gaussian variables. Then we have with high probability, max i ( u ( t ) i ) 2 median ( u ( t ) i ) 2 = Ω(log d ) .

For Adam, we can try to convert it to signed descent ( e.g. in eq. (29) and Lemma 24). The main intuition is that as long as the learning rate is small enough, the movement of the gradient will be slow and Adam will be similar to signed descent. Since the update of signed descent per step is ± η , we can prove that ∀ i ∈ [ d ] : u ( t ) i ∈ {± 1 } , which gives us max i ( u ( t ) i ) 2 median ( u ( t ) i ) 2 = 1 . Substituting into eq. (4) completes the proof.

Finally, we want to emphasize that the assumption β 2 = β 2 1 is not essential. As discussed above, as long as the learning rate is sufficiently small, Adam can be approximated to signed descent regardless of the ratio between β 1 , β 2 . Hence our result can generalize to cases with other β 1 , β 2 , for example, β 2 > β 2 1 where Zhang et al. [ZCS + 22] proved the convergence of Adam to stationary points.

D Analysis of SGD+M

Let ˜ A ( t ) , ˜ Λ ( t ) xx and ˜ g ( t ) k , k = 1 , 2 be the corresponding batch versions at time t . Let E ( t ) := W ( t ) 2 W ( t ) 1 -A , and use E ( t ) i , A i and ( W ( t ) 2 W ( t ) 1 ) i to represent the i -th coordinates of E ( t ) , A and W ( t ) 2 W ( t ) 1 , respectively. By eq. (2), the update rules of W 1 and W 2 for SGD+M are given by:

where

Based on the magnitude of W 2 and W 1 , we can intuitively divide the training procedure into 2 phases.

First phase : the first several iterations when W 1 and W 2 are 'small' so that W 2 W 1 -A ≈ -A .

Second phase : later iterations when W 2 W 1 cannot be ignored.

More formally, the boundary between the first and second phase is defined below.

Definition 1 (End of the first phase) . The end of the first phase (denoted as T 1 ) is defined as .

By Assumption 2 and the assumption that ∀ j ∈ [ d ] : A j > 0 , A j = Θ(1) , at the beginning, w.h.p., ∀ j ∈ [ d ] : ( W 2 W 1 ) j -A j < 0 . During the training, each ( W 2 W 1 ) j increases and approaches A j . We hope that by choosing a small learning rate, when ( W 2 W 1 ) j overshoots for some coordinate j , i.e. ( W 2 W 1 ) j > A j , 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 glyph[epsilon1] > 0 , denote glyph[epsilon1] 0 := 1 d 1 4 α -1 + glyph[epsilon1] log √ d glyph[epsilon1] . Define

Definition 3 (Convergence time) . For glyph[epsilon1] > 0 , we define the 'convergence time'

We can first show that after the first phase, i.e. when t = T 1 , W 1 will become an approximately rank-1 matrix, as described in the following lemma.

Lemma 2. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 d α/ 2+1 . By picking η ≤ O ( 1 d α ) , we have that when t = T 1 , ¯ L ( W ( T 1 ) ) = Θ( d ) , and that

where c ( T 1 ) ∈ R , u ( T 1 ) , v ( T 1 ) , R ( T 1 ) 2 ∈ R d and R ( T 1 ) 1 ∈ R d × d . Denote the i -th coordinate of u ( T 1 ) , v ( T 1 ) , R ( T 1 ) 2 as u ( T 1 ) i , v ( T 1 ) i , R ( T 1 ) 2 i , respectively, and the ( i, j ) -th element of R ( T 1 ) 1 as R ( T 1 ) 1 [ i, j ] . Then w.h.p.,

The following lemma tells us that this approximate rank-1 structure is preserved when T 1 ≤ t ≤ min { T 2 , T 3 } .

Lemma 3. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 d α/ 2+1 . By picking η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we have that w.h.p. for T 1 ≤ t ≤ min { T 2 , T 3 } ,

where and glyph[epsilon1] 0 is defined in Definition 2. Moreover, when t = min { T 2 , T 3 } , ¯ L ( W ( t ) ) = O ( glyph[epsilon1] 0 d ) .

The following lemma gives us a more detailed description of u ( T 1 ) .

Lemma 4. The u ( T 1 ) in Lemma 2 and 3 can be written as u ( T 1 ) = X + Y where X i , i ∈ [ d ] are i.i.d Gaussian random variables and that w.h.p. ∀ i ∈ [ d ] : | Y i | | X i | ≤ ˜ O ( 1 d 1 4 α -1 2 ) .

Now we are ready to prove the SGD+M part of Theorem 1.

D.1 Proof of the SGD+M part of Theorem 1

Define T SGD , 1 = T 1 , T SGD , 2 = min { T 2 , T 3 } . By picking η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we can apply Lemma 2 and 3 to conclude that ¯ L ( W ( T SGD , 1 ) ) = Θ( d ) and ¯ L ( W ( T SGD , 2 ) ) = O ( glyph[epsilon1] 0 d ) . For any p > 0 , by picking 0 < glyph[epsilon1] < 1 d p and α ≥ 4( p +2) , we have ¯ L ( W ( T SGD , 2 ) ) = O ( glyph[epsilon1] 0 d ) ≤ ˜ O ( 1 d p ) .

Moreover, when t ∈ [ T SGD , 1 , T SGD , 2 ] , the conditions in Lemma 31 are satisfied with δ = ˜ O ( glyph[epsilon1] 0 ) . Then we can apply Lemma 31 and get that

By Lemma 4, u ( T 1 ) = X + Y where w.h.p. ∀ i ∈ [ d ] : | Y i | | X i | ≤ ˜ O ( 1 d 1 4 α -1 2 ) . This fact yields

Here X i , i ∈ [ d ] are i.i.d Gaussian random variables by Lemma 4. To prove the concentration of median X 2 i , we borrow the Proposition 12 in Chapter 2.3 of [Ler19]. By setting K = N = d in this proposition, we have

Denote σ 2 as the variance of X i , i ∈ [ d ] . Then E [ X 2 i ] = σ 2 and Var ( X 2 i ) = 2 σ 4 . Hence

That means with high probability, median X 2 i ≤ Cσ 2 for some C > 0 . By Lemma 35 in Appendix H, we know that w.h.p.

which gives us w.h.p.

Hence we have proved that R SGDM med , 1 ( t ) , R SGDM med , 2 ( t ) ≥ Ω(log d ) .

D.2 Proof of Lemma 2

In the first phase, W 2 W 1 is 'small', and we write the update equations in the following way

where

Similarly, we have

where

The following lemma gives us an explicit formula of W ( t ) 2 .

Lemma 5. Let λ 1 < λ 2 be the two roots of the quadratic equation x 2 -2 x +1 -η 2 (1 -β ) 2 ‖ A ‖ 2 2 = 0 . Pick η < 1 -β ‖ A ‖ 2 , then we have that

where C 1 = -W (1) 2 -λ 2 W (0) 2 λ 2 -λ 1 , C 2 = W (1) 2 -λ 1 W (0) 2 λ 2 -λ 1 . r ( t ) 5 will be specified in the proof.

We can prove that in the first phase, r ( t ) 5 is 'small'. More specifically, denote its i -th coordinate as r ( t ) 5 i , and the i -th coordinate of C 2 as C 2 i . Then the following lemmas tell us that ∀ i ∈ [ d ] , ∣ ∣ ∣ r ( t ) 5 i ∣ ∣ ∣ ≤ O ( 1 d p ( α ) ) , where w.h.p. O ( 1 d p ( α ) ) glyph[lessmuch] min i ∈ [ d ] | C 2 i | .

We first have the following bounds of ∣ ∣ ∣ r ( t ) 1 i ∣ ∣ ∣ , ∣ ∣ ∣ r ( t ) 2 i ∣ ∣ ∣ and ∣ ∣ ∣ r ( t ) 5 i ∣ ∣ ∣ for i ∈ [ d ] .

Lemma 6. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 d α/ 2+1 and pick η ≤ O ( 1 d α ) . We have w.h.p. for all t ≤ T 1 , ∀ i ∈ [ d ] : ∣ ∣ ∣ r ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 d 3 2 α -1 ) , ∣ ∣ ∣ r ( t ) 2 i ∣ ∣ ∣ ≤ ˜ O ( 1 d 3 2 α -2 ) .

Lemma 7. Under conditions of Lemma 6, we have that w.h.p. for all t ≤ T 1 , ∀ i ∈ [ d ] : ∣ ∣ ∣ r ( t ) 5 i ∣ ∣ ∣ ≤ ˜ O ( 1 d 3 2 α -1 ) .

Next we prove upper and lower bounds of | C 1 i | and | C 2 i | for i ∈ [ d ] .

Lemma 8. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 d α/ 2+1 . Pick η < 1 -β ‖ A ‖ 2 , we have that i) w.h.p., ∀ i ∈ [ d ] : | C 1 i | ≤ ˜ O ( 1 d α ) , | C 2 i | ≤ ˜ O ( 1 d α ) ;

ii) C 2 can be written as C 2 := 1 2 ( C 3 + C 4 ) where C 3 i , i ∈ [ d ] are i.i.d Gaussian random variables and that w.h.p. ∀ i ∈ [ d ] : | C 4 i | | C 3 i | ≤ ˜ O ( 1 d 1 4 α -1 2 ) ;

Now we are ready to prove Lemma 2. Lemma 5 tells us that

where λ 1 = 1 -η 1 -β ‖ A ‖ 2 and λ 2 = 1 + η 1 -β ‖ A ‖ 2 .

Under the conditions of Theorem 1 and pick η ≤ O ( 1 d α ) , by Lemma 7 and 8, we know that w.h.p. ∀ t ≤ T 1 , ∀ 1 ≤ i ≤ d ,

We first prove that ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ reaches 1 d α/ 2 for some coordinate i before ∣ ∣ ∣ W ( t ) 1 [ k, j ] ∣ ∣ ∣ for ∀ k, j ∈ [ d ] . To see this, first note that

where v ( t ) T = 1 -β A and

Moreover, we have that

For t ≤ T 1 , by eq. (7), we get that w.h.p.,

For t ≤ T 1 , by Lemma 6, ∀ 1 ≤ i, j ≤ d : ∣ ∣ ∣ r ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 d 3 2 α -1 ) . Then we have that w.h.p.

Here we used ∀ i ∈ [ d ] : A i = Θ(1) by Assumption 1.

Since λ 1 = 1 -η 1 -β ‖ A ‖ 2 , we have that | C 1 i λ t 1 | ≤ | C 1 i | ≤ ˜ O ( 1 d α ) and that

Using the Gaussian tail bound and union bound, we have w.h.p. ∀ 1 ≤ i, j ≤ d : ∣ ∣ ∣ W (0) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 d 2 α ) . Combining the above bounds together yields that for t ≤ T 1 and ∀ i, j ∈ [ d ] ,

where for ∀ i, j ∈ [ d ] . ∣ ∣ ∣ R ( t ) 11 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 d α + 1 2 ) , ∣ ∣ ∣ R ( t ) 21 ,i ∣ ∣ ∣ ≤ ˜ O ( 1 d α ) and ∣ ∣ ∣ e ( t ) 1 [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ e ( t ) 2 i ∣ ∣ ∣ ≤ ˜ O ( 1 d 1 4 α -1 ) .

Further we notice that for t ≤ T 1 , we have ∀ j ∈ [ d ] ,

which yields that ∣ ∣ ∣ u ( t ) i v ( t ) j ∣ ∣ ∣ ≤ O ( 1 √ d )∣ ∣ ∣ c ( t ) u ( t ) i ∣ ∣ ∣ . Together with eq. (9) gives us that ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ reaches 1 d α/ 2 for some i ∈ [ d ] before ∣ ∣ ∣ W ( t ) 1 [ k, j ] ∣ ∣ ∣ for ∀ k, j ∈ [ d ] , i.e. T 1 = inf { t ≥ 0 : ∃ i ∈ [ d ] : ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ ≥ 1 d α 2 } .

Further, we know that at time T 1 , ∣ ∣ ∣ c ( T 1 ) u ( T 1 ) i 0 ∣ ∣ ∣ = | C 2 i 0 | λ T 1 2 = Θ ( 1 d α/ 2 ) for some i 0 ∈ [ d ] , which means w.h.p.

This is the length of the first phase. As for c ( T 1 ) u ( T 1 ) i and u ( T 1 ) i v ( T 1 ) j for other coordinates, we have that w.h.p. ∀ 1 ≤ i, j ≤ d ,

Here in ( i ) we used λ 2 = 1 + η 1 -β ‖ A ‖ 2 . Then we have at time T 1 , ∀ i, j ∈ [ d ] , ∣ ∣ ∣ R ( T 1 ) 11 [ i,j ] ∣ ∣ ∣ ∣ ∣ ∣ u ( T 1 ) i v ( T 1 ) j ∣ ∣ ∣ ≤

˜ O ( 1 d 1 4 α ) and that ∣ ∣ ∣ R ( T 1 ) 21 ,i ∣ ∣ ∣ ∣ ∣ ∣ c ( T 1 ) u ( T 1 ) i ∣ ∣ ∣ ≤ ˜ O ( 1 d 1 4 α ) . Together with eq. (9), we have the following weight structure:

where w.h.p.,

Finally, we consider the loss. Since ∀ j ∈ [ d ] : ( W ( T 1 ) 2 W ( T 1 ) 1 ) j -A j = -Θ(1) , we know that ¯ L ( W ( T 1 ) ) = Θ( d ) .

D.3 Proof of Lemma 4

Eq. (8) tells us that u ( T 1 ) = ∑ T 1 -1 τ =0 λ τ 2 C T 2 . Lemma 8 tells us that C 2 can be written as C 2 := 1 2 ( C 3 + C 4 ) where C 3 i , i ∈ [ d ] are i.i.d Gaussian random variables and that w.h.p. ∀ i ∈ [ d ] : | C 4 i | | C 3 i | ≤ ˜ O ( 1 d 1 4 α -1 2 ) . Combining these two facts together finishes the proof.

D.4 Proof of Lemma 5

Replacing t by t -1 in eq. (6), we get

Eq. (6)-(11) and substituting eq. (5) yield where r 3

For the equation x 2 -2 x + 1 -η 2 (1 -β ) 2 ‖ A ‖ 2 2 = 0 , the roots are λ 1 = 1 -η 1 -β ‖ A ‖ 2 and λ 2 = 1 + η 1 -β ‖ A ‖ 2 . We have that

We further have

where r ( t ) 5 = ∑ t τ =1 λ -τ 2 r ( τ ) 4 , C 1 = -W (1) 2 -λ 2 W (0) 2 λ 2 -λ 1 and C 2 = W (1) 2 -λ 1 W (0) 2 λ 2 -λ 1 .

D.5 Proof of Lemma 6

Let's first try to bound ∣ ∣ ∣ q ( t ) 12 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ q ( t ) 22 ,i ∣ ∣ ∣ . For any τ ≤ T 1 , we have that

and thus ∀ i ∈ [ d ] : ∣ ∣ ∣ E ( τ ) i ∣ ∣ ∣ = O (1) . Then we have for all i, j ∈ [ d ] ,

That gives us ∀ i, j ∈ [ d ] ,

Then we bound ∣ ∣ ∣ q ( t ) 13 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ q ( t ) 23 ,i ∣ ∣ ∣ . We have for ∀ i, j ∈ [ d ] ,

Finally we use Lemma 32 to bound ∣ ∣ ∣ q ( t ) 14 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ q ( t ) 24 ,i ∣ ∣ ∣ . For t ≤ T 1 , the M ( t ) 1 , M ( t ) 2 in Lemma 32 are upper bounded by 1 d α 2 . In the theorem we consider the training period before T SGD , 2 so the time T in Lemma 32 is set as T SGD , 2 . In the following sections, we will prove that T SGD , 2 ≤ O ( d α log( √ d/glyph[epsilon1] ) η ) . Then by Lemma 32, we have with probability at least 1 -1 d , for ∀ t ≤ T 1 and

∀ i, j ∈ [ d ] ,

By picking σ ≤ η 3 / 2 d α/ 2+1 , we have w.h.p. for ∀ t ≤ T 1 and ∀ i, j ∈ [ d ] , ∣ ∣ ∣ Dg ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ η ˜ O ( 1 d α 2 ) and ∣ ∣ ∣ Dg ( t ) 2 i ∣ ∣ ∣ ≤ η ˜ O ( 1 d α 2 -1 ) , which yields

Combining all the above bounds and substituting η ≤ O ( 1 d α ) gives us for ∀ t ≤ T 1 and ∀ i, j ∈ [ d ] ,

For t ≤ T 1 , we have ∀ i, j ∈ [ d ] , ∣ ∣ ∣ w ( t ) 2 i A j ∣ ∣ ∣ ≤ O ( 1 d α/ 2 ) and ∣ ∣ ∣ ∑ d j =1 A j W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ O ( 1 d α/ 2 -1 ) , which gives us ∣ ∣ ∣ r ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ O ( 1 d α/ 2 ) and ∣ ∣ ∣ r ( t ) 2 i ∣ ∣ ∣ ≤ O ( 1 d α/ 2 -1 ) . Substituting into eq. (5) and eq. (6) yields that for t ≤ T 1 and ∀ i, j ∈ [ d ] ,

Hence for t ≤ min { α log d log(1 /β ) , T 1 } , we have ∀ i, j ∈ [ d ] ,

Then we know that T 1 > α log d log(1 /β ) and also get tighter bounds of ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ for t ≤ α log d log(1 /β ) . Now we use these new bounds to analyze ∣ ∣ ∣ r ( t ) 1 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ r ( t ) 2 i ∣ ∣ ∣ again.

When t ≤ α log d log(1 /β ) , we have for all i, j ∈ [ d ] , β t +1 ∣ ∣ ∣ w ( t ) 2 i A j ∣ ∣ ∣ ≤ ∣ ∣ ∣ w ( t ) 2 i A j ∣ ∣ ∣ ≤ ˜ O ( 1 d 3 α 2 -1 ) and β t +1 ∣ ∣ ∣ ∑ d j =1 A j W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ∣ ∣ ∣ ∑ d j =1 A j W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 d 3 α 2 -1 ) . When α log d log(1 /β ) < t ≤ T 1 , we have β t +1 ≤ 1 d α , suggesting that ∀ i, j ∈ [ d ] , β t +1 ∣ ∣ ∣ w ( t ) 2 i A j ∣ ∣ ∣ ≤ 1 d α ˜ O ( 1 d α 2 ) ≤ ˜ O ( 1 d 3 α 2 ) and β t +1 ∣ ∣ ∣ ∑ d j =1 A j W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ 1 d α ˜ O ( 1 d α 2 -1 ) ≤ ˜ O ( 1 d 3 α 2 -1 ) . Substituting into (12) completes the proof.

D.6 Proof of Lemma 7

Based on the bound in Lemma 6, we have

Since λ 1 = 1 -η 1 -β ‖ A ‖ 2 , λ 2 = 1 + η 1 -β ‖ A ‖ 2 , and note that ‖ A ‖ 2 = Θ ( √ d ) , we have that

D.7 Proof of Lemma 8

For the equation x 2 -2 x + 1 -η 2 (1 -β ) 2 ‖ A ‖ 2 2 = 0 , the roots are λ 1 = 1 -η 1 -β ‖ A ‖ 2 and λ 2 = 1 + η 1 -β ‖ A ‖ 2 , which gives us

where ˜ r (0) 2 = -W (0) 2 W (0) 1 W (0) T 1 -Dg (0) 2 . Note that this is slightly different from the definition of r (0) 2 in eq. (6). Now let's bound the i -th coordinate of ˜ r (0) 2 .

In Section D.5 we have shown that w.h.p. for ∀ t ≤ T 1 and ∀ i, j ∈ [ d ] , ∣ ∣ ∣ Dg ( t ) 2 i ∣ ∣ ∣ ≤ η ˜ O ( 1 d α 2 -1 ) = ˜ O ( 1 d 3 α 2 -1 ) , which also applies to t = 0 . Using the Gaussian tail bound and union bound, w.p. at least 1 -δ , for ever 1 ≤ i, j ≤ d , we have that

Then we have that w.p. at least 1 -δ , ∀ 1 ≤ i, j ≤ d : ,

Next, we bound the i -th coordinate of W (0) 2 + 1 -β ‖ A ‖ 2 AW (0) T 1 , i.e. w (0) 2 i + 1 -β ‖ A ‖ 2 A ( W (0) 1 [ i, :] ) T . By independence under Assumption 2, we have that

Using the Gaussian tail bound and union bound, w.p. at least 1 -δ , for ever 1 ≤ i ≤ d , we have that

Since for X ∼ N (0 , σ 2 ) , we have that P ( | X | ≤ t ) ≤ 2 t √ 2 πσ , then for a fixed i ,

Then by union bound, we have that w.p. at least 1 -1 d α 4 -1 , for every 1 ≤ i ≤ d ,

Now define C 3 := W (0) 2 + 1 -β ‖ A ‖ 2 AW (0) T 1 and C 4 := 1 -β 2 ‖ A ‖ 2 ˜ r (0) 2 i . We get that C 3 i , i ∈ [ d ] are i.i.d Gaussian random variables and that C 2 = 1 2 ( C 3 + C 4 ) , where w.h.p. for all i ∈ [ d ] ,

where ( i ) follows from eq. (14) and the fact that ‖ A ‖ 2 = d . Then we get that w.h.p.

Substituting eq. (15) into eq. (13), we get that w.h.p.,

Similarly, note that

we can use the same techniques to get that i) w.p. at least 1 -δ , ∀ i ∈ [ d ] : | C 1 i | ≤ ˜ O ( 1 d α ) , ii) w.p. at least 1 -δ -1 d α 4 -1 , ∀ i ∈ [ d ] , | C 1 i | ≥ ˜ Ω ( 1 d 5 4 α ) .

D.8 Proof of Lemma 3

The proof in Section D.2 tells us that at the end of the first phase (when t = T 1 ),

Denote the i -th coordinate of u ( t ) , v ( t ) , R ( t ) 2 as u ( t ) i , v ( t ) i , R ( t ) 2 i , respectively. Denote the ( i, j ) -th element of R ( t ) 1 as R ( t ) 1 [ i, j ] . For t ≥ T 1 , we prove by induction that,

with r ( t ) 1 := η ∑ t τ =0 β t -τ ( W ( t ) T 2 E ( t ) -W ( τ ) T 2 E ( τ ) ) -η ∑ t τ =0 β t -τ Dg ( τ ) 1 , E ( t ) := W ( t ) 2 W ( t ) 1 -A , η t = η ∑ t τ =0 β t -τ and r ( t ) 2 = η ∑ t τ =0 β t -τ ( E ( t ) W ( t ) T 1 -E ( τ ) W ( τ ) T 1 ) -η ∑ t τ =0 β t -τ Dg ( τ ) 2 .

Note that the r ( t ) 1 and r ( t ) 2 here are different from those defined in Section D.2, but we abuse the notation and still use r ( t ) 1 and r ( t ) 2 to 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. (5) and eq. (6), we have that

Plugging in the inductive hypothesis yields

It implies that our lemma holds for t +1 , which completes the proof.

where

Now we analyze the error terms ∣ ∣ ∣ R ( t ) 1 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ R ( t ) 2 i ∣ ∣ ∣ . Eq. (16) tells us that c ( T 1 ) and ∀ i ∈ [ d ] , v ( T 1 ) i are all positive. We first prove by induction that for all T 1 ≤ t ≤ T 2 , c ( t ) > 0 , ∀ i ∈ [ d ] , v ( t ) i > 0 . The above discussion already proves the base case. Suppose at time t , we have c ( t ) > 0 , ∀ i ∈ [ d ] , v ( t ) i > 0 . Note that when T 1 ≤ t < T 2 , ∀ i ∈ [ d ] : E ( t ) i ≤ 0 , then for t +1 ,

Therefore by induction, we have proved that for all T 1 ≤ t ≤ T 2 , c ( t ) > 0 , ∀ i ∈ [ d ] , v ( t ) i > 0 . Now we prove that for all T 1 ≤ t ≤ T 2 ,

where

The left hand sides of the inequalities are trivial since we have proved that c ( t ) > 0 , ∀ i ∈ [ d ] , v ( t ) i > 0 for all T 1 ≤ t ≤ T 2 . Now we prove the right hand sides by induction.

The base case is already verified by the definition of δ i . Suppose eq.(18) holds for T 1 ≤ t < T 2 . Then for t +1 , using ∀ i ∈ [ d ] : E ( t ) i ≤ 0 and v ( t +1) ≥ v ( t ) , c ( t +1) ≥ c ( t ) , we can get that ∀ 1 ≤ i, j ≤ d

Similarly, we have that ∀ 1 ≤ i ≤ d

Therefore by induction, eq. (18) holds for all t in the second phase.

So far we have proved the rank 1 structure stated in Lemma 3. The remaining part of the proof is given by the following lemma, whose proof is deferred to Section D.9.

Lemma 9. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 d α/ 2+1 . By picking η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we have that w.h.p. for T 1 ≤ t ≤ min { T 2 , T 3 } ,

and that when t = min { T 2 , T 3 } , we have ∥ ∥ E ( t ) ∥ ∥ 2 2 = O ( glyph[epsilon1] 0 d ) .

D.9 Proof of Lemma 9

We first have the following lemma which describes the structure of v ( t ) for t ≥ T 1 .

Lemma 10. Under Assumption 1, 2 and 3, for t ≥ T 1 , we can write v ( t ) T as v ( t ) T = a ( t ) A + R ( t ) T v , with a ( T 1 ) = η 1 -β , R ( T 1 ) T v = [0 , 0 , ..., 0] , and

Moreover, we have that

We prove Lemma 9 by induction. Denote the i -th coordinate of R ( t ) 3 and R ( t ) v as R ( t ) 3 i and R ( t ) vi , respectively. The following lemmas constitute the inductive part.

Lemma 11. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 d α/ 2+1 and pick η ≤ O ( 1 d α ) . Consider any t such that T 1 ≤ t < min { T 2 , T 3 } . Suppose for all T 1 ≤ τ ≤ t , we have ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ ≤ O ( d 1 / 4 ) , ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ ≤ O ( 1 d 1 / 4 ) , then we have that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ r ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( η 2 d 11 / 4 ) , ∣ ∣ ∣ r ( t ) 2 i ∣ ∣ ∣ = ˜ O ( η 2 d 13 / 4 ) . Moreover, we can get that ∀ i ∈ [ d ] : glyph[epsilon1] ( t ) i = ˜ O ( η 2 d 3 4 α + 13 4 ) , where glyph[epsilon1] ( t ) i is defined in eq. (18) .

Lemma 12. Under the conditions of Lemma 11 and pick η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we have that at time t +1 ,

where glyph[epsilon1] 0 is defined in Definition 2.

Lemma 13. Under the conditions of Lemma 11 and pick η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we have that at time t +1 ,

Moreover,

Lemma 14. Under the conditions of Lemma 11 and pick η ≤ O ( glyph[epsilon1] 7 α 4 +4 ) , if we further suppose that

d ,

∥ ∥

By combining Lemma 11, 12 and 14, we can prove by induction that for all T 1 ≤ t ≤ min { T 2 , T 3 } , eq. (19) holds (which follows from Lemma 12), and

which follows from the part (A) of Lemma 14. Now the only thing to verify is the base case, i.e. when t = T 1 . More specifically, we want to prove that 1) ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( T 1 ) 2 i ∣ ∣ ∣ ≤ O ( d 1 / 4 ) , ∣ ∣ ∣ W ( T 1 ) 1 [ i, j ] ∣ ∣ ∣ ≤ O ( 1 d 1 / 4 ) and that 2) ∀ j ∈ [ d ] : v ( T 1 ) j c ( T 1 ) = Θ ( 1 √ d ) , and that 3) ∣ ∣ ∣ R ( T 1 ) 3 j ∣ ∣ ∣ A j and ∣ ∣ ∣ R ( T 1 ) vj ∣ ∣ ∣ a ( T 1 ) A j are of order ˜ O ( glyph[epsilon1] 0 ) All of them can be verified by the proof in Section D.2 and the definition of R ( t ) v , R ( t ) 3 .

So far we have proved eq. (19) in Lemma 9. Now let's prove when t = min { T 2 , T 3 } , we have that ∥ ∥ E ( t ) ∥ ∥ 2 2 = O ( glyph[epsilon1] 0 d ) .

If min { T 2 , T 3 } = T 3 , by Definition 3, we have ∥ ∥ E ( t ) ∥ ∥ 2 2 ≤ glyph[epsilon1] . If min { T 2 , T 3 } = T 2 , by Definition 2, there exists j ∈ [ d ] such that E ( t ) j = -Θ ( √ glyph[epsilon1] 0 ) . Combining with eq. (22) gives us ∀ i ∈ [ d ] : E ( t ) i = -Θ ( √ glyph[epsilon1] 0 ) . Combining these two cases, we get that when t = min { T 2 , T 3 } , ∥ ∥ E ( t ) ∥ ∥ 2 2 ≤ max { glyph[epsilon1], Θ( glyph[epsilon1] 0 d ) } = O ( glyph[epsilon1] 0 d ) .

D.10 Proof of Lemma 10

We prove this lemma by induction. The base case ( t = T 1 ) of v Suppose at time t , v ( t ) T = a ( t ) A + R ( t ) T v , then by eq. 17, we have that

( t ) is verified by eq. (16).

where d ( t ) := c ( t ) ∥ ∥ u ( T 1 ) ∥ ∥ 2 + R ( t ) T 2 u ( T 1 ) , R ( t ) T 3 := c ( t ) u ( T 1 ) T R ( t ) 1 + R ( t ) T 2 R ( t ) 1 . That gives us

Therefore we have proved by induction that for t in the second phase, v ( t ) = a ( t ) A + R ( t ) T v . The above steps also proved eq. (20).

D.11 Proof of Lemma 11

Write r ( t ) 1 = q ( t ) 11 + q ( t ) 12 where we have q ( t ) 11 = η ∑ t τ =0 β t -τ ( W ( t ) T 2 E ( t ) -W ( τ ) T 2 E ( τ ) ) , q ( t ) 12 = -η ∑ t τ =0 β t -τ Dg ( τ ) 1 . Write r ( t ) 2 = q ( t ) 21 + q ( t ) 22 where q ( t ) 22 = -η ∑ t τ =0 β t -τ Dg ( τ ) 2 , q ( t ) 21 = η ∑ t τ =0 β t -τ ( E ( t ) W ( t ) T 1 -E ( τ ) W ( τ ) T 1 ) .

Let's first bound ∣ ∣ ∣ q ( t ) 11 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ q ( t ) 21 ,i ∣ ∣ ∣ . By definition of T 2 , we know that for T 1 ≤ τ ≤ t , ∀ i ∈ [ d ] : ∣ ∣ ∣ E ( τ ) i ∣ ∣ ∣ = O (1) . Then we have for all i, j ∈ [ d ] ,

Note that

We can further get that for ∀ j ∈ [ d ] ,

Combining the above inequalities gives us ∀ i, j ∈ [ d ] ,

Next let's bound ∣ ∣ ∣ q ( t ) 12 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ q ( t ) 22 ,i ∣ ∣ ∣ . By the assumption of this lemma and the analysis before T 1 , we know that for all τ ≤ t , the M ( τ ) 1 , M ( τ ) 2 in Lemma 32 are upper bounded by O ( 1 d 1 / 4 ) and O ( d 1 / 4 ) , respectively. In the theorem we consider the training period before T SGD , 2 so the time T in

Lemma 32 is set as T SGD , 2 . In the following sections, we will prove that T SGD , 2 ≤ O ( d α log( √ d/glyph[epsilon1] ) η ) . Then by Lemma 32, we have with probability at least 1 -1 d , for ∀ τ ≤ t and ∀ i, j ∈ [ d ] ,

By picking σ ≤ η 3 / 2 d α/ 2+1 , we have ∣ ∣ ∣ Dg ( τ ) 1 [ i, j ] ∣ ∣ ∣ ≤ η ˜ O ( d 11 4 ) and ∣ ∣ ∣ Dg ( τ ) 2 i ∣ ∣ ∣ ≤ η ˜ O ( d 13 4 ) , which yields

Combining the above bounds, we get that ∀ i, j ∈ [ d ] ,

By the analysis in Section D.2, we know that at time T 1 , for some i 0 ∈ [ d ] , c ( T 1 ) ∣ ∣ ∣ u ( T 1 ) i 0 ∣ ∣ ∣ = Θ ( 1 d α 2 ) , and for ∀ i, j ∈ [ d ] , we have c ( T 1 ) ∣ ∣ ∣ u ( T 1 ) i ∣ ∣ ∣ = ˜ Ω ( 1 d 3 α 4 ) and ∣ ∣ ∣ u ( T 1 ) i ∣ ∣ ∣ v ( T 1 ) j = ˜ Ω ( 1 d 3 α 4 + 1 2 ) , which gives us ∀ i, j ∈ [ d ] ,

Hence we get the bound ∀ i ∈ [ d ] : glyph[epsilon1] ( t ) i ≤ ˜ O ( η 2 d 3 4 α + 13 4 ) .

D.12 Proof of Lemma 12

Let's first try to bound the length of min { T 2 , T 3 } . More formally, we prove that under the conditions of Lemma 11 and pick η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we have that min { T 2 , T 3 } ≤ O ( d α log( √ d/glyph[epsilon1] ) η ) . Under the conditions of Lemma 11, we know that

Combining with eq. (23), we get

Then we have

When T 1 ≤ t < T 2 , we have proved that c ( t ) is increasing over time in Section D.8, which implies that ∥ ∥ ∥ W ( t ) 2 ∥ ∥ ∥ 2 2 ≥ C ∥ ∥ ∥ W ( T 1 ) 2 ∥ ∥ ∥ 2 2 since c ( t ) u ( T 1 ) T is the leading term of W ( t ) 2 . Combining with η t ≥ η gives us

where ( i ) uses ∥ ∥ E ( T 1 ) ∥ ∥ 2 = O ( √ d ) . By picking η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) and noticing that ∥ ∥ ∥ W ( T 1 ) 2 ∥ ∥ ∥ 2 2 ≥ Ω ( 1 d α ) , we have ηd 4 ∥ ∥ ∥ W ( T 1 ) 2 ∥ ∥ ∥ 2 2 < √ glyph[epsilon1] 2 . Hence when t -T 1 ≥ Θ ( log ( √ d/glyph[epsilon1] ) η ∥ ∥ ∥ W ( T 1 ) 2 ∥ ∥ ∥ 2 2 ) , we have that ∥ ∥ E ( t ) ∥ ∥ 2 ≤ √ glyph[epsilon1] , i.e. ∥ ∥ E ( t ) ∥ ∥ 2 2 ≤ glyph[epsilon1] .

That means after at most O ( log ( √ d/glyph[epsilon1] ) η ∥ ∥ ∥ W ( T 1 ) 2 ∥ ∥ ∥ 2 2 ) steps from T 1 , either t ≥ T 2 , or we have ∥ ∥ E ( t ) ∥ ∥ 2 2 ≤ glyph[epsilon1] . In other words, min { T 2 , T 3 } ≤ T 1 + O ( log ( √ d/glyph[epsilon1] ) η ∥ ∥ ∥ W ( T 1 ) 2 ∥ ∥ ∥ 2 2 ) ≤ O ( d α log ( √ d/glyph[epsilon1] ) η ) .

Now we are ready to bound eq. 18.

Combining min { T 2 , T 3 } ≤ O ( d α log( √ d/glyph[epsilon1] ) η ) and Lemma 11 yields that for t +1 ≤ min { T 2 , T 3 } , ∀ i ∈ [ d ] ,

Lemma 2 tells us that δ i = ˜ O ( 1 d 1 4 α -1 ) . Substituting these bounds into eq. (18) completes the proof.

D.13 Proof of Lemma 13

The proof in Section D.8 tells us that for T 1 ≤ τ ≤ T 2 , c ( τ ) > 0 , ∀ j ∈ [ d ] : v ( τ ) j > 0 , which gives us 0 ≤ ∣ ∣ ∣ R ( t +1) T 2 u ( T 1 ) ∣ ∣ ∣ c ( t +1) ‖ u ( T 1 ) ‖ 2 and 0 ≤ ∣ ∣ ∣ R ( t +1) 3 j ∣ ∣ ∣ c ( t +1) ‖ u ( T 1 ) ‖ 2 v ( t +1) j . By Lemma 12, we have that

which gives us

Lemma 10 tells us that

And we have that

Therefore

Then we have that ∀ j ∈ [ d ] ,

Since t < T 2 , we have ∀ j ∈ [ d ] : ( W ( t +1) 2 W ( t +1) 1 ) j A j = O (1) , which yields

which proves eq. (21), since ∀ j ∈ [ d ] : 0 ≤ ∣ ∣ ∣ R ( t +1) 3 j ∣ ∣ ∣ c ( t +1) ‖ u ( T 1 ) ‖ 2 v ( t +1) j ≤ ˜ O ( glyph[epsilon1] 0 ) .

D.14 Proof of Lemma 14

(A) Under the conditions of Lemma 11 and pick η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we can apply the technique when proving eq. (24) to show that eq. (24) also holds at time t . Since ∣ ∣ ∣ R ( t ) vj ∣ ∣ ∣ a ( t ) A j ≤ ˜ O ( glyph[epsilon1] 0 ) , we get that

Substituting into the time t version of eq.(24) yields

That gives us

Since t < T 2 , we have E ( t ) j < - √ glyph[epsilon1] 0 . Combining with A j = Θ(1) , gives us a ( t ) c ( t ) ∥ ∥ u ( T 1 ) ∥ ∥ 2 -1 = -Ω ( √ glyph[epsilon1] 0 ) . Then we can rewrite E ( t ) j as ∀ j ∈ [ d ] ,

where ∣ ∣ ∣ e ( t ) Ej ∣ ∣ ∣ = ˜ O ( √ glyph[epsilon1] 0 ) . Hence ∀ i, j ∈ [ d ] : E ( t ) j E ( t ) i = Θ(1) .

(B) Note that we assume ∀ j ∈ [ d ] : v ( t ) j c ( t ) = Θ ( 1 √ d ) , then we have for j ∈ [ d ] ,

Then for t +1 , we have that for j ∈ [ d ] ,

(C) Combining eq. (25) and ∀ j ∈ [ d ] : A j = Θ(1) , we know that

which yields ∀ j ∈ [ d ] ,

Hence ∀ i, j ∈ [ d ] ,

where ( i ) and ( ii ) use Lemma 12.

(D) The fact that ∀ j ∈ [ d ] : ∣ ∣ ∣ R ( t +1) 3 j ∣ ∣ ∣ A j ≤ ˜ O ( glyph[epsilon1] 0 ) was already proved in Lemma 13 in eq.(21). To analyze ∣ ∣ ∣ R ( t +1) vi ∣ ∣ ∣ a ( t +1) A i , we first prove that 1 -η t c ( t ) d ( t ) > 0 .

It is not hard to prove that eq.(26) also holds for time t . Recall that d ( t ) = c ( t ) ∥ ∥ u ( T 1 ) ∥ ∥ 2 + R ( t ) T 2 u ( T 1 ) and Lemma 13 tells us that 0 ≤ ∣ ∣ ∣ R ( t ) T 2 u ( T 1 ) ∣ ∣ ∣ c ( t ) ‖ u ( T 1 ) ‖ 2 ≤ ˜ O ( glyph[epsilon1] 0 ) , then we have

Under the conditions of Lemma 11, and pick η ≤ O ( glyph[epsilon1] d 7 α 4 +4 ) , we have that 1 -η t c ( t ) d ( t ) ≥ 1 -ηc ( t ) d ( t ) > 0 .

The assumption ∀ j ∈ [ d ] : ∣ ∣ ∣ R ( t ) 3 j ∣ ∣ ∣ A j ≤ ˜ O ( glyph[epsilon1] 0 ) together with c ( t ) > 0 gives us

Combining with the assumption ∣ ∣ ∣ R ( t ) vi ∣ ∣ ∣ a ( t ) A i ≤ ˜ O ( glyph[epsilon1] 0 ) yields

E Analysis of Adam

Note that A = 1 m Y X T , Λ xx := 1 m XX T . Denote g ( t ) k := ∇ W k L ( W ( t ) ) , k = 1 , 2 . We have that g ( t ) 1 = W ( t ) T 2 ( W ( t ) 2 W ( t ) 1 -A ) , g ( t ) 2 = ( W ( t ) 2 W ( t ) 1 -A ) W ( t ) T 1 .

Let ˜ A ( t ) , ˜ Λ ( t ) xx and ˜ g ( t ) k , k = 1 , 2 be the corresponding batch versions at time t . Let E ( t ) := W ( t ) 2 W ( t ) 1 -A , and denote E ( t ) j as the j -th component of E ( t ) . We also denote ∆ w ( t ) 2 i := w ( t +1) 2 i -w ( t ) 2 i , ∆ W ( t ) 1 [ i, j ] := W ( t +1) 1 [ i, j ] -W ( t ) 1 [ i, j ] . By eq. (2), the update equations of Adam are given by

Denote the i -th coordinate of W 2 W 1 and A as ( W 2 W 1 ) i and A i , respectively. By Assumption 2 and the assumption that ∀ i ∈ [ d ] : A i > 0 , A i = Ω(1) , at the beginning, w.h.p., ∀ i ∈ [ d ] : ( W 2 W 1 ) i -A i < 0 . Based on this, we divide the training procedure into two phases (note that these two phases are different from those of SGD+M).

First phase : when the error ( W 2 W 1 ) i -A i is negative and its absolute value is big for all i ∈ [ d ] .

Second phase : when ( W 2 W 1 ) i -A i is close to zero for some coordinate i ∈ [ d ] .

More formally, we define the boundary between the two phases below.

Definition 4 (End of the first phase) . The end of the first phase (denoted as T 1 ) is defined as T 1 = inf { t > 0 : ∃ i ∈ [ d ] : E ( t ) i ≥ - √ ηd } .

In the second phase, we define some time points.

For t < T 1 , we have ∀ i ∈ [ d ] : E ( t ) i < 0 by Definition 4. For t > T 1 , some E ( t ) i may flip the sign and become positive. For certain coordinate i , we define the following 'flip time'.

Definition 6. Define T f,i := inf { t > T 1 : E ( t ) i ≥ - √ ηd } . Define T f := max i T f,i as the largest 'flip time' over all i ∈ [ d ] , i.e. the 'flip time' of the last E i which flips the sign. Moreover, denote ˜ T := min { T g , T f } .

We can first show that after a few steps in the first phase, W 1 will become an approximately rank-1 matrix, as described in the following lemma.

Lemma 15. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . By picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 , there exists t inc > 0 such that w.h.p. for t inc ≤ t < T 1 ,

Specially, when t = T 1 , we have that

where η ( T 1 -t inc ) = Θ ( 1 √ d ) and

The following lemma tells us that this approximate rank-1 structure is preserved when T 1 ≤ t ≤ ˜ T . Lemma 16. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . By picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 , we have w.h.p. for T 1 ≤ t < ˜ T ,

and that L ( W ( ˜ T ) ) ≤ ˜ O ( ηd 4 ) .

Now we are ready to prove the Adam part of Theorem 1.

E.1 Proof of the Adam part of Theorem 1

Define T Adam , 1 = t inc + 1 ηd α 2 . Note that this choice of T Adam , 1 gives η ( T Adam , 1 -t inc ) = 1 d α 2 . By picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 and β 2 = β 2 1 , we can apply Lemma 15 to get that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( T Adam , 1 ) 2 i ∣ ∣ ∣ = Θ ( 1 d α 2 ) , ∣ ∣ ∣ W ( T Adam , 1 ) 1 [ i, j ] ∣ ∣ ∣ = Θ ( 1 d α 2 ) , and therefore ∀ i ∈ [ d ] : E ( T Adam , 1 ) i = -Θ(1) and L ( W ( T Adam , 1 ) ) = Θ( d ) . Define T Adam , 2 = ˜ T . By Lemma 16, we have L ( W ( T Adam , 2 ) ) = ˜ O ( ηd 4 ) . For any p > 0 , by picking α ≥ p +4 3 , we have L ( W ( T Adam , 2 ) ) = ˜ O ( ηd 4 ) ≤ ˜ O ( 1 d p ) .

Moreover, combining Lemma 15 and 16, we get that when t ∈ [ T Adam , 1 , T Adam , 2 ] , the conditions in Lemma 31 are satisfied with δ = ˜ O ( η 1 4 + 1 d α 2 -1 4 ) . The i -th component of the u vector (denoted as u i ) is sign ( w (0) 2 i ) . That means ∀ i ∈ [ d ] : u 2 i = 1 and max i ( u i ) 2 median ( u i ) 2 = 1 . Then we can apply Lemma 31 and get that

Note that by definition, R Adam med , 1 ( t ) and R Adam med , 2 ( t ) are always larger than or equal to 1, then we have

E.2 Proof of Lemma 15

For some time t , we introduce two conditions.

Condition 1.

Condition 2.

Next prove that, under Assumption 1 and 2, by picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 , there exists t inc > 0 such that for t inc ≤ t < T 1 , the update of Adam can be approximated as that of signed descent.

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 a finite-step average, and trying to bound the stochastic error terms in eq. (28).

Lemma 17. Under Assumption 1, 2 and 3 and pick β 2 = β 2 1 . Let M ( t ) 1 := max i,j ∈ [ d ] ,τ ≤ t ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ , M ( t ) 2 := max i,j ∈ [ d ] ,τ ≤ t ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ , G ( t ) 1 := max i,j ∈ [ d ] ,τ ≤ t ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ and G ( t ) 2 := max i,j ∈ [ d ] ,τ ≤ t ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ . We have that w.h.p., for all t ≤ ˜ O ( 1 √ dη ) and ∀ i, j ∈ [ d ] , ∆ W ( t ) 1 [ i, j ] = -η t (1 -β 1 ) ∑ H τ =0 β τ 1 g ( t -τ ) 1 [ i, j ] + glyph[epsilon1] ( t ) 1 n [ i, j ] √ (1 -β 2 ) ∑ H τ =0 β τ 2 ( g ( t -τ ) 1 [ i, j ] ) 2 + glyph[epsilon1] ( t ) 1 d [ i, j ] + ξ , ∆ w ( t ) 2 i = -η t (1 -β 1 ) ∑ H τ =0 β τ 1 g ( t -τ ) 2 i + glyph[epsilon1] ( t ) 2 n,i √ (1 -β 2 ) ∑ H τ =0 β τ 2 ( g ( t -τ ) 2 i ) 2 + glyph[epsilon1] ( t ) 2 d,i + ξ , where H ≥ 1 1 -β 1 log max { G ( t ) 1 ,G ( t ) 2 , ( G ( t ) 1 ) 2 , ( G ( t ) 2 ) 2 } ηξ 2 and

Corollary 2. Under the conditions of Lemma 17 and suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Consider any t ≤ ˜ O ( 1 √ dη ) . If M ( t ) 1 , M ( t ) 2 ≤ ˜ O ( 1 √ d ) , G ( t ) 1 ≤ ˜ O ( 1 √ d ) , G ( t ) 2 ≤ ˜ O ( √ d ) , then H in Lemma 17 can be picked as 1 1 -β 1 log d ηξ 2 and we can get that ∀ i, j ∈ [ d ] , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 n,i ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 d,i ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) .

The following lemma analyzes the magnitude of weights during a short period at the beginning.

Lemma 18. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Pick ξ ≤ 1 d 3 2 α , then there exists some time point t inc ∈ ( H,T 1 ) , such that w.h.p., for t ≤ t inc, for every i, j ∈ [ d ] ,

Specifically, when t = t inc , we have sign ( w ( t inc ) 2 i ) = sign ( W ( t inc ) 1 [ i, j ] ) = sign ( w (0) 2 i ) , ∣ ∣ ∣ W ( t inc ) 1 [ i, j ] ∣ ∣ ∣ = Θ ( 1 d 3 2 α +1 ) and ∣ ∣ ∣ g ( t inc ) 1 [ i, j ] ∣ ∣ ∣ ≥ Ω( ξ ) , ∣ ∣ ∣ g ( t inc ) 2 i ∣ ∣ ∣ ≥ Ω( ξ ) . Moreover, Condition 1 and 2 are satisfied for t = t inc. The s ( t ) 1 [ i, j ] and s ( t ) 2 i in the conditions are both -sign ( w (0) 2 i ) .

The following lemma gives us lower bounds of ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ g ( t ) 2 i ∣ ∣ ∣ .

Lemma 19. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Pick ξ ≤ √ η d 3 α -1 , η ≤ O ( 1 d 3 α ) . Consider t inc in Lemma 18. We have w.h.p. for any t ∈ [ t inc , T 1 ) , and for ∀ i, j ∈ [ d ] , sign ( ∆ W ( t ) 1 [ i, j ] ) = sign ( ∆ w ( t ) 2 i ) = sign ( w (0) 2 i ) and that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ ≥ ˜ Ω ( √ η ) , ∣ ∣ ∣ g ( t ) 2 i ∣ ∣ ∣ ≥ ˜ Ω ( √ ηd ) . Moreover, we have ∀ τ ≤ t , ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) , ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) and ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) , ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ ≤ ˜ O ( √ d ) .

The following lemma shows that when t inc ≤ t < T 1 , we have ∀ i, j ∈ [ d ] : ∣ ∣ ∣ g ( t ) 2 i ∣ ∣ ∣ glyph[greatermuch] ∣ ∣ ∣ g ( t ) 2 i -g ( t -1) 2 i ∣ ∣ ∣ and that ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ glyph[greatermuch] ∣ ∣ ∣ g ( t ) 1 [ i, j ] -g ( t -1) 1 [ i, j ] ∣ ∣ ∣ .

Lemma 20. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Pick ξ ≤ √ η d 3 α -1 , η ≤ O ( 1 d 3 α ) . For t inc in Lemma 18, we have that w.h.p. for t inc ≤ t < T 1 and τ ≤ t , ∀ i, j ∈ [ d ] ,

Equipped with these lemmas, now let's prove eq. (29).

For any t ∈ [ t inc , T 1 ) , by Lemma 19, we know that M ( t ) 1 , M ( t ) 2 ≤ ˜ O ( 1 √ d ) , and that G ( t ) 1 ≤ ˜ O ( 1 √ d ) , G ( t ) 2 ≤ ˜ O ( √ d ) . At the end of the proof for this lemma, we will show that T 1 = Θ ( 1 √ dη ) . Then we can pick H := 1 1 -β 1 log d ηξ 2 and apply Lemma 17 and Corollary 2 to get that, w.h.p., for all t ∈ [ t inc , T 1 ) and ∀ i, j ∈ [ d ] , eq. (27) can be written as

where ∀ i, j ∈ [ d ] , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 n,i ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 d,i ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) .

Let's first look at the update of W ( t ) 1 [ i, j ] . For t in the first phase, we write the RHS of eq. (32) as

where

We have already shown that ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) . By Lemma 20, we have that ∀ i, j ∈ [ d ] ,

Similarly, we have ∀ i, j ∈ [ d ] ,

By Lemma 19, we know that ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ = Ω ( √ η ) . Then we have that

Therefore by Lemma 34 in Appendix H, we have

where ∣ ∣ ∣ ˜ e ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( √ η ) .

Since β ∈ (0 , 1) , we know that log β ≤ β -1 < 0 . Then our choice of H gives us H = 1 1 -β 1 log d ηξ 2 ≥ log ηξ 2 d log β 1 and H > 1 1 -β 2 log d ηξ 2 ≥ log ηξ 2 d log β 2 , which implies that β H 1 , β H 2 ≤ ηξ 2 /d . Hence for t ≥ t inc > H , η t 1 -β H +1 1 √ 1 -β H +1 2 = η √ 1 -β t +1 2 √ 1 -β H +1 2 1 -β H +1 1 1 -β t +1 1 = η (1 ±O ( η )) .

Combining all of the above yields that

where ∣ ∣ ∣ e ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( √ η ) . The proof for w ( t ) 2 i is similar.

So far we have successfully proved eq. (29). By sign ( ∆ W ( t ) 1 [ i, j ] ) = sign ( ∆ w ( t ) 2 i ) = sign ( w (0) 2 i ) in Lemma 19, we know that sign ( -g ( t ) 1 [ i, j ] ) = sign ( -g ( t ) 2 i ) = sign ( w (0) 2 i ) , which gives us

where ∣ ∣ ∣ R ( t ) 1 [ i,j ] ∣ ∣ ∣ η ( t -t inc ) = ˜ O ( √ η + ∣ ∣ ∣ W ( t inc ) 1 [ i,j ] ∣ ∣ ∣ η ( t -t inc ) ) and ∣ ∣ ∣ R ( t ) 2 i ∣ ∣ ∣ η ( t -t inc ) = ˜ O ( √ η + ∣ ∣ ∣ w ( t inc ) 2 i ∣ ∣ ∣ η ( t -t inc ) ) . Now it suffices to show that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( t inc ) 2 i ∣ ∣ ∣ ≤ O ( 1 d α ) , ∣ ∣ ∣ W ( t inc ) 1 [ i, j ] ∣ ∣ ∣ ≤ O ( 1 d α ) , which is implied by Lemma 18.

Finally to complete the proof, we show that T 1 = Θ ( 1 √ dη ) . When t = T 1 , we have ∀ j ∈ [ d ] : ∑ d i =1 w ( T 1 ) 2 i W ( T 1 ) 1 [ i, j ] = Θ(1) . Combining with the above results, we know that dη 2 ( T 1 -t inc ) 2 = Θ(1) , i.e. η ( T 1 -t inc ) = Θ ( 1 √ d ) . In Section E.5, we will prove t inc = Θ ( 1 ηd 3 2 α +1 ) . Then we have T 1 = Θ ( 1 √ dη ) .

E.3 Proof of Lemma 17

For certain t and H , we write eq. (27) as

where

and r ( t ) 1 n [ i, j ] , r ( t ) 1 d [ i, j ] , r ( t ) 2 n,i , r ( t ) 2 d,i are defined in eq. (28).

Since β 2 = β 2 1 < β 1 , then if we pick H ≥ 1 1 -β 1 log max { G ( t ) 1 ,G ( t ) 2 , ( G ( t ) 1 ) 2 , ( G ( t ) 2 ) 2 } ηξ 2 , we can get that H ≥ 1 1 -β 1 log G ( t ) 1 ηξ 2 , H ≥ 1 1 -β 2 log ( G ( t ) 1 ) 2 ηξ 2 , H ≥ 1 1 -β 1 log G ( t ) 2 ηξ 2 , H ≥ 1 1 -β 2 log ( G ( t ) 2 ) 2 ηξ 2 . Hence we can apply Lemma 33 in Appendix H to get that ∣ ∣ ∣ q ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ q ( t ) 1 d [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ q ( t ) 2 n,i ∣ ∣ ∣ , ∣ ∣ ∣ q ( t ) 2 d,i ∣ ∣ ∣ ≤ ηξ 2 .

Pick T in Lemma 32 as of order ˜ O ( 1 √ dη ) . By Lemma 32, we have with probability at least 1 -1 d , for all t ≤ T , ∀ τ ≤ t and ∀ i, j ∈ [ d ] ,

∣ ∣ ∣ Dg ( τ ) 2 i ∣ ∣ ∣ = ∣ ∣ ∣ ˜ g ( τ ) 2 i -g ( τ ) 2 i ∣ ∣ ∣ ≤ ˜ O   d 4 ( M ( t ) 1 ) 2 M ( t ) 2 σ √ d 1 / 2 η   + ˜ O   dM ( t ) 1 σ √ d 3 / 2 η   := D ( t ) 2 . Plugging into eq. (28) gives us t

E.4 Proof of Corollary 2

By M ( t ) 1 , M ( t ) 2 ≤ ˜ O ( 1 √ d ) , G ( t ) 1 ≤ ˜ O ( 1 √ d ) , G ( t ) 2 ≤ ˜ O ( √ d ) and the assumption σ ≤ η 3 / 2 ξ 2 d 13 / 4 , we get that D ( t ) 1 and D ( t ) 2 are upper bounded by D ( t ) 1 ≤ ˜ O ( d 7 / 4 ση -1 / 2 ) and D ( t ) 2 ≤ ˜ O ( d 11 / 4 ση -1 / 2 ) , which yields ∀ i, j ∈ [ d ] , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 n,i ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 d,i ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) .

E.5 Proof of Lemma 18

The proof is based on the following two lemmas.

Lemma 21. Under Assumption 1 and 2, we have that w.p. at least 1 -1 d α 2 -1 , for every 1 ≤ i ≤ d , √ π d 3 2 α ≤ ∣ ∣ ∣ w (0) 2 i ∣ ∣ ∣ ≤ √ 2 d 2 α log 2 d δ , and that w.p. at least 1 -δ for any given δ > 0 , ∣ ∣ ∣ W (0) 1 [ i, j ] ∣ ∣ ∣ ≤ √ 2 d 4 α log 2 d 2 δ .

Lemma 22. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Pick β 2 = β 2 1 , ξ ∈ (0 , 1) , η < 1 4 . Consider any time point t ≤ ˜ O ( 1 √ dη ) . If ∀ τ ≤ t, ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) , ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) and ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) , ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ ≤ ˜ O ( √ d ) , we will have

where the ˜ O notation depends on H = 1 1 -β 1 log d ηξ 2 .

Furthermore, if for certain i, j ∈ [ d ] , Condition 1 (resp. Condition 2) is satisfied, we will have

Now we prove Lemma 18. Define t d := inf { t : ∃ i, j : ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ > 1 d or ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ > 1 d } . Now we want to find a time point t inc before t d for the lemma to hold. During the period t < t d , we have ∀ j ∈ [ d ] , E j = -Θ(1) (which means t d < T 1 ) and therefore for all i, j ∈ [ d ] , ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ 1 d and ∣ ∣ ∣ g ( t ) 2 i ∣ ∣ ∣ ≤ 1 . Then we can use Lemma 22 to get that for t ≤ min { t d , 1 √ dη } , we have ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( η ) , ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ ≤ ˜ O ( η ) . Hence t d ≥ ˜ Ω ( 1 ηd ) .

Define t sign = inf { t < min { t d , 1 √ dη } : ∃ i ∈ [ d ] : ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ ≤ 1 d 3 2 α } . By Lemma 21, w.h.p. ∀ i ∈ [ d ] : ∣ ∣ ∣ w (0) 2 i ∣ ∣ ∣ ≥ √ π d 3 2 α , combining with ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ ≤ ˜ O ( η ) gives us that w.h.p., t sign ≥ √ π -1 d 3 2 α / ˜ O ( η ) = ˜ Ω ( 1 ηd 3 2 α ) .

Now let's analyze the behavior of W 1 during the period t < t sign . Consider any i, j ∈ [ d ] . By definition, sign ( w ( t ) 2 i ) = sign ( w (0) 2 i ) . Note that E ( t ) j = -Θ(1) , then we have sign ( g ( t ) 1 [ i, j ] ) = -sign ( w (0) 2 i ) and that ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ = Ω ( 1 d 3 2 α ) = Ω( ξ ) by our choice of ξ . Then we know that Condition 1 is satisfied with s ( t ) 1 [ i, j ] = -sign ( w (0) 2 i ) (for all H < t ≤ t sign), which by Lemma 22 yields sign ( ∆ W ( t ) 1 [ i, j ] ) = sign ( w (0) 2 i ) and ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Θ( η ) .

glyph[negationslash]

Lemma 21 tells us that w.h.p., ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W (0) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 d 2 α ) . For any i, j , if initially sign ( W (0) 1 [ i, j ] ) = sign ( w (0) 2 i ) , then for the following steps before t sign, we will have sign ( W ( t ) 1 [ i, j ] ) = sign ( w (0) 2 i ) . If initially sign ( W (0) 1 [ i, j ] ) = sign ( w (0) 2 i ) , then after at most t 0 = ˜ O ( 1 ηd 2 α ) steps, W 1 [ i, j ] will flip the sign. Note that t 0 = ˜ O ( 1 ηd 2 α ) is smaller than t sign .

Hence we have shown that at some time point t 0 , we have ∀ i, j ∈ [ d ] : sign ( W ( t ) 1 [ i, j ] ) = sign ( w ( t ) 2 i ) = sign ( w (0) 2 i ) . Now we analyze the period t ≥ t 0 .

When t 0 < t ≤ t sign, we still have sign ( ∆ W ( t ) 1 [ i, j ] ) = sign ( w (0) 2 i ) and ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Θ( η ) . Combining these two with the fact sign ( W ( t 0 ) 1 [ i, j ] ) = sign ( w (0) 2 i ) , we know that for all t ∈ [ t 0 , t sign ] , sign ( W ( t ) 1 [ i, j ] ) = sign ( w (0) 2 i ) and that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W ( t +1) 1 [ i, j ] ∣ ∣ ∣ = ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ + ˜ Θ( η ) . Then at certain step t inc which satisfies t inc = t 0 + ˜ Θ ( 1 ηd 3 2 α +1 ) ∈ ( H,t sign ) , we will have ∀ t inc -H ≤ τ ≤ t inc , ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = Θ ( 1 d 3 2 α +1 ) and therefore ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ = ∣ ∣ ∣ ∑ d j =1 W ( τ ) 1 [ i, j ] E ( τ ) j ∣ ∣ ∣ = ∑ d j =1 ∣ ∣ ∣ W ( τ ) 1 [ i, j ] E ( τ ) j ∣ ∣ ∣ = Θ ( 1 d 3 2 α ) = Ω( ξ ) . For t ≤ t inc, we have ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = O ( 1 d 3 2 α +1 ) .

Since t inc < t sign, we have ∣ ∣ w t inc 2 i ∣ ∣ = Ω ( 1 d 3 2 α ) . For t ≤ t inc, note that ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ ≤ ˜ O ( η ) , t inc = t 0 + ˜ Θ ( 1 ηd 3 2 α +1 ) = ˜ Θ ( 1 ηd 3 2 α +1 ) , combining with the upper bound in Lemma 21 yields

∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ ≤ ∣ ∣ ∣ w (0) 2 i ∣ ∣ ∣ + t inc ˜ O ( η ) ≤ ˜ O ( 1 d 3 2 α +1 ) ≤ O ( 1 d α ) .

Moreover, ∀ t inc -H ≤ τ ≤ t inc , ∀ i ∈ [ d ] : sign ( g ( τ ) 2 i ) = -sign ( w (0) 2 i ) . Then Condition 2 is satisfied with s ( t ) 2 i = -sign ( w (0) 2 i ) for t = t inc. In the analysis of g ( t ) 1 [ i, j ] , we have already shown that for all t ≤ t sign (and thus for t = t inc), Condition 1 is satisfied, which completes the proof.

E.6 Proof of Lemma 21

Since for X ∼ N ( 0 , σ 2 ) , we have that P ( | X | ≤ t ) ≤ 2 t √ 2 πσ , then for a fixed i ,

∣ ∣ √

Then by union bound, we have that w.p. at least 1 -1 d α 2 -1 , for every 1 ≤ i ≤ d , ∣ ∣ w (0) 2 i ∣ ∣ ≥ π d 3 2 α .

As for the upper bounds, using the Gaussian tail bound and union bound, we have w.p. at least 1 -δ ,

E.7 Proof of Lemma 22

Now we analyze the magnitude order of ∆ W ( t ) 1 [ i, j ] . The analysis of ∆ w ( t ) 2 i is similar. For t ≤ ˜ O ( 1 √ dη ) . By assumption, M ( t ) 1 , M ( t ) 2 ≤ ˜ O ( 1 √ d ) , G ( t ) 1 ≤ ˜ O ( 1 √ d ) , G ( t ) 2 ≤ ˜ O ( √ d ) , and σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Hence we can pick H := 1 1 -β 1 log d ηξ 2 and apply Lemma 17 and Corollary 2 to get that, w.h.p., for all t ≤ ˜ O ( 1 √ dη ) and ∀ i, j ∈ [ d ] , eq. (27) can be written as

where ∀ i, j ∈ [ d ] , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 n,i ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 2 d,i ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) .

On one hand, using ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) and β 2 = β 2 1 , and √ x + y ≥ √ x -√ | y | when x ≥ 0 , x + y ≥ 0 , we get from eq. (33) that

where ( i ) uses Cauchy-Schwarz inequality for the numerator.

On the other hand, when sign ( g ( t -H ) 1 [ i, j ] ) = sign ( g ( t -H +1) 1 [ i, j ] ) = ... = sign ( g ( t ) 1 [ i, j ] ) = s ( t ) 1 [ i, j ] , we have

If we further have (1 -β 1 ) ∣ ∣ ∣ ∑ H τ =0 β ( τ ) 1 g ( t -τ ) 1 [ i, j ] ∣ ∣ ∣ ≥ Ω( ξ ) , then combining with ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) < ξ we will get

Using x + y ≤ √ | x | + √ | y | , we obtain that

Together with the upper bound completes the proof.

E.8 Proof of Lemma 19

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 23. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Pick ξ ≤ min { √ η d 3 α -1 , 1 d 3 2 α } , for t inc in Lemma 18, we have that w.h.p. for all t inc ≤ t ≤ T 1 , ∀ i, j ∈ [ d ] .

Specially, at the end of the first phase ( t = T 1 ), we have ∀ i, j ∈ [ d ] , ∣ ∣ w ( T 1 ) 2 i ∣ ∣ = ˜ Θ ( 1 √ d ) and ∣ ∣ ∣ W ( T 1 ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) .

Now we go back to the proof of Lemma 19. For t inc ≤ t < T 1 , since E ( t ) j = ( W ( t ) 2 W ( t ) 1 ) j -A j = ∑ d i =1 w ( t ) 2 i W ( t ) 1 [ i, j ] -A j , we have,

Combining Lemma 23 and eq. (34) gives us ∀ j ∈ [ d ] ,

Let's first analyze g ( t ) 1 [ i, j ] . Note that

where sign ( w ( t +1) 2 i ∆ E ( t ) j ) = sign ( w (0) 2 i ) while sign ( ∆ w ( t ) 2 i E ( t ) j ) = -sign ( w (0) 2 i ) .

Now we analyze the sign of g ( t ) 1 [ i, j ] when t inc ≤ t < T 1 . Using ∣ ∣ ∣ w ( t inc +1) 2 i ∣ ∣ ∣ = ˜ O ( 1 d α ) and eq. (35), we get that ∣ ∣ ∣ w ( t inc +1) 2 i ∆ E ( t inc ) j ∣ ∣ ∣ ≤ ˜ O ( 1 d α · √ dη ) . While on the other hand, ∣ ∣ ∣ ∆ w ( t inc ) 2 i E ( t inc ) j ∣ ∣ ∣ = ˜ Θ( η ) . That means sign ( ∆ g ( t inc ) 1 [ i, j ] ) = -sign ( w (0) 2 i ) . Note that sign ( g ( t inc ) 1 [ i, j ] ) = -sign ( w ( t inc ) 2 i ) = -sign ( w (0) 2 i ) , we know that ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ will increase when t = t inc .

In the following steps, ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ will keep increasing as long as ∣ ∣ ∣ ∆ w ( t ) 2 i E ( t ) j ∣ ∣ ∣ > ∣ ∣ ∣ w ( t +1) 2 i ∆ E ( t ) j ∣ ∣ ∣ . Since ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ keep increasing while ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ remain ˜ Θ( η ) , by eq. (35), we know that the trend of ∣ ∣ ∣ ∆ E ( t ) j ∣ ∣ ∣ is to increase. On the other hand, ∣ ∣ ∣ E ( t ) j ∣ ∣ ∣ keeps decreasing since E ( t ) j < 0 while ∆ E ( t ) j > 0 . Then after some time point we will have ∣ ∣ ∣ ∆ w ( t ) 2 i E ( t ) j ∣ ∣ ∣ < ∣ ∣ ∣ w ( t +1) 2 i ∆ E ( t ) j ∣ ∣ ∣ and in the following steps ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ will have the trend to decrease. Specially, when t = T 1 -1 , we have ∣ ∣ ∣ E ( t ) j ∣ ∣ ∣ = Θ ( √ ηd ) and ∣ ∣ ∣ W ( t ) 1 ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) , ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) by Lemma 23, which gives us

.

Therefore we have proved that when t inc ≤ t < T 1 , the trend of ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ is to first increase and then decrease. In order to prove ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( √ η ) , it suffices to show that ∣ ∣ ∣ g ( t inc ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( √ η ) and ∣ ∣ ∣ g ( T 1 ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( √ η ) . When t = t inc ,

When t = T 1 , we have

As for g ( t ) 2 i , since for ∀ i ∈ [ d ] , W ( t ) 1 [ i, j ] for different j have the same sign. Combining with ∀ j ∈ [ d ] : E ( t ) j < 0 gives us

Then it suffices to show that for t inc ≤ t < T 1 , ∣ ∣ ∣ E ( t ) j W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( √ η ) , which can be proven using the same technique as above.

Finally, for ∀ τ ≤ t, ∀ i, j ∈ [ d ] , note that the upper bounds of ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ are already given in Lemma 23. As for ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ and ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ , we have ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ∣ ∣ ∣ w ( τ ) 2 i E ( τ ) j ∣ ∣ ∣ = ˜ O ( 1 √ d ) , ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ ≤ ∑ d j =1 ∣ ∣ ∣ E ( τ ) j W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( √ d ) .

E.9 Proof of Lemma 23

For any i, j ∈ [ d ] , and any t in the interval [ t inc , T 1 ] , we prove by induction that

The base case t = t inc was already proven by Lemma 18.

For t ∈ [ t inc , T 1 ) , suppose (B) and (C) hold for time t and (A) holds for all τ ∈ [ t inc , t ] . From (A), we get that ∀ τ ∈ [ t inc , t ] : ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ∣ ∣ ∣ w ( τ ) 2 i E ( τ ) j ∣ ∣ ∣ = ˜ O ( 1 √ d ) , ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ ≤ ∑ d j =1 ∣ ∣ ∣ E ( τ ) j W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( √ d ) . Since when t < T 1 , ∀ j ∈ [ d ] : E ( t ) j < 0 , from (B) we know that ∀ τ ∈ [ t -H,t ] : sign ( g ( τ ) 1 [ i, j ] ) = sign ( g ( τ ) 2 i ) = -sign ( w (0) 2 i ) . Combining with (C) tells us that Condition 1 and 2 are satisfied.

In Section E.2 we have shown that T 1 = Θ ( 1 √ dη ) . Then for t ∈ [ t inc , T 1 ) , we can use Lemma 22 to get that ∀ t inc ≤ τ ≤ t , ∀ i, j ∈ [ d ] ,

Since when t = t inc, sign ( W ( t inc ) 1 [ i, j ] ) = sign ( w ( t inc ) 2 i ) = sign ( w (0) 2 i ) . We get that for t inc ≤ τ ≤ t ,

Now for t +1 , we have ∀ i, j ∈ [ d ] ,

That means ∀ τ ∈ [ t +1 -H,t +1] : sign ( W ( τ ) 1 [ i, j ] ) = sign ( w ( τ ) 2 i ) = sign ( w (0) 2 i ) . This proves (B) for time t +1 .

On the other hand, we get that ∣ ∣ ∣ W ( t +1) 1 [ i, j ] ∣ ∣ ∣ ≥ ∣ ∣ ∣ W ( t inc ) 1 [ i, j ] ∣ ∣ ∣ = Θ ( 1 d 3 2 α +1 ) and ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ ≥ ∣ ∣ ∣ w ( t inc ) 2 i ∣ ∣ ∣ = Ω ( 1 d 3 2 α ) . Since t +1 ≤ T 1 which means ∀ j ∈ [ d ] : ∣ ∣ ∣ E ( t +1) j ∣ ∣ ∣ ≥ √ ηd . Then

This proves (C) at time t +1 .

Since t +1 ≤ T 1 which means ∀ j ∈ [ d ] : ( W ( t +1) 2 W ( t +1) 1 ) j ≤ O (1) , we obtain that

Note that ∣ ∣ ∣ W ( t inc ) 1 [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ w ( t inc ) 2 i ∣ ∣ ∣ < 1 d (since t inc < t d ), we get that ( t +1 -t inc ) ˜ Θ( η ) = O ( 1 √ d ) , which gives us ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ = ˜ O ( 1 √ d ) and ∣ ∣ ∣ W ( t +1) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) and hence (A) holds at time t +1 .

Therefore by induction, we can prove that (A), (B), (C) hold for all t inc ≤ t ≤ T 1 . Then applying Lemma 22, we get that for all t inc ≤ t ≤ T 1 , ∀ i, j ∈ [ d ] : ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Θ( η ) , ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ( η ) .

Specially, at the end of the first phase, we have ∀ j ∈ [ d ] : ( W ( t +1) 2 W ( t +1) 1 ) j = Θ(1) . Repeating the above proof techniques gives us ∣ ∣ ∣ w ( T 1 ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) and ∣ ∣ ∣ W ( T 1 ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) for ∀ i, j ∈ [ d ] .

E.10 Proof of Lemma 20

Let's first prove eq. (30).

By Lemma 19, for t inc ≤ t < T 1 , we have ∀ i, j ∈ [ d ] , ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( √ η ) , ∣ ∣ ∣ g ( t ) 2 i ∣ ∣ ∣ = ˜ Ω ( √ ηd ) . Then it suffices to show that for t inc ≤ t < T 1 , ∣ ∣ ∣ g ( t ) 1 [ i, j ] -g ( t -τ ) 1 [ i, j ] ∣ ∣ ∣ = τ ˜ O ( η ) and ∣ ∣ ∣ g ( t ) 2 i -g ( t -τ ) 2 i ∣ ∣ ∣ = τ ˜ O ( ηd ) . It suffices to show that when t < T 1 , ∣ ∣ ∣ g ( t +1) 1 [ i, j ] -g ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( η ) and ∣ ∣ ∣ g ( t +1) 2 i -g ( t ) 2 i ∣ ∣ ∣ = ˜ O ( ηd ) .

By Lemma 18 and 23, we know that when t < T 1 , ∀ i, j ∈ [ d ] , ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( η ) , ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ ≤ ˜ O ( η ) and that ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) , ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) . Then the bound ∣ ∣ ∣ ∆ E ( t ) j ∣ ∣ ∣ ≤ ˜ O ( η √ d ) in eq. (35) hold for all t < T 1 (not only t inc ≤ t < T 1 ). Substituting these bounds into eq. (36) gives us ∀ t < T 1 ,

Similarly, we have that ∣ ∣ ∣ g ( t +1) 2 i -g ( t ) 2 i ∣ ∣ ∣ = ˜ O ( ηd ) , which proves eq. (30).

Note that for a, b ∈ R :

Then eq. (31) immediately follows from eq. (30).

E.11 Proof of Lemma 16

We divide Lemma 16 into the following three lemmas. Combining them together immediately gives us the whole proof.

The first lemma below gives us the structure of W 2 in the second phase and that of W 1 under some conditions.

Lemma 24. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . By picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 , we have w.h.p. for T 1 ≤ t < ˜ T ,

and moreover

As for W 1 , if for certain i.j ∈ [ d ] and certain t ∈ [ T 1 , ˜ T ) we have ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( √ η ) , then

The second lemma below also analyzes the structure of W 1 but removes the conditions in Lemma 24.

Lemma 25. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . By picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 , we have w.h.p. for T 1 ≤ t < ˜ T , ∀ i, j ∈ [ d ] , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) and for any j ∈ [ d ] ,

The third lemma proves the convergence of Adam at time ˜ T .

Lemma 26. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . By picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 , at time ˜ T , we have that w.h.p. ∀ j ∈ [ d ] : ∣ ∣ ∣ E ( ˜ T ) j ∣ ∣ ∣ ≤ ˜ O ( d √ ηd ) , which implies ∥ ∥ ∥ E ( ˜ T ) ∥ ∥ ∥ 2 2 ≤ ˜ O ( ηd 4 ) .

E.12 Proof of Lemma 24

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 second phase.

Lemma 27. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . By picking η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 , we have w.h.p. for all T 1 ≤ t < ˜ T ,

Moreover, we have that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) .

Equipped with Lemma 27, we are ready to prove Lemma 24. We will only prove the results of w ( t ) 2 i . The proof for W ( t ) 1 [ i, j ] uses the same techniques.

Lemma 27 gives us upper bounds of ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ , as well as ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ and ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ for all i, j ∈ [ d ] . Then we know that eq.(35) still holds, which gives us ∀ j ∈ [ d ] : ∣ ∣ ∣ E ( t +1) j -E ( t ) j ∣ ∣ ∣ = ˜ O ( √ dη ) . Then we can use the same strategy in Lemma 20 to prove that ∣ ∣ ∣ g ( t +1) 2 i -g ( t ) 2 i ∣ ∣ ∣ = ˜ O ( ηd ) .

By definition, for T 1 ≤ t < ˜ T , we know that ∣ ∣ ∣ g ( t ) 2 i ∣ ∣ ∣ = Ω ( d √ η ) . Combining with the bound ∣ ∣ ∣ g ( t +1) 2 i -g ( t ) 2 i ∣ ∣ ∣ = ˜ O ( ηd ) , we know that the g ( t ) 2 i parts in eq.(30) and eq.(31) still hold. Then we can use the same strategy in Section E.2 to prove that the w ( t ) 2 i part of eq. (29) still holds, which gives us

By Lemma 15, we have that at the end of the first phase ( t = T 1 ),

Combining with ∀ i ∈ [ d ] , ∀ t ≤ T i : sign ( g ( t ) 2 i ) = -sign ( w (0) 2 i ) yields that during the second phase, for t ≤ ˜ T , we have

E.13 Proof of Lemma 27

By definition of T f , there exists j 0 ∈ [ d ] such that E ( τ ) j 0 < - √ ηd for T 1 ≤ t ≤ ˜ T . We prove by induction that during this period, ∀ i ∈ [ d ] : sign ( w ( t ) 2 i ) = sign ( W ( t ) 1 [ i, j 0 ] ) = sign ( w (0) 2 i ) and that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) .

The base case ( t = T 1 ) was already proven by Lemma 23. Now suppose for some t such that T 1 ≤ t < ˜ T , for all τ such that T 1 ≤ τ ≤ t , we have ∀ i ∈ [ d ] : sign ( w ( τ ) 2 i ) = sign ( W ( τ ) 1 [ i, j 0 ] ) = sign ( w (0) 2 i ) and that ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) , ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) . Using these bounds, we get that ∀ j ∈ [ d ] : ∣ ∣ ∣ E ( τ ) j ∣ ∣ ∣ ≤ ∑ d i =1 ∣ ∣ ∣ w ( τ ) 2 i W ( τ ) 1 [ i, j ] ∣ ∣ ∣ + | A j | = O (1) , which then yields two upper bounds ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ∣ ∣ ∣ w ( τ ) 2 i E ( τ ) j ∣ ∣ ∣ = ˜ O ( 1 √ d ) and ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ ≤ ∑ d j =1 ∣ ∣ ∣ E ( τ ) j W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( √ d ) .

By definition of T g , we know that for all T 1 ≤ τ ≤ t , ∀ i ∈ [ d ] : ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ ≥ d √ η = Ω( ξ ) and that sign ( g ( τ ) 2 i ) = -sign ( w (0) 2 i ) , which implies that Condition 2 is satisfied for ∀ i ∈ [ d ] . At the end of the proof of this lemma, we will show that ˜ T = ˜ Θ ( 1 √ dη ) . Together with the upper bound of ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ , we can apply Lemma 22 to get that w.h.p. for T 1 ≤ τ ≤ t , sign ( ∆ w ( τ ) 2 i ) = sign ( w (0) 2 i ) and ∣ ∣ ∣ ∆ w ( τ ) 2 i ∣ ∣ ∣ = ˜ Θ( η ) . Combining with the inductive hypothesis sign ( w ( τ ) 2 i ) = sign ( w (0) 2 i ) gives us that ∣ ∣ ∣ w ( τ +1) 2 i ∣ ∣ ∣ = ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ + ˜ Θ( η ) . Specially, when τ = t , we get the lower bound ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ ≥ ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) and that sign ( w ( t +1) 2 i ) = sign ( w (0) 2 i ) .

Since E ( τ ) j 0 < - √ ηd , we have that ∀ i ∈ [ d ] : ∣ ∣ ∣ g ( τ ) 1 [ i, j 0 ] ∣ ∣ ∣ = ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ ∣ ∣ ∣ E ( τ ) j 0 ∣ ∣ ∣ = ˜ Ω ( √ η ) = Ω( ξ ) and that sign ( g ( τ ) 1 [ i, j 0 ] ) = -sign ( w (0) 2 i ) . That means Condition 1 is satisfied for ∀ i ∈ [ d ] and j 0 . Using the same technique as when we deal with w ( τ ) 2 i , we get that for T 1 ≤ τ ≤ t ,

∀ i ∈ [ d ] : ∣ ∣ ∣ W ( τ +1) 1 [ i, j 0 ] ∣ ∣ ∣ = ∣ ∣ ∣ W ( τ ) 1 [ i, j 0 ] ∣ ∣ ∣ + ˜ Θ( η ) , sign ( W ( t +1) 1 [ i, j 0 ] ) = sign ( w (0) 2 i ) and that ∀ i, j ∈ [ d ] , ∣ ∣ ∣ ∆ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( η ) .

Now we analyze the magnitude order of ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ , ∣ ∣ ∣ W ( t +1) 1 [ i, j ] ∣ ∣ ∣ . Let's first analyze ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ . By Lemma 15, when t = T 1 , ∀ i, j ∈ [ d ] ,

Combining with the facts that for T 1 ≤ τ ≤ t , ∣ ∣ ∣ W ( τ +1) 1 [ i, j 0 ] ∣ ∣ ∣ = ∣ ∣ ∣ W ( τ ) 1 [ i, j 0 ] ∣ ∣ ∣ + ˜ Θ( η ) and ∣ ∣ ∣ w ( τ +1) 2 i ∣ ∣ ∣ = ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ + ˜ Θ( η ) yields ∣ ∣ ∣ W ( t +1) 1 [ i,j 0 ] ∣ ∣ ∣ ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ = ˜ Θ(1) . Since we just proved ∀ i ∈ [ d ] : sign ( w ( t +1) 2 i ) = sign ( W ( t +1) 1 [ i, j 0 ] ) = sign ( w (0) 2 i ) , we get that

which gives us that ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ = ˜ O ( 1 √ d ) . Recall that we have shown ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ ≥ ˜ Ω ( 1 √ d ) , then ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) .

Now we prove ∣ ∣ ∣ W ( t +1) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) . We have proved that T 1 ≤ τ ≤ t , ∀ i, j ∈ [ d ] , ∣ ∣ ∣ ∆ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( η ) and ∣ ∣ ∣ w ( τ +1) 2 i ∣ ∣ ∣ -∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ = ˜ Θ( η ) , then ∀ i, j ∈ [ d ] ,

where the last equality uses ∣ ∣ ∣ W ( T 1 ) 1 [ i,j ] ∣ ∣ ∣ ∣ ∣ ∣ w ( T 1 ) 2 i ∣ ∣ ∣ = 1 ± ˜ O ( √ η + 1 d α -1 / 2 ) . Since we already proved that ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) , we get ∣ ∣ ∣ W ( t +1) 1 ∣ ∣ ∣ = ˜ O ( 1 √ d ) .

Therefore by induction, for all t in the interval [ T 1 , ˜ T ) , we have ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) . From the proof we also get ∀ i ∈ [ d ] : ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ > ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ , and that ∣ ∣ ∣ ∆ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ( η ) , ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( η ) .

Now we verify that ˜ T = ˜ Θ ( 1 √ dη ) . Combining ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( ˜ T ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) and ∀ t ∈ [ T 1 , ˜ T ) , ∣ ∣ ∣ w ( t +1) 2 i ∣ ∣ ∣ -∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ( η ) , we immediately get that ˜ T -T 1 = ˜ Θ ( 1 √ dη ) . In Section E.2 we have shown that T 1 = Θ ( 1 √ dη ) , then we get ˜ T = ˜ Θ ( 1 √ dη ) .

E.14 Proof of Lemma 25

We prove this lemma by induction. The base case ( t = T 1 ) can be verified by Lemma 15. Now suppose for t in the interval [ T 1 , ˜ T ) , we have ∀ i, j ∈ [ d ] , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) .

For t ∈ [ T 1 , ˜ T ) , by the proof of Lemma 27 (Section E.13), we know that for ∀ τ ≤ t , ∀ i, j ∈ [ d ] : ∣ ∣ ∣ w ( τ ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) , ∣ ∣ ∣ W ( τ ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) and that ∣ ∣ ∣ g ( τ ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( 1 √ d ) , ∣ ∣ ∣ g ( τ ) 2 i ∣ ∣ ∣ ≤ ˜ O ( √ d ) , and that ˜ T 1 = ˜ Θ ( 1 √ dη ) . Then we can pick H := 1 1 -β 1 log d ηξ 2 and apply Lemma 17 and Corollary 2 to get that, w.h.p., for all t ∈ [ T 1 , ˜ T ) and ∀ i, j ∈ [ d ] , the update of W 1 can be written as

where ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ , ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) . By Lemma 24, we have that for 1 ≤ i, j ≤ d ,

Using the fact that for a, b ∈ R , | a 2 -b 2 | a 2 ≤ 2 | a -b | | a | + ( | a -b | | a | ) 2 , we get that

where ∣ ∣ ∣ R ( t ) gsqr , 1 [ i,j ] ∣ ∣ ∣ ( c ( t ) E ( t ) j ) 2 = ˜ O ( √ η + 1 d α -1 / 2 ) . That yields

Since ( c ( t -τ ) E ( t -τ ) j ) 2 > 0 , in eq. (38) we have that

However we cannot similarly prove that ∣ ∣ ∣ ∑ H τ =0 β τ 1 R ( t -τ ) g, 1 [ i, j ] ∣ ∣ ∣ glyph[lessmuch] ∣ ∣ ∣ ∑ H τ =0 β τ 1 c ( t -τ ) E ( t -τ ) j ∣ ∣ ∣ in eq. (37) because c ( t -τ ) E ( t -τ ) j may not have the same sign for τ = 0 , 1 , ..., H . To deal with eq.(37), we need to consider the two cases where ∣ ∣ ∣ ∑ H τ =0 β τ 1 R ( t -τ ) g, 1 [ i, j ] ∣ ∣ ∣ glyph[lessmuch] ∣ ∣ ∣ ∑ H τ =0 β τ 1 c ( t -τ ) E ( t -τ ) j ∣ ∣ ∣ or ∣ ∣ ∣ ∑ H τ =0 β τ 1 R ( t -τ ) g, 1 [ i, j ] ∣ ∣ ∣ glyph[negationslash]glyph[lessmuch] ∣ ∣ ∣ ∑ H τ =0 β τ 1 c ( t -τ ) E ( t -τ ) j ∣ ∣ ∣ .

Case 1. ∣ ∣ ∣ (1 -β 1 ) ∑ H τ =0 β τ 1 R ( t -τ ) g, 1 [ i, j ] + glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ ≤ δ ∣ ∣ ∣ (1 -β 1 ) ∑ H τ =0 β τ 1 c ( t -τ ) E ( t -τ ) j ∣ ∣ ∣ where δ = ( η 1 4 + 1 d α 2 -1 4 ) .

Note that from eq. (39) we have

Combining with ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) ≤ ˜ O ( η 1 4 + 1 d α 2 -1 4 ) 2 ξ 2 , we can apply Lemma 34 to get that

where ∣ ∣ ∣ e ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( η 1 4 + 1 d α 2 -1 4 ) . Since ∣ ∣ ∣ W ( t +1) 1 [ i, j ] -W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( η ) , we get that ∣ ∣ ∣ v ( t ) j ∣ ∣ ∣ = ˜ O ( η ) .

Case 2. ∣ ∣ ∣ (1 -β 1 ) ∑ H τ =0 β τ 1 R ( t -τ ) g, 1 [ i, j ] + glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ > δ ∣ ∣ ∣ (1 -β 1 ) ∑ H τ =0 β τ 1 c ( t -τ ) E ( t -τ ) j ∣ ∣ ∣ where δ = ( η 1 4 + 1 d α 2 -1 4 ) .

Since ∣ ∣ ∣ R ( t ) g, 1 [ i,j ] ∣ ∣ ∣ c ( t ) ∣ ∣ ∣ E ( t ) j ∣ ∣ ∣ = ˜ O ( √ η + 1 d α -1 / 2 ) , we have that

where ( i ) uses Cauchy-Schwarz inequality and β 2 = β 2 1 , ( ii ) uses eq. (38) and (39).

Combining with ∣ ∣ ∣ glyph[epsilon1] ( t ) 1 n [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( ηξ 2 ) ≤ ˜ O ( √ η + 1 d α -1 / 2 ) ( ξ -√∣ ∣ ∣ glyph[epsilon1] ( t ) 1 d [ i, j ] ∣ ∣ ∣ ) gives us

which implies

Consider certain i.j ∈ [ d ] and the period from T 1 to t . Denote T as the set of time points when Case 1 is satisfied. By Lemma 27, we know that η ( t -T 1 ) = ˜ O ( 1 √ d ) , which gives us

By the first phase analysis, we have that

where V ( T 1 ) j = O ( 1 √ d ) , ∣ ∣ ∣ R ( T 1 ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O (√ η d + 1 d α ) . Combining with the analysis of Case 1, we have that

where | R T [ i, j ] | ≤ O ( η 1 4 + 1 d α 2 -1 4 )(∣ ∣ ∣ V ( T 1 ) j ∣ ∣ ∣ + ∑ τ ∈T ∣ ∣ ∣ v ( τ ) j ∣ ∣ ∣ ) .

Since for τ ∈ T , ∣ ∣ ∣ v ( τ ) j ∣ ∣ ∣ = ˜ O ( η ) , V ( T 1 ) j = O ( 1 √ d ) and η ( t -T 1 ) = ˜ O ( 1 √ d ) , we can bound | R T [ i, j ] | by

Combining the above results together yields

By the inductive hypothesis ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) , we get that ∣ ∣ ∣ V ( t ) j ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) , which gives us ∣ ∣ ∣ R ( t ) 1 [ i,j ] ∣ ∣ ∣ ∣ ∣ ∣ V ( t ) j ∣ ∣ ∣ ≤ ˜ O ( η 1 4 + 1 d α 2 -1 4 ) .

Therefore, we have that for any j ∈ [ d ] and any i 1 , i 2 ∈ [ d ] ,

By Lemma 24, we know that ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ with different i are also roughly equal, i.e. ∣ ∣ ∣ w ( t ) 2 i 1 ∣ ∣ ∣ ∣ ∣ ∣ w ( t ) 2 i 2 ∣ ∣ ∣ = 1 ± ˜ O ( √ η + 1 d α -1 / 2 ) . Then we have for any j ∈ [ d ] ,

where k can be any index in { 1 , 2 , ..., d } and the last equality uses ∀ i ∈ [ d ] : ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) .

Now we analyze the lower bound of ∣ ∣ ∣ W ( t +1) 1 [ k, j ] ∣ ∣ ∣ . Although it may decrease during some period, we observe that once ∣ ∣ ∣ W ( t ) 1 [ k, j ] ∣ ∣ ∣ decreases to some value of order ˜ Θ ( 1 √ d ) such that ( W 2 W 1 ) ( t ) j < A j - √ ηd , i.e. E ( t ) j < - √ ηd , we can apply the technique in Section E.13 when analyzing W ( t ) 1 [ i, j 0 ] to get that ∣ ∣ ∣ W ( t ) 1 [ k, j ] ∣ ∣ ∣ will increase in the next step. This mechanism ensures a ˜ Ω ( 1 √ d ) lower bound of ∣ ∣ ∣ W ( t +1) 1 [ k, j ] ∣ ∣ ∣ . Since k, j are arbitrary, we have proved that at time t +1 , ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W ( t +1) 1 [ i, j ] ∣ ∣ ∣ ≥ ˜ Ω ( 1 √ d ) .

Therefore by induction, we conclude that when T 1 ≤ t < ˜ T , for ∀ i, j ∈ [ d ] , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) . The remaining part of this lemma has also been proved by the analysis above.

E.15 Proof of Lemma 26

Lemma 27 tells us that for any i ∈ [ d ] , ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ keeps increasing when t < ˜ T . However, the behavior of W ( t ) 1 [ i, j ] is more complicated. The following lemma tells us that ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ will increase until T f,j . After that ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ and E ( t ) j may zigzag, but E ( t ) j will not fluctuate dramatically and will be trapped in a small interval around zero.

Lemma 28. Under Assumption 1, 2 and 3, suppose σ ≤ η 3 / 2 ξ 2 d 13 / 4 . Pick η ≤ O ( 1 d 3 α ) , ξ ≤ √ η d 3 α -1 , and β 2 = β 2 1 . Consider certain coordinate j . For T 1 ≤ t < min { ˜ T, T f,j } , we have ∀ i ∈ [ d ] : ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ keeps increasing. If T f,j < ˜ T , then for T f,j ≤ t < ˜ T , we will have -˜ O ( √ ηd ) ≤ E ( t ) j ≤ ˜ O ( √ ηd ) .

glyph[negationslash]

Now we start proving Lemma 26. At time ˜ T , denote S := { j : T f,j < ˜ T } , i.e. the set of coordinates whose E j have passed its 'flip time'. By Lemma 28, we know that ∀ j ∈ S, ∣ ∣ ∣ E ( ˜ T ) j ∣ ∣ ∣ ≤ ˜ O ( √ ηd ) . If S c = φ , which means ∀ j ∈ [ d ] : ∣ ∣ ∣ E ( ˜ T ) j ∣ ∣ ∣ ≤ ˜ O ( √ ηd ) , then our lemma will immediately follow. If S c = φ , we have ˜ T = min { T g , T f } = T g and that ∀ j ∈ S c : E ( ˜ T ) j < 0 . By the definition of T g , we know that ∃ i 0 ∈ [ d ] : ∣ ∣ ∣ g ( ˜ T ) 2 i 0 ∣ ∣ ∣ ≤ O ( d √ η ) . Then

By Lemma 25, we know that when T 1 ≤ t < ˜ T , for ∀ i, j ∈ [ d ] , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) . Since the update per step ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( η ) , we know that sign ( W ( t ) 1 [ i, j ] ) remains unchanged during this period and sign ( W ( t ) 1 [ i, j ] ) = sign ( W ( T 1 ) 1 [ i, j ] ) = sign ( w (0) 2 i ) independent of j . Combining with ∀ j ∈ S c : E ( ˜ T ) j < 0 gives us that E ( ˜ T ) j W ( ˜ T ) 1 [ i 0 , j ] for different j have the same sign. Therefore for any j 0 ∈ S c ,

Note that the above inequality holds for any j 0 ∈ S c , which means ∀ j ∈ S c : ∣ ∣ ∣ E ( ˜ T ) j ∣ ∣ ∣ ≤ ˜ O ( d √ ηd ) . Combining with the fact that ∀ j ∈ S : ∣ ∣ ∣ E ( ˜ T ) j ∣ ∣ ∣ ≤ ˜ O ( d √ ηd ) completes the proof.

E.16 Proof of Lemma 28

Consider certain j ∈ [ d ] , when t < min { ˜ T, T f,j } , we have that E ( t ) j < - √ ηd . Therefore we can use the same argument as in Section E.13 to prove that ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ keeps increasing, and sign ( W ( t ) 1 [ i, j ] ) = sign ( w (0) 2 i ) for all i ∈ [ d ] .

At time the 'flip time' t = T f,j , by definition, E ( t ) j ≥ - √ ηd . After that E ( t ) j may oscillate. Now we prove that once E ( t ) j ≥ √ ηd (or E ( t ) j ≤ - √ ηd ), after a short period E ( t ) j will decrease (or increase) until E ( t ) j ≤ ˜ O ( √ ηd ) (or E ( t ) j ≥ -˜ O ( √ ηd ) ). Moreover, during this period, E ( t ) j won't change too much.

We first recall that when T 1 ≤ t < ˜ T , Lemma 27 gives us for all i ∈ [ d ] , ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) and ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) . Then eq.(35) we obtained in the first phase analysis still holds, which tells us that the change of E ( t ) j per step satisfies ∣ ∣ ∣ E ( t +1) j -E ( t ) j ∣ ∣ ∣ = ˜ O ( η √ d ) for all T 1 ≤ t < ˜ T .

We divide the analysis into two cases, based on whether these E ( t ) j ≥ √ ηd or E ( t ) j ≤ - √ ηd . By Lemma 25, we know that when T 1 ≤ t < ˜ T , ∀ i ∈ [ d ] , ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ = ˜ Ω ( 1 √ d ) . Since the update per step ∣ ∣ ∣ ∆ W ( t ) 1 [ i, j ] ∣ ∣ ∣ ≤ ˜ O ( η ) , we know that sign ( W ( t ) 1 [ i, j ] ) remains unchanged during this period and sign ( W ( t ) 1 [ i, j ] ) = sign ( W ( T 1 ) 1 [ i, j ] ) = sign ( w (0) 2 i ) .

By the analysis of w ( t ) 2 i in Lemma 24, we have for all i ∈ [ d ] , w ( t +1) 2 i = w ( t ) 2 i + sign ( w (0) 2 i ) ∆ ( t ) 2 i , where ∆ ( t ) 2 i = η ( 1 ± ˜ O ( √ η ) ) .

Case 1. Consider some time point t such that E ( t ) j ≤ - √ ηd . Note that for all i ∈ [ d ] , ∣ ∣ ∣ g ( t ) 1 [ i, j ] ∣ ∣ ∣ = ∣ ∣ ∣ w ( t ) 2 i E ( t ) j ∣ ∣ ∣ = ˜ Ω ( √ η ) and that sign ( g ( t ) 1 [ i, j ] ) = -sign ( w ( t ) 2 i ) = -sign ( w (0) 2 i ) . By Lemma 24, for all i ∈ [ d ] we have W ( t +1) 1 [ i, j ] = W ( t ) 1 [ i, j ] + sign ( w (0) 2 i ) ∆ ( t ) 1 [ i, j ] with ∆ ( t ) 1 [ i, j ] = η ( 1 ± ˜ O ( √ η ) ) . That gives us

where ( i ) is because sign ( w ( t ) 2 i ) = sign ( W ( t ) 1 [ i, j ] ) = sign ( w (0) 2 i ) . Therefore we have proved that E ( t ) j will increase in the next step. After that for τ ≥ t +1 , as long as E ( τ ) j ≤ - √ ηd , the above analysis will hold and E ( τ ) j will keep increasing until E ( τ ) j > - √ ηd or we reach ˜ T .

Case 2. Consider some time point t such that E ( t ) j ≥ √ ηd . We will prove that E ( t ) j will decrease after a short period, and during this period, the change of it is at most ˜ O ( √ ηd ) .

By similar arguments as in Case 1, we can get that W ( t +1) 1 [ i, j ] = W ( t ) 1 [ i, j ] -sign ( w (0) 2 i ) ∆ ( t ) 1 [ i, j ] , where ∆ ( t ) 1 [ i, j ] = η ( 1 ± ˜ O ( √ η ) ) , Then

where ( i ) is because sign ( w ( t ) 2 i ) = sign ( W ( t ) 1 [ i, j ] ) = sign ( w (0) 2 i ) . E ( t +1) j may not be smaller than E ( t ) j , but we will show that after at most t s steps for some t s , we will have E ( t + t s +1) j < E ( t + t s ) j . To see this, first note that by the bounds of ∆ ( t ) 1 [ i, j ] and ∆ ( t ) 2 i , we get ∆ ( t ) 1 [ i, j ] ≥ ∆ ( t ) 2 i -η ˜ O ( √ η ) . Since ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ increases by ˜ Θ( η ) per step, and ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ keeps decreasing, then we have either i) after t s steps for some t s , ∀ i ∈ [ d ] : ∣ ∣ ∣ w ( t + t s ) 2 i ∣ ∣ ∣ ≥ ∣ ∣ ∣ W ( t + t s ) 1 [ i, j ] ∣ ∣ ∣ + √ η or ii) we reach ˜ T .

For i), if E ( t + t s ) j < √ ηd , then it's already what we want. Otherwise we will have ∆ ( t + t s ) 1 [ i, j ] = η ( 1 ± ˜ O ( √ η ) ) . Hence

where the last inequality uses ∀ i, j ∈ [ d ] : ∣ ∣ ∣ W ( t + t s ) 1 [ i, j ] ∣ ∣ ∣ = ˜ O ( 1 √ d ) . Therefore E ( t + t s +1) j < E ( t + t s ) j . After that for τ ≥ t + t s +1 , as long as E ( τ ) j ≥ √ ηd , the above analysis will hold and E ( τ ) j will keep decreasing until E ( τ ) j < √ ηd or we reach ˜ T .

Now we prove that during these t s steps, the change of E j is ˜ O ( √ ηd ) . Since at each step the difference | w 2 i | - | W 1 [ i, j ] | will be enlarged by ˜ Ω( η ) , then we know that t s = √ η/ ˜ Ω( η ) = ˜ O ( 1 √ η ) .

Combining with the fact that for all T 1 ≤ τ ≤ ˜ T , ∣ ∣ ∣ E ( τ +1) j -E ( τ ) j ∣ ∣ ∣ = ˜ O ( η √ d ) gives us

Combining the above two cases, we find that if for some t , E ( t ) j ≥ √ ηd , then after at most t s steps E j will decrease and keeps decreasing until E j < √ ηd or we reach ˜ T . During these steps, E j can increase at most ˜ O ( √ ηd ) . If for some t , E ( t ) j ≤ - √ ηd , then after one step it will increase and keeps increasing until E j > √ ηd or we reach ˜ T . That means once for some coordinate j , E j overshoots, it will zigzag in a small region around zero, which is [ -˜ O ( √ ηd ) , ˜ O ( √ ηd ) ] .

F 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 a rigorous theoretical analysis on a two-layer network under Assumption 1 and 2.

F.1 Empirical Results

Let's first define the diagonal domination of the i -th coordinate at time t .

glyph[negationslash]

To measure the diagonal domination of the whole Hessian, we need to consider the distribution of r OPT diag ,i ( t ) for different i . Figure 15 shows the mean and median of r SGDM diag ,i ( t ) and r Adam diag ,i ( t ) 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 4 layers. Since the number of parameters is very large, we did the computation by random sampling. As we can see, for both r SGDM diag ,i ( t ) and r Adam diag ,i ( t ) , the trend of their mean or median is to decrease over time, although there might be some oscillation.













∮√˜√∐√]}{





















⋃⊔〈〉⊕

̂∐⌉









∮√˜√∐√]}{

















⋃⊔〈〉⊕

̂∐⌉

(a) Mean

(b) Median

F.2 Theoretical Analysis

To simplify the theoretical analysis, we consider the mean of r OPT diag ,i ( t ) over all coordinate and define

We consider a 2-layer network under Assumption 1 and 2, and have two goals in our proof:

To show that R OPT diag ( t ) after training is smaller than that before training ( t = 0 ).

Note that in our setting (see in Assumption 1), the Hessian is a ( d 2 + d ) × ( d 2 + d ) matrix. For a completely 'uniform' matrix with the same size, we have that R OPT diag ( t ) = Θ ( √ d 2 + d ) = Θ( d ) . Hence our second goal is to show that the R OPT diag ( t ) after training is on lower order than Θ( d ) .

Theorem 2. Consider the ratio R OPT diag ( t ) defined in eq. (40) . Under Assumption 1 and 2, we have that before training ( t = 0 ), with high probability,

For SGD+M defined in eq. (3) . For any p > 0 , by picking the same hyperparameters as in Theorem 1, for T SGD , 1 , T SGD , 2 mentioned in Theorem 1, we have with constant probability, for any t ∈ [ T SGD , 1 , T SGD , 2 ] ,

where the trend of q ( t ) is to decrease over time and q ( T SGD , 2 ) ≤ ˜ O ( 1 d p/ 2 -1 ) = o ( d ) .

For Adam defined in eq. (3) . For any p > 0 , by picking the same hyperparameters as in Theorem 1, for T Adam , 1 , T Adam , 2 mentioned in Theorem 1, we have with high probability, for any t ∈ [ T Adam , 1 , T Adam , 2 ] ,

where the trend of r ( t ) is to decrease over time and r ( T Adam , 2 ) ≤ ˜ O ( 1 d p -1 2 ) = o ( √ d ) .

F.3 Proof of Theorem 2

Lemma 4.3 of [Kaw16] gives us the following forms of Hessian.

For any k ∈ { 1 , 2 , ..., H +1 } , we know that ∇ vec ( W k ) ( ∇ vec ( W k ) L ( W )) equals

and for k ∈ { 2 , 3 , ..., H +1 } ,

where r = ( W H +1 . . . W 1 -A ) T , C = W H +1 W H · · · W 2

For the 2-layer linear network, write the Hessian as

Intuitively, before training the elements of W 1 and W 2 are very close to zero, and W 2 W 1 -A ≈ -A . Since the elements of A are Θ(1) , we know that the magnitudes of elements of H 21 are much bigger than those of H 11 and H 22 .

After training, for both SGD+M and Adam, W 2 W 1 -A ≈ 0 . Then H 21 ≈ ( W 2 ) T ⊗ ( W 1 ) T and the magnitudes of its elements are no longer much larger than those of H 11 and H 22 . From the formula of H 11 , we know that all the diagonal entries are nonzero, and among the d 4 -d 2 off-diagonal entries, there are only d 3 -d 2 nonzero entries, which helps us to bound R OPT diag ( t ) .

then we have that

F.3.1 Proof of eq. (41)

Let's first analyze the weights and Hessian before training ( t = 0 ). For ease of notation, we omit the superscript ( t ) .

For the i -th row where 1 ≤ i ≤ d , i.e. the i -th row of the submatrix [ H 22 H T 21 ] , we have glyph[negationslash]

glyph[negationslash]

On the other hand, for the diagonal elements, we have w.h.p.

Then we have that for 1 ≤ i ≤ d , glyph[negationslash]

For the ( id + k ) -th row where 1 ≤ i ≤ d, 1 ≤ k ≤ d , i.e. the (( i -1) d + k ) -th row of the submatrix [ H 21 H 11 ] , we have glyph[negationslash]

glyph[negationslash]

On the other hand, for the diagonal elements, we have w.h.p.

Then we have that for 1 ≤ i ≤ d, 1 ≤ k ≤ d , glyph[negationslash]

Taking the average, we obtain that before training, i.e. when t = 0 ,

F.3.2 Proof of eq. (42)

The proof is based on the lemma below.

Lemma 29. Suppose the weight matrices have the following structure:

where ∀ 1 ≤ i, j ≤ d : | R 1 [ i,j ] | | u i v j | ≤ δ, | R 2 i | | cu i | ≤ δ, δ ∈ (0 , 1) .

Then we have for 1 ≤ i ≤ d , glyph[negationslash]

and for 1 ≤ i ≤ d, 1 ≤ k ≤ d glyph[negationslash]

Now we are ready to prove eq. (42).

By the analyses in Section D.1, we know that for t ∈ [ T SGD , 1 , T SGD , 2 ] , the weights obtained by GD with momentum satisfy

where T SGD , 1 = T 1 and

Here glyph[epsilon1] 0 is defined in Definition 2. Since u ( T 1 ) doesn't depend on time t in the period ( T SGD , 1 , T SGD , 2 ] , we write u ( T 1 ) as u for ease of notation.

Hence by Lemma 29, when t ∈ [ T SGD , 1 , T SGD , 2 ] , we have for 1 ≤ i ≤ d , glyph[negationslash]

and for 1 ≤ i ≤ d, 1 ≤ k ≤ d , glyph[negationslash]

By Lemma 4, we have u = X + Y where X i , i ∈ [ d ] are i.i.d Gaussian random variables and w.h.p.,

which yields that

By the proof in Section D.8, we know that for t ∈ [ T SGD , 1 , T SGD , 2 ] , ∀ i ∈ [ d ] : v ( t ) i , c ( t ) are positive. The induction in Section D.9 further gives us that for t ∈ [ T SGD , 1 , T SGD , 2 ] , w.h.p. ∀ k ∈ [ d ] : v ( t ) k c ( t ) = Θ ( 1 √ d ) , which yields c ( t ) ‖ v ( t ) ‖ 2 = Θ(1) . Combining with eq. (47), we obtain

By the proof in Section D.8, we know that for t ∈ [ T SGD , 1 , T SGD , 2 ] , ∀ i ∈ [ d ] : v ( t ) i , c ( t ) are positive and monotonically increasing. On the other hand, the proof in Section D.2 and D.9 tells us that w.h.p. ∥ ∥ E ( t ) ∥ ∥ 2 (resp. ∀ k ∈ [ d ] , ∣ ∣ ∣ E ( t ) k ∣ ∣ ∣ ) decreases from Θ( √ d ) (resp. Θ(1) ) when t = T SGD , 1 to O ( √ glyph[epsilon1] 0 d )

(resp. O ( √ glyph[epsilon1] 0 ) ) when t = T SGD , 2 . Therefore, the trend of ‖ E ( t ) ‖ 2 u 2 i ‖ v ( t ) ‖ 2 2 and ∣ ∣ ∣ E ( t ) k ∣ ∣ ∣ ( c ( t ) ) 2 u 2 i is to decrease over time, and when t = T SGD , 2 , we have w.h.p.

Moreover, when t = T SGD , 2 , the inequality in eq. (26) becomes equality, i.e. c 2 ‖ u ‖ 2 2 = Θ ( √ d ) and ∀ j ∈ [ d ] : ‖ u ‖ 2 2 v 2 j = Θ ( 1 √ d ) .

Using u = X + Y and eq. (46), we have

which together with the second inequality in eq. (47) yields

Combining with eq. (49), we get that

Substituting eq. (48) and (50) into eq. (44) and (45) gives us glyph[negationslash]

where the trend of q ( t ) 1 i is to decrease over time and q ( T SGD , 2 ) 1 i ≤ O ( ∑ d j =1 X 2 j X 2 i · √ glyph[epsilon1] 0 ) . We also have glyph[negationslash]

where the trend of q ( t ) 2 i is to decrease over time and q ( T SGD , 2 ) 2 i ≤ O ( ∑ d j =1 X 2 j X 2 i · √ glyph[epsilon1] 0 d ) . Hence

where the trend of q ( t ) is to decrease over time and

Denote σ 2 as the variance of X i for i ∈ [ d ] . By concentration of chi-squared distribution, we know that with probability at least 1 -δ for δ > 0 ,

By Lemma 36 in Appendix H, we know that with constant probability 1 d ∑ d i =1 1 | X i | = O ( 1 σ log d ) .

Then with constant probability, 1 d ∑ d i =1 1 X 2 i ≤ 1 d ( ∑ d i =1 1 | X i | ) 2 = O ( d σ 2 log 2 d ) . Hence

Therefore with constant probability,

where the trend of q ( t ) is to decrease over time and q ( T SGD , 2 ) ≤ ˜ O ( d √ glyph[epsilon1] 0 d ) . For any p > 0 , by picking the same hyperparameters as in Theorem 1, we have glyph[epsilon1] 0 d ≤ ˜ O ( 1 d p ) and hence q ( T SGD , 2 ) ≤ ˜ O ( 1 d p/ 2 -1 ) = o ( d ) .

F.3.3 Proof of eq. (43)

By the analyses in Section E.1, we know that for t ∈ [ T Adam , 1 , T Adam , 2 ] , the weights obtained by Adam satisfy

where ∀ i ∈ [ d ] : u i = sign ( w (0) 2 i ) ∈ {± 1 } and

Hence by Lemma 29, when t ∈ [ T Adam , 1 , T Adam , 2 ] , we have for 1 ≤ i ≤ d , glyph[negationslash]

and for

glyph[negationslash]

Recall the following facts of Adam.

(A) By Lemma 15, we know that for t ∈ [ T Adam , 1 , T 1 ] (where T 1 is defined in Definition 4), w.h.p. ∀ k ∈ [ d ] : v ( t ) k = c ( t ) = η ( t -t inc ) . Specially, when t = T Adam , 1 , ∀ k ∈ [ d ] : v ( t ) k = c ( t ) = 1 d α 2 . Lemma 25 and 27 tell us that for t ∈ [ T 1 , T Adam , 2 ] w.h.p. ∀ i, j ∈ [ d ] : | W 1 [ i, j ] | = ˜ Θ ( 1 √ d ) , | w 2 i | = ˜ Θ ( 1 √ d ) , which gives us ∀ k ∈ [ d ] : ∣ ∣ ∣ v ( t ) k ∣ ∣ ∣ = ˜ Θ ( 1 √ d ) and ∣ ∣ c ( t ) ∣ ∣ = ˜ Θ ( 1 √ d ) . That means when t ∈ [ T Adam , 1 , T Adam , 2 ] , ∀ k ∈ [ d ] : ∣ ∣ ∣ v ( t ) k ∣ ∣ ∣ and ∣ ∣ c ( t ) ∣ ∣ increase from 1 d α 2 to ˜ Θ( 1 √ d ) and ∣ ∣ ∣ v ( t ) k ∣ ∣ ∣ | c ( t ) | = ˜ Θ(1) , | c ( t ) | ‖ v ( t ) ‖ 2 = ˜ Θ ( 1 √ d ) .

(B) Lemma 15 and 26 tell us that w.h.p. ∥ ∥ E ( t ) ∥ ∥ 2 (resp. ∀ k ∈ [ d ] , ∣ ∣ ∣ E ( t ) k ∣ ∣ ∣ ) decreases from Θ( d ) (resp. Θ(1) ) when t = T Adam , 1 to ˜ O ( d 2 √ η ) (resp. ˜ O ( d √ ηd ) ) when t = T Adam , 2 .

Combining (A) and (B), we get that the trend of ‖ E ( t ) ‖ 2 ‖ v ( t ) ‖ 2 2 and ∣ ∣ ∣ E ( t ) k ∣ ∣ ∣ ( c ( t ) ) 2 is to decrease over time, and when t = T Adam , 2 , we have w.h.p.

Substituting (A) and eq. (53) into eq. (51) and (52) gives us w.h.p., glyph[negationslash]

where the trend of r ( t ) 1 i is to decrease over time and r ( T Adam , 2 ) 1 i ≤ ˜ O ( d 2 √ η ) . We also have glyph[negationslash]

where the trend of r ( t ) 2 i is to decrease over time and r ( T Adam , 2 ) 2 i ≤ ˜ O ( d 2 √ ηd ) . Hence R Adam diag ( t ) = ˜ O ( √ d ) + 1 d 2 + d ∑ d i =1 r ( t ) 1 i + d d 2 + d ∑ d i =1 r ( t ) 2 i := ˜ O ( √ d ) + r ( t ) where the trend ( t )

of r is to decrease over time and

For any p > 0 , by picking the same hyperparameters as in Theorem 1, we have ηd 4 ≤ ˜ O ( 1 d p ) and hence r ( T Adam , 2 ) ≤ ˜ O ( 1 d p -1 2 ) = o ( √ d ) .

F.4 Proof of Lemma 29

By the assumed weight structure, we get that

For the i -th row where 1 ≤ i ≤ d , i.e. the i -th row of the submatrix [ H 22 H T 21 ] , by triangle inequality, we have glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

Then we have that for 1 ≤ i ≤ d , glyph[negationslash]

glyph[negationslash]

For the ( id + k ) -th row where 1 ≤ i ≤ d, 1 ≤ k ≤ d , i.e. the (( i -1) d + k ) -th row of the submatrix [ H 21 H 11 ] , by triangle inequality again, we have glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

Then we have that for 1 ≤ i ≤ d, 1 ≤ k ≤ d , glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

G Connection between diagonal of loss Hessian and weights

The partial derivative at W i of the cost function for each i is given by:

In our experiments, we were interested in the diagonal elements of the hessian. These are given by:

for each possible i, a, b . For ease in notation, define for each i , the quantities M i := W T i +1 . . . W T H +1 and N i := W T 1 . . . W T i -1 . Then we have the following lemma.

Lemma 30. The diagonal elements of the hessian of the cost function are given by:

for each possible i, a, b .

Proof. We have:

This implies that:

where the last step follows since M i and N i are not functions of W i .

Since M i := W T i +1 . . . W T H +1 , N i := W T 1 . . . W T i -1 , by defining C i := M i W H +1 W H . . . W i +1 = M i M T i and D i := W i -1 . . . W 2 W 1 N i = N T i N i we have that:

where C i and D i are not functions of W i . Now, Equation 74 in the Matrix Cookbook 13 shows us that for any matrices A and X we have:

Note that W i ∈ R d i × d i -1 , then we can apply this to obtain that:

This completes the proof.

For ease of notation, let's now drop the superscript OPT and ( t ) and write R OPT med , 1 ( t ) as R med , 1 and R OPT med , 2 ( t ) as R med , 2 . For a 2-layer linear network, H = 1 . Consider the Hessian w.r.t W 1 , we have M 1 M T 1 = W T 2 W 2 and N T 1 N 1 is an identity matrix. Under Assumption 1, we know that W 2 is a row vector, which can be denoted as W 2 = [ w 21 , w 22 , ..., w 2 d 1 ] . Then we have

Similarly, consider the Hessian w.r.t. W 2 , we have that M 1 M T 1 is an identity matrix and N T 1 N 1 = W 1 W T 1 . Therefore,

Hence we have related the uniformity of diagonal Hessian to that of weight matrices. In the detailed analysis, for both GD and Adam, we can prove that W 1 converges to an approximately rank 1 matrix. The following lemma allows us to use this rank 1 structure to compute R med , 1 and R med , 2 .

Lemma 31. Suppose W 1 ∈ R d × d and W 2 ∈ R 1 × d have the following structure:

where u ∈ R d , v ∈ R d , R 1 ∈ R d × d , R 2 ∈ R 1 × d and that

Then we have

Proof. Let's first consider R med , 1 . we have

which yields

Similarly, for R med , 2 . We have that

which yields

That means

H Auxiliary lemmas

Lemma 32. Let A = 1 m Y X T , Λ xx := 1 m XX T , g ( t ) k = ∇ W k L ( W ( t ) ) , k = 1 , 2 . Denote ˜ A ( t ) , ˜ Λ ( t ) xx and ˜ g ( t ) k , k = 1 , 2 as the corresponding batch versions at time t . Let M ( t ) 1 = max i,j ∣ ∣ ∣ W ( t ) 1 [ i, j ] ∣ ∣ ∣ and M ( t ) 2 = max i ∣ ∣ ∣ w ( t ) 2 i ∣ ∣ ∣ . Under Assumption 3, we have with probability at least 1 -1 d , for ∀ t ≤ T and ∀ i, j ∈ [ d ] ,

Proof. By Assumption 3 and Chebyshev's inequality, we have for fixed i, j ∈ [ d ] and t ≤ T ,

Applying the union bound gives us

which gives us with probability at least 1 -1 d , for ∀ t ≤ T, ∀ i, j ∈ [ d ] ,

Now we are ready to bound ˜ g ( t ) k -g ( t ) k for k = 1 , 2 and t ≤ T .

Note that for all t ≤ T and ∀ i ∈ [ d ] ,

Then we have with probability at least 1 -1 d , for all t ≤ T and ∀ i ∈ [ d ] ,

Combining with ˜ g ( t ) 1 -g ( t ) 1 = W ( t ) T 2 ( W ( t ) 2 W ( t ) 1 ( ˜ Λ ( t ) xx -Λ xx ) -( ˜ A ( t ) -A )) , we get that with probability at least 1 -1 d , for all t ≤ T and ∀ i, j ∈ [ d ] ,

Similarly, note that ˜ g ( t ) 2 i -g ( t ) 2 i = ( W ( t ) 2 W ( t ) 1 ( ˜ Λ ( t ) xx -Λ xx ) -( ˜ A ( t ) -A )) W ( t ) T 1 , we then have that with probability at least 1 -1 d , for all t ≤ T and ∀ i, j ∈ [ d ] ,

Lemma 33. Consider two sequences { a ( t ) } t ≥ 0 , { b ( t ) } t ≥ 0 , which satisfy

Suppose ∀ τ ≤ t : ∣ ∣ b ( t ) ∣ ∣ ≤ B , then for any glyph[epsilon1] > 0 , the following truncated version

with H ≥ 1 1 -β log B glyph[epsilon1] = ˜ Ω ( 1 1 -β ) satisfies

Proof. We have that

To make it less than glyph[epsilon1] , it suffices to choose H ≥ log( glyph[epsilon1] B ) / log β .

Since β ∈ (0 , 1) , we know that log β ≤ β -1 < 0 . We also have log glyph[epsilon1] B < 0 . Then it suffices to choose

Lemma 34. Suppose a, b, c, e a , e b , e c ∈ R , b > 0 , c > 0 satisfy b + e b + e c > 0 , | e a | ≤ δ | a | , | e b | ≤ δb, | e c | ≤ δ 2 c 2 with 0 < δ glyph[lessmuch] 1 , then we have

Proof. We have

Define R := q 1 + q 2 . The term | q 1 | can be bounded by

where | q 3 | can be bounded by

Here the denominator of ( i ) uses b + e b ≥ b (1 -δ ) > 0 and √ x + y ≥ √ x -√ | y | when x ≥ 0 , x + y ≥ 0 .

Now let's bound | q 4 | . If e c > 0 , we have e c = | e c | and | q 4 | ≤ √ e c √ e c = 1 since b + e b ≥ b (1 -δ ) > 0 .

If e c ≤ 0 , note that b + e b + e c > 0 , we have | e c | < b + b e ≤ b (1 + δ ) , which yields | q 4 | ≤ √ | e c | √ b = O (1) . Combining the above bounds give us | q 1 | ≤ | q 3 | + δ | q 4 | = O ( δ ) .

On the other hand, | q 2 | can be bounded by

Then | R | ≤ | q 1 | + | q 2 | = O ( δ )

Lemma 35. Suppose X 1 , X 2 , ..., X d are i.i.d Gaussian with mean 0 and variance σ 2 , then for 0 < δ < 1 e , we have with probability at least 1 -δ ,

for some C 1 , C 2 > 0 .

Proof. It suffices to assume that σ 2 = 1 and prove that w.p. at least 1 -δ , max 1 ≤ i ≤ d X 2 i ≥ C 1 log d -C 2 log log 1 δ .

First, by the lower bound of Gaussian tail, there exists α, β > 0 such that P ( | X i | > x ) = 2 P ( X i > x ) ≥ αe -βx 2 for x ≥ 0 . Then by i.i.d., we have

where the last inequality uses 1 -x ≤ e -x for x ∈ [0 , 1] . Let exp( -dαe -βx 2 ) = δ , we get that w.p. at least 1 -δ ,

Then we have w.p. at least 1 -δ ,

Lemma 36. Suppose X 1 , X 2 , ..., X d are i.i.d Gaussian with mean 0 and variance σ 2 , then we have with constant probability,

Proof. It suffices to assume that σ 2 = 1 and prove that with constant probability, 1 d ∑ d i =1 1 | X i | ≤ O (log d ) .

Consider X i for some fixed i . Since X i ∼ N (0 , 1) , we have P ( | X i | ≤ t ) ≤ 2 t √ 2 π . Then we know that with probability at least 1 -Θ ( d -1 ) , | X i | ≥ C d for some C > 0 . Then by union bound, with constant probability, ∀ i ∈ [ d ] : | X i | ≥ C d .

Now we split the interval [ C d , 1] into several subintervals I k = { i : | X i | ∈ [2 -k -1 , 2 -k ] } for k = 0 , 1 , ..., glyph[ceilingleft] log 2 d C glyph[ceilingright]1 . Let p k = P ( | X i | ∈ [2 -k -1 , 2 -k ]) , we know that |I k | ∼ Binomial ( d, p k ) and p k ≤ C 1 · 2 -k -1 . Then by the concentration of binomial variables, we have w.p. at least 1 -d -p for p > 0 , |I k | = O ( dp k + √ dp k log d +log d ) = O ( d · 2 -k -1 + √ d · 2 -k -1 log d +log d ) . Then we have

Therefore, with constant probability,

which means with constant probability, d i =1 | X i | = O (log d ) .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to understand why Adam works better than SGD for language model training from the perspective of local geometry of the training loss around the algorithm iterates.\n\nSTRENGTHS:\nThe problem that this paper considers is indeed a very important question, and this work sets out to answer the question from first principles.\nThis paper has many good intuitions. This paper also has some partial theory that backs up their empirical observation. Given their empirical observation is quite hard to theoretically analyze, I think their theory is quite valuable although it only shows a weak result for a very specific model.\n\nWEAKNESSES:\n- Many plots have too small fonts. Please adjust the font sizes for better readability.\n- In general, the plots have legends and axes that are quite hard to read, please adjust them.\n\n- In Figure 8, how can the spectrum on the right plot approximately low-rank? All singular values are greater than 0.5 and compared to the left plot, I wouldn't consider this approximate low rank. Please explain\n\n- Regarding Section A.6, the plots seem a bit inconclusive. Given how noisy the plots are, I don't see this as alignment. Do you have any intuitions behind why the layer gradients align with diagonal of Hessian?\n\n- Are the results established in Theorem 1 empirically tight? I think the experiments on this very simple setting to see if your empirical observations hold true for this very simple setting is required. In particular, for this two layer MLP, do you see lower R-value for Adam than SGD-M?\n\n- In Figure 9, I see that the observation doesn't hold for the right singular vectors.. This makes the observation quite brittle in my opinion. Could you explain why it doesn't hold for the right singular vectors?\n\n- In general, this paper lacks empirical investigation on smaller/toy models. I think to what extent this observation holds for simpler models would help readers comprehend how universal this observed phenomenon is.\n\n- Given that the language model training is done mostly with the cross entropy loss, I'm curious whether the conclusion you made in Theorem 1 similarly holds true for the logistic loss?\n\n- In the Gaussian initialization, I see the variances of the different layers are chosen quite carefully. Also the variance looks quite small. Could you justify this assumption as to why making these assumptions is practical. Also, could you explain intuition behind why these assumptions are required for the proof?\n\n- In the Theorem statement, how big are $T_1^{sgd}, T_2^{sgd}$ compared to $T_1^{adam}, T_2^{adam}$? Are they comparable? Or as you observed in the empirical results, are $T_1^{sgd}, T_2^{sgd}$ provably larger than $T_1^{adam}, T_2^{adam}$?\n\n- Could you share some intuitions as to why training behavior of vision models (CIFAR 10, IMAGENET) is so different than that of language models? I don't find the explanation that it has to do with transformer architecture quite conclusive...\n\n- Minor comment: It would be helpful if the proof sketch of Theorem 1 appears in the main text. I think it has many good intuitions.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to study the connections between an optimization algorithm and the geometric properties observed during the training of a neural network with that algorithm. The authors argue that when analyzing these connections, it is important to consider the local iterates. To address this, they propose the utilization of a metric called $R^{OPT}_{med}$, which involves the local condition number of the Hessian calculated at the iterates. The authors conduct experiments to illustrate their findings that fast convergence is often associated with low statistics, and adaptive methods such as Adam often has low statistics. They proves that for some small theoretical setting, Adam has better statistic than SGD with some probability.\n\nSTRENGTHS:\nStudying the influence of algorithms on the loss geometry is interesting. This paper provides some insights on the uniformity of diagonal geometry by the $R^{OPT}_{med}$ metric where the values found by Adam are smaller than those found by SGDM in the empirical results. \n\nThe theoretical setting shows that the diagonal of loss Hessian for Adam has good uniformity while the diagonal of loss Hessian for SGD+M is less uniform.\n\nWEAKNESSES:\nWhile there may exist a correlation between the proposed statistic and the optimization algorithm, the reasons behind the favorable behavior of these algorithms remain uncertain. Is it better due to the uniformity or is the diagonal uniformity a byproduct of Adam's inherent algorithmic properties that are unrelated to its overall effectiveness? As a result, it is more important to demonstrate the contribution of the proposed statistic to achieving successful convergence rather than focusing on the relationship between Adam and the the diagonal uniformity.\n\nWhy the large batch setting are needed to prove the correlation of Adam and diagonal uniformity effectiveness? It is not clear how the algorithms behave in other settings and how high the probability can be in the theoretical analysis.\n\nQUESTIONS:\nPlease see above.\n\n---\nI thank the authors for your rebuttal. My recommendation remains the same.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis study makes the comparison between the adaptive optimization method (Adam) and the non-adaptive one (stochastic gradient method with momentum) through a lens of ``uniformity'' of diagonal components of the Hessian denoted by $R_{\\rm med}^{\\rm OPT}(t)$. The comparison is conducted experimentally for various deep learning models such as BERT-small and theoretically for the two-layer linear neural networks under several assumptions. Both comparisons conclude that the uniformity will be smaller when using the Adam, which results in faster convergence.\n\nSTRENGTHS:\n- The experimental comparison is conducted carefully.\n- The theoretical results have a novelty, which gives new insight into analyzing the implicit bias of various optimization methods.\n\nWEAKNESSES:\n- I couldn't fully understand the motivation for employing $R_{\\rm med}^{\\rm OPT}(t)=\\frac{{\\rm max}\\{|H_{ii}^{(t)}|\\}}{{\\rm median}\\{|H_{ii}^{(t)}}|\\}$ as for the notion of uniformity. I realize this value is more stable than the condition number; as explained in footnote 1, many types of variants could be considered. Indeed, in the experiment, there are some cases where $\\frac{R_{\\rm med}^{\\rm SGD}(t)}{R_{\\rm med}^{\\rm Adam}(t)}$ is smaller than 1.\n- In the theoretical comparison (Theorem 1), there are several unclear points to me. First, the training dynamics will change by the selection of hyperparameter and affects to $R_{\\rm med}^{\\rm OPT}(t)$, but some parameters (such as $\\alpha$, $\\sigma$) can be taken in a different order. Will it be a fair comparison? Moreover, I think it will be more helpful for readers if the authors clarify the dependence of $T_{\\rm OPT,(1,2)}$ on other parameters.\n- (minor) The reference to Figure 4. seems wrong. The authors write ``see Section 4.1'' in the caption, but it is explained in Section 6.\n\nQUESTIONS:\n- Although the initialization order required in Theorem 1 seems to be very small, can the authors import the theoretical result or its insights into the practical settings? If not, what is the difficulty?\n- The setting in Section 5 requires $d_1=d_0=d$, but can the overparameterized settings (i.e., $d_1\\ge d_0$ holds) be treated in the same way?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to study the interaction between optimizers and the local loss landscape. The authors introduce the notion of $R_{\\mathrm{med}}^{\\mathrm{OPT}}$ to characterize the uniformity of the Hessian diagonal. Through extensive experiments, they show that, compared with SGD,  Adam biases the trajectory towards regions with higher Hessian diagonal uniformity, which, they argue, contributes to faster optimization. As a theoretical support, the authors demonstrate that Adam indeed possesses this property in a two-layer linear network setting.\n\nSTRENGTHS:\n1. This paper is well-written, with extensive experiments and solid theoretical analysis.\n\n2. The interaction between optimizers and local loss landscape is an important and interesting topic. The findings of this paper may inspire future studies on developing a deeper understanding of the success of adaptive methods.\n\nWEAKNESSES:\nI was wondering whether $R_{\\mathrm{med}}^{\\mathrm{OPT}}$ can characterize the degree of the loss landscape being ill-conditioned. While it is true that $R_{\\textrm{med}}^{\\textrm{OPT}}$ conveys the similar message as the condition number when the Hessian is dominated by the diagonal ($\\nabla^2 \\mathcal{L}(\\theta)\\approx\\mathrm{diag}(\\nabla^2 \\mathcal{L}(\\theta))$,  it remains unclear whether $\\nabla^2 \\mathcal{L}(\\theta)\\approx\\mathrm{diag}(\\nabla^2 \\mathcal{L}(\\theta))$ holds throughout the training trajectory.  Can the authors provide some empircial evidence?\n\nQUESTIONS:\nRegarding the comparison of different training runs, specifically run 1 (using Adam throughout), run 2 (using Adam initially and then switching to SGDM halfway), and run 3 (using SGDM throughout), Figure 1 illustrates that run 2 achieves a considerably lower loss level compared to run 3. However, I am interested in the gap in the final loss between run 1 and run 2.\n\n If this gap is acceptable, it raises a possibility: Could we design an algorithm that employs Adam during the initial phase and then seamlessly switches to SGDM halfway through the training process? Such an approach may offer the potential benefit of memory savings.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a new explanation of why adaptive gradient methods perform better than SGD with momentum though extensive experiments and theoretical analysis. The key insight is that adaptive gradient methods especially Adam bias toward solutions with uniform Hessian diagonal values, and this property may contribute to faster convergence. Theoretical analysis is conducted on a simple setting with 2 layer neural network, it is proven that Adam guarantees to converge to solution with more uniform Hessian diagonal values than SGD.\n\nSTRENGTHS:\nThe insight that Adam bias toward solutions with uniform Hessian diagonal values is an interesting and novel observation and this measure is easy to compute compared with Hessian singular value-based metrics. The observation is supplemented with comprehensive experiment results and an analysis in simplified setting of neural network, which make the work pretty solid in both theory and empirical sides.\n\nWEAKNESSES:\nTo my understanding, the weaknesses have two folds:\n\nThe experiments in the paper mainly focus on settings where Adam outperforms SGD, in which more uniformity in Hessian diagonals is observed with Adam. However, it lacks a comparison in settings where SGD outperforms Adam. It remains a question whether uniformity is still a good measure of performance in this case. In other words, would SGD iterates observe more uniformity in Hessian diagonals instead?\n\nIn the explanation of the paper, Adam biases towards solutions with more uniformity in Hessian diagonals, and this contributes to faster convergence. Thus, Adam converges faster than SGD. The bias in Adam has been well experimented with and studied. However, it is not very clear why uniformity in Hessian diagonals implies faster convergence. The paper does not provide a theory on it, and the experimental support of this claim is not strong. Additionally, intuitively, it is hard to connect uniformity with faster convergence since regions with more uniformity are allowed to have either large curvature or small curvature. The folklore intuition on a flat region encourages better optimization and generalization does not seem to hold for this measure.\n\nQUESTIONS:\nCould the authors provide more intuition on how uniformity in Hessian diagonals can induce faster optimization?\n\nLIMITATIONS:\nLimitations are adequately discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-20"}