Original: D IMENSION -R EDUCED A DAPTIVE G RADIENT
Revised: D IMENSION -R EDUCED A DAPTIVE G RADIENT

Original: M ETHOD
Revised: M ETHOD

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Adaptive gradient methods, such as Adam, have shown faster convergence speed than SGD across various kinds of network models.
Revised: Adaptive gradient methods, such as Adam, have shown faster convergence speed than SGD across various kinds of network models.

Original: However, adaptive algorithms often suffer from inferior generalization performance than SGD.
Revised: However, adaptive algorithms often suffer from inferior generalization performance than SGD.

Original: Though much effort via combining Adam and SGD have been invested to solve this issue, adaptive methods still fail to attain as good generalization as SGD.
Revised: Though much effort via combining Adam and SGD have been invested to solve this issue, adaptive methods still fail to attain as good generalization as SGD.

Original: In this work, we proposed a Dimension-Reduced Adaptive Gradient Method (DRAG) to eliminate the generalization gap.
Revised: In this work, we proposed a Dimension-Reduced Adaptive Gradient Method (DRAG) to eliminate the generalization gap.

Original: DRAG makes an elegant combination of SGD and Adam by adopting a trust-region like framework.
Revised: DRAG makes an elegant combination of SGD and Adam by adopting a trust-region like framework.

Original: We observe that 1) Adam adjusts stepsizes for each gradient coordinate , and indeed decomposes the n -dimensional gradient into n independent directions to search; 2) SGD uniformly scales gradient for all gradient coordinates and actually has only one descent direction to minimize.
Revised: We observe that 1) Adam adjusts stepsizes for each gradient coordinate, and indeed decomposes the n -dimensional gradient into n independent directions to search; 2) SGD uniformly scales gradient for all gradient coordinates and actually has only one descent direction to minimize.

Original: 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.
Revised: 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.

Original: 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 low-dimensional, e.g. k
Revised: Then per iteration, DRAG finds the best stepsizes for k descent directions by solving a trustregion subproblem whose computational overhead is negligible since the trustregion subproblem is low-dimensional, e.g. k

Original: = 2 in this work.
Revised: = 2 in this work.

Original: DRAG is compatible with the common deep learning training pipeline without introducing extra hyper-parameters and with negligible extra computation.
Revised: DRAG is compatible with the common deep learning training pipeline without introducing extra hyper-parameters and with negligible extra computation.

Original: Moreover, we prove the convergence property of DRAG for non-convex stochastic problems that often occur in deep learning training.
Revised: Moreover, we prove the convergence property of DRAG for non-convex stochastic problems that often occur in deep learning training.

Original: Experimental results on representative benchmarks testify the fast convergence speed and also superior generalization of DRAG.
Revised: Experimental results on representative benchmarks testify the fast convergence speed and also superior generalization of DRAG.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: SGD (Robbins & Monro, 1951) and its variant with momentum (Sutskever et al., 2013) are used widely in training deep neural networks.
Revised: SGD (Robbins & Monro, 1951) and its variant with momentum (Sutskever et al., 2013) are used widely in training deep neural networks.

Original: They perform well empirically and have theoretical guarantee (Szegedy et al., 2015;
Revised: They perform well empirically and have theoretical guarantee (Szegedy et al., 2015;

Original: He et al., 2016; Lee et al., 2016; Hardt et al., 2016).
Revised: He et al., 2016; Lee et al., 2016; Hardt et al., 2016).

Original: However, SGD suffers from two issues.
Revised: However, SGD suffers from two issues.

Original: It often has slow convergence speed since it adopts a single learning rate for all the gradient coordinates.
Revised: It often has slow convergence speed since it adopts a single learning rate for all the gradient coordinates.

Original: Moreover, it is also hard to tune the single learning rate (Wilson et al., 2017), since not all gradient coordinates share the same optimization properties.
Revised: Moreover, it is also hard to tune the single learning rate (Wilson et al., 2017), since not all gradient coordinates share the same optimization properties.

Original: To resolve this problem, several adaptive gradient methods have been proposed to adopt different learning rate for different gradient coordinates.
Revised: To resolve this problem, several adaptive gradient methods have been proposed to adopt different learning rate for different gradient coordinates.

Original: Typical examples of such methods include Adagrad (Duchi et al., 2011), RMSProp (Tieleman et al., 2012), and Adam (Kingma & Ba, 2014).
Revised: Typical examples of such methods include Adagrad (Duchi et al., 2011), RMSProp (Tieleman et al., 2012), and Adam (Kingma & Ba, 2014).

Original: Emprically, these methods have shown faster convergence speed and eased the burden of carefully tuning the learning rate in SGD across many kinds of networks.
Revised: Emprically, these methods have shown faster convergence speed and eased the burden of carefully tuning the learning rate in SGD across many kinds of networks.

Original: However, their generalization performance are often worse than SGD in many scenarios (Wilson et al., 2017; Zhou et al., 2020).
Revised: However, their generalization performance are often worse than SGD in many scenarios (Wilson et al., 2017; Zhou et al., 2020).

Original: Some algorithms are proposed to combine the fast convergence speed of adaptive gradient methods and good generalization performance of SGD.
Revised: Some algorithms are proposed to combine the fast convergence speed of adaptive gradient methods and good generalization performance of SGD.

Original: Instances of this type of algorithms include SWATS (Keskar & Socher, 2017) which automatically switchs from Adam to SGD, ND-Adam (Zhang et al., 2017) which utilizes vector learning rate and normalization to control direction and stepsize, and
Revised: Instances of this type of algorithms include SWATS (Keskar & Socher, 2017) which automatically switchs from Adam to SGD, ND-Adam (Zhang et al., 2017) which utilizes vector learning rate and normalization to control direction and stepsize, and

Original: AMSGrad (Reddi et al., 2018) which maintains a monotone increasing second moment.
Revised: AMSGrad (Reddi et al., 2018) which maintains a monotone increasing second moment.

Original: Unfortunately, these methods only slightly bridge the generalization gap between SGD and Adam, but does not attain as good generalization performance as SGD, needless to say the state-of-the-art performance on test set.
Revised: Unfortunately, these methods only slightly bridge the generalization gap between SGD and Adam, but does not attain as good generalization performance as SGD, needless to say the state-of-the-art performance on test set.

Original: Accordingly, these algorithms are rarely used to train deep networks in practice.
Revised: Accordingly, these algorithms are rarely used to train deep networks in practice.

Original: To combine the merits of Adam and SGD, i.e. fast convergence speed in Adam and excellent generalization in SGD, we proposed a Dimension-Reduced Adaptive Gradient Method (DRAG for short) which minimizes the loss from several descent directions to trade-off the whole space search in
Revised: To combine the merits of Adam and SGD, i.e. fast convergence speed in Adam and excellent generalization in SGD, we proposed a Dimension-Reduced Adaptive Gradient Method (DRAG for short) which minimizes the loss from several descent directions to trade-off the whole space search in

Original: Adam and the minimization along a single gradient direction in SGD.
Revised: Adam and the minimization along a single gradient direction in SGD.

Original: For Adam, adjusting stepsizes for each gradient coordinate actually transforms the n -dimensional gradient into n independent directions to optimize, in which each direction inherits one coordinate element from the gradient and sets the remaining coordinate positions as zeros.
Revised: For Adam, adjusting stepsizes for each gradient coordinate actually transforms the n -dimensional gradient into n independent directions to optimize, in which each direction inherits one coordinate element from the gradient and sets the remaining coordinate positions as zeros.

Original: In contrast, SGD only uses a single learning rate for all gradient coordinates and minimizes the loss along one descent direction.
Revised: In contrast, SGD only uses a single learning rate for all gradient coordinates and minimizes the loss along one descent direction.

Original: Though the adaptive learning rate for each coordinate shows faster convergence speed than a single learning rate for all coordinates, as shown in many works (Wilson et al., 2017; Zhou et al., 2018), it also leads to the inferior generalization in Adam, since minimizing n independent directions means searching the whole parameter space and could results in overfitting.
Revised: Though the adaptive learning rate for each coordinate shows faster convergence speed than a single learning rate for all coordinates, as shown in many works (Wilson et al., 2017; Zhou et al., 2018), it also leads to the inferior generalization in Adam, since minimizing n independent directions means searching the whole parameter space and could results in overfitting.

Original: So it is natural to trade-off the number of descent directions.
Revised: So it is natural to trade-off the number of descent directions.

