An Exponential Learning Rate Schedule for Deep Learning

• Training can be done using SGD with momentum and an exponentially increasing learning rate schedule , i.e. , learning rate increases by some ( 1 + α ) factor in every epoch for some α > 0 . ( Precise statement in the paper . ) To the best of our knowledge this is the first time such a rate schedule has been successfully used , let alone for highly successful architectures . As expected , such training rapidly blows up network weights , but the network stays wellbehaved due to normalization . • Mathematical explanation of the success of the above rate schedule : a rigorous proof that it is equivalent to the standard setting of BN + SGD + Standard Rate Tuning + Weight Decay + Momentum . This equivalence holds for other normalization layers as well , Group Normalization ( Wu & He , 2018 ) , Layer Normalization ( Ba et al. , 2016 ) , Instance Norm ( Ulyanov et al. , 2016 ) , etc . • A worked-out toy example illustrating the above linkage of hyperparameters . Using either weight decay or BN alone reaches global minimum , but convergence fails when both are used . 1 INTRODUCTION Batch Normalization ( BN ) offers significant benefits in optimization and generalization across architectures , and has become ubiquitous . Usually best performance is attained by adding weight decay and momentum in addition to BN . Usually weight decay is thought to improve generalization by controlling the norm of the parameters . However , it is fallacious to try to separately think of optimization and generalization because we are dealing with a nonconvex objective with multiple optima . Even slight changes to the training surely lead to a different trajectory in the loss landscape , potentially ending up at a different solution ! One needs trajectory analysis to have a hope of reasoning about the effects of such changes . In the presence of BN and other normalization schemes , including GroupNorm , LayerNorm , and InstanceNorm , the optimization objective is scale invariant to the parameters , which means rescaling parameters would not change the prediction , except the parameters that compute the output which do not have BN . However , Hoffer et al . ( 2018b ) shows that fixing the output layer randomly doesn ’ t harm the performance of the network . So the trainable parameters satisfy scale invariance . ( See more in Appendix E ) The current paper introduces new modes of analysis for such settings . This rigorous analysis yields the surprising conclusion that the original learning rate ( LR ) schedule and weight decay ( WD ) can be folded into a new exponential schedule for learning rate : in each iteration multiplying it by ( 1 + α ) for some α > 0 that depends upon the momentum and weight decay rate . Theorem 1.1 ( Main , Informal ) . SGD on a scale-invariant objective with initial learning rate η , weight decay factor λ , and momentum factor γ is equivalent to SGD with momentum factor γ where at iteration t , the learning rate η̃t in the new exponential learning rate schedule is defined as η̃t = α −2t−1η without weight decay ( λ̃ = 0 ) where α is a non-zero root of equation x2 − ( 1 + γ − λη ) x+ γ = 0 , ( 1 ) Specifically , when momentum γ = 0 , the above schedule can be simplified as η̃t = ( 1−λη ) −2t−1η . The above theorem requires that the product of learning rate and weight decay factor , λη , is small than ( 1−√γ ) 2 , which is almost always satisfied in practice . The rigorous and most general version of above theorem is Theorem 2.12 , which deals with multi-phase LR schedule , momentum and weight decay . There are other recently discovered exotic LR schedules , e.g . Triangular LR schedule ( Smith , 2017 ) and Cosine LR schedule ( Loshchilov & Hutter , 2016 ) , and our exponential LR schedule is an extreme example of LR schedules that become possible in presence of BN . Such an exponential increase in learning rate seems absurd at first sight and to the best of our knowledge , no deep learning success has been reported using such an idea before . It does highlight the above-mentioned viewpoint that in deep learning , optimization and regularization are not easily separated . Of course , the exponent trumps the effect of initial lr very fast ( See Figure 4 ) , which explains why training with BN and WD is not sensitive to the scale of initialization , since with BN , tuning the scale of initialization is equivalent to tuning the initial LR η while fixing the product of LR and WD , ηλ ( See Lemma 2.7 ) . Note that it is customary in BN to switch to a lower LR upon reaching a plateau in the validation loss . According to the analysis in the above theorem , this corresponds to an exponential growth with a smaller exponent , except for a transient effect when a correction term is needed for the two processes to be equivalent ( see discussion around Theorem 2.12 ) . Thus the final training algorithm is roughly as follows : Start from a convenient LR like 0.1 , and grow it at an exponential rate with a suitable exponent . When validation loss plateaus , switch to an exponential growth of LR with a lower exponent . Repeat the procedure until the training loss saturates . In Section 3 , we demonstrate on a toy example how weight decay and normalization are inseparably involved in the optimization process . With either weight decay or normalization alone , SGD will achieve zero training error . But with both turned on , SGD fails to converge to global minimum . In Section 4 , we experimentally verify our theoretical findings on CNNs and ResNets . We also construct better exponential LR schedules by incorporating the Cosine LR schedule on CIFAR10 , which opens the possibility of even more general theory of rate schedule tuning towards better performance . 1.1 RELATED WORK There have been other theoretical analyses of training models with scale-invariance . ( Cho & Lee , 2017 ) proposed to run Riemanian gradient descent on Grassmann manifold G ( 1 , n ) since the weight matrix is scaling invariant to the loss function . observed that the effective stepsize is proportional to ηw ‖wt‖2 . ( Arora et al. , 2019 ) show the gradient is always perpendicular to the current parameter vector which has the effect that norm of each scale invariant parameter group increases monotonically , which has an auto-tuning effect . ( Wu et al. , 2018 ) proposes a new adaptive learning rate schedule motivated by scale-invariance property of Weight Normalization . Previous work for understanding Batch Normalization . ( Santurkar et al. , 2018 ) suggested that the success of BNhas does not derive from reduction in Internal Covariate Shift , but by making landscape smoother . ( Kohler et al. , 2018 ) essentially shows linear model with BN could achieve exponential convergence rate assuming gaussian inputs , but their analysis is for a variant of GD with an inner optimization loop rather than GD itself . ( Bjorck et al. , 2018 ) observe that the higher learning rates enabled by BN empirically improves generalization . ( Arora et al. , 2019 ) prove that with certain mild assumption , ( S ) GD with BN finds approximate first order stationary point with any fixed learning rate . None of the above analyses incorporated weight decay , but ( Zhang et al. , 2019 ; Hoffer et al. , 2018a ; van Laarhoven , 2017 ) argued qualitatively that weight decay makes parameters have smaller norms , and thus the effective learning rate , ηw‖wt‖2 is larger . They described experiments showing this effect but didn ’ t have a closed form theoretical analysis like ours . None of the above analyses deals with momentum rigorously . 1.2 PRELIMINARIES AND NOTATIONS For batch B = { xi } Bi=1 , network parameter θ , we denote the network by fθ and the loss function at iteration t by Lt ( fθ ) = L ( fθ , Bt ) . When there ’ s no ambiguity , we also use Lt ( θ ) for convenience . We say a loss function L ( θ ) is scale invariant to its parameter θ is for any c ∈ R+ , L ( θ ) = L ( cθ ) . In practice , the source of scale invariance is usually different types of normalization layers , including Batch Normalization ( Ioffe & Szegedy , 2015 ) , Group Normalization ( Wu & He , 2018 ) , Layer Normalization ( Ba et al. , 2016 ) , Instance Norm ( Ulyanov et al. , 2016 ) , etc . Implementations of SGD with Momentum/Nesterov comes with subtle variations in literature . We adopt the variant from Sutskever et al . ( 2013 ) , also the default in PyTorch ( Paszke et al. , 2017 ) . L2 regularization ( a.k.a . Weight Decay ) is another common trick used in deep learning . Combining them together , we get the one of the mostly used optimization algorithms below . Definition 1.2 . [ SGD with Momentum and Weight Decay ] At iteration t , with randomly sampled batch Bt , update the parameters θt and momentum vt as following : θt =θt−1 − ηt−1vt ( 2 ) vt =γvt−1 +∇θ ( Lt ( θt−1 ) + λt−1 2 ‖θt−1‖2 ) , ( 3 ) where ηt is the learning rate at epoch t , γ is the momentum coefficient , and λ is the factor of weight decay . Usually , v0 is initialized to be 0 . For ease of analysis , we will use the following equivalent of Definition 1.2. θt − θt−1 ηt−1 = γ θt−1 − θt−2 ηt−2 −∇θ ( ( L ( θt−1 ) + λt−1 2 ‖θt−1‖22 ) , ( 4 ) where η−1 and θ−1 must be chosen in a way such that v0 = θ0−θ−1 η−1 is satisfied , e.g . when v0 = 0 , θ−1 = θ0 and η−1 could be arbitrary . A key source of intuition is the following simple lemma about scale-invariant networks Arora et al . ( 2019 ) . The first property ensures GD ( with momentum ) always increases the norm of the weight . ( See Lemma C.1 in Appendix C ) and the second property says that the gradients are smaller for parameteres with larger norm , thus stabilizing the trajectory from diverging to infinity . Lemma 1.3 ( Scale Invariance ) . If for any c ∈ R+ , L ( θ ) = L ( cθ ) , then ( 1 ) . 〈∇θL , θ〉 = 0 ; ( 2 ) . ∇θL ∣∣ θ=θ0 = c∇θL ∣∣ θ=cθ0 , for any c > 0 2 DERIVING EXPONENTIAL LEARNING RATE SCHEDULE As a warm-up in Section 2.1 we show that if momentum is turned off then Fixed LR + Fixed WD can be translated to an equivalent Exponential LR . In Section 2.2 we give a more general analysis on the equivalence between Fixed LR + Fixed WD + Fixed Momentum Factor and Exponential LR + Fixed Momentum Factor . While interesting , this still does completely apply to real-life deep learning where reaching full accuracy usually requires multiple phases in training where LR is fixed within a phase and reduced by some factor from one phase to the next . Section 2.3 shows how to interpret such a multi-phase LR schedule + WD + Momentum as a certain multi-phase exponential LR schedule with Momentum . 2.1 REPLACING WD BY EXPONENTIAL LR IN MOMENTUM-FREE SGD We use notation of Section 1.2 and assume LR is fixed over iterations , i.e . ηt = η0 , and γ ( momentum factor ) is set as 0 . We also use λ to denote WD factor and θ0 to denote the initial parameters . The intuition should be clear from Lemma 1.3 , which says that shrinking parameter weights by factor ρ ( where ρ < 1 ) amounts to making the gradient ρ−1 times larger without changing its direction . Thus in order to restore the ratio between original parameter and its update ( LR×Gradient ) , the easiest way would be scaling LR by ρ2 . This suggests that scaling the parameter θ by ρ at each step is equivalent to scaling the LR η by ρ−2 . To prove this formally we use the following formalism . We ’ ll refer to the vector ( θ , η ) the state of a training algorithm and study how this evolves under various combinations of parameter changes . We will think of each step in training as a mapping from one state to another . Since mappings can be composed , any finite number of steps also correspond to a mapping . The following are some basic mappings used in the proof . 1 . Run GD with WD for a step : GDρt ( θ , η ) = ( ρθ − η∇Lt ( θ ) , η ) ; 2 . Scale the parameter θ : Πc1 ( θ , η ) = ( cθ , η ) ; 3 . Scale the LR η : Πc2 ( θ , η ) = ( θ , cη ) . For example , when ρ = 1 , GD1t is vanilla GD update without WD , also abbreviated as GDt . When ρ = 1−λη0 , GD1−λη0t is GD update with WD λ and LR η0 . Here Lt is the loss function at iteration t , which is decided by the batch of the training samples Bt in tth iteration . Below is the main result of this subsection , showing our claim that GD + WD ⇔ GD+ Exp LR ( when Momentum is zero ) . It will be proved after a series of lemmas . Theorem 2.1 ( WD ⇔ Exp LR ) . For every ρ < 1 and positive integer t following holds : GDρt−1 ◦ · · · ◦ GD ρ 0 = [ Πρ t 1 ◦Π ρ2t 2 ] ◦Πρ −1 2 ◦ GDt−1 ◦Π ρ−2 2 ◦ · · · ◦ GD1 ◦Π ρ−2 2 ◦ GD0 ◦Π ρ−1 2 . With WD being λ , ρ is set as 1 − λη0 and thus the scaling factor of LR per iteration is ρ−2 = ( 1− λη0 ) −2 , except for the first iteration it ’ s ρ−1 = ( 1− λη0 ) −1 . We first show how to write GD update with WD as a composition of above defined basic maps . Lemma 2.2 . GDρt = Π ρ 2 ◦Π ρ 1 ◦ GDt ◦Π ρ−1 2 . Below we will define the proper notion of equivalence such that ( 1 ) . Πρ1 ∼ Π ρ−2 2 , which implies GDρt ∼ Π ρ−1 2 ◦ GDt ◦Π ρ−1 2 ; ( 2 ) the equivalence is preserved under future GD updates . We first extend the equivalence between weights ( same direction ) to that between states , with additional requirement that the ratio between the size of GD update and that of parameter are the same among all equivalent states , which yields the notion of Equivalent Scaling . Definition 2.3 ( Equivalent States ) . ( θ , η ) is equivalent to ( θ′ , η′ ) iff ∃c > 0 , ( θ̃ , η̃ ) = [ Πc1 ◦ Πc 2 2 ] ( θ , η ) = ( cθ , c 2η ) , which is also denoted by ( θ̃ , η̃ ) c∼ ( θ , η ) . Πc1 ◦ Πc 2 2 is called Equivalent Scaling for all c > 0 . The following lemma shows that equivalent scaling commutes with GD update with WD , implying that equivalence is preserved under GD update ( Lemma 2.4 ) . This anchors the notion of equivalence — we could insert equivalent scaling anywhere in a sequence of basic maps ( GD update , LR/parameter scaling ) , without changing the final network . Lemma 2.4 . For any constant c , ρ > 0 and t ≥ 0 , GDρt ◦ [ Πc1 ◦Πc 2 2 ] = [ Π c 1 ◦Πc 2 2 ] ◦ GD ρ t .

Summary: This work makes an interesting observation that it is possible to use exponentially growing learning rate schedule when training with neural networks with batch normalization. This paper provides both theoretical insights and empirical demonstration of this remarkable property. In detail, the authors prove that for stochastic gradient descent (SGD) with momentum, this exponential learning rate schedule is equivalent to constant learning rate + weight decay, for any scale invariant networks, including networks with Batch Normalization and other normalization methods. This paper also contains an interesting toy example where  gd converges when normalization or weight decay is used alone while not when normalization and weight decay are used together.