Title: SOFTADAM: UNIFYING SGD AND ADAM FOR BETTER STOCHASTIC GRADIENT DESCENT

Abstract: Stochastic gradient descent (SGD) and Adam are commonly used to optimize deep neural networks, but choosing one usually means making tradeoffs between speed, accuracy and stability. Here we present an intuition for why the tradeoffs exist as well as a method for unifying the two in a continuous way. This makes it possible to control the way models are trained in much greater detail. We show that for default parameters, the new algorithm equals or outperforms SGD and Adam across a range of models for image classification tasks and outperforms SGD for language modeling tasks.

1. INTRODUCTION:

One of the most common methods of training neural networks is stochastic gradient descent (SGD) (Bottou et al. (2016) ). SGD has strong theoretical guarantees, including convergence in locally non-convex optimization problems (Lee et al. (2016) ). It also shows improved generalization and stability when compared to other optimization algorithms (Smith & Le (2018)).

There have been various efforts in improving the speed and generalization of SGD. One popular modification is to use an adaptive gradient (Duchi et al. (2011) ), which scales the gradient step size to be larger in directions with consistently small gradients. Adam, an implementation that combines SGD with momentum and an adaptive step size inversely proportional to the RMS gradient, has been particularly successful at speeding up training and solving particular problems (Kingma & Ba (2014)). However, at other problems it pays a penalty in worse generalization ( Here we develop an intuition for adaptive gradient methods that allows us to unify Adam with SGD in a natural way. The new optimizer, SoftAdam, descends in a direction that mixes the SGD with Adam update steps. As such, it should be able to achieve equal or better optimization results across a variety of problems.

1.1. RELATED WORK:

Several authors have recently tried to combine Adam and SGD to get the best of both worlds. However, these have not always enabled better generalization or performance across different problems.

In one study, the optimization algorithm was switched from Adam to SGD during training based on a scale-free criterion, preventing the addition of a new hyper-parameter (Keskar & Socher (2017)). The result is that the longer the convolutional networks were trained on Adam, the worse their generalization performance compared to SGD. The best performance came from switching to SGD in this case as soon as possible.

Another recent algorithm takes the approach of clipping the large Adam updates to make them more similar to SGD as the training nears the end (Luo et al. (2019) ). However, this approach requires two new hyper-parameters: the rate at which the training is switched over, and a learning rate for both SGD and Adam.

Similar to this work, partially adaptive methods (Chen & Gu (2018) ) can allow arbitrary mixing between SGD and Adam. However, in that work the step size is not strictly smaller than the SGD step and so the same guarantees cannot be made about convergence. It is of interest to see whether there is any advantage over these methods.

Other algorithms have shown improvements on SGD by averaging weights over many steps (Polyak & Juditsky (1992) ; Zhang et al. (2019); Izmailov et al. (2018) ). These algorithms are complementary to the algorithm developed here, as they require an underlying algorithm to choose the step direction at any point.

2. SGD AND ADAM:

The fundamental idea of gradient descent is to follow the path of steepest descent to an optimum. Stochastic gradient descent enables us to optimize much larger problems by using randomly subsampled training data. The stochastic gradient descent algorithm will minimize the loss function J(θ; x), which is parameterized by θ and takes as input training data x,

g t ← [∇ θ J(θ; x t )] θ=θt-1 θ t ← θ t-1 -αg t ,

where α is a learning rate that may vary with t and x t is the training data selected for the batch at step t. The convergence rate can be improved further by using a running average of the gradient, initializing with m 0 ← 0. This method, known as momentum (Goh (2017)), may be written as,

g t ← [∇ θ J(θ; x t )] θ=θt-1 m t ← β 1 m t-1 + (1 -β 1 ) g t , θ t ← θ t-1 -αm t ,

A further development that has improved convergence, especially for LSTM and language modeling tasks, involves the second gradient as well. This specific version is known as the Adam algorithm (Kingma & Ba (2014)),

g t ← [∇ θ J(θ; x t )] θ=θt-1 m t ← β 1 m t-1 + (1 -β 1 ) g t , v t ← β 2 v t-1 + (1 -β 2 ) g 2 t , θ t ← θ t-1 - α mt √ vt ,

where mt = m t / 1 -β t 1 and vt = v t / 1 -β t 2 are unbiased estimators of the first and second moment respectively.

2.1. CONVERGENCE:

In order to analyze the convergence of these algorithms, we can consider a second-order approximation of J on its combined argument z = (θ; x) in the region of (θ t ; x t ),

J(z) = 1 2 z T H t z -b T z,

where H t is the Hessian of J(z) around z t . This gives us the gradient,

g t = ∇J(z t ) = H t z t -b,

which becomes the SGD update step,

z t+1 = z t -α (H t z -b) .

Unrolling this update step can be shown to lead to an expression for the distance from the optimal value z , in the basis of Hessian eigenvectors ξ i :

z t -z = i (z 0 • ξ i ) (1 -αλ i ) t ξ i .

We can see that the learning is stable if the learning rate α satisfies,

|1 -αλ i | < 1.

In addition, we find that the value for the learning rate that leads to the fastest overall convergence is,

EQUATION

where λ 1 and λ n are the max and min eigenvalues of H, respectively.

2.2. ADAPTIVE MOMENT ESTIMATION:

If rather than a single learning rate α, we were to use a diagonal matrix D such that the update is,

z t+1 = z t -Dg t ,

we may be able to modify the diagonal entries [d 0 , d 1 , . . . , d n ] such that faster overall convergence is achieved. For example, in the special case that the Hessian is diagonal, the convergence rate for the i-th element becomes,

r i = (1 -d i λ i ) ,

which has the solution d i = λ -1 i . In this situation, if the eigenvalues λ i are known, the algorithm can converge to the minimum in exactly one step. This corresponds with some intuition behind adaptive moment methods: that taking a step with a "constant" size in every direction toward the target will reach convergence faster than taking a step proportional to the gradient size.

Because the eigenvalues and eigenvectors not known a priori, for a practical algorithm we must rely on an approximation to find d i . One technique named AdaGrad (Duchi et al. (2011) ) prescribes the diagonal elements:

EQUATION

For building our intuition, we consider the special case where the Hessian is diagonal,

g 2 ti = (λ i z ti -b i ) 2 .

Combining this with Equation 2, we compare the AdaGrad coefficient to the optimal value for

EQUATION

As long as αλ i , this will be true when,

EQUATION

which is the sum squared distance of the parameter value z ti to its optimum z i = b i /λ i . This reveals some of the tradeoffs being made by AdaGrad. It will perform ideally when all parameters start the same distance from their targets.

This may be true on average if the initializations z 0i and optima z i can be made over the same subspaces. That is, if z i is uncorrelated to λ i , we can expect this to have good performance on average.

However, there can be significant errors in both overestimating and underestimating the eigenvalue. One would expect that for a typical problem b i and λ i might be drawn from uncorrelated distributions. In this case, large values of z i will be likely to correspond to small values of λ i . Since z 0 can only be drawn from the average distribution (no information is known at this point), the estimated λ i is more likely to be large, as the initialization will be far from the optimum. Intuitively, the gradient is large because the optimum is far from the initialization, but the algorithm mistakes this large gradient for a large eigenvalue.

On the other hand, when the parameter is initialized close to its optimum, the algorithm will mistakenly believe the eigenvalue is small, and so take relatively large steps. Although they do not affect the overall convergence much on their own (since the parameter is near its optimum), these steps can add significant noise to the process, making it difficult to accurately measure gradients and therefore find the optimum in other parameters. This problem will be significantly worse for Adam, which forgets its initialization with some decay factor β 2 . In that case, as each parameter reaches its optimum, its estimated eigenvalue λ i drops and the step size gets correspondingly increased. In fact, the overall algorithm can be divergent as each parameter reaches its optimum, as the step size will grow unbounded unless α is scheduled to decline properly or a bounding method like AMSGrad is used (Reddi et al. ( 2018)).

In addition, reviewing our earlier assumption of small , these algorithms will perform worse for small eigenvalues λ i < /α. 2019)). As such, each element might be better optimized by a learning rate that better serves both its min and max eigenvalues.

