Why are Adaptive Methods Good for Attention Models?

Introduction

Stochastic gradient descent (SGD) is the canonical algorithm for training neural networks . SGD iteratively updates model parameters in the negative gradient direction and seamlessly scales to large-scale settings. Though a well-tuned SGD outperforms adaptive methods in many tasks including ImageNet classification (see Figure 1(a)), certain tasks necessitate the use of adaptive variants of SGD (e.g., Adagrad , Adam , AMSGrad ), which employ adaptive learning rates. For instance, consider training an attention model using BERT . Figure 1(e) shows that in spite of extensive hyperparameter tuning, SGD converges much slower than Adam during BERT training.

In this work, we provide one explanation for why adaptivity can facilitate convergence with theoretical and empirical evidence. The significant hint that initializes our work comes from the distribution of the stochastic gradients. For Imagenet, the norms of the mini-batch gradients are typically quite small and well concentrated around their mean. On the other hand, the mini-batch gradient norms for BERT take a wide range of values and are sometimes much larger than their mean value. More formally, while the distribution of the stochastic gradients in Imagenet is well approximated by a Gaussian, the distribution for BERT seems to be heavy-tailed. Such observation leads us to the question: does adaptivity stabilize optimization under heavy-tailed noise?

We provide a positive answer to the above question by performing both theoretical and empirical studies of the convergence of optimization methods under heavy-tailed noise. In this setting, some of the stochastic gradients are much larger than the mean and can excessively influence the updates of SGD. This makes SGD unstable and leads to its poor performance. A natural strategy to stabilize the updates is to clip the magnitude of the stochastic gradients. We prove that indeed it is sufficient to ensure convergence even under heavy-tailed noise. Based on the analysis, we then motivate the design of a novel algorithm (ACClip) that outperforms ADAM on BERT related tasks. Specifically, we make the following contributions:

We empirically show that in tasks on which Adam outperforms SGD (BERT pretraining), the noise in stochastic gradients is heavy-tailed. On the other hand, on tasks where traditionally SGD outperforms Adam (ImageNet training), we show that the noise is well concentrated.

In section 3, we study the convergence of gradient methods under heavy-tailed noise condition where SGD’s performance degrades and its convergence might fail. We then establish (with upper and lower bounds) the convergence of clipped gradient methods under the same condition and prove that they obtain theoretically optimal rates.

Though clipping speeds up SGD, it does not close the gap between SGD and ADAM. In section 4, we motivated the a novel adaptive-threshold coordinate-wise clipping algorithm and in section 5 experimentally show that it outperforms Adam on BERT training tasks.

Adaptive step sizes. Adaptive step sizes during optimization have long been studied . More recently, Duchi et al. developed the Adagrad algorithm that benefits from the sparsity in stochastic gradients. Inspired by Adagrad, several adaptive methods have been proposed in the deep learning community . Recently, there has been a surge in interest to study the theoretical properties of these adaptive gradient methods due to , which pointed out the non-convergence of Adam and proposed an alternative algorithm, AMSGrad. Since then, many works studied different interesting aspects of adaptive methods, see . Another direction of related work is normalized gradient descent, which has been studied for quasi-convex and non-convex settings . In contrast to our work, these prior works assume standard noise distributions that might not be applicable to key modern applications such as attention models, which exhibit heavy-tailed noise. Furthermore, convergence rates of adaptive methods are mostly worse than SGD.

Noise in neural network. There has been little study of the actual stochastic gradient noise distributions in neural network training. To our knowledge, start the topic and observe heavy tailed noise in network training. Our work differs in two important ways: First, we treat the noise as a high dimensional vector, while treat deviations in each coordinate as scaler noises to estimate tail index. Hence, we observe that the example given in is well-concentrated when viewed as a random vector. This is also confirmed by . More experimental comparisons are in Appendix H. Second, we focus on convergence of optimization algorithm, the previously mentioned works focus on Langevin dynamics and escaping saddle points. The convergence rate given in is for global Holder-continuous functions, which restricts the function variations and excludes examples like quadratic functions. Our analysis instead provides the first convergence rates under the standard L-smoothness setting. Further, studies accelerated first order methods under less concentrated noise, however, there “heavy-tailedness” refers to non-sub-Gaussianity.

Heavy-tailed noise in stochastic gradients

To gain intuition about the difference between SGD and adaptive methods, we start our discussion with the study of noise distributions of stochastic gradient that arise during neural network training. In particular, we focus on noise distributions while training two popular deep learning models — BERT and ResNet. Note that BERT and ResNet are typically trained with Adam and SGD (with momentum) respectively and can thus, provide insights about difference between these optimizers.

