S OFT A DAM : U NIFYING SGD AND A DAM FOR BETTER STOCHASTIC GRADIENT DESCENT
Anonymous authors
Paper under double-blind review
A BSTRACT
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 classiﬁcation tasks and outperforms
SGD for language modeling tasks.
I NTRODUCTION
One of the most common methods of training neural networks is stochastic gradient descent (SGD) (Bottou et al.
SGD has strong theoretical guarantees, including convergence in locally non-convex optimization problems (Lee et al.
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 modiﬁcation 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 (Wilson et al.
(2017); Keskar & Socher (2017)),
and it requires additional modiﬁcations to achieve a convergence guarantee (Reddi et al.
; Li & Orabona (2018)).
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. R ELATED W ORK
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.
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.
Izmailov et al.
These algorithms are complementary to the algorithm developed here, as they require an underlying algorithm to choose the step direction at any point.
SGD AND A DAM
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 ,
[Equation]
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,
[Equation]
A further development that has improved convergence, especially for LSTM and language modeling tasks, involves the second gradient as well.
This speciﬁc version is known as the Adam algorithm (Kingma & Ba (2014)),
[Equation]
where ˆ m [Equation] and ˆ v t = v t / (cid:112) 1 − β t 2 are unbiased estimators of the ﬁrst and second moment respectively.
2. C ONVERGENCE
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 ) ,
[Equation]
where H t is the Hessian of J ( z ) around z t .
This gives us the gradient,
[Equation]
which becomes the SGD update step,
[Equation]
Unrolling this update step can be shown to lead to an expression for the distance from the optimal value z (cid:63) , in the basis of Hessian eigenvectors ξ i :
[Equation]
We can see that the learning is stable if the learning rate α satisﬁes,
[Equation]
In addition, we ﬁnd 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. A DAPTIVE M OMENT E STIMATION
If rather than a single learning rate α , we were to use a diagonal matrix D such that the update is,
[Equation]
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,
[Equation]
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 ﬁnd 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,
[Equation]
Combining this with Equation 2, we compare the AdaGrad coefﬁcient to the optimal value for d i = λ − 1 i , ﬁnding,
[Equation]
As long as (cid:15) (cid:28) αλ
i , this will be true when,
[Equation]
which is the sum squared distance of the parameter value
z ti to its optimum z (cid:63)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 0 i and optima z (cid:63)i can be made over the same subspaces.
That is, if z (cid:63)i is uncorrelated to λ i , we can expect this to have good performance on average.
However, there can be signiﬁcant 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 (cid:63)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 signiﬁcant noise to the process, making it difﬁcult to accurately measure gradients and therefore ﬁnd the optimum in other parameters.
This problem will be signiﬁcantly 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.
In addition, reviewing our earlier assumption of small (cid:15) , these algorithms will perform worse for small eigenvalues
λ
i < (cid:15)/α .
This might be especially bad in Adam where late in training when the learning rate α ∼ (cid:80) t (cid:16)
z ti − b
i
λ
i (cid:17) 2 is small.
We ﬁnally note that the Hessians in deep learning problems are not diagonal (Sagun et al.
; Li et al.
As such, each element might be better optimized by a learning rate that better serves both its min and max eigenvalues.
S MOOTHED A DAPTIVE G RADIENT
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.
[Equation]
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,
[Equation]
where ¯ v t is the mean value of v t , to write,
[Equation]
One issue with the above estimation is that its variance is very large at the beginning of training (Liu et al.
It was suggested that this is the reason that warmup is needed for Adam and shown
Input : θ 0 ∈ F initial parameters
{ α t > rameters
J
Output : θ T ∈ F m 0 , v 0 ←(cid:46)
Initialize for t ← 1 . . .
T do g t ← ∇ θ J ( z ) (cid:46) Find gradient m t ← β 1 m t − 1 + (1 − β 1 ) g t (cid:46) Update ﬁrst moment ˆ β 2 t ← min ( β 2 , 1 − 1 /t ) v t ← ˆ β 2 t v t − 1 + (cid:16) 1 − ˆ β 2 t (cid:17)
g 2 t (cid:46) second moment ¯ v t ← mean elem [ v t ] (cid:46) Average over all elements r t ← sqrt (cid:104) (1 − β 2 ) / (cid:16) 1 − ˆ β 2 t (cid:17)(cid:105) d t = α (1 + η ) (cid:16) 1 + η − ηr t + ηr t (cid:112) v t / (¯ v t + (cid:15) 2 ) (cid:17) −(cid:46)
Calculate θ t = θ t − 1 − m t d t (cid:46) Perform 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 deﬁne r t = (cid:112) n t /n ∞ and: [Equation]
This ﬁnally forms the basis for our algorithm.
3. A LGORITHM
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:
[Equation]
In addition, the second moment v t is calculated in an unbiased way using an effective ˆ β 2 ( t ) , or by abuse of notation ˆ β 2 t :
[Equation]
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 ∞ :
[Equation]
We ﬁnally note that the weight decay should be calculated separately from this update as in AdamW (Loshchilov & Hutter (2017)).
[Table]
These steps are combined and summarized in Algorithm 1.
A reference implementation for the
PyTorch package (Paszke et al.
(2017)) is provided in Appendix A.
E MPIRICAL R ESULTS
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 η = (cid:15) − 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. CIFAR-
We trained a variety of networks: 1 AlexNet (Krizhevsky et al.
VGG19 with batch normalization (Simonyan & Zisserman (2014)), ResNet-110 with bottleneck blocks (He et al.
PreResNet-56 with bottleneck blocks (He et al.
(2016)), DenseNet-BC with L=100 and k=12 (Huang et al.
on the CIFAR-10 dataset (Krizhevsky (2012)) using SGD, AdamW and SoftAdam.
For each model and optimization method, the weight decay was varied over [1e-4,2e-4,5e-4,1e3,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 ﬁnal 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 ﬁnd that SoftAdam equals or outperforms SGD in training classiﬁers on this dataset.
Due to the larger optimal weight decay constant, SoftAdam achieves lower validation loss at a higher train loss than SGD.
4. P ENN T REEBANK
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.
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 signiﬁcantly 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. IWSLT’14 G ERMAN TO E NGLISH
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 [Figure] [Figure]
[Table]
SGD.
Following Liu et al.
(2019), we use a linear learning rate schedule.
For these parameters, we again ﬁnd that SoftAdam outperforms SGD signiﬁcantly.
C ONCLUSIONS
In this paper, we have motivated and demonstrated a new optimization algorithm that naturally uniﬁes 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 classiﬁcation tasks.
It also performed signiﬁcantly better than SGD on language modeling tasks.
Together with ﬁnding 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.
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 a new algorithm that can be used to improve performance across a diverse set of problems. [Table]
A R EFERENCE P Y T ORCH I MPLEMENTATION import math
import t o r c h from t o r c h .
optim .
o p t i m i
z e r import Optimizer import numpy as np c l a s s SoftAdam ( Optimizer ) :
def i n i t ( s e l f , params , l r =1e − 1, b e t a s =( 0.9 , 0 .
9 9 9 ) , eps =1e − 8, e t a =1.0 , weight decay =0 , n e s t e r o v =False , ) :
d e f a u l t s =
d i c t ( l r = lr , b e t a s = betas , eps=eps , e t a =eta , weight decay =weight decay , n e s t
e
r
o v = nesterov , ) super ( SoftAdam , s e l f ) .
i n i t ( params , d e f a u l t s ) def 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 not None : l o s s =
c l o s u r e ( ) f
o r group in s e l f .
param groups : f o r p in group [” params ” ] : i f p .
grad i s
None : c on t in u e grad = p .
grad .
data i f grad .
i s s p a r s e : r a i s e RuntimeError ( ” SoftAdam does not support s p a r s e g r a d
i e n t s ” )
n e s t e r o v = group [” n e s t e r o v ”] s t a t e = s e l f .
s t a t e [ p ] # S t a t e i n i t i a l i z a t
i o n i f len ( s t a t e ) =
= 0: s t a t e [” s t e p ”] =
# Exponential moving average of g r a d i e n t values s t a t e [” exp avg ”]
=
t o r c h .
z e r o s l i k e ( p . data )
# Exponential moving average of # squared g r a d i e n t values s t a t e [” exp avg sq ”]
=
t o r c h .
z e r o s l i k e ( p . data ) exp avg , exp avg sq
=
( s t a t e [” exp avg ” ] , s t a t e [” exp avg sq ” ] , ) beta1 , beta2 = group [” b e t a s ”] s t a t e [” s t e p ”]
+=b e t a 2 h a t = min ( beta2 , 1.0 − 1.0 / ( s t a t e [” s t e p ” ] ) )
r b e t a
=
1 − beta2 ) /
(1 − b e t a 2 h a t ) e t a h a t 2 =
( group [” e t a ”] ∗ group [” e t a ”] ∗ r b e t a ) # Decay the f i r s t and second moment with the # running average c o e f f i
c i e n t exp avg .
mul ( beta1 ) .
add (1 − beta1 , grad )
exp avg sq .
mul ( b e t a 2 h a t ) .
addcmul ( 1 − bet a2 hat , grad , grad )
# Create temporary t e n s o r f o r the denominator denom =
exp avg sq .
mul ( e t a h a t/ ( t o r c h .
mean ( exp avg sq ) + group [” eps ”] ∗ group [” eps ”] ) ) denom .
s q r t ( ) .
add ( 1 + group [” e t a ”] − np .
s q r t ( e t a h a t 2 ) )
wd = group [” weight decay ”] ∗ group [” l r ”] p .
data .
add ( − wd , p .
data ) l r e f f = group [” l r ”] ∗ (1 + group [” e t a ” ] ) i
f n e s t
e r o v : p .
data .
addcdiv ( − l r e f f ∗ beta1 , exp avg , denom ) p .
data .
addcdiv ( − l r e f f ∗ (1 − beta1 ) , grad , denom )
e l s e : p .
data .
addcdiv ( − l r e f f , exp avg , denom ) r e t u r n l o s s