3. SMOOTHED ADAPTIVE GRADIENT:

Overall, this understanding has led us to believe that adaptive moments might effectively estimate λ i when it is large, but might be less effective when it is small. In order to incorporate this information about large eigenvalues, as well as optimize the learning rate to account for variation in the eigenvalues contributing to convergence of a particular component, we consider the an update to Eq. 1,

d i = α λ (1 + η) λ + ηλ i ,

where λ is an average eigenvalue and η is a new hyper-parameter that controls the weighting of the eigenvalue estimation. Here we have added λ to the numerator so that α does not need to absorb the changes to the RMS error as it does in Eq. 4. This also recovers the SGD convergence guarantees, since the step is always within a factor of η to an SGD step. In addition, this will allow us to recover SGD with momentum exactly in the limit η → 0. We use the adaptive gradient estimation,

λ 2 i λ2 ≈ g 2 i 2 + g 2 avg = v ti 2 + vt ,

where vt is the mean value of v t , to write,

d i = α (1 + η) 1 + η v ti / ( 2 + vt ) .

One issue with the above estimation is that its variance is very large at the beginning of training (Liu et 

T ∈ F m 0 , v 0 ← 0 Initialize moments to 0 for t ← 1 . . . T do g t ← ∇ θ J(z) Find gradient m t ← β 1 m t-1 + (1 -β 1 ) g t Update first moment β2t ← min (β 2 , 1 -1/t) v t ← β2t v t-1 + 1 -β2t g 2 t Update second moment vt ← mean elem [v t ]

Average over all elements

r t ← sqrt (1 -β 2 ) / 1 -β2t d t = α (1 + η) 1 + η -ηr t + ηr t v t / (v t + 2 ) -1

Calculate the denominator

θ t = θ t-1 -m t d t

Perform the update end for return θ T that rectifying it can make warmup unnecessary. Where v t is the average of n t elements and v ∞ the average of n ∞ , we define r t = n t /n ∞ and:

d i = α (1 + η) 1 + η -ηr t + ηr t v ti / ( 2 + vt ) .

This finally forms the basis for our algorithm.

3.1. ALGORITHM:

Our algorithm differs from Adam in a few other ways. First, the biased gradient estimate is used rather than the unbiased one. This matches the SGD implementation of momentum, and also avoids magnifying the large variance of early gradient estimates:

m t = β 1 m t-1 + (1 -β 1 ) g t .

In addition, the second moment v t is calculated in an unbiased way using an effective β2 (t), or by abuse of notation β2t :

β2t = min β 2 , 1 - 1 t , v t = β2t v t-1 + 1 -β2t g t .

This has a negligable impact on the performance of the algorithm, but makes tracking the moment over time easier since it does not need to be un-biased later. We then calculate the ratio of the number of samples n t used to calculate the moment v t to the steady state number of samples in the average n ∞ :

r t = n t n ∞ = 1 -β 2 1 -β2t .

We finally note that the weight decay should be calculated separately from this update as in AdamW (Loshchilov & Hutter (2017) ). 

4. EMPIRICAL RESULTS:

In order to test the ability of this algorithm to reach better optima, we performed testing on a variety of different deep learning problems. In these problems we keep η at the default value of 1. Because the algorithm is the same as SGDM if η = 0 and is comparable to Adam when η = -1 , getting results at least as good as those algorithms is just a matter of parameter tuning. These results are intended to show that SoftAdam performs remarkably well with a common parameter choice. For the best performance, the hyper-parameter η and learning rate schedule α should be optimized for a particular problem.

4.1. CIFAR-10:

We trained a variety of networks: For each model and optimization method, the weight decay was varied over [1e-4,2e-4,5e-4,1e-3,2e-3,5e-3]. For AdamW the learning rate was varied over [1e-4,2e-4,5e-4,1e-3,2e-3]. For each optimizer and architecture and the best result was chosen, and three runs with separate initializations were used go generate the final data. The learning schedule reduced the learning rate by a factor of 10 at 50% and 75% through the total number of epochs. The results are summarized in Table 1 . We find that SoftAdam equals or outperforms SGD in training classifiers on this dataset. Due to the larger optimal weight decay constant, SoftAdam achieves lower validation loss at a higher train loss than SGD.

4.2. PENN TREEBANK:

We trained a 3-layer LSTM with 1150 hidden units per layer on the Penn Treebank dataset (Mikolov et al. (2010) ) in the same manner as Merity et al. (2017) . For SoftAdam the weight drop was increased from 0.5 to 0.6. Results for the average of three random initializations are shown in Figure 2 (a) and are summarized in Table 2 . For these parameters, SoftAdam outperforms SGD significantly but does not quite achieve the same results as Adam. Note that for this experiment we chose Adam instead of AdamW for comparison due to its superior performance.

4.3. IWSLT'14 GERMAN TO ENGLISH:

We also trained a transformer using the fairseq package by Ott et al. (2019) on the IWSLT'14 German to English dataset. Results for each method with optimized hyperparameters are summarized in Table 3 . Note that no warmup is used for training SoftAdam, but warmup is used for AdamW and 

5. CONCLUSIONS:

In this paper, we have motivated and demonstrated a new optimization algorithm that naturally unifies SGD and Adam.

We have focused our empirical results on the default hyper-parameter setting, η = 1, and predetermined learning schedules. With these parameters, the algorithm was shown to produce optimization that is better than or equal to SGD and Adam on image classification tasks. It also performed significantly better than SGD on language modeling tasks.

Together with finding the optimal values for η, we expect a better understanding of the learning schedule to bring light to the way in which the adaptive gradient methods improve convergence. SoftAdam now also makes it possible to create a learning schedule on η, which may be another fruitful avenue of research, expanding on the work of Ward et al. (2018) .

Better understanding of how adaptive gradients improve the convergence of practical machine learning models during training will enable larger models to be trained to more accurately in less time. This paper provides a useful intuition for how that occurs and provides new algorithm that can be used to improve performance across a diverse set of problems. # S t a t e i n i t i a l i z a t i o n i f l e n ( s t a t e ) == 0 : s t a t e [ " s t e p " ] = 0 # E x p o n e n t i a l moving a v e r a g e o f g r a d i e n t v a l u e s s t a t e [ " e x p a v g " ] = t o r c h . z e r o s l i k e ( p . d a t a ) # E x p o n e n t i a l moving a v e r a g e o f # s q u a r e d g r a d i e n t v a l u e s s t a t e [ " e x p a v g s q " ] = t o r c h . z e r o s l i k e ( p . d a t a ) e x p a v g , e x p a v g s q = ( s t a t e [ " e x p a v g " ] , s t a t e [ " e x p a v g s q " ] , 

:

Implementation was based on the public repository at https://github.com/bearpaw/ pytorch-classification

Figures and Tables:

Wilson et al. (2017); Keskar & Socher (2017)), and it requires additional modifications to achieve a convergence guarantee (Reddi et al. (2018); Li & Orabona (2018)).