We first investigate the distribution of the gradient noise norm $\left\lVert g-\nabla f(x)\right\rVert$ in the aforementioned neural network models, where $g$ is the stochastic gradient computed from a minibatch sample. In particular, we fix the model at initialization without doing any updates. We then iterate through the dataset to compute the noise norm for each minibatch. Figure 1 (b) and (f) show these distributions for ResNet50 on ImageNet and BERT on the Wikipedia and books dataset at model initialization respectively. For comparison, we plot distributions of a normalized sum of squared Gaussians, a well-concentrated distribution, and a Levy-$\alpha$-stable distribution, a heavy-tailed distribution, in Figure 1 (c) and (g) respectively. We observe that the noise distribution for BERT appears heavy-tailed, while that of ResNet50 is well-concentrated. Results for noise distributions at other stages of training are displayed in Figure 2.

To support this observation, in Figure 1 (d) and (h) we further show the empirical variance of stochastic gradients with respect to the sample size used in the estimation. The results highlight that while the corresponding estimator converges for Imagenet, the empirical variance does not converge in BERT training even as the sample size approaches $10^{7}$.

From the obeservation that the noise can be heavy-tailed, we hypothesize that this is one major aspect that determines the performance of SGD and adaptive methods. In the rest of the paper, we argue and provide evidence that adaptive methods can be faster than SGD in scenarios where heavy-tailed noise distributions arise. More experiment details can be found in Section 5.

Convergence of gradient methods under heavy-tailed noise

In this section we study the performance of SGD and adaptive methods under heavy-tailed noise. More precisely, we analyze algorithms of the following form

where $x_{k}$ represent the current parameters, $\eta_{k}$ is the step size and $g_{k}$ is the stochastic (mini-batch) gradient evaluated at $x_{k}$. We show that if the stochasticity in the gradient $g_{k}$ is heavy-tailed, it is critical for the step sizes to be adaptive i.e. $\eta_{k}$ must depend on the observed gradients. We propose to use one such algorithm GClip and prove that it obtains optimal convergence rates.

The above assumption with $\alpha=2$ corresponds to the standard variance bound, but in general is weaker. It is indeed possible (e.g. Pareto or $\alpha$-stable Levy random variables) for the variance of $g(x)$ to be unbounded, while simultaneously satisfying assumption 1 for $\alpha<2$. One should note that even if the variance may not actually be infinite in practice, it might be too large to be practically useful. All our analyses and insights carry over to this setting as well.

The possibility that the variance is unbounded has a profound impact on the optimization process.

The issue is that SGD is easily influenced by a single-stochastic gradient, which could be very large and incorrect. A simple strategy to circumvent this issue is to use a biased clipped stochastic gradient estimator. This allows us to circumvent the problem of unbounded variance and ensures optimal convergence rates even under heavy-tailed noise. Our results are summarized in Table 1, and all proofs are relegated to the Appendices.

1 Convergence of Clipped Methods

A simple clipping strategy is to globally clip the norm of the update to threshold $\tau_{k}$:

We refer to this strategy as GClip (Global Clip), as opposed to coordinate-wise clipping which we discuss later. We first state the rates for smooth non-convex functions.

Suppose that $f$ is $L$-smooth and that the stochastic gradients satisfy Assumption 1 for $\alpha\in(1,2]$. Let $\{x_{k}\}$ be the iterates of GClip with parameters $\eta_{k}=\eta=\min\{\frac{1}{4L},\frac{\sigma^{\alpha}}{L\tau^{\alpha}},\frac{1}{24L\tau}\}$ and \tau_{k}=\tau=\max\{2,48^{1/(\alpha-1)}\sigma^{\alpha/(\alpha-1)},8\sigma,\big{(}\frac{f_{0}}{\sigma^{2}K}\big{)}^{\frac{\alpha}{3\alpha-2}}\big{/}L^{\frac{2\alpha-2}{3\alpha-2}},\}. Then for $F_{0}:=f(x_{0})-f^{*}$,

We prove improved rates of convergence for non-smooth strongly-convex functions in a bounded domain. Due to limited space, we relegate the definitions and assumptions to Appendix A.

Suppose that the stochastic gradients satisfy Assumption 4 for $\alpha\in(1,2]$. Let $\{x_{k}\}$ be the iterates of projected GClip (proj-GClip) with clipping parameter $\tau_{k}=Gk^{\alpha-1}$ and steps-size $\eta_{k}=\frac{4}{\mu(k+1)}$. Define the output to be a $k$-weighted combination of the iterates: $\bar{x}_{k}=\sum_{j=1}^{k}jx_{j-1}/(\sum_{j=1}^{k}j)\,.$ Then the output $\bar{x}_{k}$ satisfies:

