Dimension-Reduced Adaptive Gradient Method
LI JINGYANG @
U
NUS
EDU
SEA
Abstract
Adaptive gradient methods, such as Adam, have shown faster convergence speed than SGD across various kinds of network models at the expense of inferior generalization performance.
In this work, we proposed a Dimension-Reduced Adaptive Gradient Method (DRAG) to eliminate the generalization gap.
DRAG makes an elegant combination of SGD and Adam by adopting a trustregion like framework.
We observe that 1) Adam adjusts stepsizes for each gradient coordinate according to some loss curvature, and indeed decomposes the n -dimensional gradient into n standard basis directions to search; 2) SGD uniformly scales gradient for all gradient coordinates and actually has only one descent direction to minimize.
Accordingly, DRAG reduces the high degree of freedom of Adam and also improves the flexibility of SGD via optimizing the loss along k ( ≪ n ) descent directions, e.g. the gradient direction and momentum direction used in this work.
Then per iteration, DRAG finds the best stepsizes for k descent directions by solving a trust-region subproblem whose computational overhead is negligible since the trust-region subproblem is lowdimensional, e.g. k
= 2 in this work.
DRAG is compatible with the common deep learning training pipeline without introducing extra hyper-parameters and with negligible extra computation.
Moreover, we prove the convergence property of DRAG for non-convex stochastic problems that often occur in deep learning training.
Experimental results on representative benchmarks testify the fast convergence speed and also superior generalization of DRAG.
1. Introduction
Training neural networks can be seen as solving the following non-convex optimization problem
[Equation]
where f is the loss function and x ∈ R n is the variable.
Among all optimizers, Adam [6] is one of the most popular algorithm to solve problem (1).
At each training iteration, Adam maintains an exponential moving average (EMA) of first and second moments of stochastic gradient v t and u t as
[Equation]
where β 1 , β 2 ∈ [0 , 1] are constant and g t − 1 :
=
˜ ∇ f
( x t − 1 ) is the stochastic gradient.
It adaptively scales the learning rates for each gradient coordinate, and actually minimizes the loss function along n descent directions
[Equation]
where ˆ v t , ˆ u t are bias-corrected v t , u t , e i is the standard basis vector with 1 for dimension i and 0 for all other dimensions.
Specifically, Adam adopts a stepsize of η √ ˆ u t,i
[Equation]
for the i -th gradient descent direction − ˆ v t,i e i .
While adaptive stepsize boosts the convergence of Adam, it weakens the generalization performance due to noise and overfitting.
In contrast, SGD generalizes well because it uses a single stepsize for all gradient coordinates and indeed optimizes the loss function only along the gradient direction.
One interpretation for their different generalization performance is that Adam’s update direction no longer falls into the subspace spanned by all stochastic gradients span { g 0 , · · · , g t } [17, 23], while SGD do.
Actually, Wilson et al.
[17] proved that on a binary classification problem,
SGD converges to the max-margin solution because its update at each step is linear combination of stochastic gradients, while adaptive gradient methods converge to solutions that generalize poorly because adaptivity makes the algorithm susceptible to noises and therefore causes overfitting.
To overcome the issue just mentioned, motivated by DRSOM [22], we proposed DRAG algorithm to optimize the loss function in (1) from the gradient direction and the momentum direction.
It maintains flexibility in the update direction while inheriting the generalization capacity of SGD.
At each step, it searches for the optimal stepsizes along these two directions by solving a twodimensional trust-region subproblem.
Therefore, from the optimization perspective, it conducts the optimal update within the two-dimensional subspace spanned by gradient direction and momentum direction.
Moreover, while DRAG adopts the trust-region framework, it is compatible with the dominant deep learning training pipeline without introducing extra hyperparameters.
2. Method
As described in Algorithm 1, at each iteration DRAG first computes stochastic gradient g t − 1 , and use it to update the first moment v t and second moment u t of stochastic gradient like Adam.
Then, we introduce the bias-corrected second moment ˆ u t to approximate the Hessian.
In this way, DRAG constructs the trust-region subproblem in line 9 of Algorithm 1.
While solving this trust-region subproblem in high-dimensional parameter space is computational expensive, DRAG solves it in the two-dimensional subspace spanned by bias-corrected first moment direction ˆ v t and momentum direction
d t − 1 , making the computational overhead negligible.
Here we intuitively set the trustregion radius as η ∥ ˆ v t ∥ , and the benefits of this setting is described in Section 2.1.
After calculating the solution α 1 t and α 2 t of the subproblem, we get an optimal update p = − α 1 t
ˆ v t + α 2 t d t − 1 in the two-dimensional subspace.
Finally, we follow [9] and conduct a decoupled weight decay step.
This is the overall framework of our DRAG.
The only extra computational overhead of DRAG compared with Adam is solving the twodimensional trust-region subproblem in line 9 of Algorithm 1.
The trust-region subproblem can be
Input
3: for
1 , · · · , m do 4:
Compute stochastic gradient g t − 1 = ˜
[Equation]
f ( x t − 1 ) .
[Equation]
[Equation]
[Equation]
8: d t − 1 = x t − 1 − x t − 2 if t ≥ 2 else d t − 1 = 0 .
9: ( α 1 t , α 2 t ) = argmin p
[Equation]
⟨ ˆ
[Equation]
t , p
[Equation]
+ 12 ⟨ p , H t p ⟩ |
[Equation]
p ∥
[Equation]
η
[Equation]
ˆ
[Equation]
t ∥ , p
[Equation] [Equation]
t
[Equation] [Equation]
t − 1 (cid:9) .
10: x t = x t − 1 − α 1 t ˆ
[Equation]
t [Equation] α 2 t d t −11: x t = x t − ηγ x t − 1 (Conduct weight decay) 12: end for 13: Output: x 1 , · · · , x m formally formulated as follows:
[Equation]
(cid:2)
αα 2 (cid:3) (cid:34) ˆ v Tt H t
ˆ v t − ˆ v Tt H t
d t −− ˆ v Tt H t d t −d Tt − 1 H t d t −(cid:35) (cid:20) αα(cid:21) s.t. [Equation] where H t = diag( √ ˆ u t + ν ) as defined in Algorithm 1.
This two-dimensional subproblem can be solved efficiently by using its global minimal condition.
In Appendix A, we transform this subproblem into a standard trust-region subproblem, and then an ϵ -global primal-dual solution satisfying
KKT condition can be found in O (log log( 1 ϵ )) time [10].
See more details in Appendix A.
2.1. Benefits of our algorithm Flexibility of update
As in Algorithm 1, DRAG updates the variable x along EMA of gradient direction ˆ v t and momentum direction d t − 1 .
This update direction choice acts as a trade-off between the whole space search of Adam and one direction search of SGD.
Moreover, the update of DRAG lies in the subspace span { ˆ v t , d t − 1 } ∈ span { g 0 , · · · , g t − 1 } .
This means that the parameter update direction is always a combination of stochastic gradients.
According to Wilson et al.
[17], this property makes DRAG always converge to the max-margin solution of the binary classification problem, which has the best generalization capacity.
This helps to explain DRAG’s excellent generalization performance in practice.
Optimal stepsizes
DRAG solves the dimension-reduced subproblem at each training epoch and finds the best update along the gradient direction and momentum direction.
This optimal update is evaluated by the quadratic approximation to the loss function, where the Hessian is approximated by second moment √ ˆ u t and gradient is approximated by first moment ˆ v t .
Since DRAG conducts optimal update along gradient and momentum direction within the learning rate we set, it converges faster than SGD on training dataset and is comparable with adaptive gradient methods.
Heuristic trust-region radius
We set the trust-region radius for the subproblem as η ∥ ˆ v t ∥ .
The intuition is that when gradient is large, we hope our algorithm can make a larger step to minimize the loss function significantly.
While when gradient is small, we hope our method to be stable and don’t change the parameters too much.
This heuristic design not only frees us from changing the radius at each step as trust region method does, but also make our algorithm compatible well with dominant deep learning training pipeline without introducing extra hyperparameters.
2.2. Convergence Analysis
For the analysis of stochastic non-convex algorithm, we follow the works Guo et al.
[2], Zhuang et al.
[24] and make the following necessary assumption.
Assumption 1 For non-convex problem min x ∈ R n f ( x ) , we assume the loss f ( x ) satisfies
• f is L -Lipschitz smooth.
The gradient estimation g is unbiased, namely E [ g t ] = ∇ f ( x t ) , and its variance can be bounded as E [ ∥ g t − ∇ f ( x t ) ∥ 2 ] ≤ σ 2 .
Then we can derive the convergence of our proposed algorithm and also provide its stochastic gradient complexity to find an ϵ -approximate first-order stationary point.
Theorem 1 Suppose Assumption 1 holds.
Let β t = β and η t =
η for all t .
Assume there exist constants α, G > 0 , such that α ≤ min t α 1 t and α 1 t ≤ ηG , | α 2 t | ≤ ηG .
In addition, η ≤ min (cid:26)2 LG , (cid:16) (1 − β ) 2 α 8 GL(cid:17) 13 , (cid:16) α96 G(cid:17) 14 , (cid:0) α 48 LG 2 (cid:1) 14 , (cid:0) α 192 L 2 G 3 (cid:1)(cid:27) .
Then, if 1 − β ≤ ϵ 2 3 C 2 σ 2 and T ≥ max
[Equation] [Equation]
(cid:111)
, DRAG can achieve
[Equation]
where C 1 = 4 ( f ( x 0 ) − f ( x ∗ )) , C 2 = 4 ηGα and C 3 = 2 ηG E [ ∥∇ f ( x 0 ) − (1 − β
[Equation]
g 0 ∥ 2 ] α .
Remark 1 Theorem 1 with its proof in Appendix C demonstrates that by properly selecting constant trust-region radius η t and constant momentum parameter β t (correspond to β 1 in Algorithm 1),
DRAG can converge to an ϵ -approximate first-order stationary point of the non-convex stochastic problem with stochastic gradient complexity O
( ϵ − 4 ) .
Note that the assumptions on α 1 t and α 2 t are mild with the design of DRAG, see details in Appendix B.
The complexity of DRAG is of the same order as the lower bound provided by Arjevani et al.
A similar complexity has also been obtained in, for example, LAMB [20], Adam-family [2].
In the analysis of DRAG, we only need a unbiased and variance-bounded stochastic gradient, without any large mini-batch sizes requirement as in LARS [19] and LAMB [20].
In addition, some previous works [8, 11, 14, 21] require the momentum parameter β t to be very close or decreasing to zero.
In contrast, DRAG requires β t to be close to one, which is more consistent with the practice.
Proof of Theorem 1 and more convergence analysis of DRAG can be found in Appendix C.
[Figure]
3. Experiments
We conduct experiments on several representative benchmarks, including VGG [15], ResNet [3],
DenseNet [5] on CIFAR10, CIFAR100 dataset [7], and LSTM [4] on the Penn Treebank dataset [12].
We compare our algorithm DRAG with some popular deep learning optimizers, including SGD [13],
Adam [6], AdamW [9], AdaBound [11], AdaBelief [24], RAdam [8], Yogi [21].
Experimental results show that DRAG has faster convergence speed compared with SGD and it achieves state-ofthe-art generalization performance.
We also conduct ablation study to show 1) two search directions (DRAG) performs better than one direction and multiple directions and 2) DRAG is robust to different learning rate schedules.
Details of the ablation study is in Appendix D.
3.1. CNNs on image classification
We conducted experiments for VGG16 with Batch Normalization, ResNet34, and DenseNet121 on
CIFAR10 an CIFAR100 dataset.
The experimental setting is borrowed from AdaBelief [24] and we also use their default setting for all the hyperparameters.
For DRAG, we choose its learning rate to be the same as in SGD, which is 0.1, and weight decay factor is 0.0015 for CIFAR10 and 0.
for CIFAR100.
Other hyperparameters of DRAG is the same as the default setting ( β 1 = 0 .
9 , β 2 = 0 .
999 , ϵ = 10 − 8 ).
As Figure 1 shows, DRAG has convergence speed comparable with adaptive gradient methods and it attains the best generalization performance.
To be specific, DRAG obtains more than 0.5% generalization accuracy gain over AdaBelief [24] on most tasks.
The detailed test accuracy is summarized in Table 1.
The possible reasons for this improvement on the convergence speed and generalization capacity is 1) DRAG searches for the optimal update along two directions and thus converges faster, 2)
DRAG confines the search of update within the two-dimensional subspace spanned by gradient and momentum direction to avoid overfitting and alleviating the influence of noises, therefore it generalizes better.
[Table] [Table]
3.2. LSTMs on language modeling
We experiment with LSTM on the Penn Treebank dataset and record the perplexity (lower is better).
We follow the exact experimental setting in Adabelief [24] and use their default hyperparameters except for SGD.
For SGD, we use the same hyperparameters as DRAG to make a fair comparison between the two.
For SGD and DRAG, we set their learning rate as 25, 75, 75 for 1,2,3-layer LSTM and weight decay factor as 2 . 5 × 10 − 6 .
SGD’s generalization performance in our setting is better than the results provided by Zhuang et al.
From Table 2, we can see that DRAG attains more than 0.5 less perplexity than other optimizers.
The good generalization performance may be due to
DRAG’s two-direction search.
The gradient direction inherits SGD’s good generalization property and the extra momentum direction further improves its performance.
4. Conclusion
In this paper we propose the DRAG algorithm, which finds the optimal update of the parameters along gradient and momentum directions at each iteration.
Compared with Adam, DRAG reduces the flexibility of update direction from searching in the whole parameter space to updating in a two-dimensional subspace, therefore is less susceptible to overfitting and has better generalization performance.
Compared with SGD, DRAG inherits the gradient update direction and also update along an extra momentum direction, thus it has faster convergence speed and comparable generalization capacity.
Our algorithm can be further generalized to any number of search directions and any choice of Hessian approximation.
Appendix A. Solve the trust-region subproblem Recall the trust region subproblem
[Equation]
where α :=
[Equation]
and
[Equation]
t
[Equation]
In order to solve this trust region subproblem
[Equation]
we transform it into a standard trust region subproblem with L 2 -norm constraint
[Equation]
When matrix G t is positive definite
[Equation]
we have
[Equation]
So we let y =
L Tt α , then α
= L − T t y and the subproblem becomes
[Equation]
In this way, the trust region subproblem is transformed to a standard spherical constrained quadratic optimization problem and it can be solved efficiently [18].
When | G t | = 0 , this means ˆ v t is linearly dependent with d t − 1 .
In this case, we solve the one-dimensional subproblem as described in Section 2.
Appendix B. Mild assumptions on α 1 t , α 2 t
The trust-region subproblem to be solved in Algorithm 1 has global optimality condition [10] given by
[Equation]
By its construction, we know that G t is positive semidefinite.
In practice, numerical issues sometimes make it indefinite, leaving the trust-region subproblem insoluble.
Thus, we make an adjustment to G t
[Equation]
where λ min is the smallest eigenvalue of G t .
In this way, when | G t | ̸ = 0 , we have
[Equation]
which means
[Equation]
With the common additional assumption that stochastic gradient g t =
˜ ∇ f
( x t ) has bounded L ∞ norm, i.e. ∥ g t ∥ ∞ ≤ G ∞ , then ˆ v t as an moving average of g t also has bounded norm ∥
ˆ v t ∥ .
Therefore, we can see that | α 1 t η | , | α 2 t η | are upper bounded by a constant.
When | G t | = 0 , which means d t − 1 is parallel with ˆ v t .
Then we only need to find the optimal update within the trust-region along gradient direction ˆ v t .
In this case, we manually set α 2 t = 0 in our implementation of DRAG, and then α 1 satisfies | α 1 | ≤ η .
From discussions above, we can see the assumption that | α 1 t η | , | α 2 t η | are upper bounded in Theorem 1 and Theorem 2 is satisfied given the common assumption that stochastic gradient g t =
˜ ∇ f
( x t ) has bounded L ∞ norm.
For the simplicity of notations, we directly make assumptions for α 1 t and α 2 t in Theorem 1 and Theorem 2.
For the assumption that α 1 t is positive and α 1 t η is lower bounded by a constant, we give an explanation here by intuition and empirical results.
Gradient direction is what we considered the most important update direction locally, because by the training pipeline of neural networks, stochastic gradients of training parameters are the new information we gain at each iteration.
Thus, we consider the update should at least move towards the gradient descent direction rather than move towards the gradient ascent direction.
Moreover, from the observations of α 1 t under all the experimental settings, α 1 t is always positive and α 1 t η is always larger than 0.1.
Therefore, this assumption on α 1 t is reasonable based on common sense and holds true in practice.
Appendix C. Convergence analysis in non-convex stochastic optimization
To clarify Assumption 1, we give the following definitions.
Definition 1 For a differentiable function f , x is said to be an ϵ -approximate first-order stationary point if it satisfies ∥∇ f ( x ) ∥ ≤ ϵ .
Definition
For a differentiable funtion f ( x ) , it is called L -Lipschitz smooth if it statisfies ∥∇ f ( x ) − ∇ f ( y ) ∥
≤ L ∥ x − y ∥ for a constant L > 0 and any x , y in domain of f .
Except for Theorem 1, we also have the following result which establishes an O (log T/ √ T ) sub-linear convergence rate for DRAG.
Theorem 2 Suppose Assumption 1 holds.
Assume there exist constants δ, G > 0 , such that 0 < δ ≤ α 1 t η t ≤ G , | α 2 t | η t ≤ G .
Set η t =
c η √ t +2 , 1 − β t =
Cc η √ t +1 , for any c η and C satisfying C ≥ L (cid:113) 8 Gδ , and c η ≤ (cid:26)
[Equation]
(cid:27) .
Then there exist two constant C 1 and Cwhich are independent with T , such that
[Equation]
Given a tolerance ϵ > 0 , if T ≥ ˜ O ( 1 ϵ 4 ) , we have
[Equation]
Remark 2 Theorem 2 establishes an O (log T/ √ T ) sub-linear convergence rate for DRAG by choosing a decreasing η t and 1 − β t with the order O (1 / √ t ) .
Similar sub-linear convergence rates are also established by Zou et al.
[25] for Adam and Guo et al.
[2] for Adam-type optimizers.
While Zou et al.
[25] has restrictions on the second moment momentum parameter β 2 .
In Theorem 2, we only need β t (corresponds β 1 in Algorithm 1) to increase to one.
One key ingredient in our analysis is an existing variance recursion of the stochastic estimator based on moving average, which is given by the following lemma.
Lemma 3 (Variance Recursion [16]) Suppose Assumption 1 holds, then we have
[Equation]
where E t [ · ] denotes the conditional expectation with respect to all randomness before g t .
Before proving Theorem 1, we need to prove the following auxiliary lemma.
Lemma 4 Suppose Assumption 1 holds.
Assume there exist α, η, δ, G > 0 , such that α ≤ min t α 1 t , max t η t ≤ η , and 0 < δ ≤ αη ≤ α 1 t η t ≤ G , | α 2 t | η t ≤ G , ( δ, G ) are constants independent with t .
In addition, η ≤ min (cid:26)
[Equation]
Then there exist positive constants C 1 , C 2 and C 3 , which are all independent with T , such that the following estimation holds:
[Equation]
Proof Since F is L -smooth
[Equation]
we have
[Equation]
By Lemma 3, we can obtain
T (cid:88)
[Equation]
Taking expectation for both sides of (5) and taking summation among t =
[Equation]
..., T − 1 , combining with (6)
[Equation]
we have
[Equation]
By AM-GM inequality,
[Equation]
Combining all together, we have
T − 1 (cid:88)
[Equation]
where x ∗ is one of the global minimizer of F .
Since α 1 t ≤ ηG ≤ 12 L , we have α 1 t (1 − Lα 1 t )≥ α 1 t 8 .
By the conditions for η and α , we have α 1 t 16 ≥ α
[Equation]
96 ≥
[Equation]
96 ≥
[Equation]
≥ α 22 ,t +1 η 2 t +1 α 1 ,
[Equation]
+ , α 1 t 96 ≥ α 96 ≥ Lη 4 G 2≥ Lα 22 ,t +1 η 2 t +1, and α 1 t 96 ≥ α 96 ≥ 2 L 2 {+η 5 G 3 ≥ L 
+} α 1 ,t +1 α 22 ,t +1 η 2 t +1 1 − Lα 1 ,t + .
By ∥ d t ∥ ≤ η t ∥ v t ∥ .
Since v 0 = 0 , we have
[Equation] [Equation]
[Equation]
Combining (8) and (9), we can obtain
T − 1 (cid:88)
[Equation]
Dividing the above two inequalities by αT 4 and αT 32 respectively
[Equation]
we have
[Equation]
which completes the proof by letting C
= [Equation] f [Equation]
x [Equation]
Proof of Theorem
Proof By the selections of α and η t in Theorem [Equation] let δ =
α/η .
By Lemma [Equation] we have
[Equation]
The conditions 1 − β ≤ ϵ 2 3 C 2 σ 2 and T ≥ max (cid:110) 3 C 1 αϵ 2 , 3 C 3 (1 − β ) ϵ(cid:111) lead to C 1 Tα ≤ ϵ 2 3 , C 2 (1 − β ) σ 2 ≤ ϵ 2 3 ,
C 3 T (1 − β ) ≤ ϵ 2 3 .
This completes the proof.
Proof of Theorem
Proof From (5) in Lemma 4, we have f
[Equation]
By Lemma 3, we have
[Equation]
Taking expectation and summation for t = 1 , ..., T , we get
T − 1 (cid:88) t
[Equation]
[Equation]
Note that 1 − β t +1 =
Cη t , so
α 1 t 2 ≤ G 2 η t = G 2 C (1 − β t +1 ) .
Taking expectation for both sides of (10) and taking summation among t = 0 , ..., T − 1 , combining with (11), we can obtain
E [ f ( x T ) − f ( x 0 )]
[Equation]
From (7) and (12), we can get
[Equation]
[Equation]
[Equation]
By the conditions for c η and C , we have α 1 t ≤ η t G ≤ 12 L , α 1 t (1 − Lα 1 t )≥ α 1 t 8 .
By similar arguments in the proof of Lemma 4, we have α 1 t 16 ≥ δη t 16 ≥ η 2 t +1 GL 2 2 η t C=
GL 2 η 2 t +1 2(1 − β t +1 ) C , α 1 t 96 ≥ δη t 96 ≥ α 22 ,t +1 η 2 t +1 α 1 ,t + , α 1 t 96 ≥ δη t 96 ≥ Lα 22 ,t +1 η 2 t +1, and α 1 t 96 ≥ δη t 96 ≥ 2 L 2 η 5 t +1 G 3 ≥ L 2 α 1 ,t +1 α 22 ,t +1 η 2 t +1 1 − Lα 1 ,t + . By ∥ d t ∥ ≤ η t ∥ v t ∥ .
Since v 0 = 0 , we can get
T
[Equation]
≤ −
[Equation]
Combining (13) and (14)
[Equation]
we can obtain
T − 1 (cid:88)
[Equation]
Then
[Equation]
the final assertion can be obtained by (cid:80)
Tt =1 1 t +
=
O (log T ) .
This completes the proof
[Equation] Appendix D. Ablation Study Different search directions
We compare the performance of algorithms that solve the trust-region subproblem in one-dimensional, two-dimensional (DRAG), and three-dimensional subspaces as described in Section 2.
As show in Table 3, DRAG generalizes better than its one search direction and three search direction counterparts.
The reason is that DRAG updates in more directions than the one search direction counterpart while its subproblem can be solved more accurately than the three direction counterpart, since low-dimensional subproblem can be solved with less numerical errors in single precision arithmetic by GPU.
[Table]
Robustness to learning rate schedule
DRAG is robust to different choices of learning rate schedule.
Except for letting the learning rate decay at epoch 150 as in Section 3.1, we also conduct experiments on decaying the learning rate at epoch 120 and adopting cosine annealing learning rate schedule.
The only change of hyperparameter setting from Section 3.1 is we increase the learning rate of DRAG from 0.1 to 0.12 in cosine annealing schedule.
The intuition is that when the trustregion radius is decreased during the training process, we need a larger initial radius to converge to a better local minima.
We compared DRAG’s test performance with other optimizers with VGG16 on
CIFAR10, details are presented in Table
[Table]
which shows that DRAG enjoys the best generalization performance for all the learning rate schedules
[Table]
[Table]
Decay at 150 epoch 94.