AlexNet (Krizhevsky et al. (2012)), VGG19 with batch normalization (Simonyan & Zisserman (2014)), ResNet-110 with bottleneck blocks (He et al. (2015)), PreResNet-56 with bottleneck blocks (He et al. (2016)), DenseNet-BC with L=100 and k=12 (Huang et al. (2016)) on the CIFAR-10 dataset (Krizhevsky (2012)) using SGD, AdamW and SoftAdam.

adamw (Mean) ━ sgd (Mean) ━ softadam (Mean) (b1) VGG19 Train Loss (b2) VGG19 Val. Loss (b3) VGG19 Val. Acc. adamw (Mean) ━ sgd (Mean) ━ softadam (Mean) adamw (Mean) ━ sgd (Mean) ━ softadam (Mean) adamw (Mean) ━ sgd (Mean) ━ softadam (Mean) (e1) DenseNet Train Loss (e2) DenseNet Val. Loss (e3) DenseNet Val. Acc.

Figure 1: Training results for CIFAR-10. Architectures are (a) AlexNet, (b) VGG19 (BN), (c) PreResNet-56, (d) ResNet-110, (e) DenseNet-100 (BC) using SoftAdam, AdamW and SGD. Each line indicates the average of 3 training runs with random initializations.

Figure 2: Language model validation perplexity during training for SoftAdam, Adam/AdamW and SGD.

