Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks

Introduction

It is a commonly accepted fact that a large initial learning rate is required to successfully train a deep network even though it slows down optimization of the train loss. Modern state-of-the-art architectures typically start with a large learning rate and anneal it at a point when the model’s fit to the training data plateaus . Meanwhile, models trained using only small learning rates have been found to generalize poorly despite enjoying faster optimization of the training loss.

A number of papers have proposed explanations for this phenomenon, such as sharpness of the local minima , the time it takes to move from initialization , and the scale of SGD noise . However, we still have a limited understanding of a surprising and striking part of the large learning rate phenomenon: from looking at the section of the accuracy curve before annealing, it would appear that a small learning rate model should outperform the large learning rate model in both training and test error. Concretely, in Fig. 1, the model trained with small learning rate outperforms the large learning rate until epoch 60 when the learning rate is first annealed. Only after annealing does the large learning rate visibly outperform the small learning rate in terms of generalization.

In this paper, we propose to theoretically explain this phenomenon via the concept of learning order of the model, i.e., the rates at which it learns different types of examples. This is not a typical concept in the generalization literature — learning order is a training-time property of the model, but most analyses only consider post-training properties such as the classifier’s complexity , or the algorithm’s output stability . We will construct a simple distribution for which the learning order of a two-layer network trained under large and small initial learning rates determines its generalization.

Informally, consider a distribution over training examples consisting of two types of patterns (“pattern” refers to a grouping of features). The first type consists of a set of easy-to-generalize (i.e., discrete) patterns of low cardinality that is difficult to fit using a low-complexity classifier, but easily learnable via complex classifiers such as neural networks. The second type of pattern will be learnable by a low-complexity classifier, but are inherently noisy so it is difficult for the classifier to generalize. In our case, the second type of pattern requires more samples to correctly learn than the first type. Suppose we have the following split of examples in our dataset:

The following informal theorems characterize the learning order and generalization of the large and small initial learning rate models. They are a dramatic simplification of our Theorems 3.4 and 3.5 meant only to highlight the intuitions behind our results.

There is a dataset with size $N$ of the form (1.1) such that with a large initial learning rate and noisy gradient updates, a two layer network will:

1) initially only learn hard-to-generalize, easy-to-fit patterns from the $0.8N$ examples containing such patterns.

2) learn easy-to-generalize, hard-to-fit patterns only after the learning rate is annealed.

Thus, the model learns hard-to-generalize, easily fit patterns with an effective sample size of $0.8N$ and still learns all easy-to-generalize, hard to fit patterns correctly with $0.2N$ samples.

In the same setting as above, with small initial learning rate the network will:

1) quickly learn all easy-to-generalize, hard-to-fit patterns.

2) ignore hard-to-generalize, easily fit patterns from the $0.6N$ examples containing both pattern types, and only learn them from the $0.2N$ examples containing only hard-to-generalize patterns.

Thus, the model learns hard-to-generalize, easily fit patterns with a smaller effective sample size of $0.2N$ and will perform relatively worse on these patterns at test time.

Together, these two theorems can justify the phenomenon observed in Figure 1 as follows: in a real-world network, the large learning rate model first learns hard-to-generalize, easier-to-fit patterns and is unable to memorize easy-to-generalize, hard-to-fit patterns, leading to a plateau in accuracy. Once the learning rate is annealed, it is able to fit these patterns, explaining the sudden spike in both train and test accuracy. On the other hand, because of the low amount of SGD noise present in easy-to-generalize, hard-to-fit patterns, the small learning rate model quickly overfits to them before fully learning the hard-to-generalize patterns, resulting in poor test error on the latter type of pattern.

Both intuitively and in our analysis, the non-convexity of neural nets is crucial for the learning-order effect to occur. Strongly convex problems have a unique minimum, so what happens during training does not affect the final result. On the other hand, we show the non-convexity causes the learning order to highly influence the characteristics of the solutions found by the algorithm.

In Section E.1, we propose a mitigation strategy inspired by our analysis. In the same setting as Theorems 1.1 and 1.2, we consider training a model with small initial learning rate while adding noise before the activations which gets reduced by some constant factor at some particular epoch in training. We show that this algorithm provides the same theoretical guarantees as the large initial learning rate, and we empirically demonstrate the effectiveness of this strategy in Section 7. In Section 7 we also empirically validate Theorems 1.1 and 1.2 by adding an artificial memorizable patch to CIFAR-10 images, in a manner inspired by (1.1).

The question of training with larger batch sizes is closely tied with learning rate, and many papers have empirically studied large batch/small LR phenomena , particularly focusing on vision tasks using SGD as the optimizer.While these papers are framed as a study of large-batch training, a number of them explicitly acknowledge the connection between large batch size and small learning rate. Keskar et al. argue that training with a large batch size or small learning rate results in sharp local minima. Hoffer et al. propose training the network for longer and with larger learning rate as a way to train with a larger batch size. Wen et al. propose adding Fisher noise to simulate the regularization effect of small batch size.

Adaptive gradient methods are a popular method for deep learning that adaptively choose different step sizes for different parameters. One motivation for these methods is reducing the need to tune learning rates . However, these methods have been observed to hurt generalization performance , and modern architectures often achieve the best results via SGD and hand-tuned learning rates . Wilson et al. construct a toy example for which ADAM generalizes provably worse than SGD. Additionally, there are several alternative learning rate schedules proposed for SGD, such as warm-restarts and . Ge et al. analyze the exponentially decaying learning rate and show that its final iterate achieves optimal error in stochastic optimization settings, but they only analyze convex settings.

There are also several recent works on implicit regularization of gradient descent that establish convergence to some idealized solution under particular choices of learning rate . In contrast to our analysis, the generalization guarantees from these works would depend only on the complexity of the final output and not on the order of learning.

Other recent papers have also studied the order in which deep networks learn certain types of examples. Mangalam and Prabhu and Nakkiran et al. experimentally demonstrate that deep networks may first fit examples learnable by “simpler” classifiers. For our construction, we prove that the neural net with large learning rate follows this behavior, initially learning a classifier on linearly separable examples and learning the remaining examples after annealing. However, the phenomenon that we analyze is also more nuanced: with a small learning rate, we prove that the model first learns a complex classifier on low-noise examples which are not linearly separable.

Finally, our proof techniques and intuitions are related to recent literature on global convergence of gradient descent for over-parametrized networks . These works show that gradient descent learns a fixed kernel related to the initialization under sufficient over-parameterization. In our analysis, the underlying kernel is changing over time. The amount of noise due to SGD governs the space of possible learned kernels, and as a result, regularizes the order of learning.

Setup and Notations

We formally introduce our data distribution, which contains examples supported on two types of components: a $\mathcal{P}$ component meant to model hard-to-generalize, easier-to-fit patterns, and a $\mathcal{Q}$ component meant to model easy-to-generalize, hard-to-fit patterns (see the discussion in our introduction). Formally, we assume that the label $y$ has a uniform distribution over $\{-1,1\}$, and the data $x$ is generated as

where $\mathcal{P}_{-1},\mathcal{P}_{1}$ are assumed to be two half Gaussian distributions with a margin $\gamma_{0}$ between them:

Memorizing 𝒬𝒬\mathcal{Q} with a two-layer net