Original: To this end, motivated by DRSOM (Zhang et al., 2022), we update the parameters along the gradient direction and momentum direction through a trust-region-like approach, which greatly reduces the high adaptivity of Adam while adding flexibility to SGD.
Revised: To this end, motivated by DRSOM (Zhang et al., 2022), we update the parameters along the gradient direction and momentum direction through a trust-region-like approach, which greatly reduces the high adaptivity of Adam while adding flexibility to SGD.

Original: At each iteration, DRAG searches for the optimal update along the gradient and the momentum which are widely used in accelerated algorithm (Polyak, 1964; Nesterov, 2003) by solving a two-dimensional trust-region subproblem to find the best stepsizes for these two directions.
Revised: At each iteration, DRAG searches for the optimal update along the gradient and the momentum which are widely used in accelerated algorithm (Polyak, 1964; Nesterov, 2003) by solving a two-dimensional trust-region subproblem to find the best stepsizes for these two directions.

Original: For the trust-region subproblem, we use a quadratic approximation to estimate the loss with the Hessian matrix estimated by the second moment in Adam which is a diagonal matrix and can greatly reduce the computational cost.
Revised: For the trust-region subproblem, we use a quadratic approximation to estimate the loss with the Hessian matrix estimated by the second moment in Adam which is a diagonal matrix and can greatly reduce the computational cost.

Original: Moreover, we heuristically design a simple and effective trust-region radius for the subproblem.
Revised: Moreover, we heuristically design a simple and effective trust-region radius for the subproblem.

Original: Despite the delicate design of our algorithm, we also theoretically prove that on non-convex problems, our DRAG can converge and enjoys a stochastic gradient complexity of O ( ϵ − 4 ) to find an ϵ -approximate first-order stationary point.
Revised: Despite the delicate design of our algorithm, we also theoretically prove that on non-convex problems, our DRAG can converge and enjoys a stochastic gradient complexity of O ( ϵ − 4 ) to find an ϵ -approximate first-order stationary point.

Original: To summarize, our contributions are as follows:
Revised: To summarize, our contributions are as follows:

Original: • We proposed the DRAG algorithm to optimize the loss from several descent directions for balancing the whole space search in Adam and the optimization along a single gradient direction in SGD.
Revised: • We proposed the DRAG algorithm to optimize the loss from several descent directions for balancing the whole space search in Adam and the optimization along a single gradient direction in SGD.

Original: Moreover, we formulate the optimum stepsize search for these descent directions into a low-dimensional trust region problem whose computational cost is negligible when compared with the vanilla cost in adaptive gradient algorithms.
Revised: Moreover, we formulate the optimum stepsize search for these descent directions into a low-dimensional trust region problem whose computational cost is negligible when compared with the vanilla cost in adaptive gradient algorithms.

Original: • We theoretically prove that to find an ϵ -approximate stationary point on non-convex stochastic problems, DRAG has the stochastic gradient complexity of O ( ϵ − 4 ) which matches the lower bound Ω( ϵ − 4 ) in (Arjevani et al., 2022) under the same non-convex optimization setting.
Revised: • We theoretically prove that to find an ϵ -approximate stationary point on non-convex stochastic problems, DRAG has the stochastic gradient complexity of O ( ϵ − 4 ) which matches the lower bound Ω( ϵ − 4 ) in (Arjevani et al., 2022) under the same non-convex optimization setting.

Original: • Experimental results show that on several representative benchmarks, our DRAG method can achieve faster convergence speed than SGD, and also state-of-the-art generalization performance.
Revised: • Experimental results show that on several representative benchmarks, our DRAG method can achieve faster convergence speed than SGD, and also state-of-the-art generalization performance.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Adaptive gradient methods, e.g. Adam (Kingma & Ba, 2014), Adagrad (Duchi et al., 2011), and
Revised: Adaptive gradient methods, e.g. Adam (Kingma & Ba, 2014), Adagrad (Duchi et al., 2011), and

Original: RMSprop (Tieleman et al., 2012), adopt different stepsizes for different gradient coordinates so as to boost training process.
Revised: RMSprop (Tieleman et al., 2012), adopt different stepsizes for different gradient coordinates so as to boost training process.

Original: Although Adam and its variants are used widely for training deep neural networks, their poor generalization performance makes SGD still dominant in some areas, such as training CNNs (He et al., 2016).
Revised: Although Adam and its variants are used widely for training deep neural networks, their poor generalization performance makes SGD still dominant in some areas, such as training CNNs (He et al., 2016).

Original: Many researchers tried to improve the generalization capacity of
Revised: Many researchers tried to improve the generalization capacity of

Original: Adam, including SWATS (Keskar & Socher, 2017) which conducts an automatic switch from Adam to SGD training strategy, ND-Adam (Zhang et al., 2017) which controls the stepsizes and update direction in a more precise way, AMSGrad (Reddi et al., 2019) which ensures monotone increasing second moment, and Padam (Chen et al., 2021) which introduces a partially adaptive parameter to control the adaptivity of stepsizes.
Revised: Adam, including SWATS (Keskar & Socher, 2017) which conducts an automatic switch from Adam to SGD training strategy, ND-Adam (Zhang et al., 2017) which controls the stepsizes and update direction in a more precise way, AMSGrad (Reddi et al., 2019) which ensures monotone increasing second moment, and Padam (Chen et al., 2021) which introduces a partially adaptive parameter to control the adaptivity of stepsizes.

Original: Our solution to improve the generalization performance is to confine the update of parameters to a two-dimensional subspace of the parameter space.
Revised: Our solution to improve the generalization performance is to confine the update of parameters to a two-dimensional subspace of the parameter space.

Original: Specifically, we solve a trust-region subproblem to determine the optimal stepsizes along the gradient direction and momentum direction at each iteration.
Revised: Specifically, we solve a trust-region subproblem to determine the optimal stepsizes along the gradient direction and momentum direction at each iteration.

Original: The idea of utilizing gradient and momentum direction to update the variable traces back to Polyak’s heavy ball method (Polyak, 1964)
Revised: The idea of utilizing gradient and momentum direction to update the variable traces back to Polyak’s heavy ball method (Polyak, 1964)

Original: [Equation]
Revised: [Equation]

Original: SGD (with momentum) (Sutskever et al., 2013) can also be written in this form after replacing the gradient direction ∇ f ( x t − 1 ) with stochastic gradient direction g t − 1 .
Revised: SGD (with momentum) (Sutskever et al., 2013) can also be written in this form after replacing the gradient direction ∇ f ( x t − 1 ) with stochastic gradient direction g t − 1 .

Original: Unlike SGD which adopts constant stepsizes, DRAG adopts the stepsizes as the solution of a two-dimensional quadratic model with a spherical constraint, it adaptively learns the optimal stepsize for each direction.
Revised: Unlike SGD which adopts constant stepsizes, DRAG adopts the stepsizes as the solution of a two-dimensional quadratic model with a spherical constraint, it adaptively learns the optimal stepsize for each direction.

Original: Trust-region method has been widely used in optimization, and some works tries to use it for machine learning (Martens et al., 2010; Dudar et al., 2017; Erway et al., 2020).
Revised: Trust-region method has been widely used in optimization, and some works tries to use it for machine learning (Martens et al., 2010; Dudar et al., 2017; Erway et al., 2020).

Original: The closest related work is the Dimension Reduced Second Order Method (DRSOM) proposed by Zhang et al.
Revised: The closest related work is the Dimension Reduced Second Order Method (DRSOM) proposed by Zhang et al.

Original: It uses a finite difference method to approximate the Hessian-vector product arisen in the trust-region subproblem and achieves higher convergence rate than first order methods.
Revised: It uses a finite difference method to approximate the Hessian-vector product arisen in the trust-region subproblem and achieves higher convergence rate than first order methods.

Original: However, difference method is inaccurate in high-dimensional problem, and it introduces huge extra computational overhead.
Revised: However, difference method is inaccurate in high-dimensional problem, and it introduces huge extra computational overhead.

Original: Our method DRAG, on the other hand, approximates the Hessian directly by the second moment of stochastic gradient directly.
Revised: Our method DRAG, on the other hand, approximates the Hessian directly by the second moment of stochastic gradient directly.

Original: It reduces extra forward-backward loops to calculate the Hessian-vector product in DRSOM and captures curvature information more accurately.
Revised: It reduces extra forward-backward loops to calculate the Hessian-vector product in DRSOM and captures curvature information more accurately.

Original: Moreover, DRAG adopts a heuristic trust-region radius, making the algorithm compatible with dominant deep learning training pipeline without introducing extra hyperparameters.
Revised: Moreover, DRAG adopts a heuristic trust-region radius, making the algorithm compatible with dominant deep learning training pipeline without introducing extra hyperparameters.