The rates of convergence for the strongly convex and non-convex cases in Theorem 4 and Theorem 2 exactly match those of the usual SGD rates ($\mathcal{O}(1/\sqrt{k})$ for convex and $\mathcal{O}(k^{-\frac{1}{4}})$ for non-convex) when $\alpha=2$ and gracefully degrade for $\alpha\in(1,2]$. As we will next show, both the strongly convex rates and non-convex rates of GClip are in fact optimal for every $\alpha\in(1,2]$.

2 Theoretic lower bounds

For any $\alpha\in(1,2]$ and any (possibly randomized) algorithm $\mathcal{A}$, there exists a problem $f$ which is 1-strongly convex and 1-smooth ($\mu=1$ and $L=1$), and stochastic gradients which satisfy Assumptions 4 with $G\leq 1$ such that the output $x_{k}$ of the algorithm $\mathcal{A}$ after processing $k$ stochastic gradients has an error

Given any $\alpha\in(1,2]$, smoothness constant $L$, and (possibly randomized) algorithm $\mathcal{A}$, there exists a constant $c_{1}$ and an $L$-smooth function $f$ with stochastic gradients satisfying Assumption 1 for any given $\sigma\geq c_{1}\sqrt{(f(0)-f^{*})L}$ such that the output $x_{k}$ of the algorithm $\mathcal{A}$ after processing $k$ stochastic gradients has an error

Theorem 6, proven in Appendix G, extends the recent work of [2, Theorem 1] to heavy-tailed noise. Here, the lower-bound matches the upper-bound in Theorem 2 up to constants, proving its optimality.

Faster Optimization with Adaptive Coordinate-wise Clipping

The previous section showed that adaptive step sizes (which depend on the gradients) are essential for convergence under heavy-tailed noise, and also showed that GClip provides the optimal rates. There are of course other adaptive methods such as Adam which employs not only the current gradients but also all past gradients to adaptively set coordinate-wise step-sizes. In this section, we study why coordinate-wise clipping may yield even faster convergence than GClip, and show how to modify GClip to design an Adaptive Coordinate-wise Clipping algorithm (ACClip).

The first technique we use is applying coordinate-wise clipping instead of global clipping. We had previously assumed a global bound on the $\alpha$-moment of the norm (or variance) of the stochastic gradient is bounded by $\sigma$. However, $\sigma$ might be hiding some dimension dependence $d$. We show a more fine-grained model of the noise in order to tease out this dependence.

Suppose we run GClip under Assumption 2 to obtain the sequence $\{x_{k}\}$. If $f$ is $\mu$-strongly convex, with appropriate step-sizes and averaging, the output $\bar{x}_{k}$ satisfies

Thus, the convergence of GClip can have a strong dependence on $d$, which for large-scale problems might be problematic. We show next that using coordinate-wise clipping removes this dependency:

Suppose we run CClip under the Assumption of 2 with $\tau_{k}=Bk^{\alpha-1}$ to obtain the sequence $\{x_{k}\}$. Then, if $f$ is $\mu$-strongly convex, with appropriate step-sizes and averaging, the output $\bar{x}_{k}$ satisfies

Note that $\lVert B\rVert_{2}\leq\lVert B\rVert_{\alpha}$. CClip has a worst-case convergence independent of $d$ under the coordinate-wise noise model. Similar comparison between GClip and CClip can be done for non-convex conditions too, but we skip for conciseness. Though we only compare upper-bounds here, when the noise across coordinates is independent the upper bounds may be tight (see Lemma 12).

2 Online moment estimation

We now present the second technique that is motivated by our observation in Figure 2. There, the distribution of gradient noise at the beginning of different epochs is shown during training for BERT with Wikipedia (top) as well as ResNet with ImageNet (bottom). The result highlights that the noise distribution is not only heavy-tailed, but also non-stationary during BERT training and becomes increasingly more concentrated. In contrast, for the ResNet model the noise distribution remains mostly unchanged.

Since the scale of the noise changes drastically during training for BERT model and our theoretical analysis suggest that we should clip proportional to the noise level, we propose to use an exponential moving average estimator to estimate the moment and clip the gradient accordingly (line 4,5 of Alg 1). This, combined with the momentum term leads to our proposed ACClip algorithm in Algorithm 1. On a high level, the algorithm applies clipping to the momentum term, where the clipping threshold is proportional to the estimated moment using an exponential moving average. From our experiment, we found the conservative choice of $\alpha=1$ leads to the best performance.

Experiments