It is easy for a two-layer relu network to memorize the labels of $x_{2}$ using two neurons with weights $w,v$ such that $\langle w,z\rangle<0$, $\langle w,z-\zeta\rangle>0$ an $\langle v,z\rangle<0$, $\langle v,z+\zeta\rangle>0$. In particular, we can verify that $-\langle w,x_{2}\rangle_{+}-\langle v,x_{2}\rangle_{+}$ will output a negative value for $x_{2}\in\{z-\zeta,z+\zeta\}$ and a zero value for $x_{2}=z$. Thus choosing a small enough $\rho>0$, the classifier $-\langle w,x_{2}\rangle_{+}-\langle v,x_{2}\rangle_{+}+\rho$ gives the correct sign for the label $y$.

We assume that we have a training dataset with $N$ examples $\{(x^{(1)},y^{(1)}),\cdots,(x^{(N)},y^{(N)})\}$ drawn i.i.d from the distribution described above. We use $p$ and $q$ to denote the empirical fraction of data points that are drawn from equation (2.2) and (2.3).

Two-layer neural network model

Here we slightly abuse the notation to use $W$ to denote both a matrix of $2d$ columns with last $d$ columns being zero, or a matrix of $d$ columns. We also extend our theorem to other $U$ such as a two layer convolution network in Section E.

Training objective

We consider a regularized training objective $\widehat{L}_{\lambda}(u,U)=\widehat{L}(u,U)+\frac{\lambda}{2}\|U\|_{F}^{2}$. For the simplicity of derivation, the second layer weight vector $u$ is random initialized and fixed throughout this paper. Thus with slight abuse of notation the training objective can be written as $\widehat{L}_{\lambda}(U)=\widehat{L}(u,U)+\frac{\lambda}{2}\|U\|_{F}^{2}$.

Notations

Main Results

The training algorithm that we consider is stochastic gradient descent with spherical Gaussian noise. We remark that we analyze this algorithm as a simplification of the minibatch SGD noise encountered when training real-world networks. There are a number of works theoretically characterizing this particular noise distribution , and we leave analysis of this setting to future work.

We initialize $U_{0}$ to have i.i.d. entries from a Gaussian distribution with variance $\tau_{0}^{2}$, and at each iteration of gradient descent we add spherical Gaussian noise with coordinate-wise variance $\tau_{\xi}^{2}$ to the gradient updates. That is, the learning algorithm for the model is

where $\gamma_{t}$ denotes the learning rate at time $t$. We will analyze two algorithms:

Algorithm 1 (L-S): The learning rate is $\eta_{1}$ for $t_{0}$ iterations until the training loss drops below the threshold $\varepsilon_{1}+q\log 2$. Then we anneal the learning rate to $\gamma_{t}=\eta_{2}$ (which is assumed to be much smaller than $\eta_{1}$) and run until the training loss drops to $\varepsilon_{2}$.

Algorithm 2 (S): We used a fixed learning rate of $\eta_{2}$ and stop at training loss $\varepsilon_{2}^{\prime}\leq\varepsilon_{2}$.

For the convenience of the analysis, we make the following assumption that we choose $\tau_{0}$ in a way such that the contribution of the noises in the system stabilize at the initialization:Let $\tau_{0}^{\prime}$ be the solution to (3.3) holding $\tau_{\xi},\eta_{1},\lambda$ fixed. If the standard deviation of the initialization is chosen to be smaller than $\tau_{0}^{\prime}$, then standard deviation of the noise will grow to $\tau_{0}^{\prime}$. Otherwise if the initialization is chosen to be larger, the contribution of the noise will decrease to the level of $\tau_{0}^{\prime}$ due to regularization. In typical analysis of SGD with spherical noises, often as long as either the noise or the learning rate is small enough, the proof goes through. However, here we will make explicit use of the large learning rate or the large noise to show better generalization performance.

After fixing $\lambda$ and $\tau_{\xi}$, we choose initialization $\tau_{0}$ and large learning rate $\eta_{1}$ so that

As a technical assumption for our proofs, we will also require $\eta_{1}\lesssim\varepsilon_{1}$.

We also require sufficient over-parametrization.

We assume throughout the paper that $\tau_{0}=1/\textup{poly}\left(\frac{d}{\varepsilon}\right)$ and $m\geq\textup{poly}\left(\frac{d}{\varepsilon\tau_{0}}\right)$ where poly is a sufficiently large constant degree polynomial. We note that we can choose $\tau_{0}$ arbitrarily small, so long as it is fixed before we choose $m$.

As we will see soon, the precise relation between $N,d$ implies that the level of over-parameterization is polynomial in $N,\epsilon$, which fits with the conditions assumed in prior works, such as .

Throughout this paper, we assume the following dependencies between the parameters. We assume that $N,d\rightarrow\infty$ with a relationship $\frac{N}{d}=\frac{1}{\kappa^{2}}$ where $\kappa\in(0,1)$ is a small value.Or in a non-asymptotic language, we assume that $N,d$ are sufficiently large compared to $\kappa$: $N,d\gg\textup{poly}(\kappa)$ We set $r=d^{-3/4}$, $p_{0}=\kappa^{2}/2$, and $q_{0}=\Theta(1)$. The regularizer will be chosen to be $\lambda=d^{-5/4}$. All of these choices of hyper-parameters can be relaxed, but for simplicity of exposition we only work this setting.

We note that under our assumptions, for sufficiently large $N$, $p\approx p_{0}$ and $q\approx q_{0}$ up to constant multiplicative factors. Thus we will mostly work with $p$ and $q$ (the empirical fractions) in the rest of the paper. We also note that our parameter choice satisfies $(rd)^{-1},d\lambda,\lambda/r\leq\kappa^{O(1)}$ and $\lambda\leq r^{2}/(\kappa^{2}q^{3}p^{2})$, which are a few conditions that we frequently use in the technical part of the paper.

Now we present our main theorems regarding the generalization of models trained with the L-S and S algorithms. The final generalization error of the model trained with the L-S algorithm will end up a factor $O(\kappa)=O(p^{1/2})$ smaller than the generalization error of the model trained with S algorithm.

Under Assumption 3.1, 3.2, and 3.3, there exists a universal constant $0<c<1/16$ such that Algorithm 1 (L-S) with annealing at loss $\varepsilon_{1}+q\log 2$ for $\varepsilon_{1}\in\left(d^{-c},\kappa^{2}p^{2}q^{3}\right)$ and stopping criterion $\varepsilon_{2}=\sqrt{\varepsilon_{1}/q}$ satisfies the following:

It anneals the learning rate within $\widetilde{O}\left(\frac{d}{\eta_{1}\varepsilon_{1}}\right)$ iterations.

It stops at at most $t=\widetilde{O}\left(\frac{d}{\eta_{1}\varepsilon_{1}}+\frac{1}{\eta_{2}r\varepsilon_{1}^{3}}\right)$. With probability at least 0.99, the solution $U_{t}$ has test (classification) error and test loss at most ${O}\left(p\kappa\log\frac{1}{\varepsilon_{1}}\right)$.

Roughly, the learning order and generalization of the L-S model is as follows: before annealing the learning rate, the model only learns an effective classifier for $\mathcal{P}$ on the $\approx(1-q)N$ samples in $\mathcal{M}_{1}$ as the large learning rate creates too much noise to effectively learn $\mathcal{Q}$ (Lemma 4.1 and Lemma 4.2). After the learning rate is annealed, the model memorizes $\mathcal{Q}$ and correctly classifies examples with only a $\mathcal{Q}$ component during test time (formally shown in Lemmas 4.3 and 4.4). For examples with only $\mathcal{P}$ component, the generalization error is (ignoring log factors and other technicalities) $p\sqrt{\frac{d}{N}}=O(p\kappa)$ via standard Rademacher complexity. The full analysis of the L-S algorithm is clarified in Section 4.