Original:  M ETHOD
Revised:  M ETHOD

Original: Training neural networks can be seen as solving the following non-convex optimization problem
Revised: Training neural networks can be seen as solving the following non-convex optimization problem

Original: [Equation]
Revised: [Equation]

Original: where f is the loss function and x ∈ R n is the variable.
Revised: where f is the loss function and x ∈ R n is the variable.

Original: Among all optimizers, Adam (Kingma & Ba, 2014) is one of the most popular algorithm to solve problem (1).
Revised: Among all optimizers, Adam (Kingma & Ba, 2014) is one of the most popular algorithm to solve problem (1).

Original: At each training iteration,
Revised: At each training iteration,

Original: Adam maintains an exponential moving average (EMA) of first and second moments of stochastic gradient v t and u t as
Revised: Adam maintains an exponential moving average (EMA) of first and second moments of stochastic gradient v t and u t as

Original: [Equation]
Revised: [Equation]

Original: where β 1 , β 2 ∈ [0 , 1] are constant and g t − 1 :
Revised: where β 1 , β 2 ∈ [0 , 1] are constant and g t − 1 :

Original: =
Revised: =

Original: ˜ ∇ f
Revised: ˜ ∇ f

Original: ( x t − 1 ) is the stochastic gradient.
Revised: ( x t − 1 ) is the stochastic gradient.

Original: It adaptively scales the learning rates for each gradient coordinate, and actually minimizes the loss function along n descent directions
Revised: It adaptively scales the learning rates for each gradient coordinate, and actually minimizes the loss function along n descent directions

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: Specifically, Adam adopts a stepsize of [Equation] for the i -th gradient descent direction ˆ v t,i e i .
Revised: Specifically, Adam adopts a stepsize of η √ ˆ u t,i + ν for the i -th gradient descent direction ˆ v t,i e i .

Original: While adaptive stepsize boosts the convergence of Adam, it weakens the generalization performance due to noise and overfitting.
Revised: While adaptive stepsize boosts the convergence of Adam, it weakens the generalization performance due to noise and overfitting.

Original: 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.
Revised: 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.

Original: 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 } (Wilson et al., 2017; Zhang et al., 2017), while SGD do.
Revised: 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 } (Wilson et al., 2017; Zhang et al., 2017), while SGD do.

Original: Actually, Wilson et al.
Revised: Actually, Wilson et al.

Original: (2017) 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.
Revised: (2017) 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.

Original: To overcome the issue just mentioned, our DRAG algorithm optimizes the loss function in (1) from the gradient direction and the momentum direction.
Revised: To overcome the issue just mentioned, our DRAG algorithm optimizes the loss function in (1) from the gradient direction and the momentum direction.

Original: It maintains flexibility in the update direction while inheriting the generalization capacity of SGD.
Revised: It maintains flexibility in the update direction while inheriting the generalization capacity of SGD.

Original: At each step, it searches for the optimal stepsizes along these two directions by solving a two-dimensional trust-region subproblem.
Revised: At each step, it searches for the optimal stepsizes along these two directions by solving a two-dimensional trust-region subproblem.

Original: Therefore, from the optimization perspective, it conducts the optimal update within the two-dimensional subspace spanned by gradient direction and momentum direction.
Revised: Therefore, from the optimization perspective, it conducts the optimal update within the two-dimensional subspace spanned by gradient direction and momentum direction.

Original: Moreover, while DRAG adopts the trust-region framework, it is compatible with the dominant deep learning training pipeline without introducing extra hyperparameters.
Revised: Moreover, while DRAG adopts the trust-region framework, it is compatible with the dominant deep learning training pipeline without introducing extra hyperparameters.

Original: 3. D ETAILS OF THE ALGORITHM
Revised: 3. D ETAILS OF THE ALGORITHM

Original: Details of our algorithm are described in Algorithm 1.
Revised: Details of our algorithm are described in Algorithm 1.

Original: At each training epoch, 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.
Revised: At each training epoch, 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.

Original: Then, we introduce the bias-corrected second moment ˆ u t to approximate the Hessian.
Revised: Then, we introduce the bias-corrected second moment ˆ u t to approximate the Hessian.

Original: In this way, DRAG constructs the trust-region subproblem in line 9 of Algorithm 1.
Revised: In this way, DRAG constructs the trust-region subproblem in line 9 of Algorithm 1.

Original: 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
Revised: 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

Original: d t − 1 , making the computational overhead negligible.
Revised: d t − 1 , making the computational overhead negligible.

Original: Here we intuitively set the trust-region radius as η ∥ ˆ v t ∥ , and the benefits of this setting is described in
Revised: Here we intuitively set the trust-region radius as η ∥ ˆ v t ∥ , and the benefits of this setting is described in

Original: Section 3.2.
Revised: Section 3.2.

Original: After calculating the solution α 1 t and α 2 t of the subproblem, we get an optimal update p = − α 1 t
Revised: After calculating the solution α 1 t and α 2 t of the subproblem, we get an optimal update p = − α 1 t

Original: ˆ v t + α 2 t d t − 1 in the two-dimensional subspace.
Revised: ˆ v t + α 2 t d t − 1 in the two-dimensional subspace.

Original: Finally, we follow (Loshchilov & Hutter, 2018) and conduct a decoupled weight decay step.
Revised: Finally, we follow (Loshchilov & Hutter, 2018) and conduct a decoupled weight decay step.

Original: This is the overall framework of our DRAG.
Revised: This is the overall framework of our DRAG.

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 8: d t − 1 = x t − 1 − x t − 2 if t ≥ 2 else d t − 1 = 0 .
Revised: 8: d t − 1 = x t − 1 − x t − 2 if t ≥ 2 else d t − 1 = 0 .

Original: 9: ( α 1 t , α 2 t ) = argmin p
Revised: 9: ( α 1 t , α 2 t ) = argmin p

Original: [Equation]
Revised: [Equation]

Original: ⟨ ˆ v t , p ⟩ +
Revised: ⟨ ˆ v t , p ⟩ +

Original: [Equation]
Revised: [Equation]

Original: p , H t p
Revised: p , H t p

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: t ∥ , p
Revised: t ∥ , p

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: t
Revised: t

Original: [Equation]
Revised: [Equation]

Original: t
Revised: t

Original: [Equation]
Revised: [Equation]

Original: 1 (cid:9) .
Revised: 1 (cid:9) .

Original: 10: x t = x t − 1 − α 1 t ˆ v t + α 2 t d t −11: x t = x t − ηγ x t − 1 (Conduct weight decay) 12: end for 13: Output: x 1 , · · · , x m
Revised: 10: x t = x t − 1 − α 1 t ˆ v t + α 2 t d t −11: x t = x t − ηγ x t − 1 (Conduct weight decay) 12: end for 13: Output: x 1 , · · · , x m

Original: The only extra computational overhead of DRAG compared with Adam is solving the twodimensional trust-region subproblem in line 9 of Algorithm 1.
Revised: The only extra computational overhead of DRAG compared with Adam is solving the twodimensional trust-region subproblem in line 9 of Algorithm 1.

Original: The trust-region subproblem can be formally formulated as follows:
Revised: The trust-region subproblem can be formally formulated as follows:

Original: [Equation]
Revised: [Equation]

Original: where H t = diag( √ ˆ u t + ν ) as defined in Algorithm 1.
Revised: where H t = diag( √ ˆ u t + ν ) as defined in Algorithm 1.

Original: This two-dimensional subproblem can be solved efficiently by using its global minimal condition.
Revised: This two-dimensional subproblem can be solved efficiently by using its global minimal condition.

Original: In Appendix A, we transform this subproblem into a standard trust-region subproblem, and then an ϵ -global primal-dual solution satisfying
Revised: In Appendix A, we transform this subproblem into a standard trust-region subproblem, and then an ϵ -global primal-dual solution satisfying

Original: KKT condition can be found in O (log log( 1 ϵ )) time (Luenberger et al., 1984).
Revised: KKT condition can be found in O (log log( 1 ϵ )) time (Luenberger et al., 1984).

Original: See more details in
Revised: See more details in

Original: Appendix A.
Revised: Appendix A.

Original: If we set p in line 9 of Algorithm 1 as
Revised: If we set p in line 9 of Algorithm 1 as

Original: [Equation]
Revised: [Equation]