In this section, we first verify the effect of coordinate-wise clipping and moment estimation introduced in Section 4. We then perform extensive evaluations of ACClip on BERT pre-training and fine-tuning tasks and demonstrate its advantage over Adam in Section 5.2. For completeness, an experiment on ImageNet is included in Appendix I. Finally, we start with a few more experiments on the noise distribution in neural network training.

In this section we instantiate the argument in Section 4 with a set of experiments. As seen in Figure 3(b), global clipping improves the vanilla SGD algorithm but is still far from the ADAM baseline. We apply two techniques (coordinate-wise clipping and online moment estimation) onto the clipped SGD algorithm analyzed in Section 3. We use a set of experiments on Transformer-XL training to demonstrate the effect of each technique.

We train a 6-layer Transformer-XL model on PTB dataset as a proof of concept. Our main experiments will be on BERT pretraining and finetuning described in the next subsection 5.2. We use adapt the author’s github repohttps://github.com/kimiyoung/transformer-xl/tree/master/pytorch, and replace the number of layers of the base model by 6. We then select the PTB data as input and set the maximum target length to be 128. The results are shown in Figure 3(b).

Observations

From Figure 3(b), we can tell that global clipping (orange curve) indeed speeds up vanilla SGD but is still much worse compared to the ADAM baseline provided by the code base. After replacing global clipping with coordinate-wise clipping, we see that the performance is already comparable to the ADAM baseline. Finally, after using the moment estimation to determine the clipping threshold, we are able to achieve faster convergence than ADAM.

2 Performance of ACClip for BERT pre-training and fine-tuning

We now evaluate the empirical performance of our proposed ACClip algorithm on BERT pre-training as well fine-tuning using the SQUAD v1.1 dataset. As a baseline, we use Adam optimizer and the same training setup as in the BERT paper . For ACClip, we set $\tau=1,\textit{learning rate}=\text{1e-4},\beta_{1}=0.9,\beta_{2}=0.99,\epsilon=\text{1e-5}$ and $\textit{weight decay}=\text{1e-5}$. We compare both setups on BERT models of three different sizes, BERTbase with 6 and 12 layers as well as BERTlarge with 24 layers.

Figure 3(b) and 3(c) shows the loss for pretraining BERTbase using SGD with momentum, GClip, Adam and ACClip. The learning rates and hyperparameters for each method have been extensively tuned to provide best performance on validation set. However, even after extensive tuning, there remains a large gap between (clipped) SGD momentum and adaptive methods. Furthermore, clipped SGD achieves faster convergence as well as lower final loss compared to standard SGD. Lastly, the proposed optimizer ACClip achieves a lower loss than the Adam. Table 2 further shows that ACClip achieves lower loss and higher masked-LM accuracy for all model sizes.

Next, we evaluate ACClip on the SQUAD v1.1 fine-tuning task. We again follow the procedure outlined in and present the results on the Dev set in Table 3. Both for F1 as well as for exact match, the proposed algorithm outperforms Adam on all model sizes. The experimental results on BERT pretraining and fine-tuning indicate the effectiveness of the proposed algorithm.

3 Noise Patterns in BERT and ImageNet Training

In our initial analysis in Figure 1, we observe that training an attention model on Wikipedia leads to heavy-tailed noise whereas training a ResNet on ImageNet data leads to well-concentrated noise. Here, we aim to disentangle the effect that model architecture and training data have on the shape of gradient noise. To this end, we measure the distribution of the gradient noise norm in an Attention and a ResNet model on both Wikipedia and synthetic Gaussian data. We used BERTbase as the Attention model, and the ResNet is constructed by removing the self-attention modules within the transformer blocks. Gaussian synthetic data is generated by replacing the token embedding layer with normalized Gaussian input. The resulting noise histograms are shown in Figure 4. The figure shows that the Attention model leads to heavy-tailed noise independently of input data. For the ResNet model, we observe that Gaussian input leads to Gaussian noise, whereas Wikipedia data leads to be heavy-tailed noise. We thus conclude that the heavy-tailed noise pattern results from both the model architecture as well as the data distribution.

Discussion

One immediate extension from this work is to view RMSProp as a clipping algorithm and prove its convergence under shifting noise. The update for RMSProp and ACClip with $\beta_{1}=0$ can be written with effective step-sizes $h_{\text{rms}}$ and $h_{\text{clip}}$ respectively as below:

Given any set of parameters for RMSProp, if we set the parameters for ACClip as

In summary, our work theoretically and empirically ties the advantage of adaptive methods over SGD to the heavy-tailed nature of gradient noise. A careful analysis of the noise and its impact yielded two insights: that clipping is an excellent strategy to deal with heavy-tailed noise, and that the ACClip yields state of the art performance for training attention models. Our results add to a growing body of work which demonstrate the importance of the structure of the noise in understanding neural network training. We believe additional such investigations into the source of the heavy tailed-ness, as well as a characterization of the noise can lead to further insights with significant impact on practice.