Let $\varepsilon_{2}$ be chosen in Theorem 3.4. Under Assumption 3.1, 3.2 and 3.3, there exists a universal constant $c>0$ such that w.h.p, Algorithm 2 with any $\eta_{2}\leq\eta_{1}d^{-c}$ and any stopping criterion $\varepsilon_{2}^{\prime}\in(d^{-c},\varepsilon_{2}]$, achieves training loss $\varepsilon_{2}^{\prime}$ in at most $\widetilde{O}\left(\frac{d}{\eta_{2}\varepsilon_{2}^{\prime}}\right)$ iterations, and both the test error and the test loss of the obtained solution are at least $\Omega(p)$.

We explain this lower bound as follows: the S algorithm will quickly memorize the $\mathcal{Q}$ component which is low noise and ignore the $\mathcal{P}$ component for the $\approx 1-p-q$ examples with both $\mathcal{P}$ and $\mathcal{Q}$ components (shown in Lemma 5.2). Thus, it only learns $\mathcal{P}$ on $\approx pN$ examples. It obtains a small margin on these examples and therefore misclassifies a constant fraction of $\mathcal{P}$-only examples at test time. This results in the lower bound of $\Omega(p)$. We formalize the analysis in Section 5.

It will be fruitful for our analysis to separately consider the gradient signal and Gaussian noise components of the weight matrix $U_{t}$. We will decompose the weight matrix $U_{t}$ as follows: $U_{t}=\overline{U}_{t}+\widetilde{U}_{t}$. In this formula, $\overline{U}_{t}$ denotes the signals from all the gradient updates accumulated over time, and $\widetilde{U}_{t}$ refers to the noise accumulated over time:

Note that when the learning rate $\gamma_{t}$ is always $\eta$, the formula simplifies to $\overline{U}_{t}=\sum_{s=1}^{t}\eta(1-\eta\lambda)^{t-s}\nabla\widehat{L}(U_{s-1})$ and $\widetilde{U}_{t}=(1-\eta\lambda)^{t}U_{0}+\sum_{s=1}^{t}\eta(1-\eta\lambda)^{t-s}\xi_{s-1}$. The decoupling and our particular choice of initialization satisfies that the noise updates in the system stabilize at initialization, so the marginal distribution of $\widetilde{U}_{t}$ is always the same as the initialization. Another nice aspect of the signal-noise decomposition is as follows: we use tools from to show that if the signal term $\overline{U}$ is small, then using only the noise component $\widetilde{U}$ to compute the activations roughly preserves the output of the network. This facilitates our analysis of the network dynamics. See Section A.1 for full details.

Decomposition of Network Outputs

For convenience, we will explicitly decompose the model prediction at each time into two components, each of which operates on one pattern: we have $N_{U_{t}}(u,U_{t};x)=g_{t}(x)+r_{t}(x)$,

In other words, the network $g_{t}$ acts on the $\mathcal{Q}$ component of examples, and the network $r_{t}$ acts on the $\mathcal{P}$ component of examples.

Characterization of Algorithm 1 (L-S)

We characterize the behavior of algorithm L-S with large initial learning rate. We provide proof sketches in Section 6.1 with full proofs in Section C.

The following lemma bounds the rate of convergence to the point where the loss gets annealed. It also bounds the total gradient signal accumulated by this point.

In the setting of Theorem 3.4, at some time step $t_{0}\leq\widetilde{O}\left(\frac{d}{\eta_{1}\varepsilon_{1}}\right)$, the training loss $\widehat{L}(U_{t_{0}})$ becomes smaller than $q\log 2+\epsilon_{1}$. Moreover, we have $\|\overline{U}_{t_{0}}\|_{F}^{2}={O}\left(d\log^{2}\frac{1}{\varepsilon_{1}}\right)$.

Our proof of Lemma 4.1 views the SGD dynamics as optimization with respect to the neural tangent kernel induced by the activation patterns where the kernel is rapidly changing due to the noise terms $\xi$. This is in contrast to the standard NTK regime, where the activation patterns are assumed to be stable . Our analysis extends the NTK techniques to deal with a sequence of changing kernels which share a common optimal classifier (see Section 6.1 and Theorem 6.2 for additional details).

The next lemma says that with large initial learning rate, the function $g_{t}$ does not learn anything meaningful for the $\mathcal{Q}$ component before the $\frac{1}{\eta_{1}\lambda}$-timestep. Note that by our choice of parameters $1/\lambda\gg d$ and Lemma 4.1, we anneal at the time step $\widetilde{O}\left(\frac{d}{\eta_{1}\varepsilon_{1}}\right)\leq\frac{1}{\eta_{1}\lambda}$. Therefore, the function has not learned anything meaningful about the memorizable pattern on distribution $\mathcal{Q}$ before we anneal.

In the setting of Theorem 3.4, w.h.p., for every $t\leq\frac{1}{\eta_{1}\lambda}$,

After iteration $t_{0}$, we decrease the learning rate to $\eta_{2}$. The following lemma bounds how fast the loss converges after annealing.

In the setting of Theorem 3.4, there exists $t=\widetilde{O}\left(\frac{1}{\varepsilon_{1}^{3}\eta_{2}r}\right)$, such that after $t_{0}+t$ iterations, we have that

Moreover, $\|\overline{U}_{t_{0}+t}-\overline{U}_{t_{0}}\|_{F}^{2}\leq\widetilde{O}\left(\frac{1}{\varepsilon_{1}^{2}r}\right)\leq O(d)$.

The following lemma bounds the training loss on the example subsets $\mathcal{M}_{1}$, $\bar{\mathcal{M}}_{1}$.

In the setting of Lemma 4.3 using the same $t=\widetilde{O}\left(\frac{1}{\varepsilon_{1}^{3}\eta_{2}r}\right)$, the average training losses on the subsets $\mathcal{M}_{1}$ and $\bar{\mathcal{M}}_{1}$ are both good in the sense that

Intuitively, low training loss of $g_{t_{0}+t}$ on $\bar{\mathcal{M}}_{1}$ immediately implies good generalization on examples containing patterns from $\mathcal{Q}$. Meanwhile, the classifier for $\mathcal{P}$, $r_{t_{0}+t}$, has low loss on $(1-q)N$ examples. Then the test error bound follows from standard Rademacher complexity tools applied to these $(1-q)N$ examples.

Characterization of Algorithm 2 (S)

We present our small learning rate lemmas, with proofs sketches in Section 6.2 and full proofs in Section D.

The below lemma shows that the algorithm will converge to small training error too quickly. In particular, the norm of $W_{t}$ is not large enough to produce a large margin solution for those $x$ such that $x_{2}=0$.

Lower bound on the generalization error

The following important lemma states that our classifier for $\mathcal{P}$ does not learn much from the examples in $\mathcal{M}_{2}$. Intuitively, under a small learning rate, the classifier will already learn so quickly from the $\mathcal{Q}$ component of these examples that it will not learn from the $\mathcal{P}$ component of examples in $\mathcal{M}_{1}\cap\mathcal{M}_{2}$. We make this precise by showing that the magnitude of the gradients on $\mathcal{M}_{2}$ is small.