Original: our algorithm DRAG can be generalized to solve the subproblem with any k search directions.
Revised: our algorithm DRAG can be generalized to solve the subproblem with any k search directions.

Original: Although according to experimental results, one search direction or multiple search directions usually perform worse than two search directions (DRAG), they can give us some intuitions on the flexibility of our method and the optimality of its update in the subspace.
Revised: Although according to experimental results, one search direction or multiple search directions usually perform worse than two search directions (DRAG), they can give us some intuitions on the flexibility of our method and the optimality of its update in the subspace.

Original: We omit ν in the discussions below for the simplicity of notation.
Revised: We omit ν in the discussions below for the simplicity of notation.

Original: One-dimensional subspace Suppose we update the variable in the one-dimensional subspace spanned by gradient direction
Revised: One-dimensional subspace Suppose we update the variable in the one-dimensional subspace spanned by gradient direction

Original: [Equation]
Revised: [Equation]

Original: where α 1 t is calculated by solving the trust-region subproblem
Revised: where α 1 t is calculated by solving the trust-region subproblem

Original: [Equation]
Revised: [Equation]

Original: In this case, the subproblem has an explicit solution α 1 t = min { η, ⟨ ˆ v t , ˆ v t ⟩ ⟨ ˆ v t , diag( √ ˆ u t )ˆ v t ⟩ } .
Revised: In this case, the subproblem has an explicit solution α 1 t = min { η, ⟨ ˆ v t , ˆ v t ⟩ ⟨ ˆ v t , diag( √ ˆ u t )ˆ v t ⟩ } .

Original: As we can see from the solution, unlike SGD which adopts the learning rate set externally as the stepsize, our algorithm update the parameter with an adaptive stepsize within the learning rate.
Revised: As we can see from the solution, unlike SGD which adopts the learning rate set externally as the stepsize, our algorithm update the parameter with an adaptive stepsize within the learning rate.

Original: Furthermore, this adaptive stepsize is optimal according to the quadratic approximation of the loss function.
Revised: Furthermore, this adaptive stepsize is optimal according to the quadratic approximation of the loss function.

Original: Full-dimensional parameter space Suppose we update the variable along n directions in the whole parameter space
Revised: Full-dimensional parameter space Suppose we update the variable along n directions in the whole parameter space

Original: [Equation]
Revised: [Equation]

Original: Then trust-region subproblem becomes
Revised: Then trust-region subproblem becomes

Original: [Equation]
Revised: [Equation]

Original: Under this scenario, the subproblem has solution as
Revised: Under this scenario, the subproblem has solution as

Original: [Equation]
Revised: [Equation]

Original: where λ ≥ 0 satisfies λ ( ∥ p ∥ − η ∥ ˆ v t ∥ ) = 0 .
Revised: where λ ≥ 0 satisfies λ ( ∥ p ∥ − η ∥ ˆ v t ∥ ) = 0 .

Original: When ∥ ˆ v t √ u t ∥ ≤ η ∥ ˆ v t ∥ , λ = 0 and the solution is the same as Adam’s update direction (2).
Revised: When ∥ ˆ v t √ u t ∥ ≤ η ∥ ˆ v t ∥ , λ = 0 and the solution is the same as Adam’s update direction (2).

Original: From the form of solution, we can see all gradient coordinates have adaptive stepsizes, which means the method optimizes the loss function along n directions in the whole parameter space.
Revised: From the form of solution, we can see all gradient coordinates have adaptive stepsizes, which means the method optimizes the loss function along n directions in the whole parameter space.

Original: Also, this update is optimal with respect to the quadratic approximation of the loss function.
Revised: Also, this update is optimal with respect to the quadratic approximation of the loss function.

Original: 3. B ENEFITS OF OUR ALGORITHM
Revised: 3. B ENEFITS OF OUR ALGORITHM

Original: 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 .
Revised: 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 .

Original: This update direction choice acts as a trade-off between the whole space search of Adam and one direction search of SGD.
Revised: This update direction choice acts as a trade-off between the whole space search of Adam and one direction search of SGD.

Original: Specifically, Adam adjusts stepsizes for each gradient coordinate as x t = x t − 1 − (cid:80) ni
Revised: Specifically, Adam adjusts stepsizes for each gradient coordinate as x t = x t − 1 − (cid:80) ni

Original: = η
Revised: = η

Original: √ ˆ u t,i + ν (ˆ v t,i e i ) , while SGD uniformly scales each coordinate of the gradient as x t = x t − 1 − η ˆ v t .
Revised: √ ˆ u t,i + ν (ˆ v t,i e i ) , while SGD uniformly scales each coordinate of the gradient as x t = x t − 1 − η ˆ v t .

Original: DRAG, on the other hand, search along two important directions as x t = x
Revised: DRAG, on the other hand, search along two important directions as x t = x

Original: t − 1 − α 1 t ˆ v t + α 2 t d t .
Revised: t − 1 − α 1 t ˆ v t + α 2 t d t .

Original: This choice maintains the flexibility of update direction while alleviating overfitting and excessive noises.
Revised: This choice maintains the flexibility of update direction while alleviating overfitting and excessive noises.

Original: Moreover, the update of DRAG lies in the subspace span { ˆ v t , d t − 1 } ∈ span { g 0 , · · · , g t − 1 } .
Revised: Moreover, the update of DRAG lies in the subspace span { ˆ v t , d t − 1 } ∈ span { g 0 , · · · , g t − 1 } .

Original: This means that the parameter update direction is always a combination of stochastic gradients.
Revised: This means that the parameter update direction is always a combination of stochastic gradients.

Original: According to Wilson et al.
Revised: According to Wilson et al.

Original: (2017), this property makes DRAG always converge to the max-margin solution of the binary classification problem, which has the best generalization capacity.
Revised: (2017), this property makes DRAG always converge to the max-margin solution of the binary classification problem, which has the best generalization capacity.

Original: This helps to explain
Revised: This helps to explain

Original: DRAG’s excellent generalization performance in practice.
Revised: DRAG’s excellent generalization performance in practice.

Original: Optimal stepsizes DRAG solves the dimension-reduced subproblem at each training epoch and finds the best update along the gradient direction and momentum direction.
Revised: Optimal stepsizes DRAG solves the dimension-reduced subproblem at each training epoch and finds the best update along the gradient direction and momentum direction.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: Heuristic trust-region radius We set the trust-region radius for the subproblem as η ∥ ˆ v t ∥ .
Revised: Heuristic trust-region radius We set the trust-region radius for the subproblem as η ∥ ˆ v t ∥ .

Original: The intuition is that when gradient is large, we hope our algorithm can make a larger step to minimize the loss function significantly.
Revised: The intuition is that when gradient is large, we hope our algorithm can make a larger step to minimize the loss function significantly.

Original: While when gradient is small, we hope our method to be stable and don’t change the parameters too much.
Revised: While when gradient is small, we hope our method to be stable and don’t change the parameters too much.

Original: 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.
Revised: 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.

Original: Other deep learning optimizers that adopt trust-region like framework, such as the Hessian-free optimization method (Martens et al., 2010), L-SSR1-TR (Erway et al., 2020), DRSOM (Zhang et al., 2022) introduced extra hyperparameters and may incur high extra computational overhead.
Revised: Other deep learning optimizers that adopt trust-region like framework, such as the Hessian-free optimization method (Martens et al., 2010), L-SSR1-TR (Erway et al., 2020), DRSOM (Zhang et al., 2022) introduced extra hyperparameters and may incur high extra computational overhead.

Original: From our knowledge, this is the first time that a trust-region like method is well-compatible with dominant deep learning training setting without extra hyperparameters.
Revised: From our knowledge, this is the first time that a trust-region like method is well-compatible with dominant deep learning training setting without extra hyperparameters.

Original:  C ONVERGENCE ANALYSIS IN NON - CONVEX STOCHASTIC OPTIMIZATION
Revised:  C ONVERGENCE ANALYSIS IN NON - CONVEX STOCHASTIC OPTIMIZATION

Original: For the analysis of stochastic non-convex algorithm, we follow the works Zhuang et al.
Revised: For the analysis of stochastic non-convex algorithm, we follow the works Zhuang et al.

Original: Guo et al.
Revised: Guo et al.

Original: and make the following necessary definitions and also mild assumptions.
Revised: and make the following necessary definitions and also mild assumptions.

Original: Definition 1.
Revised: Definition 1.

Original: For a differentiable function f , x is said to be an ϵ -approximate first-order stationary point if it satisfies ∥∇ f ( x ) ∥ ≤ ϵ .
Revised: For a differentiable function f , x is said to be an ϵ -approximate first-order stationary point if it satisfies ∥∇ f ( x ) ∥ ≤ ϵ .