Broader impact

We study convergence rates of gradient methods under a more relaxed noise condition. The result under this setting reaches conclusions that are closer to practice compared to results under the standard setting. Hence, our work provides one way to bridge the theory-practice gap and can facilitate more future works in this direction.

References

Appendix A Additional definitions and assumptions

Here we describe some of the formal assumptions which were previously skipped.

We define the standard notion of smoothness.

$f$ is $L$-smooth, i.e. there exist positive constants $L$ such that $\forall x,y$, $f(y)\leq f(x)+\langle\nabla f(x),y-x\rangle+\frac{L}{2}\lVert y-x\rVert^{2}\,.$

Note that we only need the smoothness assumption for non-convex functions.

A.2 Assumptions in the strongly convex setting

For strongly-convex optimization, instead of bounding the noise, we assume that the stochastic oracle has bounded moment.

Note that the above assumption implies a uniform bound on gradient norm. Such bound is necessary for nonsmooth strongly convex problems, as one can no longer factor out the gradient norm using the smoothness assumption. See for example, .

$f$ is $\mu$-strongly convex, if there exist positive constants $\mu$ such that $\forall x,y$,

The strong convexity assumption and the bounded gradient assumption implies that the domain is bounded, which we state explicitly below,

We look for a solution $x$ within a bounded convex set $\mathcal{X}$.

We didn’t upper bound the domain diameter as it is not used explicitly in the proof. To ensure all updates are within a domain, we use the projected version of (GClip) defined as follows:

The projection operator $x=\operatorname{proj}_{\mathcal{X}}(y)$ finds the point $x\in\mathcal{X}$ that has the least distance to $y$.

Appendix B Effect of global clipping on variance and bias

We focus on (GClip) under stochastic gradients which satisfy Assumption 1.

By the fact that $\hat{g}(x)\leq\tau$, we get

As we increase the clipping parameter $\tau$, note that the variance (the first term in Lemma 9) increases while the bias (which is the second term) decreases. This way, we can carefully trade-off the variance of our estimator against its bias, thereby ensuring convergence of the algorithm.

Appendix C Non-convex Rates (Proof of Theorem 2)

The lemma in the previous section can be readily used in the nonsmooth strongly convex setting. However, we need a variant of Lemma 9 in the smooth case.

For any $g(x)$ suppose that assumption 1 holds with $\alpha\in(1,2]$. If $\|\nabla f(x)\|\leq\tau/2$, then the estimator $\hat{g}:=\min\{1,\tau/\|g_{k}\|\}g_{k}$ from (GClip) with global clipping parameter $\tau\geq 0$ satisfies:

The expectation is taken with respect to the randomness in noise. The second last inequality follows by the fact that $\|\nabla f(x)-\hat{g}(x)\|<2\tau$.

Next, we need a subprocedure at the end proof of Lemma 2 from .

At each iteration, we consider two cases, either $\|\nabla f(x_{k})\|<\tau/2$ or $\|\nabla f(x_{k})\|\geq\tau/2$.

Here the last step used the AM-GM inequality. Then, taking expectation in both sides and using Lemma 10 gives

In the last step we used $\{\eta_{k}=\eta\leq\frac{1}{4L}\}$.

Recall $\hat{g}_{k}=\min\{1,\tau/\|g_{k}\|\}g_{k}$ and parameter choices $\eta_{k}=\eta=\min\{\frac{1}{4L},\frac{1}{L\tau^{\alpha}},\frac{1}{24L\tau}\}$ and $\tau_{k}=\tau=\max\{2,48^{1/(\alpha-1)}\sigma^{\alpha/(\alpha-1)},8\sigma,\sigma K^{\frac{1}{3\alpha-2}}\}$. We use $\nabla f$ as a shorthand for $\nabla f(x_{k})$:

The second line follows by Lemma 11 and (C). The third line follows by that $\tau\geq 2$, and $\|\nabla f\|\geq\tau/2$ imply $\tfrac{p}{2}\|\nabla f\|^{2}\geq p\|\nabla f\|/3.$ Then, by $\tau\geq 48^{1/(\alpha-1)}\sigma^{\alpha/(\alpha-1)}$, we have $\tfrac{\sigma^{\alpha}}{(\tau/4)^{\alpha-1}}\leq\tfrac{1}{12}$. By $\tau\geq 8\sigma$, $\tfrac{8}{3}\sigma\leq\tau/3\leq\|\nabla f\|/6$.

The last inequality above follows by $\frac{1}{24L\tau}$.