-1, b e t a s = ( 0 . 9 , 0 . 9 9 9 ) , e p s =1e -8 , e t a = 1 . 0 , w e i g h t d e c a y =0 , n e s t e r o v = F a l s e , ) : d e f a u l t s = d i c t ( l r = l r , b e t a s = b e t a s , e p s = eps , e t a = e t a , w e i g h t d e c a y = w e i g h t d e c a y , n e s t e r o v = n e s t e r o v , ) s u p e r ( SoftAdam , s e l f ) . i n i t ( params , d e f a u l t s ) d e f s t e p ( s e l f , c l o s u r e =None ) : l o s s = None i f c l o s u r e i s n o t None : l o s s = c l o s u r e ( ) f o r g r o u p i n s e l f . p a r a m g r o u p s : f o r p i n g r o u p [ " p a r a m s " ] : i f p . g r a d i s None : c o n t i n u e g r a d = p . g r a d . d a t a i f g r a d . i s s p a r s e : r a i s e R u n t i m e E r r o r ( " SoftAdam d o e s n o t s u p p o r t s p a r s e g r a d i e n t s " ) n e s t e r o v = g r o u p [ " n e s t e r o v " ] s t a t e = s e l f . s t a t e [ p ]

) b e t a 1 , b e t a 2 = g r o u p [ " b e t a s " ] s t a t e [ " s t e p " ] += 1 b e t a 2 h a t = min ( b e t a 2 , 1 . 0 -1 . 0 / ( s t a t e [ " s t e p " ] ) ) r b e t a = ( 1b e t a 2 ) / ( 1b e t a 2 h a t ) e t a h a t 2 = ( g r o u p [ " e t a " ] * g r o u p [ " e t a " ] * r b e t a ) # Decay t h e f i r s t and s e c o n d moment w i t h t h e # r u n n i n g a v e r a g e c o e f f i c i e n t e x p a v g . mul ( b e t a 1 ) . a d d ( 1b e t a 1 , g r a d ) e x p a v g s q . mul ( b e t a 2 h a t ) . a d d c m u l ( 1b e t a 2 h a t , g r a d , g r a d ) # C r e a t e t e m p o r a r y t e n s o r f o r t h e d e n o m i n a t o r denom = e x p a v g s q . mul ( e t a h a t 2 / ( t o r c h . mean ( e x p a v g s q ) + g r o u p [ " e p s " ] * g r o u p [ " e p s " ] ) ) denom . s q r t ( ) . a d d ( 1 + g r o u p [ " e t a " ]np . s q r t ( e t a h a t 2 ) ) wd = g r o u p [ " w e i g h t d e c a y " ] * g r o u p [ " l r " ] p . d a t a . a d d (-wd , p . d a t a ) l r e f f = g r o u p [ " l r " ] * ( 1 + g r o u p [ " e t a " ] ) i f n e s t e r o v : p . d a t a . a d d c d i v ( l r e f f * b e t a 1 , e x p a v g , denom ) p . d a t a . a d d c d i v ( l r e f f * ( 1b e t a 1 ) , g r a d , denom ) e l s e : p . d a t a . a d d c d i v (-l r e f f , e x p a v g , denom ) r e t u r n l o s s