Original: Definition 2.
Revised: Definition 2.

Original: For a differentiable funtion f ( x ) , it is called L -Lipschitz smooth if it statisfies ∥∇ f ( x ) − ∇ f ( y ) ∥
Revised: For a differentiable funtion f ( x ) , it is called L -Lipschitz smooth if it statisfies ∥∇ f ( x ) − ∇ f ( y ) ∥

Original: ≤ L ∥ x − y ∥ for a constant L > 0 and any x , y in domain of f .
Revised: ≤ L ∥ x − y ∥ for a constant L > 0 and any x , y in domain of f .

Original: Based on these definitions, we have the following assumption.
Revised: Based on these definitions, we have the following assumption.

Original: Assumption 1.
Revised: Assumption 1.

Original: For non-convex problem min x ∈ R n f ( x ) , we assume the loss f ( x ) satisfies
Revised: For non-convex problem min x ∈ R n f ( x ) , we assume the loss f ( x ) satisfies

Original: • f is L -Lipschitz smooth.
Revised: • f is L -Lipschitz smooth.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: Theorem 1.
Revised: Theorem 1.

Original: Suppose Assumption 1 holds.
Revised: Suppose Assumption 1 holds.

Original: Let β t = β and η t =
Revised: Let β t = β and η t =

Original: η for all t .
Revised: η for all t .

Original: Assume there exist constants α, G > 0 , such that α ≤ min t α 1 t and α 1 t ≤ ηG , | α 2 t | ≤ ηG .
Revised: Assume there exist constants α, G > 0 , such that α ≤ min t α 1 t and α 1 t ≤ ηG , | α 2 t | ≤ ηG .

Original: In addition, η ≤ min
Revised: In addition, η ≤ min

Original: (cid:26)2 LG , (cid:16) (1 − β ) 2 α 8 GL 2 (cid:17) 13 , (cid:16) α 2 96 G 2 (cid:17) 14 , (cid:0) α 48 LG 2 (cid:1) 14 , (cid:0) α 192 L 2 G 3 (cid:1)
Revised: (cid:26)2 LG , (cid:16) (1 − β ) 2 α 8 GL 2 (cid:17) 13 , (cid:16) α 2 96 G 2 (cid:17) 14 , (cid:0) α 48 LG 2 (cid:1) 14 , (cid:0) α 192 L 2 G 3 (cid:1)

Original: (cid:27) .
Revised: (cid:27) .

Original: Then, if 1 − β ≤ ϵ 2 3 C 2 σ 2 and T ≥ max (cid:110) 3 C 1 αϵ 2 , 3 C
Revised: Then, if 1 − β ≤ ϵ 2 3 C 2 σ 2 and T ≥ max (cid:110) 3 C 1 αϵ 2 , 3 C

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: β ) ϵ 2 (cid:111) , DRAG can achieve
Revised: β ) ϵ 2 (cid:111) , DRAG can achieve

Original: [Equation]
Revised: [Equation]

Original: where C 1 = 4 ( f ( x 0 ) − f ( x ∗ )) , C 2 = 4 ηGα and C 3 = 2 ηG E [ ∥∇ f ( x 0 ) − (1 − β 0 ) g 

Revised: where C 1 = 4 ( f ( x 0 ) − f ( x ∗ )) , C 2 = 4 ηGα and C 3 = 2 ηG E [ ∥∇ f ( x 0 ) − (1 − β 0 ) g

Original: ∥ 2 ] α .
Revised: ∥ 2 ] α .

Original: Remark 1.
Revised: Remark 1.

Original: 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),
Revised: 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),

Original: DRAG can converge to an ϵ -approximate first-order stationary point of the non-convex stochastic problem with stochastic gradient complexity O
Revised: DRAG can converge to an ϵ -approximate first-order stationary point of the non-convex stochastic problem with stochastic gradient complexity O

Original: ( ϵ − 4 ) .
Revised: ( ϵ − 4 ) .

Original: Note that the assumptions for α 1 t and α 2 t are satisfied naturally with the design of DRAG, see details in Appendix B.
Revised: Note that the assumptions for α 1 t and α 2 t are satisfied naturally with the design of DRAG, see details in Appendix B.

Original: The complexity of DRAG is of the same order as the lower bound provided by Arjevani et al.
Revised: The complexity of DRAG is of the same order as the lower bound provided by Arjevani et al.

Original: A similar complexity has also been obtained in, for example, LAMB (You et al., 2019), Adam-family (Guo et al., 2021).
Revised: A similar complexity has also been obtained in, for example, LAMB (You et al., 2019), Adam-family (Guo et al., 2021).

Original: 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 (You et al., 2017) and LAMB (You et al., 2019).
Revised: 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 (You et al., 2017) and LAMB (You et al., 2019).

Original: In addition, some previous works (Luo et al., 2018; Zaheer et al., 2018; Liu et al., 2019; Shi et al., 2020) require the momentum parameter β t to be very close or decreasing to zero.
Revised: In addition, some previous works (Luo et al., 2018; Zaheer et al., 2018; Liu et al., 2019; Shi et al., 2020) require the momentum parameter β t to be very close or decreasing to zero.

Original: In contrast, DRAG requires β t to be close to one, which is more consistent with the practice.
Revised: In contrast, DRAG requires β t to be close to one, which is more consistent with the practice.

Original: Theorem 2.
Revised: Theorem 2.

Original: Suppose Assumption 1 holds.
Revised: Suppose Assumption 1 holds.

Original: Assume there exist constants δ, G > 0 , such that 0 < δ ≤ α 1 t η t ≤ G , | α 2 t | η t ≤ G .
Revised: Assume there exist constants δ, G > 0 , such that 0 < δ ≤ α 1 t η t ≤ G , | α 2 t | η t ≤ G .

Original: Set η t =
Revised: Set η t =

Original: c η √ t +2 , 1 − β t =
Revised: c η √ t +2 , 1 − β t =

Original: Cc η √ t +1 , for any c η and C satisfying C ≥ L (cid:113) 8 Gδ , and c η ≤ (cid:26)
Revised: Cc η √ t +1 , for any c η and C satisfying C ≥ L (cid:113) 8 Gδ , and c η ≤ (cid:26)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: (cid:27) .
Revised: (cid:27) .

Original: Then there exist two constant C 1 and Cwhich are independent with T , such that
Revised: Then there exist two constant C 1 and Cwhich are independent with T , such that

Original: [Equation]
Revised: [Equation]

Original: Given a tolerance ϵ > 0 , if T ≥ ˜ O ( 1 ϵ 4 ) , we have
Revised: Given a tolerance ϵ > 0 , if T ≥ ˜ O ( 1 ϵ 4 ) , we have

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: Remark 2.
Revised: Remark 2.

Original: Theorem 2 with its proof in Appendix D 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 ) .
Revised: Theorem 2 with its proof in Appendix D 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 ) .

Original: Similar sub-linear convergence rates are also established by Zou et al.
Revised: Similar sub-linear convergence rates are also established by Zou et al.

Original: for Adam and Guo et al.
Revised: for Adam and Guo et al.

Original: (2021) for Adam-type optimizers.
Revised: (2021) for Adam-type optimizers.

Original: While Zou et al.
Revised: While Zou et al.

Original: (2019) has restrictions on the second moment momentum parameter β 2 .
Revised: (2019) has restrictions on the second moment momentum parameter β 2 .

Original: In Theorem 2, we only need β t (corresponds β 1 in Algorithm 1) to increase to one.
Revised: In Theorem 2, we only need β t (corresponds β 1 in Algorithm 1) to increase to one.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We conduct experiments on several representative benchmarks, including VGG (Simonyan & Zisserman, 2014), ResNet (He et al., 2016), DenseNet (Huang et al., 2017) on CIFAR10, CIFAR dataset (Krizhevsky et al., 2009), ResNet18 on ImageNet (Deng et al., 2009), and LSTM (Hochreiter & Schmidhuber, 1997) on the Penn Treebank dataset (Marcinkiewicz, 1994).
Revised: We conduct experiments on several representative benchmarks, including VGG (Simonyan & Zisserman, 2014), ResNet (He et al., 2016), DenseNet (Huang et al., 2017) on CIFAR10, CIFAR dataset (Krizhevsky et al., 2009), ResNet18 on ImageNet (Deng et al., 2009), and LSTM (Hochreiter & Schmidhuber, 1997) on the Penn Treebank dataset (Marcinkiewicz, 1994).

