Momentum Centering and Asynchronous Update for Adaptive Gradient Methods

T ) for the stochastic non-convex case , which matches the oracle rate and outperforms the O ( logT√ T ) rate of RMSProp and Adam . We validate ACProp in extensive empiri- cal studies : ACProp outperforms both SGD and other adaptive optimizers in image classification with CNN , and outperforms well-tuned adaptive optimizers in the training of various GAN models , reinforcement learning and transformers . To sum up , ACProp has good theoretical properties including weak convergence condition and optimal convergence rate , and strong empirical performance including good generalization like SGD and training stability like Adam . We provide the implementation at https : //github.com/juntang-zhuang/ACProp-Optimizer . 1 Introduction . Deep neural networks are typically trained with first-order gradient optimizers due to their computational efficiency and good empirical performance [ 1 ] . Current first-order gradient optimizers can be broadly categorized into the stochastic gradient descent ( SGD ) [ 2 ] family and the adaptive family . The SGD family uses a global learning rate for all parameters , and includes variants such as Nesterov-accelerated SGD [ 3 ] , SGD with momentum [ 4 ] and the heavy-ball method [ 5 ] . Compared with the adaptive family , SGD optimizers typically generalize better but converge slower , and are the default for vision tasks such as image classification [ 6 ] , object detection [ 7 ] and segmentation [ 8 ] . The adaptive family uses element-wise learning rate , and the representatives include AdaGrad [ 9 ] , AdaDelta [ 10 ] , RMSProp [ 11 ] , Adam [ 12 ] and its variants such as AdamW [ 13 ] , AMSGrad [ 14 ] AdaBound [ 15 ] , AdaShift [ 16 ] , RAdam [ 17 ] and AdaBelief [ 18 ] . Compared with the SGD family , the adaptive optimizers typically converge faster and are more stable , hence are the default for generative adversarial networks ( GANs ) [ 19 ] , transformers [ 20 ] , and deep reinforcement learning [ 21 ] . We broadly categorize adaptive optimizers according to different criteria , as in Table . 1 . ( a ) Centered v.s . uncentered Most optimizers such as Adam and AdaDelta uses uncentered second momentum in the denominator ; RMSProp-center [ 11 ] , SDProp [ 22 ] and AdaBelief [ 18 ] use square root of centered 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . second momentum in the denominator . AdaBelief [ 18 ] is shown to achieve good generalization like the SGD family , fast convergence like the adaptive family , and training stability in complex settings such as GANs . ( b ) Sync vs async The synchronous optimizers typically use gradient gt in both numerator and denominator , which leads to correlation between numerator and denominator ; most existing optimizers belong to this category . The asynchronous optimizers decorrelate numerator and denominator ( e.g . by using gt as numerator and use { g0 , ... gt−1 } in denominator for the t-th update ) , and is shown to have weaker convergence conditions than synchronous optimizers [ 16 ] . We propose Asynchronous Centering Prop ( ACProp ) , which combines centering of second momentum with the asynchronous update . We show that ACProp has both good theoretical properties and strong empirical performance . Our contributions are summarized as below : • Convergence condition ( a ) Async vs Sync We show that for the example by Reddi et al . ( 2018 ) , asynchronous optimizers ( AdaShift , ACProp ) converge for any valid hyper-parameters , while synchronous optimizers ( Adam , RMSProp et al . ) could diverge if the hyper-paramaters are not carefully chosen . ( b ) Async-Center vs Async-Uncenter Within the asynchronous optimizers family , by example of an online convex problem with sparse gradients , we show that Async-Center ( ACProp ) has weaker conditions for convergence than Async-Uncenter ( AdaShift ) . • Convergence rate We demonstrate that ACProp achieves a convergence rate of O ( 1√ T ) for stochastic non-convex problems , matching the oracle of first-order optimizers [ 23 ] , and outperforms the O ( logT√ T ) rate of Adam and RMSProp . • Empirical performance We validate performance of ACProp in experiments : on image classification tasks , ACProp outperforms SGD and AdaBelief , and demonstrates good generalization performance ; in experiments with transformer , reinforcement learning and various GAN models , ACProp outperforms well-tuned Adam , demonstrating high stability . ACProp often outperforms AdaBelief , and achieves good generalization like SGD and training stability like Adam . 2 Overview of algorithms . 2.1 Notations . x , xt ∈ Rd : x is a d−dimensional parameter to be optimized , and xt is the value at step t. f ( x ) , f∗ ∈ R : f ( x ) is the scalar-valued function to be minimized , with optimal ( minimal ) f∗ . αt , ∈ R : αt is the learning rate at step t. is a small number to avoid division by 0. gt ∈ Rd : The noisy observation of gradient∇f ( xt ) at step t. β1 , β2 ∈ R : Constants for exponential moving average , 0 ≤ β1 , β2 < 1. mt ∈ Rd : mt = β1mt−1 + ( 1− β1 ) gt . The Exponential Moving Average ( EMA ) of observed gradient at step t. ∆gt ∈ Rd : ∆gt = gt −mt . The difference between observed gradient gt and EMA of gt . vt ∈ Rd : vt = β2vt−1 + ( 1− β2 ) g2t . The EMA of g2t . st ∈ Rd : st = β2st−1 + ( 1− β2 ) ( ∆gt ) 2 . The EMA of ( ∆gt ) 2 . 2.2 Algorithms . In this section , we summarize the AdaBelief [ 18 ] method in Algo . 1 and ACProp in Algo . 2 . For the ease of notations , all operations in Algo . 1 and Algo . 2 are element-wise , and we omit the bias-correction step of mt and st for simplicity . ΠF represents the projection onto feasible set F . We first introduce the notion of “ sync ( async ) ” and “ center ( uncenter ) ” . ( a ) Sync vs Async The update on parameter xt can be generally split into a numerator ( e.g . mt , gt ) and a denominator ( e.g . Algorithm 1 : AdaBelief Initialize x0 , m0 ← 0 , s0 ← 0 , t← 0 While xt not converged t← t+ 1 gt ← ∇xft ( xt−1 ) mt ← β1mt−1 + ( 1− β1 ) gt st ← β2st−1+ ( 1−β2 ) ( gt−mt ) 2 xt ← ∏ F , √st ( xt−1 − α√st+ mt ) Algorithm 2 : ACProp Initialize x0 , m0 ← 0 , s0 ← 0 , t← 0 While xt not converged t← t+ 1 gt ← ∇xft ( xt−1 ) mt ← β1mt−1 + ( 1− β1 ) gt xt ← ∏ F , √st−1 ( xt−1 − α√st−1+ gt ) st ← β2st−1+ ( 1−β2 ) ( gt−mt ) 2 √ st , √ vt ) . We call it “ sync ” if the denominator depends on gt , such as in Adam and RMSProp ; and call it “ async ” if the denominator is independent of gt , for example , denominator uses information up to step t−1 for the t-th step . ( b ) Center vs Uncenter The “ uncentered ” update uses vt , the exponential moving average ( EMA ) of g2t ; while the “ centered ” update uses st , the EMA of ( gt −mt ) 2 . Adam ( Sync-Uncenter ) The Adam optimizer [ 12 ] stores the EMA of the gradient in mt , and stores the EMA of g2t in vt. For each step of the update , Adam performs element-wise division between mt and √ vt . Therefore , the term αt 1√vt can be viewed as the element-wise learning rate . Note that β1 and β2 are two scalars controlling the smoothness of the EMA for the first and second moment , respectively . When β1 = 0 , Adam reduces to RMSProp [ 24 ] . AdaBelief ( Sync-Center ) AdaBelief optimizer [ 18 ] is summarized in Algo . 1 . Compared with Adam , the key difference is that it replaces the uncentered second moment vt ( EMA of g2t ) by an estimate of the centered second moment st ( EMA of ( gt −mt ) 2 ) . The intuition is to view mt as an estimate of the expected gradient : if the observation gt deviates much from the prediction mt , then it takes a small step ; if the observation gt is close to the prediction mt , then it takes a large step . AdaShift ( Async-Uncenter ) AdaShift [ 16 ] performs temporal decorrelation between numerator and denominator . It uses information of { gt−n , ... gt } for the numerator , and uses { g0 , ... gt−n−1 } for the denominator , where n is the “ delay step ” controlling where to split sequence { gi } ti=0 . The numerator is independent of denominator because each gi is only used in either numerator or denominator . ACProp ( Async-Center ) Our proposed ACProp is the asynchronous version of AdaBelief and is summarized in Algo . 2 . Compared to AdaBelief , the key difference is that ACProp uses st−1 in the denominator for step t , while AdaBelief uses st . Note that st depends on gt , while st−1 uses history up to step t − 1 . This modification is important to ensure that E ( gt/ √ st−1|g0 , ... gt−1 ) = ( Egt ) / √ st−1 . It ’ s also possible to use a delay step larger than 1 similar to AdaShift , for example , use EMA ( { gi } ti=t−n ) as numerator , and EMA ( { ( gi −mi ) 2 } t−n−1 i=0 ) for denominator . 3 Analyze the conditions for convergence . We analyze the convergence conditions for different methods in this section . We first analyze the counter example by Reddi et al . ( 2018 ) and show that async-optimizers ( AdaShift , ACProp ) always converge ∀β1 , β2 ∈ ( 0 , 1 ) , while sync-optimizers ( Adam , AdaBelief , RMSProp et al . ) would diverge if ( β1 , β2 ) are not carefully chosen ; hence , async-optimizers have weaker convergence conditions than sync-optimizers . Next , we compare async-uncenter ( AdaShift ) with async-center ( ACProp ) and show that momentum centering further weakens the convergence condition for sparse-gradient problems . Therefore , ACProp has weaker convergence conditions than AdaShift and other sync-optimizers . 3.1 Sync vs Async We show that for the example in [ 14 ] , async-optimizers ( ACProp , AdaShift ) have weaker convergence conditions than sync-optimizers ( Adam , RMSProp , AdaBelief ) . Lemma 3.1 ( Thm.1 in [ 14 ] ) . There exists an online convex optimization problem where syncoptimizers ( e.g . Adam , RMSProp ) have non-zero average regret , and one example is ft ( x ) = { Px , if t % P = 1 −x , Otherwise x ∈ [ −1 , 1 ] , P ∈ N , P ≥ 3 ( 1 ) Lemma 3.2 ( [ 25 ] ) . For problem ( 1 ) with any fixed P , there ’ s a threshold of β2 above which RMSProp converges . gence and divergence area for sync-optimizers ( Adam , RMSProp , AdaBelief ) , while async-optimizers ( ACProp , AdaShift ) always converge . Lemma 3.3 . For the problem defined by Eq . ( 1 ) , using learning rate schedule of αt = α0√t , asyncoptimizers ( ACProp and AdaShift with n = 1 ) always converge ∀β1 , β2 ∈ ( 0 , 1 ) , ∀P ∈ N , P ≥ 3 . The proof is in the appendix . Note that for AdaShift , proof for the always-convergence property only holds when n = 1 ; larger n could cause divergence ( e.g . n = 10 causes divergence as in Fig . 2 ) . The always-convergence property of ACProp and AdaShift comes from the un-biased stepsize , while the stepsize for sync-optimizers are biased due to correlation between numerator and denominator . Taking RMSProp as example of sync-optimizer , the update is −αt gt√vt = −αt gt√ βt2g 2 0+ ... +β2g 2 t−1+g 2 t . Note that gt is used both in the numerator and denominator , hence a large gt does not necessarily generate a large stepsize . For the example in Eq . ( 1 ) , the optimizer observes a gradient of −1 for P − 1 times and a gradient of P once ; due to the biased stepsize in sync-optimizers , the gradient of P does not generate a sufficiently large stepsize to compensate for the effect of wrong gradients −1 , hence cause non-convergence . For async-optimizers , gt is not used in the denominator , therefore , the stepsize is not biased and async-optimizers has the always-convergence property . Remark Reddi et al . ( 2018 ) proposed AMSGrad to track the element-wise maximum of vt in order to achieve the always-convergence property . However , tracking the maximum in the denominator will in general generate a small stepsize , which often harms empirical performance . We demonstrate this through experiments in later sections in Fig . 6 . 3.2 Async-Uncenter vs Async-Center In the last section , we demonstrated that async-optimizers have weaker convergence conditions than sync-optimizers . In this section , within the async-optimizer family , we analyze the effect of centering second momentum . We show that compared with async-uncenter ( AdaShift ) , async-center ( ACProp ) has weaker convergence conditions . We consider the following online convex problem : ft ( x ) = P/2× x , t % P == 1 −x , t % P == P − 2 0 , otherwise P > 3 , P ∈ N , x ∈ [ 0 , 1 ] . ( 2 ) Initial point is x0 = 0.5 . Optimal point is x∗ = 0 . We have the following results : Lemma 3.4 . For the problem defined by Eq . ( 2 ) , consider the hyper-parameter tuple ( β1 , β2 , P ) , there exists cases where ACProp converges but AdaShift with n = 1 diverges , but not vice versa . We provide the proof in the appendix . Lemma . 3.4 implies that ACProp has a larger area of convergence than AdaShift , hence the centering of second momentum further weakens the convergence conditions . We first validate this claim with numerical experiments in Fig . 3 ; for sanity check , we plot the trajectories of different optimizers in Fig . 4 . We observe that the convergence of AdaShift is influenced by delay step n , and there ’ s no good criterion to select a good value of n , since Fig . 2 requires a small n for convergence in problem ( 1 ) , while Fig . 4 requires a large n for convergence in problem ( 2 ) . ACProp has a larger area of convergence , indicating that both async update and second momentum centering helps weaken the convergence conditions . We provide an intuitive explanation on why momentum centering helps convergence . Due to the periodicity of the problem , the optimizer behaves almost periodically as t → ∞ . Within each period , the optimizer observes one positive gradient P/2 and one negative gradient -1 . As in Fig . 5 , between observing non-zero gradients , the gradient is always 0 . Within each period , ACprop will perform a positive update P/ ( 2 √ s+ ) and a negative update −1/ √ s− , where s+ ( s− ) is the value of denominator before observing positive ( negative ) gradient . Similar notations for v+ and v− in AdaShift . A net update in the correct direction requires P 2 √ s+ > 1√ s− , ( or s+/s− < P 2/4 ) . When observing 0 gradient , for AdaShift , vt = β2vt−1 + ( 1− β2 ) 02 ; for ACProp , st = β2st−1 + ( 1 − β2 ) ( 0 − mt ) 2 where mt 6= 0 . Therefore , v− decays exponentially to 0 , but s− decays to a non-zero constant , hence s + s− < v+ v− , hence ACProp is easier to satisfy s +/s− < P 2/4 and converge .

Summary: This paper proposes an adaptive optimizer, ACProp, which combines centering of second momentum and asynchronous update. The authors show that ACProp has a convergence rate of $O(\frac{1}{\sqrt{T}})$ in a stochastic non-convex setting, outperforming the  $O(\frac{\log T}{\sqrt{T}})$ of RMSProp and Adam in the same setting. Experiments on a range of tasks, including image classification, training GANs, reinforcement learning and machine translation, show the effectiveness of ACProp.