The SoftAdam Algorithm Input: θ 0 ∈ F: initial parameters, {α t > 0} T t=1 : learning rate, α wd , β 1 , β 2 , η, : other hyperparameters, J t (θ): loss function Output: θ

Table content: <table/>

Top-1 validation accuracy (%) on CIFAR-10 for different AdamW, SGD and SoftAdam. SoftAdam is restricted to η = 1.

Table content: <table><tr><td colspan="2">ARCHITECTURE ADAMW</td><td>SGD</td><td>SOFTADAM</td></tr><tr><td>AlexNet</td><td colspan="3">74.55 ± 0.37 77.10 ± 0.50 79.30 ± 0.25</td></tr><tr><td>VGG19 (BN)</td><td colspan="3">92.50 ± 0.25 93.52 ± 0.18 93.89 ± 0.12</td></tr><tr><td>PreResNet-56</td><td colspan="3">91.93 ± 0.18 94.01 ± 0.05 94.02 ± 0.03</td></tr><tr><td>ResNet-110</td><td colspan="3">92.11 ± 0.04 94.54 ± 0.29 94.75 ± 0.15</td></tr><tr><td>DenseNet-100</td><td colspan="3">93.25 ± 0.19 95.04 ± 0.08 95.00 ± 0.11</td></tr><tr><td colspan="4">These steps are combined and summarized in Algorithm 1. A reference implementation for the</td></tr><tr><td colspan="3">PyTorch package (Paszke et al. (2017)) is provided in Appendix A.</td></tr></table>