Original: We compare our algorithm DRAG with some popular deep learning optimizers, including SGD (Robbins & Monro, 1951), Adam (Kingma & Ba, 2014), AdamW (Loshchilov & Hutter, 2018), AdaBound (Luo et al., 2018), AdaBelief (Zhuang et al., 2020), RAdam (Liu et al., 2019), Yogi (Zaheer et al., 2018), and
Revised: We compare our algorithm DRAG with some popular deep learning optimizers, including SGD (Robbins & Monro, 1951), Adam (Kingma & Ba, 2014), AdamW (Loshchilov & Hutter, 2018), AdaBound (Luo et al., 2018), AdaBelief (Zhuang et al., 2020), RAdam (Liu et al., 2019), Yogi (Zaheer et al., 2018), and

Original: Padam (Chen et al., 2021).
Revised: Padam (Chen et al., 2021).

Original: Experimental results show that DRAG has faster convergence speed compared with SGD and it achieves state-of-the-art generalization performance.
Revised: Experimental results show that DRAG has faster convergence speed compared with SGD and it achieves state-of-the-art generalization performance.

Original: 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.
Revised: 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.

Original: At the end of ablation study, we give some advice for practitioners to use DRAG.
Revised: At the end of ablation study, we give some advice for practitioners to use DRAG.

Original: 5. CNN S ON IMAGE CLASSIFICATION
Revised: 5. CNN S ON IMAGE CLASSIFICATION

Original: We conducted experiments for VGG16 with Batch Normalization, ResNet34, and DenseNet121 on
Revised: We conducted experiments for VGG16 with Batch Normalization, ResNet34, and DenseNet121 on

Original: CIFAR10 an CIFAR100 dataset.
Revised: CIFAR10 an CIFAR100 dataset.

Original: The experimental setting is borrowed from AdaBelief (Zhuang et al., 2020) and we also use their default setting for all the hyperparameters.
Revised: The experimental setting is borrowed from AdaBelief (Zhuang et al., 2020) and we also use their default setting for all the hyperparameters.

Original: 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
Revised: 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

Original: CIFAR10 and 0.0025 for CIFAR100.
Revised: CIFAR10 and 0.0025 for CIFAR100.

Original: Other hyperparameters of DRAG is the same as the default setting ( β 1 = 0 .
Revised: Other hyperparameters of DRAG is the same as the default setting ( β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: 999 , ϵ = 10 − 8 ).
Revised: 999 , ϵ = 10 − 8 ).

Original: As Figure 1 shows, DRAG has convergence speed comparable with adaptive gradient methods and it attains the best generalization performance.
Revised: As Figure 1 shows, DRAG has convergence speed comparable with adaptive gradient methods and it attains the best generalization performance.

Original: To be specific, DRAG obtains more than 0.5% generalization accuracy gain over AdaBelief (Zhuang et al., 2020) on most tasks.
Revised: To be specific, DRAG obtains more than 0.5% generalization accuracy gain over AdaBelief (Zhuang et al., 2020) on most tasks.

Original: The detailed test accuracy is summarized in Table 1.
Revised: The detailed test accuracy is summarized in Table 1.

Original: We also train ResNet18 on ImageNet under the official training setting in He et al.
Revised: We also train ResNet18 on ImageNet under the official training setting in He et al.

Original: and compare the top-1 test accuracy of DRAG with the best result of other optimizers in the literature.
Revised: and compare the top-1 test accuracy of DRAG with the best result of other optimizers in the literature.

Original: Experimental results show DRAG has the best generalization capacity.
Revised: Experimental results show DRAG has the best generalization capacity.

Original: Details are in Table 2.
Revised: Details are in Table 2.

Original: 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)
Revised: 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)

Original: 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.
Revised: 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.

Original: 5. LSTM S ON LANGUAGE MODELING
Revised: 5. LSTM S ON LANGUAGE MODELING

Original: We experimented with LSTM on the Penn Treebank dataset and record the perplexity (lower is better).
Revised: We experimented with LSTM on the Penn Treebank dataset and record the perplexity (lower is better).

Original: We follow the exact experimental setting in Adabelief (Zhuang et al., 2020) and use their
Revised: We follow the exact experimental setting in Adabelief (Zhuang et al., 2020) and use their

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: default hyperparameters except for SGD.
Revised: default hyperparameters except for SGD.

Original: For SGD, we use the same hyperparameters as DRAG to make a fair comparison between the two.
Revised: For SGD, we use the same hyperparameters as DRAG to make a fair comparison between the two.

Original: 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 .
Revised: 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 .

Original: SGD’s generalization performance in our setting is better than the results provided by Zhuang et al.
Revised: SGD’s generalization performance in our setting is better than the results provided by Zhuang et al.

Original: As show in
Revised: As show in

Original: Figure 2, for 1-layer, 2-layer, and 3-layer LSTM, DRAG’s convergence speed is faster than SGD and comparable to adaptive gradient methods.
Revised: Figure 2, for 1-layer, 2-layer, and 3-layer LSTM, DRAG’s convergence speed is faster than SGD and comparable to adaptive gradient methods.

Original: From Table 3, we can see that DRAG attains more than 0.5 less perplexity than other optimizers.
Revised: From Table 3, we can see that DRAG attains more than 0.5 less perplexity than other optimizers.

Original: The fast convergence speed may be attributed to the optimal update DRAG takes and the good generalization performance may be due to DRAG’s twodirection search.
Revised: The fast convergence speed may be attributed to the optimal update DRAG takes and the good generalization performance may be due to DRAG’s twodirection search.

Original: The gradient direction inherits SGD’s good generalization property and the extra momentum direction further improves its performance.
Revised: The gradient direction inherits SGD’s good generalization property and the extra momentum direction further improves its performance.

Original: 5. A BLATION STUDY
Revised: 5. A BLATION STUDY

Original: 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 3.1.
Revised: 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 3.1.

Original: As show in Table 4 in Appendix E, DRAG generalizes better than its one search direction and three search direction counterparts.
Revised: As show in Table 4 in Appendix E, DRAG generalizes better than its one search direction and three search direction counterparts.

Original: 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.
Revised: 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.

Original: Robustness to learning rate schedule DRAG is robust to different choices of learning rate schedule.
Revised: Robustness to learning rate schedule DRAG is robust to different choices of learning rate schedule.

Original: Except for letting the learning rate decay at epoch 150 as in Section 5.1, we also conduct experiments on decaying the learning rate at epoch 120 and adopting cosine annealing learning rate schedule.
Revised: Except for letting the learning rate decay at epoch 150 as in Section 5.1, we also conduct experiments on decaying the learning rate at epoch 120 and adopting cosine annealing learning rate schedule.

Original: The only change of hyperparameter setting from Section 5.1 is we increase the learning rate of DRAG from 0.1 to 0.12 in cosine annealing schedule.
Revised: The only change of hyperparameter setting from Section 5.1 is we increase the learning rate of DRAG from 0.1 to 0.12 in cosine annealing schedule.

Original: 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.
Revised: 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.

Original: We compared DRAG’s test performance with other optimizers with VGG on CIFAR10, details are presented in Table 5 in Appendix E, which shows that DRAG enjoys the best generalization performance for all the learning rate schedules.
Revised: We compared DRAG’s test performance with other optimizers with VGG on CIFAR10, details are presented in Table 5 in Appendix E, which shows that DRAG enjoys the best generalization performance for all the learning rate schedules.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: For practitioners, any task that can use SGD can use DRAG to achieve faster convergence and comparable generalization performance with negligible extra computational overhead.
Revised: For practitioners, any task that can use SGD can use DRAG to achieve faster convergence and comparable generalization performance with negligible extra computational overhead.

Original: The user only needs to set the learning rate the same as in SGD or slightly larger.
Revised: The user only needs to set the learning rate the same as in SGD or slightly larger.

Original: For a new task, if one values good generalization performance, one can always use DRAG instead of SGD to enjoy easier hyperparameter tunning.
Revised: For a new task, if one values good generalization performance, one can always use DRAG instead of SGD to enjoy easier hyperparameter tunning.

Original: DRAG is more robust than SGD when large learning rate is used.
Revised: DRAG is more robust than SGD when large learning rate is used.

Original: For instance, when training VGG16 on CIFAR10 dataset, setting the learning rate to 0.5 still allows
Revised: For instance, when training VGG16 on CIFAR10 dataset, setting the learning rate to 0.5 still allows