The above lemma implies that $W$ does not learn much from examples in $\mathcal{M}_{2}$, and therefore must overfit to the $pN$ examples in $\bar{\mathcal{M}}_{2}$. As $pN\leq d/2$ by our choice of parameters, we will not have enough samples to learn the $d$-dimensional distribution $\mathcal{P}$. The following lemma formalizes the intuition that the margin will be poor on samples from $\mathcal{P}$.

As the margin is poor, the predictions will be heavily influenced by noise. We use this intuition to prove the classification lower bound for Theorem 3.5.

Proof Sketches

We first introduce notations that will be useful in these proofs. We will explicitly decouple the noise in the weights from the signal by abstracting the loss as a function of only the signal portion $\overline{U}_{t}$ of the weights. Let us define the following:

Now the proof of Lemma 4.1 relies on the following two results, which we state below and prove in Section C.1. The first says that there is a common target for the signal part of the network that is a good solution for all of the $K_{t}$.

In the setting of Lemma 4.1, there exists a solution $U^{\star}$ satisfying a) $\|U^{\star}\|_{F}^{2}\leq{O}\left(d\log^{2}\frac{1}{\varepsilon_{1}}\right)$ and b) for every $t\geq 0$

Now the second statement is a general one proving that gradient descent on a sequence of convex, but changing, functions will still find a optimum provided these functions share the same solution.

$K_{t}$’s are $L$-Lipschitz, i.e., $\|\nabla K_{t}(z)\|_{2}\leq L,\forall z,t$

For every $\mu>0$, we have that for $\lambda R^{2}\leq\frac{1}{100}\mu$ and $\eta\leq\frac{\mu}{100(\lambda^{2}R^{2}+L^{2})}$, $\eta T>\frac{R^{2}}{\mu}$, there is a $t^{\star}\in[T]$ such that:

Furthermore, the iterates satisfy $\|z_{t}-z^{\star}\|_{2}\leq R$ for all $t\leq t^{\star}$.

Combining these two statements leads to the proof of Lemma 4.1.

We can apply Theorem 6.2 with $K_{t}$ defined in (6.2) and $z^{\star}=U^{\star}$ defined in Lemma 6.1, using $R=O\left(d\log^{2}\frac{1}{\varepsilon_{1}}\right)$. We note that $\eta_{1}$ satisfies the conditions of Theorem 6.2 by our parameter choices, which completes the proof. ∎

To prove Lemma 4.2, we will essentially argue in Section C.2 that the change in activations caused by the noise will prevent the model from learning $\mathcal{Q}$ with a large learning rate. This is because the examples in $\mathcal{Q}$ require a very specific configuration of activation patterns to learn correctly, and the noise will prevent the model from maintaining this configuration.

Now after we anneal the learning rate, in order to conclude Lemmas 4.3 and 4.4, the following must hold: 1) the network learns the $\mathcal{Q}$ component of the distribution and 2) the network does not forget the $\mathcal{P}$ component that it previously learned. To prove the latter, we rely on the following lemma stating that the activations do not change much with a small learning rate:

The activation patterns do not change much after annealing the learning rate: for every $t_{0},t\leq\frac{1}{\eta_{2}\lambda}$, for any $x$ and for any row $[U_{t}]_{i}$ of the weight matrix $U$, we have that

Moreover, for all $i\in[m]$, $\left\|[\overline{U}_{t}]_{i}\right\|_{2}\leq\frac{1}{\lambda\sqrt{m}}$, it holds that w.h.p. for every $x$:

We prove the above lemma in Section C.3. Now to complete the proof of Lemma 4.3, we will construct a target solution for all timesteps after annealing the learning rate based on the activations at time $t_{0}$ (as they do not change by much in subsequent time steps because of Lemma 6.3) and reapply Theorem 6.2. Finally, to prove Lemma 4.4, we use the fact that the $W_{t}$ component of the solution does not change by much, and therefore the loss on $\mathcal{M}_{1}$ is still low.

2 Proof Sketches for Small Learning Rate

The proof of Lemma 5.1 proceeds similarly as the proof of Lemma 4.3: we will show the existence of a target solution of $K_{t}$ for all iterations, and use Theorem 6.2 to prove convergence to this target solution.

The next two statements argue that $\rho_{t}$ can be large only in a limited number of time steps. As the training loss converges quickly with small learning rate, this will be used to argue that the $\mathcal{P}$ components of examples in $\mathcal{M}_{2}$ provide a very limited signal to $W_{t}$. The proofs of these statements are in Section D.2.

We first show the following lemma that says that if $\rho_{t}$ is large (which means the loss is large as well), then the total gradient norm has to be big. This lemma holds because there is little noise in the $\mathcal{Q}$ component of the distribution, and therefore the gradient of $V_{t}$ will be large if $\rho_{t}$ is large.

For every $t\leq\frac{1}{\eta_{2}\lambda}$, we have that if $\rho_{t}=\Omega\left(\frac{1}{N}\right)$, then w.h.p.

Now we use the above lemma to bound the number of times when $\rho_{t}$ is large.

In the setting of Lemma 5.2, let $\mathcal{T}$ be the set of iterations where $\rho_{t}\geq\varepsilon_{2}^{\prime 2}\varepsilon_{3}^{2}$, where $\varepsilon_{3}$ is defined in Lemma 5.2. Then w.h.p, $|\mathcal{T}|\lesssim\frac{1}{r\varepsilon_{2}^{\prime 8}\varepsilon_{3}^{8}\eta_{2}}.$

Now if $\rho_{t}$ is small, the gradient accumulated on $W_{t}$ from examples in $\mathcal{M}_{2}$ must be small. We formalize this argument in our proof of Lemma 5.2 in Section D.2.

Lemma 5.3 will then follow by explicitly decomposing $\overline{W}_{t}$ into a component in $\text{span}\{x_{1}^{(i)}\}_{i\in\bar{\mathcal{M}}_{2}}$ and some remainder, which is shown to be small by Lemma 5.2. This is presented in the below lemma, which is proved in Section D.3.

There exists real numbers $\{\alpha_{k}\}_{k\in\bar{\mathcal{M}}_{2}}$ such that for every $j\in[m]$, we have

with $\|\overline{W}_{t}^{\prime}\|_{F}\leq\widetilde{O}\left(\varepsilon_{3}\sqrt{d}\right)$.

This allows us to conclude Lemma 5.3 via computations carried out in Section D.3.

Finally, to complete the proof of Theorem 3.5, we will argue in Section B.2 that a classifier $r_{t}$ of the form given by (5.2) cannot have small generalization error because it will be too heavily influenced by the noise in $x_{1}$.

Experiments

Our theory suggests that adding noise to the network could be an effective strategy to regularize a small learning rate in practice. We test this empirically by adding small Gaussian noise during training before every activation layer in a WideResNet16 architecture, as our analysis highlights pre-activation noise as a key regularization mechanism of SGD. The noise level is annealed over time. We demonstrate on CIFAR-10 images without data augmentation that this regularization can indeed counteract the negative effects of small learning rate, as we report a 4.72% increase in validation accuracy when adding noise to a small learning rate. Full details are in Section G.1.