Rearrange and sum the terms above for some fixed step-size and threshold $\{\tau_{k}=\tau\}$ to get

Since we use a stepsize $\eta\leq\frac{1}{L\tau^{\alpha}}$, we can simplify $T_{2}$ as

Denote $F_{0}=f(x_{0})-f^{\star}$ to ease notation. Then, adding $T_{2}$ back to $T_{1}$ and using a threshold $\tau\geq\sigma K^{\frac{1}{3\alpha-2}}$ we get

This proves the statement of the theorem. ∎

Appendix D Strongly-Convex Rates (Proof of Theorem 4)

The first inequality follows by the nonexpansivity of projections onto convex sets.

After taking expectation and apply the inequality from Lemma 9, we get

Then take $\eta_{k}=\frac{4}{\mu(k+1)},\tau_{k}=Gk^{\frac{1}{\alpha}}$ and multiply both side by $k$, we get

Notice that $\sum_{k=1}^{K}k^{\frac{2-\alpha}{\alpha}}\leq\int_{0}^{K+1}k^{\frac{2-\alpha}{\alpha}}dk\leq(K+1)^{2/\alpha}$. Sum over $k$ and we get

Devide both side by $\frac{K(K+1)}{2}$ and we get

Notice that for $K\geq 1$, $K^{-1}\leq 2(K+1)^{-1}$. We have

The theorem then follows by Jensen’s inequality. ∎

Appendix E Effect of coordinate-wise moment bound

We now examine how the rates would change if we replace Assumption 4 with Assumption 2.

We now look at(GClip) under assumption 2.

The proof of both the convex and non-convex rates following directly from the following Lemma.

For any $g(x)$ suppose that assumption 2 with $\alpha\in(1,2]$. Then suppose we have a constant upper-bound

Note that the function $(\cdot)^{\alpha/2}$ is concave for $\alpha\in(1,2]$. Using Jensen’s inequality we can rewrite as:

Since the right hand-side can be as large as $d^{\frac{\alpha}{2}-1}\lVert B\rVert_{\alpha}^{\alpha}$, we have our first inequality. On the other hand, we also have an upper bound below:

where $\lVert B\rVert^{\alpha}_{\alpha}=\sum_{i=1}^{d}B_{i}^{\alpha}$. Thus, we have shown that

We know that Jensen’s inequality is tight when all the co-ordinates have equal values. This means that if the noise across the coordinates is linearly correlated the lower bound is tighter, whereas the upper bound is tighter if the coordinates depend upon each other in a more complicated manner or are independent of each other. ∎

Substituting this bound on $G$ in Theorems 4 and 2 gives us our corollaries.

E.2 Convergence of CClip (Proof of Theorem 8)

The proof relies on the key lemma which captures the bias-variance trade off under the new noise-assumption and coordinate-wise clipping.

For any $g(x)$ suppose that assumption 2 with $\alpha\in(1,2]$ holds. Denote $g_{i}$ to be $i_{th}$ component of $g(x)$, $\nabla f(x)_{i}$ to be $i_{th}$ component of $\nabla f(x)$. Then the estimator $\hat{g}(x)=[\hat{g}_{1};\cdots;\hat{g}_{d}]$ from (CClip) with clipping parameter $\tau=[\tau_{1};\tau_{2};\cdots;\tau_{d}]$ satisfies:

Apply Lemma 9 to the one dimensional case in each coordinate. ∎

After taking expectation and apply the inequality from Lemma 9, we get

Then take $\eta_{k}=\frac{4}{\mu(k+1)},\tau_{i}=B_{i}k^{\frac{1}{\alpha}}$ and multiply both side by $k$, we get

Notice that $\sum_{k=1}^{K}k^{\frac{2-\alpha}{\alpha}}\leq\int_{0}^{K+1}k^{\frac{2-\alpha}{\alpha}}dk\leq(K+1)^{2/\alpha}$. Sum over $k$ and we get

Devide both side by $\frac{K(K+1)}{2}$ and we get

Notice that for $K\geq 1$, $K^{-1}\leq 2(K+1)^{-1}$. We have

The theorem then follows by Jensen’s inequality. ∎∎

Appendix F Lower Bound (Proof of Theorem 5)

We consider the following simple one-dimensional function class parameterized by $b$:

Also suppose that for $\alpha\in(1,2]$ and $b\in[0,1/2]$ the stochastic gradients are of the form:

Note that the function class (3) has $\mu=1$ and optimum value $f_{b}(b)=0$, and the $\alpha$-moment of the noise in (4) satisfies $G=B\leq 1$. Thus, we want to prove the following:

