Robustness to Unbounded Smoothness of Generalized SignSGD
Anonymous Author(s)
AffiliationAddress
Abstract
Traditional analyses in non-convex optimization typically rely on the smoothnessassumption, namely requiring the gradients to be Lipschitz.
However, recentevidence shows that this smoothness condition does not capture the properties ofsome deep learning objective functions, including the ones involving Recurrent
Neural Networks and LSTMs.
Instead, they satisfy a much more relaxed condition,with potentially unbounded smoothness.
Under this relaxed assumption, it has beentheoretically and empirically shown that the gradient-clipped SGD has an advantageover the vanilla one.
In this paper, we show that clipping is not indispensablefor Adam-type algorithms in tackling such scenarios: we theoretically provethat a generalized SignSGD algorithm can obtain similar convergence rates as
SGD with clipping but does not need explicit clipping at all.
This family ofalgorithms on one end recovers SignSGD and on the other end closely resembles thepopular Adam algorithm.
Our analysis underlines the critical role that momentumplays in analyzing SignSGD-type and Adam-type algorithms: it not only reducesthe effects of noise, thus removing the need for large mini-batch in previousanalyses of SignSGD-type algorithms, but it also substantially reduces the effectsof unbounded smoothness and gradient norms.
We also compare these algorithmswith popular optimizers on a set of deep learning tasks, observing that we canmatch the performance of Adam while beating the others.
Introduction
Recent years have witnessed a surge in non-convex machine learning models, with a focus on deepneural networks [26].
DNNs have achieved tremendous progress in a variety of tasks, includingcomputer vision [25, 17, 22], natural language processing [10, 46] and a lot more.
Despite their hugeempirical success, the theoretical analyses of non-convex optimization [20] prove to be fundamentallymore challenging than the established convex optimization theory [4].
Among the numerous literature,many of them assume smoothness of the objective function, namely requiring the gradients to be
Lipschitz.
Under this scenario, past works have succeeded in proving the convergence rates for anumber of algorithms, e.g., Stochastic Gradient Descent [13], AdaGrad [48, 29], and STORM [7].
Nevertheless, it was recently observed that the smoothness assumption does not capture the trainingof LSTMs [19]: the Hessian can grow with the size of the gradients [52].
Inspired by this, Zhang etal.
[52] proposed a relaxed smoothness assumption, named ( L 0 , L 1 ) smoothness:
[Equation]
They also showed that the well-known gradient clipping technique can ensure SGD’s convergencein such scenarios.
Later, their results were improved to show that SGD with clipping can be madeunaffected by the L 1 in (1) and is able to recover the optimal convergence rate of SGD under theoriginal smoothness setting [51, 21].
[Figure] [Figure]
Nevertheless, the ( L 0 , L 1 ) condition has not yet been empirically verified beyond LSTMs.
Therefore,our first contribution lies in studying the applicability and generalization of the ( L 0 , L 1 ) condition.
In particular, we have empirically verified that the popular Transformer [46] model also seems tosatisfy this assumption, see Figure 1.
Yet, we noticed that different coordinates, especially when theyare in different layers of the model, exhibit very distinct L 0 and L 1 values as shown in Figure 3.
Hence,we propose to refine the ( L 0 , L 1 ) assumption in (1) to a coordinate-wise version (Assumption 2) andconsider this to better capture the loss surface when training deep neural networks like Transformers.
Given that we assume (a generalization) of the ( L 0 , L 1 ) assumption, it would be natural to use someclipping procedure.
However, we found out that the use of clipping on Adam [24], while carried outin common practice [e.g., 49], has no effect on the training and testing performance on optimizinga large transformer model as shown in Figure 2.
In retrospect, this might not be surprising: It isknown that Adam has an implicit clipping behavior due to the normalization by the estimated secondmoment of the gradients.
Indeed, Adam can be interpreted as a variant of SignSGD [2].
Inspired by this, our second contribution is to propose and analyze a generalized SignSGD algorithmunder the relaxed smoothness assumption.
It is parameterized in such a way that it on one endrecovers SignSGD while on the other end closely resembles Adam.
Apart from the convergence rates,we also located the critical role the momentum plays in analyzing Adam-type algorithms: it not onlyreduces the effects of noise but also gives an exponential decaying effect on the unbounded gradientnorms and smoothness.
This can partly explain the phenomenon that clipping does not help Adam.
The structure of this paper is as follows.
Section 2 discusses related works and how our paper buildsupon and distinguishes from them.
The settings and assumptions are carried out in Section 3.
We willintroduce formally the generalized SignSGD algorithm and its analysis in Section 4, with a detaileddiscussion on the bounds and the role of momentum.
The experimental results are shown in Section 5,comparing our algorithm with some popular competitors in deep learning tasks.
Finally, we drawsome conclusions and discuss the limitations of our work in Section 6.
Notations We will use [ d ] to denote the sequence [1 , 2 , . . .
, d ] and use bold letters to represent vectors,e.g., u ∈ R d .
The j th coordinate of a vector u is u j .
Throughout this paper, we study the Euclideanspace R d with the inner product ⟨· , ·⟩ .
The gradient of F at x is denoted by ∇ F ( x ) .
We use I ( · )to denote the indicator function, ∥ u ∥ p to denote the p -norm: ∥ u ∥ p :=
(cid:16)(cid:80)
dj
=1 | u j | p (cid:17) 1 /p and ∥ u ∥ ∞ themaximum norm: ∥ u ∥ ∞ := max {| u 1 | ,..., | u d |} .
We also denote by (cid:80) jk =
i x k =
when i > j .
Related Works
Adaptive Gradient Methods Adaptive gradient methods [32, 11, 24, 18, 40] are popular optimizersfor training deep neural networks.
The traditional analysis of adaptive gradient methods is providingregret bounds under the online convex optimization framework [11, 24, 40].
Recently, there aresome analysis of adaptive gradient methods for nonconvex smooth functions [6, 5, 50, 9, 55].
Zouet al.
[54] introduces an intriguing connection between Adam [24] and SignGD [3] when traininga two-layer neural network in the deterministic setting, where SignGD is an algorithm followingthe negative gradient sign direction to perform the update.
However, these works cannot be directlyextended to nonconvex functions with unbounded smoothness in the stochastic setting.
To the best ofour knowledge, this work is the first one establishing guarantees for coordinate-wise type optimizerslike generalized SignSGD as well as Adam-type updates under a relaxed smoothness condition.
Gradient Clipping The algorithm and analysis of gradient clipping can be traced back to [1, 44, 12]under the assumption that the function is convex and rapidly growing.
Hazan et al.
[16] consideredgradient clipping in quasi-convex optimization.
Mai and Johansson [30] showed the stability andconvergence of stochastic gradient clipping algorithms for convex problems without the smoothnesscondition.
Gradient clipping is a standard technique in training deep neural networks [35, 36]such as RNNs and LSTMs.
The theoretical analysis of gradient clipping for nonconvex models ispioneered by [52], in which the authors analyzed the convergence of gradient clipping under therelaxed smoothness assumption rather than the standard smoothness assumption.
Zhang et al.
[51]further improved the convergence rate bound under the same assumption as in [52].
Gradient clippingis also used when there is a heavy tail noise in the stochastic gradient to establish high probabilityconvergence rates [8, 14, 53].
Cutkosky and Mehta [7] proved that normalized momentum improvesnormalized SGD under a second-order smoothness condition.
A close algorithm is the one in [21]which employs gradient normalization, momentum, and no gradient clipping to tackle the ( L 0 , L 1 )condition (1) and control noise.
Yet, their algorithm normalizes each coordinate with the same scaleunlike popular optimizers such as Adam [24].
Moreover, we observe empirically that normalized
SGD with momentum performs worse than Adam.
Motivated by this, we propose a coordinate-wiseoptimization algorithm which requires new analysis tools compared with [21].
Employ m 2 t to compute v t in Adam Designed to combine the advantages of Adagrad [11] and
RMSProp [45], the update of Adam [24] employs the ratio between the exponential moving averageof the stochastic gradient ( m t ) and the exponential moving average of the squared stochastic gradient( v t ).
Many variants of Adam have been proposed ever since.
Among them, one idea is to use m 2 tto compute v t instead of g 2 t .
The intuition is that m t represents a better update direction than g tand can thus better capture the second-moment information.
Reddi et al.
[39] adopted this change toprove the convergence of Adam in a federated learning setting; yet, they only consider the smoothsetting and require a large ϵ to obtain convergence in contrast to the original Adam.
Later, Wang etal.
[47] explored this idea in more detail, but their analyses are still restricted to the smooth setting.
Settings and Preliminaries
In this paper, we focus on the following stochastic optimization problem:
[Equation]
where ξ is a random variable representing a randomly selected data sample or random noise followingan unknown distribution D .
We will use the following assumptions.
Assumption 1.
F : R d → R is differentiable and bounded from below with infimum F ∗ .
Assumption 2.
We say that a twice differentiable function F ( x ) is ( L 0 , L 1 ) -smooth
coordinatewisely, if for any x , y ∈ R d for which ∥ x − y ∥ 2 ≤ 1 ∥ L 1 ∥ ∞ , we have for any j ∈ [ d ] that
[Equation]
We will denote L 0 :=
[ L 0 , 1 , L 0 , 2 , . . .
, L 0 ,d ] T and L 1 :=
[ L 1 , 1 , L 1 , 2 , . . .
, L 1 ,d ] T .
The original ( L 0 , L 1 ) smoothness assumption (1) in [52] was proposed as a generalization of themore common smoothness assumption, which says that the gradient should be Lipschitz.
Indeed, when L 1 ,j are zero, we recover the smoothness assumption.
In contrast, when L 1 ,j are non-zero,the smoothness of the function is potentially unbounded .
Yet, [52] works with norms and appliesto the global scale, while ours is more fine-grained and applies to each coordinate separately.
Onemotivation for this assumption comes from [Remark 2.3, 51] where they noted that (1) can be relaxedto an assumption on gradient differences: ∃ K 0 , K 1 > 0 s.t. for all x , y ∈ R d with ∥ x − y ∥ 2 ≤ 1 K 1 , [Equation]
[Equation] Indeed, our Assumption 2 implies (3) when L 0 ,j = L 0 and L 1 ,j = L 1 for all j ∈ [ d ] , up to constants(See Lemma 12 in the Appendix).
Note that the 1 √ d factor in ours is exactly for easy comparisonwith (3).
The reason we turn to the current coordinate-wise version is that we observed a vast varianceacross different layers in training Transformer models: (1) is still true globally (Figure 1), but eachlayer or even each coordinate satisfies has a very different ( L 0 , L 1 ) pair (Figure 3).
The smoothnessassumption has been generalized in orthogonal directions in other work [41, 3, 23].
One merit of Assumption 2 is that it gives us the following descent lemma.
Lemma 1.
Let F be ( L 0 , L 1 ) -smooth coordinate-wisely.
Then, for any x , y ∈ R d for which∥ x − y ∥ 2 ≤ 1 ∥ L 1 ∥ ∞ , we have
[Equation]
∥ y − x ∥ | y j − x j | .
Our last assumption is common in the literature studying the ( L 0 , L 1 ) smooth condition [52, 51].
Assumption 3.
For each j ∈ [ d ] , there exists σ j > 0 such that for all x ∈ R d and ξ ∼ D , the noisesatisfies (cid:12)(cid:12)(cid:12) [ ∇ f ( x , ξ )] j − ∂F∂x j ( x ) (cid:12)(cid:12)(cid:12) ≤ σ j with probability 1 .
We will denote σ :=
[ σ 1 , σ 2 , . .
, σ d ] T .
A Generalized SignSGD Algorithm
In this section, we present in Algorithm 1 a generalized SignSGD algorithm.
This algorithmencompasses a variety of optimization algorithms.
At first sight, it seems very similar to Adam.
Indeed, if we employ g 2 t in computing v t instead of m 2 t ,then it is exactly Adam, except for the bias correction terms.
We would like to clarify that the ideaof this change has been explored before, as detailed in Section 2.
In this paper, the motivation foradopting this idea comes from the known effect of momentum on reducing the influence of noises [7].
Indeed, in our analysis the difference between m t and ∇ F ( x t ) is much more controllable thanbetween g t and ∇ F ( x t ) .
Thus, we consider employing m t in computing v t a better choice.
On the other end, the careful reader might observe that Algorithm 1 recovers the SignSGD with
Momentum algorithm, also called SIGNUM in [3], when setting β 2 = 0 .
Sign-based algorithmsare naturally suited to distributed learning [28] and the idea dated back to at least Rprop [42].
Theconvergence to a stationary point (with ℓ 1 norm) under a coordinate-wise smoothness conditionhas been established for SignSGD with/without the momentum in [3] though they necessitate largemini-batches to control the variance of the noise.
Yet, we are more interested in their property of theupdate size being bounded without the need for explicit clipping.
Note that both SignSGD and Adam are good candidates for optimization algorithms whose updatemust be bounded on functions that satisfy the ( L 0 , L 1 ) condition.
Indeed, SignSGD can be seen asan extreme form of gradient clipping.
On the other hand, as said in the introduction, Adam does notseem to require gradient clipping at all when used to train the large Transformer model in Figure 2.
Hence, we expect our algorithm, a generalization of SignSGD and a close resemblance to Adam, canenjoy the merits of both and be robust to the unbounded smoothness in the ( L 0 , L 1 ) scenario.
In thenext section, we will formalize this claim by presenting the theoretical analysis of Algorithm 1.
4.1 Theoretical Convergence Analysis
Theorem 1.
Under Assumptions 1, 2, and 3, assume M j :=
sup (cid:110)(cid:12)(cid:12)(cid:12) ∂F∂x
[Equation] ( x [Equation] (cid:12)(cid:12)(cid:12) : F ( x ) ≤ F ( x 1 ) (cid:111) isfinite for each j ∈ [ d ] , let ∆ be any upper bound on F ( x 1 ) − F ∗ , α = min (cid:18)
[Equation]
[Equation] [Figure]
[Equation]
[Figure] [Equation] [Figure] [Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
Algorithm 1 guarantees, with probability at least 1 − δ , that
[Equation]
Furthermore, for the case when β 2 = 0 , we have the following refined guarantee:
[Equation]
Here, M j denotes the maximum absolute value of the partial derivative of F for coordinate j amongthe sub-level set of F ( x 1 ) , namely any point x with F ( x ) ≤ F ( x 1 ) .
In other words, we assumegradients to be bounded in the sub-level set of F ( x 1 ) ; yet, we do not make any restriction on gradientsoutside of this set.
We believe this is not a strong assumption, for example, when the sub-levelset of F ( x 1 ) is bounded, then by the assumed continuity of gradients it trivially holds.
Also, wejust require an upper bound and it can even be exponentially large as we have an exponentiallydecaying coefficient to counteract it: notice how the term ∥ M ∥ 1 is multiplied by a term that decaysexponentially with T .
Better still, when β 2 = 0 , we no longer even need this assumption and thealgorithm is entirely free of the influence of ∥ M ∥ 1 .
To see why this is good, we show a refined lowerbound of Gradient Descent under the relaxed smoothness scenario below which is originally in [52].
Theorem 2.
Fix ϵ > 0 , L 0 > 0 , L 1 > 0 , M ≥ max( L 0 L 1 , ϵ ) , and x 0 ∈ R .
Pick any constant learningrate η for GD, with the knowledge of the above constants.
Then, there exists a 1-d ( L 0 , L 1 ) -smoothfunction, bounded from below by f ∗ (finite), and such that sup {| f ′ ( x ) | : f ( x ) ≤ f ( x 0 ) } ≤ M onwhich the number of iterations T of GD with learning rate η to guarantee | f ′ ( x T ) | <
ϵ is at least
[Equation]
Theorem 2 shows that in the relaxed smoothness setting, GD with any constant step size will sufferfrom a linear term depending on L 1 M .
On a side note, it is a fixed version of the lower bound in[52]: we provide in Appendix an explanation of errors in their lower bound and our corrected proof.
Compared with GD, our algorithm only has an exponentially decaying dependence on L 1 M .
Weconsider this to be substantial merit of our algorithm.
Furthermore, when β 2 = 0 in which case werecover the SignSGD with Momentum algorithm, we can even show that it completely removes theeffects of the unbounded gradient norms.
Also notice that in such case we actually no longer need theassumption of M j := sup (cid:110)(cid:12)(cid:12)(cid:12) ∂F∂x j
( x ) (cid:12)(cid:12)(cid:12) : F ( x ) ≤ F ( x 1 ) (cid:111) being finite for each j ∈
[ d ] anymore, andthe ∥
L 1 ∥ ∞ term does not appear in the final bound anymore.
We also would like to point out that this bound closely resembles the one achieved by SGD withgradient clipping algorithm [51] except that we consider the coordinate-wise setting: take the settingof β 2 = 0 for example, we need at most O (cid:16) ∆ max (cid:110) ∥ σ ∥
∥ L 0 ∥ 1 ϵ, d 2 ∥ σ ∥ 21 ∥ L 1 ∥ 4 ∞ ∥ L 0 ∥, d ∥ L 1 ∥ 2 ∞ ∥ L 0 ∥ 1 (cid:111)(cid:17) to get apoint x with ∥∇ F ( x ) ∥ 1 ≤ ϵ with high probability.
Remark
The almost surely bounded assumption 3 can be relaxed to sub-gaussian noise, usingstandard extensions of Freedman inequality [e.g., 15].
Remark 2 When β 2 = 0 , we can prove an average-iterate complexity bound (see Proof of Theorem 1for β 2 = 0 in Appendix A.3); yet, we use the min form for consistency between the two cases.
Remark {+
 Our bound is incomparable with+} the one in [51, Theorem 3.2].
Yet, as we said, if
L 0 ,j = L 0 and L 1 ,j = L 1 for all j ∈ [ d ] , then the function satisfies (3). In this case, assuming thenoise vector and the gradient vector to be dense to be able to compare the ℓ 1 -norm and the ℓ 2 -norm,we recover the same bound of [51, Theorem 3.2] in terms of dependencies on L 1 , L 0 , and T . Instead,in the more general case when L 0 ,j and L 1 ,j are not uniform vectors, our bound allows a finer controlof the unbounded smoothness.
The proof of the theorem is highly technical and it uses recent advancements in the analysis ofmomentum methods [7], key techniques to deal with the ( L 0 , L 1 ) assumption [51], as well as a noveland essential inductive argument to control the norm of past gradients.
We want to stress that thedifficulty mainly comes from analyzing Adam-type updates when β 2 > 0 , while for the other case ofβ 2 =
the proof is significantly simpler.
The full proof is in the Appendix, but here we present aproof sketch that underlines the main steps.
First, we list some key lemmas we used but move theirproofs to the appendix due to space constraints.
Lemma 4.
With notations in Algorithm 1, for τ ≤ ¯ τ = [Equation] − β
[Equation] √ d [Equation]
L 1 ∥ ∞ , we have ∥ x t − τ − x t ∥ 2 ≤ 1 ∥ L 1 ∥ ∞ .
Lemma 4 limits our focus to the most recent ¯ τ steps on which Assumption 2 and Lemma 1 can apply.
Lemma 6.
Assume Assumption 3.
With the notation of Algorithm 1, let j ∈ [ d ] and β 1 ≤ 1 .
Then,with probability at least 1 − 3 δ , for any t 0 ∈ [ t ] , we have
[Equation]
Lemma 6 is the major tool we use to handle the noise we incur during drawing stochastic gradients.
It is derived based on Lemma 12 in [8].
Lemma 9.
With the notation of Algorithm 1 and under the assumptions of Theorem 1, if (cid:12)(cid:12)(cid:12) ∂F∂x j
( x τ ) (cid:12)(cid:12)(cid:12) ≤
M j holds for all τ ≤ t and j ∈ [ d ] , and D > 0 , then, with probability at least 1 − 3 tδ we have that,
[Equation]
< 5 B j D or | m t,j | √ v t,j ≥ ρD 5 √ 1 − β 2 , where B j ≜ ηL 0 ,j √ 1 − β 2 (1 − β 1 ) + β ¯ τ 1 ( M j + σ j ) + (1 − β 1 ) E j and D ≜ 1 − 2 η √ d ∥
L 1 ∥ ∞ √ 1 − β 2 (1 − β 1 ) .
Lemma 9 is similar to Lemma A.2 in [54] which considered the deterministic and smooth setting;in contrast, our proof is much more challenging in that we need to tackle both the noise and theunbounded smoothness.
With this lemma, we know that either the true gradient is small or that theupdate of our Algorithm 1 can be lower bounded.
Lemma 11.
Under Assumptions 1, 2, and 3, using the hyperparameters in Theorem 1, denoting 216 α = 1 − β 1 and ϵ t = m t − ∇ F ( x t ) , for all t and j ∈ [ d ] we have, with probability at least 1 − 3 δ ,| [Equation] t +1 ,j |
[Equation]
t
[Equation]
[Equation]
Lemma 11 shows how the use of momentum can help control the noise by choosing β 1 wisely.
It isadapted from the proof of Theorem 2 in [8] but with the added difficulty of unbounded smoothness.
Proof sketch of Theorem 1.
Observing the formula of setting β 1 , we can see that when ∥ σ ∥ 1 ≤ 220 (cid:112) ∥
L 0 ∥ 1 ∆ / √ T , β 1 = 0 .
As β 2 < β 1 , Algorithm 1 reduces to SignSGD.
In this case, the keycomponent is Lemma 11 using which we are able to show that (cid:80) Tt =1 (cid:12)(cid:12)(cid:12) m t,j − ∂F∂x j
( x t ) (cid:12)(cid:12)(cid:12) can be controlled as C 1 (cid:80)
Tt =1 (cid:12)(cid:12)(cid:12) ∂F∂x j
( x t ) (cid:12)(cid:12)(cid:12) + C 2 .
The summation of true gradients over time can then beoffsetted by choosing η and β 1 wisely when we invoke the descent lemma 1.
The rest is standard.
Now for the other case in which ∥ σ ∥ 1 > (cid:112) ∥
L 0 ∥ 1 ∆ / √ T , we take a different route.
First, notice that Assumption 2 and the Descent Lemma 1 only hold when two points are not too faraway.
Thus, we need to restrict our attention to the recent updates (Lemma 4), beyond which wewould have no control.
This means we want the influence of those updates too long ago to not havea big effect on the current one.
To make this happen, one natural idea is to use a bounded gradientassumption, then with the use of exponential averaging, their effect would be quickly reduced.
Yet,assuming directly that all gradients are bounded would trivialize the ( L 0 , L 1 ) assumption.
Thus,we pose a much weaker condition, assuming that M j := sup (cid:110)(cid:12)(cid:12)(cid:12) ∂F∂x j
( x ) (cid:12)(cid:12)(cid:12) : F ( x ) ≤ F ( x 1 ) (cid:111) beingfinite for each j ∈ [ d ] .
Then, we prove that M j will provide an upper bound to all the true gradientsthe algorithm see.
We prove it using induction, analyzing separately the case that either the truegradient is already very small and we have reached the proximity of a stationary point, or the objectivefunction is monotonically non-increasing and the gradient remains bounded.
Having controlled the past gradients, we prove in Lemma 9 that the update of Algorithm 1 is eithervery small that we can pass or having a constant lower bound that we can use in the Descent Lemma 1.
Also, considering that this is the stochastic setting, noise typically slows down convergence or caneven cause the algorithm to diverge if the hyperparameters are not chosen wisely.
To handle this, weinvoke Freedman’s inequality to show that the addition of adjacent stochastic noise almost cancelsout each other and the absolute value of the sum remains controlled (Lemma 6).
Yet, we still need another block to handle the difference between the true gradient and the momentumas we are updating in the direction of the momentum instead of the true gradient.
Turns our that wecan prove that sign ( m t,j )
= sign (cid:16) ∂F ∂x j ( x t ) (cid:17) when (cid:12) (cid:12)(cid:12) ∂F
∂x j ( x t ) (cid:12) (cid:12)(cid:12) is not too small.
As before, in the case
(cid:12)(cid:12)(cid:12) ∂F∂x j ( x t ) (cid:12)(cid:12)(cid:12)
is small, we have converged on that coordinate.
Combining all these blocks together, we are able to arrive at the final results.
Experiments
We conducted our experiments using PyTorch [37] on Nvidia V100 GPUs.
5.1 Comparison with Other Optimizers
To validate the efficacy of our Algorithm 1, we compare it with Adam [24], SGD [43], SGD Momentum Normalized [21], SGDClipGrad, and SGDClipMomentum.
The latter two are from Algorithm1 in [51] where SGDClipGrad corresponds to the case where ν
= 0 and SGDClipMomentumcorresponds to the case when ν = 1 .
[Table]
Training Unless otherwise specified, we use grid-search to fine-tune the initial learning rate for alloptimizers
[Table]
as well as the clipping threshold for SGDClipGrad and SGDClipMomentum
[Table]
and β 2 for
Adam and our algorithm, to select the one giving the best validation performance on a separatedvalidation set.
We then employ the best performing hyperparameters to train the model over alltraining data and report the testing performance.
The testing is repeated with random seeds 5 times toeliminate the influence of stochasticity.
For more details, please refer to Section A.4.
Resnet for Image Classification on CIFAR-10 We employ the 20-layer Residual Network model [17]to do image classification on the CIFAR-10 dataset.
Images are normalized per channel using themeans and standard deviations computed from all training images.
We adopt the data augmentationtechnique following [27] (for training only): 4 pixels are padded on each side of an image and a 32 ×32 crop is randomly sampled from the padded image or its horizontal flip.
The mini-batch size is 128and we train all algorithms for 164 epochs.
We do not employ any learning rate decay schedule inorder to focus on the comparison of the optimizers themselves.
We fixed the weight decay value to be0 . 0001 and the momentum parameter ( β 1 ) to be 0 .
Figure 4 and Table 1 report the training andtesting performance for each algorithm, showing that ours is among the best of all.
LSTM for Language Modeling on Penn Treebank
We adopt a 3-layer AWD-LSTM [33] to dolanguage modeling on the Penn Treebank (PTB) dataset [31](word level).
The mini-batch size is 40and we trained each algorithm for 750 epochs.
Apart from the hyperparameters we stated above, wefurther fine-tuned the weight decay value for all algorithms noticing its significant influence on theperformance.
We choose the set of hyperparameters that give the smallest final validation perplexity.
We report the results in Figure 5 and Table 1.
It can be seen that we can match the performance of
Adam while beating the others.
5.2 Transformers Observe ( L 0 , L 1 ) -smoothness
For Figure 1 which verifies the original form (1) of the ( L 0 , L 1 ) condition using the norm, wefollowed the method in Section H.3 of [52].
Specifically, given x t and x t +1 , denote d := x t +1 − x t .
We estimate the smoothness at x t by
[Equation]
where { δ 1 , δ 2 , . . .
, δ N } denotes the sample locations and we use { 16 , 26 , 36 , 46 , 56 } .
For Figure 3 verifying the coordinate-wise version (2) of the ( L 0 , L 1 ) condition, note that theequation is symmetric in that if we just swap x and y it shall still holds.
Thus, during plotting, wecompare
[Equation]
Figure 1(a) is on training a 2 -layer Transformer Encoder to do language modeling on the Wikitext-2dataset.
The implementation, settings, and parameter choices follow this.
1 We only plot the first 5training epochs.
Figure 1(b) and 3 are on training a 6 -layer Transformer [46] to do machine translationon the WMT’16 Multimodal Machine Translation Task German-English dataset.
The implementationof the transformer is forked from here
and we also follow their default settings.
The mini-batch sizeis
and we trained for 400 epochs using Adam and report the whole training trajectory.
5.3 Clipping does not Affect Adam’s Performance
We compare clipping and non-clipping for Adam optimizer on the Wikitext-103 (103 million tokens,180MB)
[34] language modeling task, with a 16-layer GPT-2 transformer model [38].
This GPT-2model has an input length of 256 tokens, 410-dimension word embedding, 16 Attention layers with10 Attention heads and 2100 hidden dimensions.
Model size is 201.58 MB.
The vocabulary size is28996.
We use the hyper-parameter settings prescribed in [49]: batch size 256, warm up learning ratefrom 0 to 2 .
5 × 10 − 4 in the first 64000 samples (i.e., 250 iterations) and then cosine-anneal learningrate to zero, on top of an Adam optimizer.
It takes about 40 hours to train 200 epochs on 8 V
GPUs.
We use clipping threshold max_norm 0.25 for the entire model as prescribed in the literature[49].
We also count that with this clipping scheme, clipping occurs in every single batch.
As we cansee from Figure 2, neither training loss (2.79 vs 2.76) nor perplexity score (27.92 vs 27.97) differsmuch in the clipping and non-clipping case, which is consistent with our theory that Adam naturallyachieves gradients clipping effect.
Conclusion and Limitations
Smoothness has been a widely adopted condition for proving convergence rates of algorithms inthe non-convex optimization scenario.
Yet, it has been found that this assumption does not capturelosses when employing some deep learning models including RNNs and LSTMs.
In light of this, arelaxed smoothness assumption was proposed that aligns well with the practice.
We observed thatthe loss surface of training using Transformers also exhibits this relaxed smoothness.
Under thisassumption, SGD with clipped gradient has been proven to work well.
However, we found thatclipping is not necessary for achieving convergence in such a setting.
Indeed, we showed that ageneralized SignSGD algorithm does not require explicit clipping but can almost guarantee the samebound as SGD with clipping.
In the analyses, we identified the key effect of using momentum inanalyzing Adam-type algorithms, that it reduces both the noise and the unbounded gradient norms.
Finally, we conducted a variety of deep learning tasks showing that our algorithm can match Adam’sperformance while exceeding others.
Limitations The current work is in no way a perfect one and there are many directions worthexploring beyond it.
First of all, though our algorithm could be seen as a close resemblance to theoriginal Adam algorithm, they are still not equal.
Considering the huge popularity of Adam andits established effectivity in practice, it is worth studying whether Adam in its original form canconverge in the relaxed smooth setting.
Second, while our Theorem 1 are upper bounds and cannot bedirectly compared between the two cases of β 2 , it does suggests that β 2 = 0 minimizes the worst-caseconvergence rate.
However, it still does not fully explain the phenomenon that a choice of β 2 close to1 yields better performance in using our Algorithm 1 as well as Adam in practice.
Third, despite thereare lower bounds showing that, for example, GD with a constant step size can be arbitrarily worsethan GD with clipping, it would be more meaningful to study whether the relaxed smooth conditionis inherently more difficult, possibly by establishing a lower bound for all first-order optimizationalgorithms.
Finally, we did show that Transformers observe the relaxed smoothness condition, but weconsider it more beneficial to research in-depth what properties or structures make a model satisfysuch conditions.
https://pytorch.org/tutorials/beginner/transformer_tutorial.html
https://github.com/jadore801120/attention-is-all-you-need-pytorch
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
This
paper is mainly a theoretical one analyzing machine learning algorithms, we do not
forsee any potential negative societal impacts of our work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 3.
(b) Did you include complete proofs of all theoretical results?
[Yes] See the Appendix.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main ex-
perimental results (either in the supplemental material or as a URL)?
[Yes]
In the
supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Section 5.1 and A.4.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Section
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes] Included in the code folder of the
supplemental materials.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[No] Data use follows the permission of corresponding licenses.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[No]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]