We will also empirically demonstrate that the choice of large vs. small initial learning rate can indeed invert the learning order of different example types. We add a memorizable 7 $\times$ 7 pixel patch to a subset of CIFAR-10 images following the scenario presented in (1.1), such that around 20% of images have no patch, 16% of images contain only a patch, and 64% contain both CIFAR-10 data and patch. We generate the patches so that they are not easily separable, as in our constructed $\mathcal{Q}$, but they are low in variation and therefore easy to memorize. Precise details on producing the data, including a visualization of the patch, are in Section G.2. We train on the modified dataset using WideResNet16 using 3 methods: large learning rate with annealing at the 30th epoch, small initial learning rate, and small learning rate with noise annealed at the 30th epoch.

Figure 3 depicts the validation accuracy vs. epoch on clean (no patch) and patch-only images. From the plots, it is apparent that the small learning rate picks up the signal in the patch very quickly, whereas the other two methods only memorize the patch after annealing.

From the validation accuracy on clean images, we can deduce that the small learning rate method is indeed learning the CIFAR images using a small fraction of all the available data, as the validation accuracy of a small LR model when training on the full dataset is around 83%, but the validation on clean data after training with the patch is 70%. We provide additional arguments in Section G.2. Our code for these experiments is online at the following link: https://github.com/cwein3/large-lr-code.

Conclusion

In this work, we show that the order in which a neural net learns to fit different types of patterns plays a crucial role in generalization. To demonstrate this, we construct a distribution on which models trained with large learning rates generalize provably better than those trained with small learning rates due to learning order. Our analysis reveals that more SGD noise, or larger learning rate, biases the model towards learning “generalizing” kernels rather than “memorizing” kernels. We confirm on articifially modified CIFAR-10 data that the scale of the learning rate can indeed influence learning order and generalization. Inspired by these findings, we propose a mitigation strategy that injects noise before the activations and works both theoretically for our construction and empirically. The design of better algorithms for regularizing learning order is an exciting question for future work.

Acknowledgements

CW acknowledges support from a NSF Graduate Research Fellowship.

References

Appendix A Basic Properties and Toolbox

In this section, we collect a few basic properties of the neural networks we are studying. In section F, we provide two lemmas on Gaussian random variables and perturbation theory of the matrices.

Let $[\nabla\widehat{L}(U)]_{i}$ be the $i$-th row of $\nabla\widehat{L}(U)$. We have that $\|[\nabla\widehat{L}(U)]_{i}\|_{2}\lesssim 1/\sqrt{m}$.

For any $t$, if $\gamma_{s}=\eta$ for every $s\leq t$, then we have that $\|[\overline{U}_{t}]_{i}\|_{2}\lesssim\min\{\frac{1}{\sqrt{m}\lambda},\eta t/\sqrt{m}\}$ and $\|\overline{U}_{t}\|_{F}\lesssim\frac{1}{\lambda}$.

By equation (3.4) and Proposition A.2, we have that

Hence $\|[\widetilde{U}x]_{+}\|_{2}\leq\|\widetilde{U}x\|_{2}\lesssim\tau\sqrt{m}\|x\|_{2}$. Now, since each $u_{i}$ is i.i.d. uniform $\{-m^{-1/2},m^{1/2}\}$, using the randomness of $u_{i}$ we know that w.h.p.

Here in the last inequality we applied Lemma A.8. The second statement follows from Proposition A.4 and triangle inequality. ∎

We have the following Rademacher complexity bound:

In this section, we collect useful statements which will help with decoupling the signal $\overline{U}$ from the noise $\widetilde{U}$ in our analysis. First, we observe that if the noise updates in the system stabilize at initialization, the marginal distribution of $U_{t}$ is always the same as the initialization.

Under Assumption 3.1, suppose we run Algorithm 1. Then for any $t$ before annealing the learning rate, $\widetilde{U}_{t}$ has marginal distribution $\mathcal{N}(0,\tau_{0}^{2}I_{m\times m}\otimes I_{d\times d})$. In other words, each entry of $\widetilde{U}_{t}$ follows $\mathcal{N}(0,\tau_{0}^{2})$ and they are independent with each others.

One nice aspect of the signal-noise decomposition is as follows: we use tools from to show that if the signal term $\overline{U}$ is small, then using only the noise component $\widetilde{U}$ to compute the activations roughly preserves the output of the network. This facilitates our analysis of the network dynamics.

As we will often apply (A.11) with $\|\overline{U}\|_{F}\lesssim\frac{1}{\lambda}$, for notational simplicity we denote throughout the paper $\varepsilon_{s}=\left(\frac{1}{\lambda\tau_{0}}\right)^{4/3}m^{-1/3}$. By our choice of $m\geq\textup{poly}(d/\tau_{0})$ we know that $\varepsilon_{s}\leq d^{-\Theta(1)}$.

Appendix B Proof of Main Theorems

We start with the following lemma that shows that if $g$ has small training error on $\bar{\mathcal{M}}_{1}$, then the output of $g$ on $x_{2}$ is large compared to $\|x_{2}\|$. This is because for the loss to be low, $g$ must have a good margin on $x_{2}$. However, as the norm of $x_{2}$ is roughly uniform in $$, the examples with small norm will force$g$ to have larger output.

W.h.p. for every $t\geq 0$ and every $\delta\geq\frac{1}{\sqrt{qN}}$, as long as $\widehat{L}_{\bar{\mathcal{M}}_{1}}(g_{t_{0}+t})\leq\delta$, we have that: for every $(x,y)$,

We use $\bar{\mathcal{M}}_{1}^{(1)}$ to denote the set of all $x_{2}^{(i)}\in\bar{\mathcal{M}}_{1}$ such that $x_{2}^{(i)}=\alpha(z-\zeta)$. Similarly, we use $\bar{\mathcal{M}}_{1}^{(2)}$ to denote the set of all $x_{2}^{(i)}\in\bar{\mathcal{M}}_{1}$ such that $x_{2}^{(i)}=\alpha(z+\zeta)$, and use $\bar{\mathcal{M}}_{1}^{(3)}$ to denote the set of all $x_{2}^{(i)}\in\bar{\mathcal{M}}_{1}$ such that $x_{2}^{(i)}=\alpha z$.

Let $g_{t_{0}+t}(z+\zeta)=\rho_{1},g_{t_{0}+t}(z-\zeta)=\rho_{2},g_{t+t_{0}}(z)=\rho_{3}$. By the positive homogeneity of ReLU, we know that for every $x_{2}\in\bar{\mathcal{M}}_{1}^{(i)}$, it holds:

Since $\widehat{L}_{\bar{\mathcal{M}}_{1}}(g_{t_{0}+t})\leq\delta$, it holds that w.h.p. for every $i\in$,

Our proof of Theorem 3.4 now amounts to carefully checking that all examples in $\mathcal{M}_{2}$ are classified correctly, and the classifier $r_{t_{0}+t}$ will generalize well on $\bar{\mathcal{M}}_{2}$.

By Lemma 4.4, we know that for $t=\widetilde{O}\left(\frac{1}{\varepsilon_{1}^{3}\eta_{2}r}\right)$ we have $\widehat{L}_{\bar{\mathcal{M}}_{1}}(g_{t_{0}+t})=O(\sqrt{\varepsilon_{1}/q^{3}})$. Thus applying Lemma B.1, we obtain that as long as $\varepsilon_{1}\geq\frac{1}{\sqrt{N}}$ (which is implied by Assumption 3.3)