For any $\alpha\in(1,2]$ there exists a distribution $\chi_{b}$ such that the stochastic gradients satisfy (4). Further, for any (possibly randomized) algorithm $\mathcal{A}$, define $\mathcal{A}_{k}\left(f_{b}+\chi_{b}\right)$ to be the output of the algorithm $\mathcal{A}$ after $k$ queries to the stochastic gradient $g(x)$. Then:

Our lower bound construction is inspired by Theorem 2 of . Let $\mathcal{A}_{k}(f_{b}+\chi_{b})$ denote the output of any possibly randomized algorithm $\mathcal{A}$ after processing $k$ stochastic gradients of the function $f_{b}$ (with noise drawn i.i.d. from distribution $\chi_{b}$). Similarly, let $\mathcal{D}_{k}(f_{b}+\chi_{b})$ denote the output of a deterministic algorithm after processing the $k$ stochastic gradients. Then from Yao’s minimax principle we know that for any fixed distribution $\mathcal{B}$ over $[0,1/2]$,

In this rest of the proof, we will try obtain a lower bound for the right hand side above.

We now describe our construction of the three quantities to be defined: the problem distribution $\mathcal{B}$, the noise distribution $\chi_{b}$, and the bijective mapping $h(\cdot)$. All of our definitions are parameterized by $\alpha\in(1,2]$ (which is given as input) and by $\epsilon\in(0,1/8]$ (which represents the desired target accuracy). We will pick $\epsilon$ to be a fixed constant which depends on the problem parameters (e.g. $k$) and should be thought of as being small.

Problem distribution: $\mathcal{B}$ picks $b_{0}=2\epsilon$ or $b_{1}=\epsilon$ at random i.e. $\nu\in\{0,1\}$ is chosen by an unbiased coin toss and then we pick

Noise distribution: Define a constant $\gamma=(4\epsilon)^{1/(\alpha-1)}$ and $p_{\nu}=(\gamma^{\alpha}-2\nu\gamma\epsilon)$. Simple computations verify that $\gamma\in(0,1/2]$ and that

Then, for a given $\nu\in\{0,1\}$ the stochastic gradient $g(x)$ is defined as

To see that we have the correct gradient in expectation verify that

Next to bound the $\alpha$ moment of $g(x)$ we see that

The above inequality used the bounds that $\alpha\geq 1$, $x\in[0,1/2]$, and $\gamma\in(0,1/2]$. Thus $g(x)$ defined in (6) satisfies condition (4).

Bijective mapping: Note that here the only unknown variable is $\nu$ which only affects $p_{\nu}$. Thus the mapping is bijective as long as the frequencies of the events are preserved. Hence given a stochastic gradient $g(x_{i})$ the mapping we use is:

Given the definitions above, the output of algorithm $\mathcal{D}_{k}$ is thus simply a function of $k$ i.i.d. samples drawn from the Bernoulli distribution with parameter $p_{\nu}$ (which is denoted by $\operatorname{Bern}(p_{\nu})$). We now show how achieving a small optimization error implies being able to guess the value of $\nu$.

Suppose we are given problem and noise distributions defined as in (5) and (6), and an bijective mapping $h(\cdot)$ as in (7). Further suppose that there is a deterministic algorithm $\mathcal{D}_{k}$ whose output after processing $k$ stochastic gradients satisfies

Suppose that we are given access to $k$ samples of $\operatorname{Bern}(p_{\nu})$. Use these $k$ samples as the input $h(f_{b}+\chi_{b}))$ to the procedure $\mathcal{D}_{k}$ (this is valid as previously discussed), and let the output of $\mathcal{D}_{k}$ be $x^{(\nu)}_{k}$. The assumption in the lemma states that

Then, using Markov’s inequality (and then taking square-roots on both sides) gives

Next using Pinsker’s inequality we can upper bound the right hand side as:

where $\lvert\cdot\rvert_{TV}$ denotes the total-variation distance and $\operatorname{KL}(\cdot,\cdot)$ denotes the KL-divergence. Recall two properties of KL-divergence: i) for a product measures defined over the same measurable space $(p_{1},\dots,p_{k})$ and $(q_{1},\dots,q_{k})$,

Recalling that $\alpha\in(1,2]$ gives us the statement of the lemma. ∎

Given Lemmas 15 and 16, this implies that for the above choice of $\epsilon$,

This finishes the proof of the theorem. Note that the readability of the proof was prioritized over optimality and it is possible to obtain significantly better constants. ∎

Appendix G Non-convex Lower Bound (Proof of Theorem 6)