Original: DRAG to attain over 90 percent test accuracy, but SGD diverge and fail in the training process.
Revised: DRAG to attain over 90 percent test accuracy, but SGD diverge and fail in the training process.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper we propose the DRAG algorithm, which finds the optimal update of the parameters along gradient and momentum directions at each iteration.
Revised: In this paper we propose the DRAG algorithm, which finds the optimal update of the parameters along gradient and momentum directions at each iteration.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Theoretically we prove that DRAG has the same order of stochastic gradient complexity as the lower bound for non-convex stochastic optimization (Arjevani et al., 2022).
Revised: Theoretically we prove that DRAG has the same order of stochastic gradient complexity as the lower bound for non-convex stochastic optimization (Arjevani et al., 2022).

Original: Experimentally we show that DRAG has faster convergence speed compared with SGD and it attains state-of-the-art generalization performance.
Revised: Experimentally we show that DRAG has faster convergence speed compared with SGD and it attains state-of-the-art generalization performance.

Original: Our algorithm can be further generalized to any number of search directions and any choice of Hessian approximation.
Revised: Our algorithm can be further generalized to any number of search directions and any choice of Hessian approximation.

Original: A S OLVE THE TRUST - REGION SUBPROBLEM
Revised: A S OLVE THE TRUST - REGION SUBPROBLEM

Original: Recall the trust region subproblem
Revised: Recall the trust region subproblem

Original: [Equation]
Revised: [Equation]

Original: where α :=
Revised: where α :=

Original: [Equation]
Revised: [Equation]

Original: , and G t :=
Revised: , and G t :=

Original: [Equation]
Revised: [Equation]

Original: In order to solve this trust region subproblem, we transform it into a standard trust region subproblem with L 2 -norm constraint.
Revised: In order to solve this trust region subproblem, we transform it into a standard trust region subproblem with L 2 -norm constraint.

Original: When matrix G t is positive definite, we have
Revised: When matrix G t is positive definite, we have

Original: [Equation]
Revised: [Equation]

Original: So we let y =
Revised: So we let y =

Original: L Tt α , then α
Revised: L Tt α , then α

Original: = L − T t y and the subproblem becomes
Revised: = L − T t y and the subproblem becomes

Original: [Equation]
Revised: [Equation]

Original: In this way, the trust region subproblem is transformed to a standard spherical constrained quadratic optimization problem and it can be solved efficiently (Wright et al., 1999).
Revised: In this way, the trust region subproblem is transformed to a standard spherical constrained quadratic optimization problem and it can be solved efficiently (Wright et al., 1999).

Original: When | G t | = 0 , this means ˆ v t is linearly dependent with d t − 1 .
Revised: When | G t | = 0 , this means ˆ v t is linearly dependent with d t − 1 .

Original: In this case, we solve the onedimensional subproblem as described in Section 3.1.
Revised: In this case, we solve the onedimensional subproblem as described in Section 3.1.

Original: B H OW α 1 t , α 2 t SATISFY THE ASSUMPTIONS NATURALLY
Revised: B H OW α 1 t , α 2 t SATISFY THE ASSUMPTIONS NATURALLY

Original: The trust-region subproblem to be solved in Algorithm 1 has global optimality condition (Luenberger et al., 1984) given by
Revised: The trust-region subproblem to be solved in Algorithm 1 has global optimality condition (Luenberger et al., 1984) given by

Original: [Equation]
Revised: [Equation]

Original: By its construction, we know that G t is positive semidefinite.
Revised: By its construction, we know that G t is positive semidefinite.

Original: In practice, numerical issues sometimes make it indefinite, leaving the trust-region subproblem insoluble.
Revised: In practice, numerical issues sometimes make it indefinite, leaving the trust-region subproblem insoluble.

Original: Thus, we make an adjustment to G t
Revised: Thus, we make an adjustment to G t

Original: [Equation]
Revised: [Equation]

Original: where λ min is the smallest eigenvalue of G t .
Revised: where λ min is the smallest eigenvalue of G t .

Original: In this way, when | G t | ̸ = 0 , we have
Revised: In this way, when | G t | ̸ = 0 , we have

Original: [Equation]
Revised: [Equation]

Original: which means
Revised: which means

Original: [Equation]
Revised: [Equation]

Original: With the common additional assumption that stochastic gradient g t =
Revised: With the common additional assumption that stochastic gradient g t =

Original: ˜ ∇ f
Revised: ˜ ∇ f

