How Does Adaptive Optimization Impact Local Neural Network Geometry?

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.

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

From a theoretical perspective, adaptive methods are well understood in the traditional context of convex optimization. For instance, Duchi et al. 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 .

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 . For instance, Cohen et al. 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 2/(step size). The viewpoint from convex optimization, where a loss function has some (potentially) non-uniform but fixed underlying geometry that we must adapt to, is thus insufficient for neural networks, since the choice of optimization algorithm can actually interact with and influence the observed geometry significantly.

To provide another example of this interactive phenomenon, we consider the following experiment. On the same network training loss function $f$, we run stochastic gradient descent with momentum (SGD+M) and Adam to obtain two different trajectories. We select an iterate $x_{\text{Adam}}$ from the Adam trajectory and an iterate $x_{\text{SGD}}$ from the SGD trajectory, such that $f(x_{\text{Adam}})=f(x_{\text{SGD}})$. We then run SGD+M twice, once from $x_{\text{Adam}}$ and once from $x_{\text{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(a), there is a noticeable difference in performance, and running SGD+M from $x_{\text{Adam}}$ achieves lower loss than running SGD+M from $x_{\text{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)?

This statistic thus measures the uniformity of the diagonal of the Hessian, where a smaller value of $R_{\text{med}}^{\text{OPT}}(t)$ implies that the Hessian has a more uniform diagonal. It can also be viewed as a stableConsider the case where one parameter has little impact on the loss, then the second derivative w.r.t. this parameter is almost zero, making $\frac{\max\{|H^{(t)}_{ii}|\}_{i=1}^{d}}{\min\{|H^{(t)}_{ii}|\}_{i=1}^{d}}$ infinity. So we consider median which is more stable. variant of the condition number. Instead of eigenvalues, we choose diagonal entries because adaptive methods used in practice are coordinate-wise, which can be viewed as the diagonal scaling approaches.Recall that the main theoretical bound in the original Adagrad paper is in terms of the diagonal scaling. Hence we believe the diagonal of Hessian is more relevant than the spectrum. As a supplementary result, in Appendix E, 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 etc. could yield to faster convergence rate for (S)GD . As for (diagonal) adaptive methods (e.g. Adagrad), they were original designed to adapt to the nonuniform diagonal geometry. Intuitively, a smaller $R_{\text{med}}^{\text{OPT}}(t)$, which implies more uniform diagonal geometry, could lead to faster convergence.

Armed with this statistic, we make the following contributions:

On a wide variety of neural network transformer architectures and language modeling datasets, we conduct experiments to compare how $R_{\text{med}}^{\text{Adam}}(t)$ and $R_{\text{med}}^{\text{SGDM}}(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_{\text{med}}^{\text{Adam}}(t)$ values that are significantly smaller than the $R_{\text{med}}^{\text{SGDM}}(t)$ values found by SGD+M. We show a simple example of this phenomenon in Figure 1(b). 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. As an aside, we observe that larger improvements in performance of Adam over SGD+M are correlated with larger gaps between $R_{\text{med}}^{\text{Adam}}(t)$ and $R_{\text{med}}^{\text{SGDM}}(t)$. This suggests that a region where the Hessian diagonal is more uniform is also a region that is more amenable to rapid optimization.

We complement our empirical results with a theoretical analysis of this phenomenon in the simplified setting of large batch Adam and SGD+M, on 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_{\text{med}}^{\text{Adam}}(t)=1\pm o(1)$ but $R_{\text{med}}^{\text{SGDM}}(t)=\Omega(\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.

Related work

Existing analyses of adaptive methods. The vast majority of prior theoretical work on adaptive methods has focused on the blackbox setting . 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, . 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 , 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. find a so-called catapult effect of initial learning rate on the training trajectory of SGD and related loss curvature. Cohen et al. 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. leave open the problem of whether similar interactive phenomena occur in algorithms that are not GD, including adaptive methods.

Overview of results and setup

As is mentioned in Section 2, existing work on adaptive algorithms has mainly focused on black-box analysis assuming some global worst-case parameters. However, these global bounds can be extremely bad in complicated deep learning models, as is discussed in Section 1. To see this, we initialized a transformer modelhttps://pytorch.org/tutorials/beginner/transformer_tutorial.html with default initialization in Pytorch but chose a large gainThis refers to the gain parameter in some commonly used initialization functions of Pytorch, e.g. torch.nn.init.xavier_uniform_()., and computed the smoothness parameter (denoted as $l$) and the condition number (denoted as $\kappa$) 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 $\kappa$ ($l\geq 10^{7}$ and $\kappa\geq 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 that when two optimization algorithms are used, they will have the same geometry in local trajectory. In particular, although in theory, adaptive algorithms can yield to a convergence rate with better dependency on certain local geometry of the function comparing to SGD (with momentum), it could still be the case that the local geometry along the trajectory of adaptive algorithm can be much worse than that of SGD (with momentum).

That motivates us to study the local geometry, especially that obtained by adaptive methods comparing 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)?

2 Overview of the experiments

As is discussed in Section 1, we consider $R_{\text{med}}^{\text{OPT}}(t)$ defined in eq. (1) as a measurement of the uniformity of the diagonal of the loss Hessian. We conduct experiments on different NLP tasks to examine $R_{\text{med}}^{\text{OPT}}(t)$, as in language models, adaptive methods have shown significantly faster convergence than SGD (with momentum). The details of these experiments will be shown in Section 4. To explore potential different patterns of different layers, 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:

When we train the neural network using Adam, the uniformity of diagonal geometry, measured by $R_{\text{med}}^{\text{OPT}}(t)$ is smaller than that when we train using SGD+M from the same initialization, except for first several layers.

Table 1 shows a typical example of $R_{\text{med}}^{\text{Adam}}(t)$ compared to $R_{\text{med}}^{\text{SGDM}}(t)$ on a sentence classification task using BERT-small (see Section 4.1 for details). We repeated the experiments for 12 times starting from the same initialization. Table 1 shows the averaged $R_{\text{med}}^{\text{Adam}}(t)$ and $R_{\text{med}}^{\text{SGDM}}(t)$ in some randomly selected layers (except for the first several). We also report the averaged $\frac{R_{\text{med}}^{\text{SGDM}}(t)}{R_{\text{med}}^{\text{Adam}}(t)}$ and their standard deviations in the brackets.$R_{\text{med}}^{\text{SGDM}}(t)$ values in Table 1 for most layers are roughly 1.4 to 2 times $R_{\text{med}}^{\text{Adam}}(t)$ in corresponding layers. In practice, it can be considered significant because it might imply 1.4 to 2 times faster convergence. Figure 2 shows the corresponding training losses of one in these 12 experiments.

To understand this phenomenon in a more principled point of view, we also provide a formal proof of the statement in a simplified setting: large batch Adam and SGD+M on a two-layer linear network. Although simple, the choice of two-layer linear network to understand learning dynamics is common in prior works (e.g. ). Section 3.3 below describes the theoretical setup.

3 Setup of the theoretical analysis

Let $[d]=\{1,2,...,d\}$. We use $\|\cdot\|_{2}$ to denote the $l_{2}$ norm of a vector, and $\|\cdot\|_{F}$ to denote the Frobenius norm of a matrix. Let $\langle\cdot,\cdot\rangle$ be the Euclidean inner product between vectors or matrices. Let $\mathcal{N}(\mu,\sigma^{2})$ be the one-dimensional Gaussian distribution with mean $\mu$ and variance $\sigma^{2}$. For a scalar (vector, matrix) $A$ which evolves over time, we use $A^{(t)}$ to denote its value at time $t$.

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

$\forall i,j:w_{2i}^{(0)}\sim\mathcal{N}(0,\frac{1}{d^{2\alpha}}),W_{1}^{(0)}[i,j]\sim\mathcal{N}(0,\frac{1}{d^{4\alpha}})$ are independently initialized with sufficiently large $\alpha>0$.

where $\eta$ is the learning rate, $\beta,\beta_{1},\beta_{2}$ are momentum parameters, and $\xi$ is for numerical stability. All operations on vectors are element-wise.

The uniformity of diagonal geometry

As is mentioned in Section 3.2, we computed $R_{\text{med}}^{\text{OPT}}(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.

We fine-tuned BERT-small on the IMDB dataset : the task is to classify whether movie reviews are positive or negative.https://huggingface.co/docs/transformers/v4.16.2/en/training The momentum parameter $\beta$ in SGD was set as 0.9. The two momentum parameters $(\beta_{1},\beta_{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_{\text{med}}^{\text{Adam}}(t)$ and $R_{\text{med}}^{\text{SGDM}}(t)$.

We trained a Seq2Seq network that uses Transformer to solve a machine translation task on Multi30k (CC BY-NC-SA 4.0): this task is to train a German to English translation model.https://pytorch.org/tutorials/beginner/translation_transformer.html The momentum parameter $\beta$ in SGD was set as 0.9. The two momentum parameters $(\beta_{1},\beta_{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(a) shows the training losses for one among them. Table 2(a) shows the averaged $R_{\text{med}}^{\text{Adam}}(t)$, $R_{\text{med}}^{\text{SGDM}}(t)$ and $\frac{R_{\text{med}}^{\text{SGDM}}(t)}{R_{\text{med}}^{\text{Adam}}(t)}$ (with standard deviation in the brackets) in some randomly selected layers.

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(b) shows the training losses for one among them. Table 2(b) shows the averaged $R_{\text{med}}^{\text{Adam}}(t)$, $R_{\text{med}}^{\text{SGDM}}(t)$ and $\frac{R_{\text{med}}^{\text{SGDM}}(t)}{R_{\text{med}}^{\text{Adam}}(t)}$ (with standard deviation in the brackets) of the same 10 layers as in Table 2(a).To prevent $R_{\text{med}}^{\text{OPT}}(t)$ from getting too large due to tiny median, we added an additional term $0.001\max\{|H^{(t)}_{ii}|\}_{i=1}^{d}$ to the denominator of eq. (1) when computing.

3 How the (adaptive) gradient aligns with diagonal of loss Hessian

In this section we present the uniformity of diagonal geometry of adaptive methods from another perspective. Denote $H_{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_{\text{adapt},i}$ to represent the $i$-th component of the adaptive gradient. Through experiments on language models, we find that $|g_{\text{adapt},i}|$ for different $i$ are quite uniform and do not align with $|H_{ii}|$ as the true gradient $|g_{i}|$ does.

4 Summarization of the empirical results and discussion

Overall, through extensive experiments on language models, we demonstrate that starting from the same initialization, the $R_{\text{med}}^{\text{OPT}}(t)$ values found by Adam are smaller than those found by SGD+M, except for the first several layers. This suggests that Adam is biased towards a region with more uniform diagonal Hessian than SGD+M.

We observe that on random dataset, SGD+M plateaus after about 400 steps and thus converges much slower when compared to Adam than on real dataset (see Figure 3(a) and Figure 3(b)). On the other hand, the gaps of $R_{\text{med}}^{\text{SGDM}}(t)$ and $R_{\text{med}}^{\text{Adam}}(t)$ are more significant on random data than on real data (see Table 2(a) and Table 2(b)) 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_{\text{med}}^{\text{OPT}}(t)$ are gradually closed after switching (see Figure 8 and Table 8). Hence we find a positive correlation between fast convergence and the uniformity of diagonal of loss Hessian, suggesting that a region with more uniform diagonal of Hessian is also a region that is more amenable to fast optimization. In Appendix A we study other adaptive algorithms (Adagrad, RMSprop and AMSGrad) and get similar observation: all these adaptive methods converge faster than SGD or SGD+M and also bias the trajectory to a region with smaller $R_{\text{med}}^{\text{OPT}}(t)$, suggesting that the uniformity of diagonal Hessian might be a universal mechanism (partially) explaining the faster optimization of adaptive algorithms than SGD (with momentum).

Considering the fact that our comparison between $R_{\text{med}}^{\text{Adam}}(t)$ and $R_{\text{med}}^{\text{SGDM}}(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, thus giving the appearance that it induces better geometry. To rule out this possibility, in Appendix A.3 we add a comparison of our measurements $R^{\text{Adam}}_{\text{med}}(t)$ and $R^{\text{SGDM}}_{\text{med}}(t^{\prime})$, where $t,t^{\prime}$ are picked such that $t$th Adam iterate and $t^{\prime}$th SGD+M iterate have the same training loss. The results (in Table 7) show that $R^{\text{Adam}}_{\text{med}}(t)<R^{\text{SGDM}}_{\text{med}}(t^{\prime})$ 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).

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 13(a) and Table 9 suggest that in this case, the relationship between $R_{\text{med}}^{\text{OPT}}(t)$ and convergence speed still holds: Adam converges faster than SGD+M and in most of the layers except for the first several, $R_{\text{med}}^{\text{Adam}}(t)$ values are smaller than $R_{\text{med}}^{\text{SGDM}}(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_{\text{med}}^{\text{Adam}}(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.

Although in this paper we focus on language models where Adam shows significant fast convergence, we also add supplementary results in Appendix A.8 on image tasks where SGD+M performs better. On a residual network trained on CIFAR-10, we observed that Adam did not converge faster than SGD+M (see Figure 13(b)) and in the meantime, $R_{\text{med}}^{\text{Adam}}(t)$ values were no longer smaller than $R_{\text{med}}^{\text{SGDM}}(t)$ during training (see Table 10). This reveals the connection between the local diagonal geometry and the convergence speed from another perspective. That is, when the diagonal of Hessian of Adam is not more uniform than SGD+M, its convergence speed is not better, either. In summary, all the observations on language and image tasks together suggest a positive correlation between the uniformity of diagonal Hessian and fast optimization.

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.

Since the weights and Hessians in different layers may have different magnitudes, we compute the $R_{\text{med}}^{\text{OPT}}(t)$ layer by layer. We denote $R_{\text{med},k}^{\text{SGDM}}(t)$ (resp. $R_{\text{med},k}^{\text{Adam}}(t)$) as the $R_{\text{med}}^{\text{OPT}}(t)$ found by SGD+M (resp. Adam) w.r.t. $W_{k}$ at time $t$ where $k=1,2$.

Under Assumption 1, 2 and 3, consider the weights $\left\{W^{(t)}_{\text{SGD}}\right\}_{t\geq 0}$ (resp. $\left\{W^{(t)}_{\text{Adam}}\right\}_{t\geq 0}$) obtained by SGD+M (resp. Adam) defined in (3).

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

The low rank structure of weight matrices and uniformity of leading singular vectors

The proof sketch in Appendix B highlights one crucial intuition of Theorem 1: After $T_{\text{SGD},1}$ (resp. $T_{\text{Adam},1}$) steps, $W_{1}$ of SGD+M (resp. Adam) becomes an approximately rank-1 matrix. Consider the left singular vector $\boldsymbol{u}:=[u_{1},u_{2},...,u_{d}]^{T}$ which corresponds to the leading singular value $\sigma_{1}$. We can show that the distribution of $u_{1}^{2},u_{2}^{2},...,u_{d}^{2}$ 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 F for more details.

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 have approximately low rank structures and for 10 of them, $R_{u}$ values (defined in (B)) obtained by Adam are smaller than those found by SGD+M. Figure 5 shows the result on one typical layer. Results of more layers can be found in Appendix A.5.

Although in multi-layer nonlinear neural networks, the connection between diagonal of loss Hessian and the weight matrices is more complicated and $R_{\text{med},2}^{\text{OPT}}(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.

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. In fact, people conjecture that the key-value-query structure in language models can be approximated by a three-layer linear network. Hence the generalization to multi-layer networks might provide more connection to real deep models and could be an interesting and challenging future direction. Moreover, it is also possible to relax our large-batch assumption (Assumption 3) and prove similar results in the general stochastic setting.

References

Appendix A More experiments of the uniformity of diagonal geometry

In this section, we present the $R_{\text{med}}^{\text{OPT}}(t)$ values defined in eq. (1) obtained by SGD and Adagrad on a language modeling taskhttps://pytorch.org/tutorials/beginner/transformer_tutorial.html. 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 both Wikitext-2 (CC BY-SA 3.0) and random dataset (generating random integers as 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.

Figure 6(a) shows the training losses on real dataset (wikitext-2). Table 3 (resp. Table 4) shows the $R_{\text{med}}^{\text{OPT}}(t)$ for Adagrad and SGD under uniform (resp. normal) initialization with different gains.

A.1.2 Experiments on random dataset

Figure 6(b) shows the training losses on random dataset and Table 5 shows the $R_{\text{med}}^{\text{OPT}}(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 are conducted on the translation task described in Section 4.1. The learning rates we used were 0.000025 for RMSprop, 0.0005 for AMSGrad and 0.03 for SGD+M. Both RMSprop and SGD+M used momentum parameter 0.9. The two momentum parameters $(\beta_{1},\beta_{2})$ of AMSGrad are $(0.9,0.98)$. Figure 7 shows the training losses and Table 6 shows the corresponding $R_{\text{med}}^{\text{OPT}}(t)$.

A.3 Comparison conditioned on the same loss

In this section, we compare $R_{\text{med}}^{\text{SGDM}}(t)$ and $R_{\text{med}}^{\text{Adam}}(t)$ conditioned on the same training loss. More precisely, we make comparison of $R^{\text{Adam}}_{\text{med}}(t)$ and $R^{\text{SGDM}}_{\text{med}}(t^{\prime})$, where $t,t^{\prime}$ are picked such that $t$th Adam iterate and $t^{\prime}$th SGD+M iterate have the same training loss. The details of the tasks are described in in Section 4.1. Table 7 shows the results of $R^{\text{Adam}}_{\text{med}}(t)$ and $R^{\text{SGDM}}_{\text{med}}(t^{\prime})$ in some layers.

A.4 Experiments of switching from SGD to Adam

In this section we describe another learning schedule: the “Adam after SGD” schedule, where we switched from SGD to Adam in the middle to see whether the loss and $R_{\text{med}}^{\text{OPT}}(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 8 and Table 8, both the loss gap and the gap of $R_{\text{med}}^{\text{OPT}}(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 9 shows the examples of layers with or without the low rank structure.

We then studied the uniformity of leading singular vectors of these 12 layers, i.e. computed $R_{u}$ and $R_{v}$ defined in (B) and the second remark of Section 6. The observation is that for 10 out of these 12 layers, $R_{u}$ values of Adam are smaller those of SGD, which implies the uniformity of leading left singular vectors of Adam. However, we did not observe significant uniformity for Adam in terms of leading right singular vectors ($R_{v}$). The second remark of Section 6 discusses possible reasons.

Figure 10 shows how $R_{u}$ and $R_{v}$ changed over time in some layers.

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

In this section, we present more empirical results on how the (adaptive) gradient aligns with diagonal of loss Hessian. The detailed setup is described in Section 4.3.

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 11.

A.6.2 SGD with momentum vs. Adam

Figure 4 presents the comparison between Adam and SGD with momentum on the sentence classification task using BERT-small. Here we add more results of the comparison of these two algorithms on the translation task described in Section 4.1. Again, we select one layer as an example, as is shown in Figure 12.

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 $\beta$ in SGD was set as 0.9. The two momentum parameters $(\beta_{1},\beta_{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 13(a) shows the training losses and Table 9 shows the values of $R_{\text{med}}^{\text{Adam}}(t)$, $R_{\text{med}}^{\text{SGDM}}(t)$ and $\frac{R_{\text{med}}^{\text{SGDM}}(t)}{R_{\text{med}}^{\text{Adam}}(t)}$ in some randomly selected layers.

A.8 Results on image tasks

We trained a ResNetWe borrowed the implementation here https://pytorch-tutorial.readthedocs.io/en/latest/tutorial/chapter03_intermediate/3_2_2_cnn_resnet_cifar10/ and replace the “layers” array with . on CIFAR-10 dataset and compared the convergence speed and $R_{\text{med}}^{\text{OPT}}(t)$ of SGD+M and Adam. The momentum parameter $\beta$ in SGD was set as 0.9. The two momentum parameters $(\beta_{1},\beta_{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 13(b) shows the training losses and Table 10 shows the values of $R_{\text{med}}^{\text{Adam}}(t)$, $R_{\text{med}}^{\text{SGDM}}(t)$ and $\frac{R_{\text{med}}^{\text{SGDM}}(t)}{R_{\text{med}}^{\text{Adam}}(t)}$.

Appendix B 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 F, C and D.

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_{2i},w_{22},...,w_{2d}]$. 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

Using the rank 1 structure, we can further simplify $R_{\text{med},1}^{\text{OPT}}(t)$ and $R_{\text{med},2}^{\text{OPT}}(t)$ by

The final step is the detailed analysis of $\boldsymbol{u}^{(t)}$.

Appendix C Analysis of SGD+M

Note that $A=\frac{1}{m}YX^{T}$, $\Lambda_{xx}:=\frac{1}{m}XX^{T}$. Denote $g_{k}^{(t)}:=\nabla_{W_{k}}L(W^{(t)}),k=1,2$. We have that

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\approx-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.

The end of the first phase (denoted as $T_{1}$) is defined as $T_{1}:=\inf\left\{t\geq 0:\exists i,j\in[d]:\left|w_{2i}^{(t)}\right|\geq\frac{1}{d^{\frac{\alpha}{2}}}\text{or }\left|W_{1}^{(t)}[i,j]\right|\geq\frac{1}{d^{\frac{\alpha}{2}}}\right\}$.

By Assumption 2 and the assumption that $\forall j\in[d]:A_{j}>0,A_{j}=\Theta(1)$, at the beginning, w.h.p., $\forall j\in[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”.

For $\epsilon>0$, denote $\epsilon_{0}:=\frac{1}{d^{\frac{1}{4}\alpha-1}}+\epsilon\log\sqrt{\frac{d}{\epsilon}}$. Define $T_{2}:=\inf\left\{t\geq 0:\exists j\in[d]:\left(W_{2}^{(t)}W_{1}^{(t)}\right)_{j}-A_{j}\geq-\sqrt{\epsilon_{0}}\right\}$.

For $\epsilon>0$, we define the “convergence time” $T_{3}:=\inf\left\{t\geq 0:\left\|E^{(t)}\right\|^{2}_{2}\leq\epsilon\right\}$.

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.

Under Assumption 1, 2 and 3, suppose $\sigma\leq\frac{\eta^{3/2}}{d^{\alpha/2+1}}$. By picking $\eta\leq\mathcal{O}\left(\frac{1}{d^{\alpha}}\right)$, we have that when $t=T_{1}$, $\bar{L}\left(W^{(T_{1})}\right)=\Theta(d)$, and that

The following lemma tells us that this approximate rank-1 structure is preserved when $T_{1}\leq t\leq\min\{T_{2},T_{3}\}$.

Under Assumption 1, 2 and 3, suppose $\sigma\leq\frac{\eta^{3/2}}{d^{\alpha/2+1}}$. By picking $\eta\leq\mathcal{O}\left(\frac{\epsilon}{d^{\frac{7\alpha}{4}+4}}\right)$, we have that w.h.p. for $T_{1}\leq t\leq\min\{T_{2},T_{3}\}$,

and $\epsilon_{0}$ is defined in Definition 2. Moreover, when $t=\min\{T_{2},T_{3}\}$, $\bar{L}\left(W^{(t)}\right)=\mathcal{O}(\epsilon_{0}d)$.

The following lemma gives us a more detailed description of $\boldsymbol{u}^{(T_{1})}$.

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

Here $X_{i},i\in[d]$ are i.i.d Gaussian random variables by Lemma 3. To prove the concentration of $\text{median }X_{i}^{2}$, we borrow the Proposition 12 in Chapter 2.3 of . By setting $K=N=d$ in this proposition, we have

That means with high probability, $\text{median }X_{i}^{2}\leq C\sigma^{2}$ for some $C>0$. By Lemma 34 in Appendix G, we know that w.h.p.

Hence we have proved that $R_{\text{med},1}^{\text{SGDM}}(t),R_{\text{med},2}^{\text{SGDM}}(t)\geq\Omega(\log d)$.

C.2 Proof of Lemma 1

In the first phase, $W_{2}W_{1}$ is “small”, and we write the update equations in the following way

The following lemma gives us an explicit formula of $W_{2}^{(t)}$.

Let $\lambda_{1}<\lambda_{2}$ be the two roots of the quadratic equation $x^{2}-2x+1-\frac{\eta^{2}}{(1-\beta)^{2}}\|A\|_{2}^{2}=0$. Pick $\eta<\frac{1-\beta}{\|A\|_{2}}$, then we have that

where $C_{1}=-\frac{W_{2}^{(1)}-\lambda_{2}W_{2}^{(0)}}{\lambda_{2}-\lambda_{1}}$, $C_{2}=\frac{W_{2}^{(1)}-\lambda_{1}W_{2}^{(0)}}{\lambda_{2}-\lambda_{1}}$. $r_{5}^{(t)}$ will be specified in the proof.

We can prove that in the first phase, $r_{5}^{(t)}$ is “small”. More specifically, denote its $i$-th coordinate as $r_{5i}^{(t)}$, and the $i$-th coordinate of $C_{2}$ as $C_{2i}$. Then the following lemmas tell us that $\forall i\in[d],\left|r_{5i}^{(t)}\right|\leq\mathcal{O}\left(\frac{1}{d^{p(\alpha)}}\right)$, where w.h.p. $\mathcal{O}\left(\frac{1}{d^{p(\alpha)}}\right)\ll\min_{i\in[d]}|C_{2i}|$.

We first have the following bounds of $\left|r_{1i}^{(t)}\right|,\left|r_{2i}^{(t)}\right|$ and $\left|r_{5i}^{(t)}\right|$ for $i\in[d]$.

Next we prove upper and lower bounds of $|C_{1i}|$ and $|C_{2i}|$ for $i\in[d]$.

Now we are ready to prove Lemma 1. Lemma 4 tells us that

where $\lambda_{1}=1-\frac{\eta}{1-\beta}\|A\|_{2}$ and $\lambda_{2}=1+\frac{\eta}{1-\beta}\|A\|_{2}$.

Under the conditions of Theorem 1 and pick $\eta\leq\mathcal{O}\left(\frac{1}{d^{\alpha}}\right)$, by Lemma 6 and 7, we know that w.h.p. $\forall t\leq T_{1},\forall 1\leq i\leq d$,

We first prove that $\left|w_{2i}^{(t)}\right|$ reaches $\frac{1}{d^{\alpha/2}}$ for some coordinate $i$ before $\left|W_{1}^{(t)}[k,j]\right|$ for $\forall k,j\in[d]$. To see this, first note that

where $\boldsymbol{v}^{(t)T}=\frac{\eta}{1-\beta}A$ and

For $t\leq T_{1}$, by eq. (7), we get that w.h.p.,

Here we used $\forall i\in[d]:A_{i}=\Theta(1)$ by Assumption 1.

Further we notice that for $t\leq T_{1}$, we have $\forall j\in[d]$,

which yields that $\left|u_{i}^{(t)}v_{j}^{(t)}\right|\leq\mathcal{O}\left(\frac{1}{\sqrt{d}}\right)\left|c^{(t)}u_{i}^{(t)}\right|$. Together with eq. (9) gives us that $\left|w_{2i}^{(t)}\right|$ reaches $\frac{1}{d^{\alpha/2}}$ for some $i\in[d]$ before $\left|W_{1}^{(t)}[k,j]\right|$ for $\forall k,j\in[d]$, i.e. $T_{1}=\inf\left\{t\geq 0:\exists i\in[d]:\left|w_{2i}^{(t)}\right|\geq\frac{1}{d^{\frac{\alpha}{2}}}\right\}$.

Further, we know that at time $T_{1}$, $\left|c^{(T_{1})}u_{i_{0}}^{(T_{1})}\right|=|C_{2i_{0}}|\lambda_{2}^{T_{1}}=\Theta\left(\frac{1}{d^{\alpha/2}}\right)$ for some $i_{0}\in[d]$, which means w.h.p.

This is the length of the first phase. As for $c^{(T_{1})}u_{i}^{(T_{1})}$ and $u_{i}^{(T_{1})}v_{j}^{(T_{1})}$ for other coordinates, we have that w.h.p. $\forall 1\leq i,j\leq d$,

Finally, we consider the loss. Since $\forall j\in[d]:\left(W_{2}^{(T_{1})}W^{(T_{1})}_{1}\right)_{j}-A_{j}=-\Theta(1)$, we know that $\bar{L}\left(W^{(T_{1})}\right)=\Theta(d)$.

C.3 Proof of Lemma 3

C.4 Proof of Lemma 4

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

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

where $r_{3}(t):=\frac{\eta^{2}}{(1-\beta)^{2}}Ar_{1}^{(t-1)T}+\frac{\eta}{1-\beta}\left(r_{2}^{(t)}-r_{2}^{(t-1)}\right)$.

For the equation $x^{2}-2x+1-\frac{\eta^{2}}{(1-\beta)^{2}}\|A\|_{2}^{2}=0$, the roots are $\lambda_{1}=1-\frac{\eta}{1-\beta}\|A\|_{2}$ and $\lambda_{2}=1+\frac{\eta}{1-\beta}\|A\|_{2}$. We have that

where $r_{5}^{(t)}=\sum_{\tau=1}^{t}\lambda_{2}^{-\tau}r_{4}^{(\tau)}$, $C_{1}=-\frac{W_{2}^{(1)}-\lambda_{2}W_{2}^{(0)}}{\lambda_{2}-\lambda_{1}}$ and $C_{2}=\frac{W_{2}^{(1)}-\lambda_{1}W_{2}^{(0)}}{\lambda_{2}-\lambda_{1}}$.

C.5 Proof of Lemma 5

Write $r_{1}^{(t)}=-\beta^{t+1}W_{2}^{(t)T}A+q_{12}^{(t)}+q_{13}^{(t)}+q_{14}^{(t)}$ where $q_{12}^{(t)}=(1-\beta)\sum_{\tau=0}^{t}\beta^{t-\tau}\left(W_{2}^{(\tau)T}-W_{2}^{(t)T}\right)A$, $q_{13}^{(t)}=-(1-\beta)\sum_{\tau=0}^{t}\beta^{t-\tau}W_{2}^{(\tau)T}W_{2}^{(\tau)}W_{1}^{(\tau)}$ and $q_{14}^{(t)}=-(1-\beta)\sum_{\tau=0}^{t}\beta^{t-\tau}Dg_{1}^{(\tau)}$. And write $r_{2}^{(t)}=-\beta^{t+1}AW_{1}^{(t)T}+q_{22}^{(t)}+q_{23}^{(t)}+q_{24}^{(t)}$, where $q_{22}^{(t)}=(1-\beta)\sum_{\tau=0}^{t}\beta^{t-\tau}A\left(W_{1}^{(\tau)T}-W_{1}^{(t)T}\right)$, $q_{23}^{(t)}=-(1-\beta)\sum_{\tau=0}^{t}\beta^{t-\tau}W_{2}^{(\tau)}W_{1}^{(\tau)}W_{1}^{(\tau)T}$ and $q_{24}^{(t)}=-(1-\beta)\sum_{\tau=0}^{t}\beta^{t-\tau}Dg_{2}^{(\tau)}$.

Let’s first try to bound $\left|q_{12}^{(t)}[i,j]\right|$ and $\left|q_{22,i}^{(t)}\right|$. For any $\tau\leq T_{1}$, we have that

and thus $\forall i\in[d]:\left|E_{i}^{(\tau)}\right|=\mathcal{O}(1)$. Then we have for all $i,j\in[d]$,

Then we bound $\left|q_{13}^{(t)}[i,j]\right|$ and $\left|q_{23,i}^{(t)}\right|$. We have for $\forall i,j\in[d]$,

Finally we use Lemma 31 to bound $\left|q_{14}^{(t)}[i,j]\right|$ and $\left|q_{24,i}^{(t)}\right|$. For $t\leq T_{1}$, the $M_{1}^{(t)},M_{2}^{(t)}$ in Lemma 31 are upper bounded by $\frac{1}{d^{\frac{\alpha}{2}}}$. In the theorem we consider the training period before $T_{\text{SGD},2}$ so the time $T$ in Lemma 31 is set as $T_{\text{SGD},2}$. In the following sections, we will prove that $T_{\text{SGD},2}\leq\mathcal{O}\left(\frac{d^{\alpha}\log(\sqrt{d/\epsilon})}{\eta}\right)$. Then by Lemma 31, we have with probability at least $1-\frac{1}{d}$, for $\forall t\leq T_{1}$ and $\forall i,j\in[d]$,

Combining all the above bounds and substituting $\eta\leq\mathcal{O}\left(\frac{1}{d^{\alpha}}\right)$ gives us for $\forall t\leq T_{1}$ and $\forall i,j\in[d]$,

For $t\leq T_{1}$, we have $\forall i,j\in[d]$, $\left|w_{2i}^{(t)}A_{j}\right|\leq\mathcal{O}\left(\frac{1}{d^{\alpha/2}}\right)$ and $\left|\sum_{j=1}^{d}A_{j}W_{1}^{(t)}[i,j]\right|\leq\mathcal{O}\left(\frac{1}{d^{\alpha/2-1}}\right)$, which gives us $\left|r_{1}^{(t)}[i,j]\right|\leq\mathcal{O}\left(\frac{1}{d^{\alpha/2}}\right)$ and $\left|r_{2i}^{(t)}\right|\leq\mathcal{O}\left(\frac{1}{d^{\alpha/2-1}}\right)$. Substituting into eq. (5) and eq. (6) yields that for $t\leq T_{1}$ and $\forall i,j\in[d]$,

Hence for $t\leq\min\left\{\frac{\alpha\log d}{\log(1/\beta)},T_{1}\right\}$, we have $\forall i,j\in[d]$,

Then we know that $T_{1}>\frac{\alpha\log d}{\log(1/\beta)}$ and also get tighter bounds of $\left|W_{1}^{(t)}[i,j]\right|,\left|w_{2i}^{(t)}\right|$ for $t\leq\frac{\alpha\log d}{\log(1/\beta)}$. Now we use these new bounds to analyze $\left|r_{1}^{(t)}[i,j]\right|$ and $\left|r_{2i}^{(t)}\right|$ again.

C.6 Proof of Lemma 6

Since $\lambda_{1}=1-\frac{\eta}{1-\beta}\|A\|_{2},\lambda_{2}=1+\frac{\eta}{1-\beta}\|A\|_{2}$, and note that $\|A\|_{2}=\Theta\left(\sqrt{d}\right)$, we have that

C.7 Proof of Lemma 7

For the equation $x^{2}-2x+1-\frac{\eta^{2}}{(1-\beta)^{2}}\|A\|_{2}^{2}=0$, the roots are $\lambda_{1}=1-\frac{\eta}{1-\beta}\|A\|_{2}$ and $\lambda_{2}=1+\frac{\eta}{1-\beta}\|A\|_{2}$, which gives us

Then we have that w.p. at least $1-\delta$, $\forall 1\leq i,j\leq d:$,

Next, we bound the $i$-th coordinate of $W_{2}^{(0)}+\frac{1-\beta}{\|A\|_{2}}AW_{1}^{(0)T}$, i.e. $w_{2i}^{(0)}+\frac{1-\beta}{\|A\|_{2}}A\left(W_{1}^{(0)}[i,:]\right)^{T}$.

By independence under Assumption 2, we have that

Using the Gaussian tail bound and union bound, w.p. at least $1-\delta$, for ever $1\leq i\leq d$, we have that

Since for $X\sim\mathcal{N}(0,\sigma^{2})$, we have that $P(|X|\leq t)\leq\frac{2t}{\sqrt{2\pi}\sigma}$, then for a fixed $i$,

Then by union bound, we have that w.p. at least $1-\frac{1}{d^{\frac{\alpha}{4}-1}}$, for every $1\leq i\leq d$,

where $(i)$ follows from eq. (14) and the fact that $\|A\|_{2}=\sqrt{d}$. Then we get that w.h.p.

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

C.8 Proof of Lemma 2

The proof in Section C.2 tells us that at the end of the first phase (when $t=T_{1}$),

Denote the $i$-th coordinate of $\boldsymbol{u}^{(t)},\boldsymbol{v}^{(t)},R_{2}^{(t)}$ as $u_{i}^{(t)},v_{i}^{(t)},R_{2i}^{(t)}$, respectively. Denote the $(i,j)$-th element of $R_{1}^{(t)}$ as $R_{1}^{(t)}[i,j]$. For $t\geq T_{1}$, we prove by induction that,

with $E^{(t)}:=W_{2}^{(t)}W_{1}^{(t)}-A$, $\eta_{t}=\eta\sum_{\tau=0}^{t}\beta^{t-\tau}$, $r_{1}^{(t)}:=\eta\sum_{\tau=0}^{t}\beta^{t-\tau}\left(W_{2}^{(t)T}E^{(t)}-W_{2}^{(\tau)T}E^{(\tau)}\right)-\eta\sum_{\tau=0}^{t}\beta^{t-\tau}Dg_{1}^{(\tau)}$ and $r_{2}^{(t)}=\eta\sum_{\tau=0}^{t}\beta^{t-\tau}\left(E^{(t)}W_{1}^{(t)T}-E^{(\tau)}W_{1}^{(\tau)T}\right)-\eta\sum_{\tau=0}^{t}\beta^{t-\tau}Dg_{2}^{(\tau)}$. Note that the $r_{1}^{(t)}$ and $r_{2}^{(t)}$ here are different from those defined in Section C.2, but we abuse the notation and still use $r_{1}^{(t)}$ and $r_{2}^{(t)}$ 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.

Now we analyze the error terms $\left|R_{1}^{(t)}[i,j]\right|$ and $\left|R_{2i}^{(t)}\right|$. Eq. (16) tells us that $c^{(T_{1})}$ and $\forall i\in[d],v_{i}^{(T_{1})}$ are all positive. We first prove by induction that for all $T_{1}\leq t\leq T_{2}$, $c^{(t)}>0,\forall i\in[d],v_{i}^{(t)}>0$.

The above discussion already proves the base case. Suppose at time $t$, we have $c^{(t)}>0,\forall i\in[d],v_{i}^{(t)}>0$. Note that when $T_{1}\leq t<T_{2}$, $\forall i\in[d]:E_{i}^{(t)}\leq 0$, then for $t+1$,

Therefore by induction, we have proved that for all $T_{1}\leq t\leq T_{2}$, $c^{(t)}>0,\forall i\in[d],v_{i}^{(t)}>0$.

Now we prove that for all $T_{1}\leq t\leq T_{2}$,

The left hand sides of the inequalities are trivial since we have proved that $c^{(t)}>0,\forall i\in[d],v_{i}^{(t)}>0$ for all $T_{1}\leq t\leq T_{2}$. Now we prove the right hand sides by induction.

The base case is already verified by the definition of $\delta_{i}$. Suppose eq.(18) holds for $T_{1}\leq t<T_{2}$. Then for $t+1$, using $\forall i\in[d]:E_{i}^{(t)}\leq 0$ and $v^{(t+1)}\geq v^{(t)},c^{(t+1)}\geq c^{(t)}$, we can get that $\forall 1\leq i,j\leq d$

Similarly, we have that $\forall 1\leq i\leq 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 2. The remaining part of the proof is given by the following lemma, whose proof is deferred to Section C.9.

Under Assumption 1, 2 and 3, suppose $\sigma\leq\frac{\eta^{3/2}}{d^{\alpha/2+1}}$. By picking $\eta\leq\mathcal{O}\left(\frac{\epsilon}{d^{\frac{7\alpha}{4}+4}}\right)$, we have that w.h.p. for $T_{1}\leq t\leq\min\{T_{2},T_{3}\}$,

and that when $t=\min\{T_{2},T_{3}\}$, we have $\left\|E^{(t)}\right\|^{2}_{2}=\mathcal{O}(\epsilon_{0}d)$.

C.9 Proof of Lemma 8

We first have the following lemma which describes the structure of $\boldsymbol{v}^{(t)}$ for $t\geq T_{1}$.

Under Assumption 1, 2 and 3, for $t\geq T_{1}$, we can write $\boldsymbol{v}^{(t)T}$ as $\boldsymbol{v}^{(t)T}=a^{(t)}A+R_{v}^{(t)T}$, with $a^{(T_{1})}=\frac{\eta}{1-\beta},R_{v}^{(T_{1})T}=[0,0,...,0]$, and

We prove Lemma 8 by induction. Denote the $i$-th coordinate of $R_{3}^{(t)}$ and $R_{v}^{(t)}$ as $R_{3i}^{(t)}$ and $R_{vi}^{(t)}$, respectively. The following lemmas constitute the inductive part.

Under the conditions of Lemma 10 and pick $\eta\leq\mathcal{O}\left(\frac{\epsilon}{d^{\frac{7\alpha}{4}+4}}\right)$, we have that at time $t+1$,

where $\epsilon_{0}$ is defined in Definition 2.

Under the conditions of Lemma 10 and pick $\eta\leq\mathcal{O}\left(\frac{\epsilon}{d^{\frac{7\alpha}{4}+4}}\right)$, we have that at time $t+1$,

$\forall i,j\in[d]:\frac{E_{j}^{(t)}}{E_{i}^{(t)}}=\Theta(1)$,

$\forall j\in[d]:\frac{v_{j}^{(t+1)}}{c^{(t+1)}}=\Theta\left(\frac{1}{\sqrt{d}}\right)$,

$\forall i,j\in[d]:\left|w_{2i}^{(t+1)}\right|\leq\mathcal{O}\left(d^{1/4}\right),\left|W_{1}^{(t+1)}[i,j]\right|\leq\mathcal{O}\left(\frac{1}{d^{1/4}}\right)$,

By combining Lemma 10, 11 and 13, we can prove by induction that for all $T_{1}\leq t\leq\min\{T_{2},T_{3}\}$, eq. (19) holds (which follows from Lemma 11), and

So far we have proved eq. (19) in Lemma 8. Now let’s prove when $t=\min\{T_{2},T_{3}\}$, we have that $\left\|E^{(t)}\right\|^{2}_{2}=\mathcal{O}(\epsilon_{0}d)$.

If $\min\{T_{2},T_{3}\}=T_{3}$, by Definition 3, we have $\left\|E^{(t)}\right\|_{2}^{2}\leq\epsilon$. If $\min\{T_{2},T_{3}\}=T_{2}$, by Definition 2, there exists $j\in[d]$ such that $E_{j}^{(t)}=-\Theta\left(\sqrt{\epsilon_{0}}\right)$. Combining with eq. (22) gives us $\forall i\in[d]:E_{i}^{(t)}=-\Theta\left(\sqrt{\epsilon_{0}}\right)$. Combining these two cases, we get that when $t=\min\{T_{2},T_{3}\}$, $\left\|E^{(t)}\right\|^{2}_{2}\leq\max\{\epsilon,\Theta\left(\epsilon_{0}d\right)\}=\mathcal{O}\left(\epsilon_{0}d\right)$.

C.10 Proof of Lemma 9

We prove this lemma by induction. The base case ($t=T_{1}$) of $\boldsymbol{v}^{(t)}$ is verified by eq. (16).

Suppose at time $t$, $\boldsymbol{v}^{(t)T}=a^{(t)}A+R_{v}^{(t)T}$, then by eq. 17, we have that

where $d^{(t)}:=c^{(t)}\left\|\boldsymbol{u}^{(T_{1})}\right\|^{2}+R_{2}^{(t)T}\boldsymbol{u}^{(T_{1})},\quad R_{3}^{(t)T}:=c^{(t)}\boldsymbol{u}^{(T_{1})T}R_{1}^{(t)}+R_{2}^{(t)T}R_{1}^{(t)}$. That gives us

Therefore we have proved by induction that for $t$ in the second phase, $\boldsymbol{v}^{(t)}=a^{(t)}A+R_{v}^{(t)T}$. The above steps also proved eq. (20).

C.11 Proof of Lemma 10

Write $r_{1}^{(t)}=q_{11}^{(t)}+q_{12}^{(t)}$ where we have $q_{11}^{(t)}=\eta\sum_{\tau=0}^{t}\beta^{t-\tau}\left(W_{2}^{(t)T}E^{(t)}-W_{2}^{(\tau)T}E^{(\tau)}\right)$, $q_{12}^{(t)}=-\eta\sum_{\tau=0}^{t}\beta^{t-\tau}Dg_{1}^{(\tau)}$. Write $r_{2}^{(t)}=q_{21}^{(t)}+q_{22}^{(t)}$ where $q_{22}^{(t)}=-\eta\sum_{\tau=0}^{t}\beta^{t-\tau}Dg_{2}^{(\tau)}$, $q_{21}^{(t)}=\eta\sum_{\tau=0}^{t}\beta^{t-\tau}\left(E^{(t)}W_{1}^{(t)T}-E^{(\tau)}W_{1}^{(\tau)T}\right)$.