LSTM validation perplexity, test perplexity, and test loss on the Penn Treebank dataset for SoftAdam, Adam and SGD. SGD. Following Liu et al. (2019), we use a linear learning rate schedule. For these parameters, we again find that SoftAdam outperforms SGD significantly.

Table content: <table><tr><td colspan="2">OPTIMIZER VAL. PPL.</td></tr><tr><td>SoftAdam</td><td>63.63 ± 0.19</td></tr><tr><td>SGD</td><td>66.64 ± 0.85</td></tr><tr><td>Adam</td><td>62.03 ± 0.34</td></tr></table>

Best results for training IWSLT'14 De-En with different optimization algorithms.

Table content: <table><tr><td colspan="3">OPTIMIZER VAL. PPL. BLEU SCORE</td></tr><tr><td>SGD</td><td>4.14</td><td>31.43</td></tr><tr><td>SoftAdam</td><td>4.02</td><td>32.76</td></tr><tr><td>AdamW</td><td>3.87</td><td>34.44</td></tr></table>

Peer review (iclr2020):
The paper tends to explain how the tradeoffs between convergence speed and convergence performance are made by different optimization methods. Moreover, the paper modifies Adam-like updating rules and proposes a novel optimization methods, SoftAdam. Finally, the paper performs numerical experiments on traditional image classification tasks as well as language modeling tasks. Pros The paper involves the language modeling tasks in empirical results besides traditional image classification tasks, which helps to explain the convergence performance of optimization methods in a wider range of applications. Cons 1. The writing of this paper is not well organised. Section 1 lacks detailed description of the main idea and the proposed optimization methods, which actually confuses the reader. Section 2 describes too much details of SGD and Adam, and lacks a clear *intuition* which readers exactly expect. 2. The notation in the paper is little confusing. In the update rule of z_{t+1} in Section 2, what is meaning of z? In Section 3, n_t and n_infty are used before a proper defination, and what is relationship between alpha and alpha_t in the implementation of the SoftAdam? 3. The motivation of the proposed method is weak. Such a weak motivation is mainly because of the insufficient *intuition* in Section 2. The author mentions *the Hessians in deep learning problems are not diagonal*, but does not provide further explanation on why more importance should be lay on serving both max and min eigenvalues. 4. There are also several minor problems on the numerical results. Firstly, why the colors of *softAdam* and *sgd* are switched several times in Figure 1? Secondly, the figural result in Figure 1 and he numerical results on language processing models both lack a display of the confidence range.