Original: ( 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 ∥
Revised: ( 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 ∥

Original: ˆ v t ∥ .
Revised: ˆ v t ∥ .

Original: Therefore, we can see that | α 1 t η | , | α 2 t η | are upper bounded by a constant.
Revised: Therefore, we can see that | α 1 t η | , | α 2 t η | are upper bounded by a constant.

Original: When | G t | = 0 , which means d t − 1 is parallel with ˆ v t .
Revised: When | G t | = 0 , which means d t − 1 is parallel with ˆ v t .

Original: Then we only need to find the optimal update within the trust-region along gradient direction ˆ v t .
Revised: Then we only need to find the optimal update within the trust-region along gradient direction ˆ v t .

Original: In this case, we manually set α 2 t = 0 in our implementation of DRAG, and then α 1 satisfies | α 1 | ≤ η .
Revised: In this case, we manually set α 2 t = 0 in our implementation of DRAG, and then α 1 satisfies | α 1 | ≤ η .

Original: 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 =
Revised: 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 =

Original: ˜ ∇ f
Revised: ˜ ∇ f

Original: ( x t ) has bounded L ∞ norm.
Revised: ( x t ) has bounded L ∞ norm.

Original: For the simplicity of notations, we directly make assumptions for α 1 t and α 2 t in Theorem 1 and Theorem 2.
Revised: For the simplicity of notations, we directly make assumptions for α 1 t and α 2 t in Theorem 1 and Theorem 2.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Thus, we consider the update should at least move towards the gradient descent direction rather than move towards the gradient ascent direction.
Revised: Thus, we consider the update should at least move towards the gradient descent direction rather than move towards the gradient ascent direction.

Original: 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.
Revised: 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.

Original: Therefore, this assumption on α 1 t is reasonable based on common sense and holds true in practice.
Revised: Therefore, this assumption on α 1 t is reasonable based on common sense and holds true in practice.

Original: C P ROOF OF T HEOREM
Revised: C P ROOF OF T HEOREM

Original: 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.
Revised: 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.

Original: Lemma 3 (Variance Recursion (Wang et al., 2017)) .
Revised: Lemma 3 (Variance Recursion (Wang et al., 2017)) .

Original: Suppose Assumption 1 holds, then we have
Revised: Suppose Assumption 1 holds, then we have

Original: E t [
Revised: E t [

Original: [Equation]
Revised: [Equation]

Original: where E t [ · ] denotes the conditional expectation with respect to all randomness before g t .
Revised: where E t [ · ] denotes the conditional expectation with respect to all randomness before g t .

Original: Before proving Theorem 1, we need to prove the following auxiliary lemma.
Revised: Before proving Theorem 1, we need to prove the following auxiliary lemma.

Original: Lemma 4.
Revised: Lemma 4.

Original: Suppose Assumption 1 holds.
Revised: Suppose Assumption 1 holds.

Original: 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 .
Revised: 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 .

Original: In addition, η ≤ min (cid:26) 1 2 LG , 1 − β 2 L (cid:113) δ 2 G ,
Revised: In addition, η ≤ min (cid:26) 1 2 LG , 1 − β 2 L (cid:113) δ 2 G ,

Original: [Equation]
Revised: [Equation]

Original: Then there exist positive constants C 1 , C 2 and C 3 , which are all independent with T , such that the following estimation holds:
Revised: Then there exist positive constants C 1 , C 2 and C 3 , which are all independent with T , such that the following estimation holds:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Since F is L -smooth, we have
Revised: Since F is L -smooth, we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: By Lemma 3, we can obtain
Revised: By Lemma 3, we can obtain

Original: T (cid:88)
Revised: T (cid:88)

Original: [Equation]
Revised: [Equation]

Original: Taking expectation for both sides of (5) and taking summation among t = 0 , ..., T − 1 , combining with (6)
Revised: Taking expectation for both sides of (5) and taking summation among t = 0 , ..., T − 1 , combining with (6)

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: By AM-GM inequality,
Revised: By AM-GM inequality,

Original: [Equation]
Revised: [Equation]

Original: Combining all together, we have
Revised: Combining all together, we have

Original: [Equation]
Revised: [Equation]

Original: where x ∗ is one of the global minimizer of F .
Revised: where x ∗ is one of the global minimizer of F .

Original: Since α 1 t ≤ ηG ≤ 12 L , we have α 1 t (1 − Lα 1 t )≥ α 1 t 8 .
Revised: Since α 1 t ≤ ηG ≤ 12 L , we have α 1 t (1 − Lα 1 t )≥ α 1 t 8 .

Original: By the conditions for η and α , we have α 1 t 16 ≥ α 16 ≥ η 3 GL 2 2(1 − β ) 2 ≥ ηGL 2 η 2 t +1 2(1 − β ) 2 , α 1 t 96 ≥ α 96 ≥ η 4 G 2 α ≥ α 22 ,t +1 η 2 t +1 α 1 ,t +1 , α 1 t 96 ≥ α 96 ≥
Revised: By the conditions for η and α , we have α 1 t 16 ≥ α 16 ≥ η 3 GL 2 2(1 − β ) 2 ≥ ηGL 2 η 2 t +1 2(1 − β ) 2 , α 1 t 96 ≥ α 96 ≥ η 4 G 2 α ≥ α 22 ,t +1 η 2 t +1 α 1 ,t +1 , α 1 t 96 ≥ α 96 ≥

Original: Lη 4 G 2≥ Lα 22 ,t +1 η 2 t
Revised: Lη 4 G 2≥ Lα 22 ,t +1 η 2 t

Original: +1, and α 1 t 96 ≥ α 96 ≥ 2 L 2 η 5 G 3 ≥ L 

Revised: +1, and α 1 t 96 ≥ α 96 ≥ 2 L 2 η 5 G 3 ≥ L

Original: α 1 ,t +1 α 22 ,t +1 η 2 t +1 1 − Lα 1 ,t +
Revised: α 1 ,t +1 α 22 ,t +1 η 2 t +1 1 − Lα 1 ,t +

Original: [Equation]
Revised: [Equation]

Original: By ∥ d t ∥ ≤ η t ∥ v t ∥ .
Revised: By ∥ d t ∥ ≤ η t ∥ v t ∥ .

Original: Since v 0 = 0 , we have
Revised: Since v 0 = 0 , we have

Original: T
Revised: T

Original: [Equation]
Revised: [Equation]

Original: 1 (cid:88)
Revised: 1 (cid:88)

Original: [Equation]
Revised: [Equation]

Original: Combining (8) and (9)
Revised: Combining (8) and (9)

Original: [Equation]
Revised: [Equation]

Original: we can obtain
Revised: we can obtain

Original: T − 1 (cid:88)
Revised: T − 1 (cid:88)

Original: [Equation]
Revised: [Equation]

Original: Dividing the above two inequalities by αT 4 and αT 32 respectively, we have
Revised: Dividing the above two inequalities by αT 4 and αT 32 respectively, we have

Original: [Equation]
Revised: [Equation]

Original: which completes the proof by letting C 

Revised: which completes the proof by letting C

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: Proof of Theorem
Revised: Proof of Theorem

Original: Proof.
Revised: Proof.

Original: By the selections of α and η t in Theorem 1, let δ =
Revised: By the selections of α and η t in Theorem 1, let δ =

Original: α/η .
Revised: α/η .

Original: By Lemma 4, we have
Revised: By Lemma 4, we have

Original: [Equation]
Revised: [Equation]

Original: The conditions 1 − β ≤ ϵ 2 3 C 2 σ 2 and T ≥ max (cid:110) 3 C 1 αϵ 2 , 3 C 3 (1 − β ) ϵ 2 (cid:111) lead to C 1 Tα ≤ ϵ 2 3 , C 2 (1 − β ) σ 2 ≤ ϵ 2 3 ,
Revised: The conditions 1 − β ≤ ϵ 2 3 C 2 σ 2 and T ≥ max (cid:110) 3 C 1 αϵ 2 , 3 C 3 (1 − β ) ϵ 2 (cid:111) lead to C 1 Tα ≤ ϵ 2 3 , C 2 (1 − β ) σ 2 ≤ ϵ 2 3 ,

Original: C 3 T (1 − β ) ≤ ϵ 2 3 .
Revised: C 3 T (1 − β ) ≤ ϵ 2 3 .

Original: This completes the proof.
Revised: This completes the proof.

Original: D P ROOF OF T HEOREM
Revised: D P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: From (5) in Lemma 4, we have
Revised: From (5) in Lemma 4, we have

Original: [Equation]
Revised: [Equation]

Original: By Lemma 3, we have
Revised: By Lemma 3, we have

Original: [Equation]
Revised: [Equation]

Original: Taking expectation and summation for t = 1 , ..., T , we get
Revised: Taking expectation and summation for t = 1 , ..., T , we get

Original: T − 1 (cid:88) t =
Revised: T − 1 (cid:88) t =

Original: [Equation]
Revised: [Equation]

Original: Note that 1 − β t +1 =
Revised: Note that 1 − β t +1 =

Original: Cη t , so
Revised: Cη t , so

Original: α 1 t 2 ≤ G 2 η t = G 2 C (1 − β t +1 ) .
Revised: α 1 t 2 ≤ G 2 η t = G 2 C (1 − β t +1 ) .

Original: Taking expectation for both sides of (10) and taking summation among t = 0 , ..., T − 1 , combining with (11), we can obtain
Revised: Taking expectation for both sides of (10) and taking summation among t = 0 , ..., T − 1 , combining with (11), we can obtain

Original: E [ f
Revised: E [ f

Original: [Equation]
Revised: [Equation]

Original: x T ) − f ( x 0 )]
Revised: x T ) − f ( x 0 )]

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: From (7) and (12), we can get
Revised: From (7) and (12), we can get

Original: T − 1 (cid:88)
Revised: T − 1 (cid:88)

Original: [Equation]
Revised: [Equation]

Original: By the conditions for c η and C , we have α 1 t ≤ η t G ≤ 12 L , α
Revised: By the conditions for c η and C , we have α 1 t ≤ η t G ≤ 12 L , α

Original: [Equation]
Revised: [Equation]

Original: 1 − Lα
Revised: 1 − Lα

Original: [Equation]
Revised: [Equation]

Original: ≥ α 1 t 8 .
Revised: ≥ α 1 t 8 .

Original: By similar arguments in the proof of Lemma 4, we have α 1 t 16 ≥
Revised: By similar arguments in the proof of Lemma 4, we have α 1 t 16 ≥

Original: [Equation]
Revised: [Equation]

Original: ≥
Revised: ≥

Original: [Equation]
Revised: [Equation]

Original: η
Revised: η

Original: [Equation]
Revised: [Equation]

Original: α 22 ,t +1 η 2 t +1 α 1 ,t +1 , α
Revised: α 22 ,t +1 η 2 t +1 α 1 ,t +1 , α

Original: [Equation]
Revised: [Equation]

Original: t 96 ≥ δη t 96 ≥ Lα
Revised: t 96 ≥ δη t 96 ≥ Lα

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: , and α 1 t 96 ≥ δη
Revised: , and α 1 t 96 ≥ δη

Original: [Equation]
Revised: [Equation]

Original: 96 ≥
Revised: 96 ≥

Original: [Equation]
Revised: [Equation]

Original: L
Revised: L

Original: [Equation]
Revised: [Equation]

Original: By ∥ d t ∥ ≤ η t ∥ v t ∥ .
Revised: By ∥ d t ∥ ≤ η t ∥ v t ∥ .

Original: Since v 0 = 0 , we can get
Revised: Since v 0 = 0 , we can get

Original: T − 1 (cid:88) t =
Revised: T − 1 (cid:88) t =

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Combining (13) and (14)
Revised: Combining (13) and (14)

Original: [Equation]
Revised: [Equation]

Original: we can obtain
Revised: we can obtain

Original: T − 1 (cid:88)
Revised: T − 1 (cid:88)

Original: [Equation]
Revised: [Equation]

Original: Gσ
Revised: Gσ

Original: C (1 − β t ) 2 .
Revised: C (1 − β t ) 2 .

Original: Then, the final assertion can be obtained by (cid:80)
Revised: Then, the final assertion can be obtained by (cid:80)

Original: Tt =1 1 t +

Revised: Tt =1 1 t +

Original: =
Revised: =

Original: O (log T ) .
Revised: O (log T ) .

Original: This completes the proof.
Revised: This completes the proof.

Original: E E XPERIMENTAL R ESULTS FOR ABLATION STUDY
Revised: E E XPERIMENTAL R ESULTS FOR ABLATION STUDY

Original: [Table] [Table]
Revised: [Table] [Table]