The proof is based on the proof of Theorem 1 in . The only difference is that we assume bounded $\alpha-$moment of the stochastic oracle instead of bounded variance as in the original proof. We refer readers to for more backgrounds and intuitions. For convenience, we study the stochastic setting ($K=1$ in ) instead of batched setting. We denote a $d-$dimensional vector $x$ as, $x=[x_{(1)};...;x_{(d)}].$ Let $\text{support}(x)$ denote the set of coordinates where $x$ is nonzero, i.e.

Denote $\text{prog}_{\beta}(x)$ as the highest index whose entry is $\beta-$far from zero.

Note that the function $\text{prog}_{\beta}(\,\cdot\,)$ is decreasing in $\beta$. The function we use to prove the theorem is the same as in . We denote

The above function satisfies the following important properties,

The function $f_{d}$ satisfies the following properties,

$f_{d}$ is $L_{0}$-smooth, where $L_{0}=152$.

For all $x$, $\|\nabla f_{d}(x)\|_{\infty}\leq 23$.

For all $x$, $\text{prog}_{0}(\nabla f_{d}(x))\leq\text{prog}_{\frac{1}{2}}(x)+1$

For all $x$, if $\text{prog}_{1}(x)<d$, then $\|\nabla f_{d}(x)\|^{2}\geq 1$.

We also define the stochastic oracle $g_{d}(x)$ as below

where $z\sim\text{Bernoulli}(p)$. The stochasticity of $g_{d}(x)$ is only in the $(\text{prog}_{\frac{1}{4}}(x)+1)$th coordinate. It is easy to see that $g_{d}(x)$ is a probability-$p$ zero chain as in [2, Definition 2] i.e. it satisfies

The second claim is because $\text{prog}_{\beta}(\,\cdot\,)$ is decreasing in $\beta$ and

The first claim is because if $z=0$, then we explicitly set the $(\text{prog}_{\frac{1}{4}}(x)+1)$th coordinate to 0. The stochastic gradient additionally has bounded $\alpha$-moment as we next show.

The stochastic oracle above is an unbiased estimator of the true gradient, and for any $\alpha\in(1,2]$

The unbiased-ness is easy to verify. For the bounded $\alpha$-moment, observe that only the $(\text{prog}_{\frac{1}{4}}+1)$-th coordinate is noisy and differs by a factor of $(\frac{z}{p}-1)$. Hence, we have

The first inequality followed from Jensen’s inequality and the convexity of $\lVert\,\cdot\,\rVert^{\alpha}$ for $\alpha\in(1,2]$:

Now we are ready to prove Theorem 6. Given accuracy parameter $\epsilon$, suboptimality $\Delta=f(0)-f^{*}$, smoothness constant $L$, and bounded $\alpha-$moment $G^{\alpha}$, we define

where $\lambda=\frac{304\epsilon}{L}$ and $d=\lfloor\frac{\Delta L}{7296\epsilon^{2}}\rfloor$. Then,

Let $x_{k}$ be the output of any zero-respecting algorithm $\mathcal{A}$. By [2, Lemma 1], we know that with probability at least 1/2, $\text{prog}_{1}(x_{k})\leq\text{prog}_{0}(x_{k})<d$ for all $k\leq\frac{(d-1)}{2p}$. Now applying Lemma 17.5, we have that for all $k\leq\frac{(d-1)}{2p}$:

Appendix H A Comparison with [25]

We are not the first to study the heavy-tailed noise behavior in neural network training. The novel work by Simsekli et al. studies the noise behavior of AlexNet on Cifar 10 and observed that the noise does not seem to come from Gaussian distribution. However, in our AlexNet training with ImageNet data, we observe that the noise histogram looks Gaussian as in Figure 5(a, b). We believe the difference results from that in , the authors treat the noise in each coordinate as an independent scaler noise, as described in the original work on applying tail index estimator. We on the other hand, consider each the noise as a high dimensional random vector computed from a minibatch. We are also able to observe heavy tailed noise if we fix a single minibatch and plot the noise in each dimension, as shown in Figure 5(c). The fact that noise is well concentrated on Cifar is also observed by Panigrahi et al. .

Furthermore, we used the tail index estimator presented in to estimate the tail index of noise norm distribution. Though some assumptions of the estimator are not satisfied (in our case, the symmetry assumption; in , the symmetry assumption and independence assumption), we think it can be an indicator for measuring the “heaviness” of the tail distribution.

Appendix I ACClip in ImageNet Training

For completeness, we test ACClip on ImageNet training with ResNet50. After hyperparameter tuning for all algorithms, ACClip is able to achieve better performance compared to ADAM, but worse performance compared to SGD. This is as expected because the noise distribution in ImageNet + ResNet50 training is well concentrated. The validation accuracy for SGD, ADAM, ACClip are $0.754,0.716,0.730$ respectively.