This implies that for $x_{1}\sim\mathcal{D}_{x_{1}}$, applying Lemma A.8 gives us

Moreover, applying Lemma A.6 on $r_{t_{0}+t}$ with $\|W_{t_{0}+t}\|_{F}^{2}\leq\|W_{t_{0}}\|_{F}^{2}+\|W_{t_{0}+t}-W_{t_{0}}\|_{F}^{2}\lesssim\left(d\log^{2}\frac{1}{\varepsilon}\right)$ by Lemma 4.2 and Lemma 4.3, we have that

where we used the fact that $\varepsilon_{1}\leq\kappa^{2}p^{2}q^{3}$.

Here the last step uses the definition of $\varepsilon_{1}$ that $\varepsilon_{1}\leq\kappa^{2}p^{2}q^{3}$. ∎

B.2 Proof of Theorem 3.5

We will prove Theorem 3.5 using Lemma 5.3 by roughly arguing that the predictions made by $r_{t}$ will be heavily influenced by a vector $\alpha$ in the low rank span of examples from $\bar{\mathcal{M}}_{2}$. With high probability, this vector $\alpha$ will be noisy and not align well with the ground truth $w^{\star}$, leading to mispredictions.

Recall that $\varepsilon_{2}^{\prime}$ denotes the stopping criterion used in Theorem 3.5 and $\varepsilon_{3}=d^{-1/32}\frac{1}{\varepsilon_{2}^{\prime 2}}$. Using Lemma 5.3, we know that w.h.p.

By Lemma F.2 we know that w.h.p. over the randomness of $x_{1}^{(i)}$’s, for $\alpha\in\text{span}\{x_{1}^{(i)}\}_{i\in\bar{\mathcal{M}}_{2}}$ we have as long as $Np\leq d/2$: $\frac{\langle\alpha,w^{\star}\rangle}{\|\alpha\|_{2}\|w^{\star}\|_{2}}\leq 0.9$. For every randomly chosen $x_{1}$, we can also write $x_{1}=\gamma w^{\star}+\beta$ where $\beta\bot w^{\star}$ so $\beta$ is independent of $\gamma$, hence

Note that $\langle\alpha,\beta\rangle\sim\mathcal{N}(0,\sigma^{2}\|\alpha\|_{2}^{2}/d)$ with $\sigma\geq 0.1$, and with probability at least $0.1$, $\gamma\leq 2\|\alpha\|_{2}/\sqrt{d}$. This implies that with probability at least $\Omega(1)$ over a randomly chosen $x_{1}$ we can have:

For $\beta$, we know that with probability at least $\Omega(1)$, we have:

Moreover, since $\beta$ is independent of $\gamma$, we know that with probability $\Omega(1)$ both events can happen, in which case:

Thus, since $\|\alpha\|_{2}=\Omega(\sqrt{Np})$ by Lemma 5.3, we know that as long as

However, since $\langle w^{\star},x_{1}\rangle<0$, we know that either $r_{t}(x_{1})<0$, which results in $r_{t}(-x_{1})<0$ but $\langle w^{\star},-x_{1}\rangle>0$. So when $x_{2}=0$, the network classifies $(-x_{1},0)$ incorrectly. On the other hand, we have when $r_{t}(x_{1})>0$ the network will classify $(x_{1},0)$ incorrectly. Since $\langle w^{\star},x_{1}\rangle<0$ and $r_{t}(x_{1})\geq r_{t}(-x_{1})$ holds with probability $\Omega(1)$, this shows that the test error is at least $\Omega(p)$. ∎

Appendix C Proofs for Large Learning Rate Lemmas

To prove Lemma 4.1, we will show that the network will learn all examples with $\mathcal{P}$ component while the learning rate is large. The key to the proof is that although the large learning rate noise only allows the network to search over coarse kernels, $\mathcal{P}$ is still learnable by these kernels because of its linearly-separable structure. To make this precise, we decompose the weights $U_{t}$ Into the signal and noise components, and show that there exists a fixed “target” signal matrix which will classify $\mathcal{P}$ correctly no matter the noise matrix.

Recall our definitions of $f_{t}(B;x)$, $K_{t}(B)$ in (6.1) and (6.2), and that

Recall that Lemma 6.1 leverages the linearly-separable structure of $\mathcal{P}$ to find a “target” signal matrix that correctly classifies $\mathcal{P}$ w.h.p over the noise matrix. We state its proof below.

By proposition A.3, $\|\overline{U}_{t}\|_{F}\leq O\left(\frac{1}{\lambda}\right)$. We apply Lemma A.8 as follows: by Proposition A.7, $\widetilde{U}_{t}$’s entry has marginal distribution $\mathcal{N}(0,\tau_{0}^{2})$ and therefore the column of $\widetilde{U}_{t}$ has distribution $\mathcal{N}(0,\tau_{0}^{2}I_{m\times m})$. Since w.h.p. $\|x\|_{2}\lesssim\sqrt{\log d}$, the coupling Lemma A.8 gives

On the other hand, we also have by Proposition A.5, using the fact that $\max_{i}\|[\overline{U}_{i}]\|_{2}\lesssim\frac{1}{\sqrt{m}\lambda}$, w.h.p.

Here in the last inequality we used the fact that the network is sufficiently over-parameterized so that $\varepsilon_{s}=\widetilde{O}(\tau_{0}\lambda)$.

Using (C.4), noting that our choice of $m,\lambda,\tau_{0}$ satisfies $\tau_{0}\log d=o(\varepsilon_{1})$, we conclude

For the term $N_{U_{t}}(u,U^{*};x)$, we know that

Note that entries of $\widetilde{W}_{t}x_{1}$ are i.i.d. random Bernoulli($1/2$), thus we know that w.h.p.

Thus, by our choice that $m^{-1/3}=O(\varepsilon_{1})$ and $\sqrt{d}\varepsilon_{s}=O(\varepsilon_{1})$,

By definition of $w^{\star}$, we know that

Now we wish to argue that even though the noise matrix is changing, gradient descent will still find the fixed target signal matrix $U^{\star}$. This leverages the fact that once we fix the activation patterns, we can view each step of the optimization as gradient descent with respect to a convex, but changing, function. Below we provide a proof of Theorem 6.2, which allows for optimization of this changing function.

For the sake of contradiction, we assume that $K_{t}(z_{t})\geq c^{\star}+\mu$ for all $t\leq T$. Using the definition of $K_{t}^{\lambda}$, we have that the update rule of $z_{t}$ can be written as

Assuming that $\|z_{t}-z^{\star}\|_{2}\leq R$, we have that as long as $\lambda R^{2}\leq\frac{1}{100}\mu$ and $\eta\leq\frac{\mu}{100(\lambda^{2}R^{2}+L^{2})}$, we have:

C.2 Proof of Lemma 4.2

In this section, we will often consider the activation patterns on the inputs $z,z-\zeta,z+\zeta$ at various time steps. For convenience, we have the following definition:

For any $s$, and vector $w$, let $\mathcal{E}^{w}_{s}\triangleq\{i\in[m]:[\widetilde{V}_{s}]_{i}w\geq 0\}$ denote the set of neurons that have positive pre-activation on the input $w$ (with weights $\widetilde{V}_{s}$), and $\bar{\mathcal{E}}^{w}_{s}\triangleq\{i\in[m]:[\widetilde{V}_{s}]_{i}w<0\}$ be the set of neurons with negative pre-activations on the input $w$. (We will mostly be interested in the quantities $\mathcal{E}^{z-\zeta},\bar{\mathcal{E}}^{z-\zeta},\mathcal{E}^{z+\zeta},\bar{\mathcal{E}}^{z+\zeta}$ and their intersections.)

Let $Q_{t}\triangleq\operatorname*{\text{diag}}(v)\overline{V}_{t}$. Then, we have that

Towards bounding the terms in equation (C.25), we will need to reason about the activations patterns of $z,z-\zeta,z+\zeta$ at various time steps. We first show that the activation patterns of $z-\zeta$ and $z+\zeta$ have to agree in most of neurons except an $\approx r$ fraction of them. This will be useful to show that the second term of the RHS of equation (C.25) is small.

In the setting of Lemma C.2, w.h.p over the randomness of the initialization and all the randomness in the algorithm, for every $t\leq\textup{poly}(d),i\in[m]$, $i\in\mathcal{E}^{z-\zeta}_{t}\oplus\mathcal{E}^{z+\zeta}_{t}$ implies that $|[\widetilde{V}_{t}]_{i}z|\lesssim\tau_{0}r\sqrt{\log d}$. Moreover, the size of the set $\mathcal{E}^{z-\zeta}_{t}\oplus\mathcal{E}^{z+\zeta}_{t}$ is bounded by

Recall that $\|\zeta\|_{2}=r$ and by Proposition A.7 $[\widetilde{V}_{t}]_{i}$ has distribution $\mathcal{N}(0,\tau_{0}^{2}I_{d\times d})$. Therefore, by standard Gaussian concentration and union bound, with high probability over the randomness of the initialization and the algorithm, for all $t\leq\textup{poly}(d)$,

Moreover, note that $\Pr\left[[|\widetilde{V}_{t}]_{i}z|\leq\tau_{0}r\sqrt{\log d}\right]\lesssim r\sqrt{\log d}$. By the independence between $[\widetilde{V}_{t}]_{i}$’s and standard concentration inequalities (Bernstein inequality), we have that with high probability, there are at most $rm\sqrt{\log d}+\log d$ entries $i\in[m]$ satisfying $|[\widetilde{V}_{t}]_{i}z|\leq\tau_{0}r\sqrt{\log d}$. Together with the first part of the lemma, and that $m$ is sufficiently large so that $rm\sqrt{\log d}+\log d\lesssim rm\sqrt{\log d}$, we complete the proof of equation (C.26). ∎

We use the lemma above to conclude that the second term in the decomposition (C.25) is at most on the order of $r^{2}/\lambda$.

In the setting of Lemma C.2, we have that

By the definition of our algorithm, before annealing the learning rate, we have

Using Proposition A.3 and that $|v_{i}|=\frac{1}{\sqrt{m}}$, we have that $\|[Q_{t}]_{i}\|_{2}\lesssim\frac{1}{\lambda m}$. It follows that

Equation above and equation (C.30) complete the proof. ∎

The following lemma decomposes $Q$ into a sum of the contribution of the gradient from all the previous steps.

In the setting of Lemma C.2, let $\Delta Q_{t}\triangleq\operatorname*{\text{diag}}(v)\nabla_{V}\widehat{L}(U_{t})$. ($\Delta Q_{t}$ can be viewed as the raw change of $Q_{t}$ at the time step $t$ without considering the effect of the regularizer.) We have that

Using the fact that $\|z\|_{2}\leq 1$ we complete the proof. ∎

Define the analog of $\mathcal{E}_{s}^{w}$ with $V_{t}$ to compute the activation pattern: for any $s$, and vector $w$, let $\mathcal{G}^{w}_{s}\triangleq\{i\in[m]:[V_{s}]_{i}w\geq 0\}$ and define $\bar{\mathcal{G}}^{w}_{s}\triangleq\{i\in[m]:[V_{s}]_{i}w<0\}$ similarly.

Suppose at some iteration $s$, $z-\zeta$ and $z+\zeta$ have the same activation pattern at neuron $i$ and $j$ in the sense that $i,j\in\mathcal{G}^{z-\zeta}_{s}\cap\mathcal{G}^{z+\zeta}_{s}$, or $i,j\in\bar{\mathcal{G}}^{z-\zeta}_{s}\cap\bar{\mathcal{G}}^{z+\zeta}_{s}$. Then the corresponding gradient update at that iteration for the weight vectors associated with $i$ and $j$ are the same up to a potential sign flip:

Moreover, suppose we have that $i,j$ satisfy that $[\widetilde{V}_{s}]_{i}x\gtrsim\tau_{0}r\sqrt{\log d}$ and $[\widetilde{V}_{s}]_{j}x\gtrsim\tau_{0}r\sqrt{\log d}$ (or $[\widetilde{V}_{s}]_{i}x\lesssim-\tau_{0}r\sqrt{\log d}$ and $[\widetilde{V}_{s}]_{j}x\lesssim-\tau_{0}r\sqrt{\log d}$) for $x\in\{z-\zeta,z+\zeta\}$, then the same conclusion holds for $i$ and $j$.

Note that by definition, $[\Delta Q_{s}]_{i}=v_{i}[\nabla_{V}\widehat{L}(U_{s})]_{i}$, and thus it suffices to prove that $v_{i}[\nabla_{V}\widehat{L}(U_{s})]_{i}=v_{j}[\nabla_{V}\widehat{L}(U_{s})]_{j}$. By Proposition A.1, we have that

Note that $x_{2}$ can only take (a positive scaling of) four values $z-\zeta,z,z+\zeta,0$. We claim that for every choice of these four values, for the $i,j$ satisfying the condition of the lemma, we have

Note that the equation above together with $v_{i}^{2}=v_{j}^{2}=1$ suffices to complete the proof.

Now to prove the second part of the lemma, suppose $i,j$ satisfy that $[\widetilde{V}_{s}]_{i}x\gtrsim\tau_{0}r\sqrt{\log d}$ and $[\widetilde{V}_{s}]_{j}x\gtrsim\tau_{0}r\sqrt{\log d}$ for $x\in\{z-\zeta,z+\zeta\}$. Using $\|[\widetilde{V}_{s}]_{i}\|_{2}\leq\frac{1}{\lambda\sqrt{m}}$ from Proposition A.3, we have that $[V_{s}]_{i}z\geq[\widetilde{V}_{s}]_{i}z-|[\overline{V}_{s}]_{i}z|\gtrsim\tau_{0}r\sqrt{\log d}-O(\frac{1}{\lambda\sqrt{m}})\geq\tau_{0}r\sqrt{\log d}$ where used the assumption that $1/\lambda=\textup{poly}(d)$ and $m=\textup{poly}(d/\tau_{0})$. Therefore, we conclude that $i,j\in\mathcal{G}^{z-\zeta}_{s}\cap\mathcal{G}^{z+\zeta}_{s}$. Now by the first lemma of the lemma we complete the proof. ∎

Now we are ready to bound the first term on the RHS of equation C.25, which is the crux of the proofs in this section. The key here is to get a bound that scales quadratically in $r$.

In the setting of Lemma C.2, let $\Delta Q_{s}$ be defined in Proposition C.5. Then, we have that

As a direct corollary of the equation above and Proposition C.5, we have that

By the set operations and the facts that $\mathcal{E}^{z-\zeta}_{t}\cap\mathcal{E}^{z+\zeta}_{t}\subset\mathcal{E}^{z}_{t}$ and that $\mathcal{E}^{z}_{t}\subset\mathcal{E}^{z-\zeta}_{t}\cup\mathcal{E}^{z+\zeta}_{t}$, we have that

where the $\lesssim,\gtrsim$ notations hide universal constants that make the first conclusion of Proposition C.3 true. By the second part of Proposition C.3 (or more directly equation (C.28)), we have that $\mathcal{F}^{+}_{s}\subset\mathcal{E}^{z-\zeta}_{s}\cap\mathcal{E}^{z+\zeta}_{s}$, and $\mathcal{F}^{-}_{s}\subset\bar{\mathcal{E}}^{z-\zeta}_{s}\cap\bar{\mathcal{E}}^{z+\zeta}_{s}$. By Proposition C.6, we have that for any $i,j\in\mathcal{F}^{-}_{s}$, $[\Delta Q_{s}]_{i}=[\Delta Q_{s}]_{j}$. For notational simplicity, let $A=\mathcal{E}^{z+\zeta}_{t}\backslash\mathcal{E}^{z}_{t}$ and $B=\mathcal{E}^{z}_{t}\backslash\mathcal{E}^{z-\zeta}_{t}$. Therefore it follows that

where in the last inequality we use that for any $i,j\in\mathcal{F}^{-}_{s}$, $[\Delta Q_{s}]_{i}=[\Delta Q_{s}]_{j}$, and the fact that $\|[\Delta Q_{s}]_{i}\|_{2}=\frac{1}{\sqrt{m}}\|[\nabla_{V}\widehat{L}(U_{s})]_{i}\|_{2}\leq 1/m$ (by Proposition A.2.)

Note that the distribution of $([\widetilde{V}_{s}]_{i},[\widetilde{V}_{t}]_{i}$’s are independent across the choice of $i$. Thus we will compute $\Pr[i\in\mathcal{E}^{z+\zeta}_{t},i\notin\mathcal{E}^{z}_{t},i\in\mathcal{F}^{+}_{s}]-\Pr[i\in\mathcal{E}^{z}_{t},i\notin\mathcal{E}^{z-\zeta}_{t},i\in\mathcal{F}^{+}_{s}]$ and then apply concentration concentration inequality for the sum. Note that the event here depends on three quantities $[\widetilde{V}_{s}]_{i}z$, $[\widetilde{V}_{t}]_{i}z$, and $[\widetilde{V}_{t}]_{i}\zeta$. First of all, $[\widetilde{V}_{t}]_{i}\zeta$ is independent of these other two because $\zeta$ is orthogonal to $z$ and $[\widetilde{V}_{t}]_{i}$ and $[\widetilde{V}_{s}]_{i}$ have spherical covariance matrices.

By the definition of $\widetilde{V}_{s},\widetilde{V}_{t}$, we can express their relationship by writing $[\widetilde{V}_{t}]_{i}z=(1-\eta_{1}\lambda)^{t-s}[\widetilde{V}_{s}]_{i}z+[\Xi_{t,s}]_{i}z$, where $\Xi_{t,s}=\eta_{1}\sum_{j\in[t-s]}(1-\eta_{1}\lambda)^{t-s-j}\xi_{s+j}$. Recall that by proposition A.7, we have $[\widetilde{V}_{s}]_{i}z\sim\mathcal{N}(0,\tau_{0}^{2})$ and $[\Xi_{t,s}]_{i}z$ are two independent Gaussians. Let $\sigma_{t,s}$ be the variance of $[\Xi_{t,s}]_{i}z$. We compute $\sigma_{t,s}$ by observing that

Note that $\zeta^{\top}z=0$, thus $[\widetilde{V}_{s}]_{i}z$ is independent of $[\widetilde{V}_{t}]_{i}\zeta$ conditioned on $[\widetilde{V}_{t}]_{i}z$, for every $s\leq t$ . For notational simplicity, let $Y_{1}=[\widetilde{V}_{s}]_{i}z$, $Y_{2}=[\widetilde{V}_{t}]_{i}z$, and $Y_{3}=[\widetilde{V}_{t}]_{i}\zeta$, and $\kappa=O(\tau_{0}r\sqrt{\log d})$ where the big O notation hide the same constant factor used in defining $\mathcal{F}^{+}_{s}$ in equation (C.40). Let $Y_{4}=[\Xi_{t,s}]_{i}z=Y_{1}-\beta Y_{2}$ where $\beta=\eta_{1}(1-\eta_{1}\lambda)^{t-s}\gtrsim 1$ (because $t\leq 1/(\eta_{1}\lambda)$). Note that by the calculation above, $Y_{4}$ has standard deviation $\sigma_{s,t}$ which is bounded from below by $\tau_{0}\sqrt{\lambda\eta_{1}(s-t)}$. Then, we have that

Now by equation (C.42) and standard concentration inequality, and the fact that $m$ is sufficiently large, we have that with high probability,

Using standard concentration inequality and the fact that $m$ is sufficiently large, we have that with high probability,

We can also prove the same bound for $|B\cap\mathcal{F}^{c}_{s}|$ analogously. Using equation (C.41) and the several equations above, we conclude that

where the last step uses that the condition that $t\leq 1/(\eta_{1}\lambda)$.

Now combining the Propositions above we are ready to prove Lemma 4.2.

Using triangle inequality, Proposition A.8, and equation (A.6) of Proposition A.5, we have that for any $x$ of norm $O(1)$,

C.3 Proof of Lemma 6.3

The proof of Lemma 6.3 relies on the fact that a smaller learning rate preserves the noise generated from the timestep before annealing. This allows us to reason that the new activations are similar to the original before reducing the learning rate.

where $\Xi_{t}:=\eta_{2}\sum_{j\leq t}(1-\lambda\eta_{2})^{t-j}\xi_{t_{0}+j}$. By properties of a sum of Independent Gaussians, we have $[\Xi_{t}]_{i}\sim\mathcal{N}(0,\sigma_{t}^{2}I)$ where $\sigma_{t}$ is the standard deviation of each entry of $\Xi_{t}$. We also have that $\Xi_{t}$ is independent of $\widetilde{U}_{t_{0}}$. Moreover, for every $t\leq\frac{1}{\eta_{2}\lambda}$, the standard deviation $\sigma_{t}$ can be bounded by

(Note that since $\eta_{2}\ll\eta_{1}$, we should expect that the standard deviations satisfy $\sigma_{t}\ll\sigma_{0}$. That is, the additional randomness introduced in the pre-activation is small.)

On the other hand, for every $t\leq\frac{1}{\eta_{2}\lambda}$, the contribution of $\widetilde{U}_{t_{0}}$ to $U_{t+t_{0}}$ is still present because the entry of $(1-\eta_{2}\lambda)^{t}[\widetilde{U}_{t_{0}}]_{i}$ has variance at least on the order of the variance of the entries of $[\widetilde{U}_{t_{0}}]_{i}$, which is $\gtrsim\tau_{0}^{2}$. This also implies that the variance of the entries of $\widetilde{U}_{t_{0}+t}$ is lower bounded by the variance of $(1-\eta_{2}\lambda)^{t}[\widetilde{U}_{t_{0}}]_{i}$. This in turn is lower bounded by $\tau_{0}^{2}$ up to constant factor.