A spherical analysis of Adam with Batch Normalization

A SPHERICAL ANALYSIS OF ADAM WITH BATCH NORMALIZATION . Anonymous authors Paper under double-blind review Batch Normalization ( BN ) is a prominent deep learning technique . In spite of its apparent simplicity , its implications over optimization are yet to be fully understood . While previous studies mostly focus on the interaction between BN and stochastic gradient descent ( SGD ) , we develop a geometric perspective which allows us to precisely characterize the relation between BN and Adam . More precisely , we leverage the radial invariance of groups of parameters , such as filters for convolutional neural networks , to translate the optimization steps on the L2 unit hypersphere . This formulation and the associated geometric interpretation shed new light on the training dynamics . Firstly , we use it to derive the first effective learning rate expression of Adam . Then we show that , in the presence of BN layers , performing SGD alone is actually equivalent to a variant of Adam constrained to the unit hypersphere . Finally , our analysis outlines phenomena that previous variants of Adam act on and we experimentally validate their importance in the optimization process . 1 INTRODUCTION The optimization process of deep neural networks is still poorly understood . Their training involves minimizing a high-dimensional non-convex function , which has been proved to be a NP-hard problem ( Blum & Rivest , 1989 ) . Yet , elementary gradient-based methods show good results in practice . To improve the quality of reached minima , numerous methods have stemmed in the last years and become common practices . One of the most prominent is Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) , which improves significantly both the optimization stability and the prediction performance ; it is now used in most deep learning architectures . However , the interaction of BN with optimization and its link to regularization remain open research topics . Previous studies highlighted mechanisms of the interaction between BN and SGD , both empirically ( Santurkar et al. , 2018 ) and theoretically ( Arora et al. , 2019 ; Bjorck et al. , 2018 ; Hoffer et al. , 2018b ) . None of them studied the interaction between BN and one of the most common adaptive schemes for Neural Networks ( NN ) , Adam ( Kingma & Ba , 2015 ) , except van Laarhoven ( 2017 ) , which tackled it only in the asymptotic regime . In this work , we provide an extensive analysis of the relation between BN and Adam during the whole training procedure . One of the key effects of BN is to make NNs invariant to positive scalings of groups of parameters . The core idea of this paper is precisely to focus on these groups of radially-invariant parameters and analyze their optimization projected on the L2 unit hypersphere ( see Fig . 1 ) , which is topologically equivalent to the quotient manifold of the parameter space by the scaling action . One could directly optimize parameters on the hypersphere as Cho & Lee ( 2017 ) , yet , most optimization methods are still performed successfully in the original parameter space . Here we propose to study an optimization scheme for a given group of radially-invariant parameters through its image scheme on the unit hypersphere . This geometric perspective sheds light on the interaction between normalization layers and Adam , and also outlines an interesting link between standard SGD and a variant of Adam adapted and constrained to the unit hypersphere : AdamG ( Cho & Lee , 2017 ) . We believe this kind of analysis is an important step towards a better understanding of the effect of BN on NN optimization . Please note that , although our discussion and experiments focus on BN , our analysis could be applied to any radially-invariant model . The paper is organized as follows . In Section 2 , we introduce our spherical framework to study the optimization of radially-invariant models . We also define a generic optimization scheme that encompasses methods such as SGD with momentum ( SGD-M ) and Adam . We then derive its image step on the unit hypersphere , leading to definitions and expressions of effective learning rate and effective learning direction . This new definition is explicit and has a clear interpretation , whereas the definition of van Laarhoven ( 2017 ) is asymptotic and the definitions of Arora et al . ( 2019 ) and of Hoffer et al . ( 2018b ) are variational . In Section 3 , we leverage the tools of our spherical framework to demonstrate that in presence of BN layers , SGD has an adaptive behaviour . Formally , we show that SGD is equivalent to AdamG , a variant of Adam adapted and constrained to the hypersphere , without momentum . In Section 4 , we analyze the effective learning direction for Adam . The spherical framework highlights phenomena that previous variants of Adam ( Loshchilov & Hutter , 2017 ; Cho & Lee , 2017 ) act on . We perform an empirical study of these phenomena and show that they play a significant role in the training of convolutional neural networks ( CNNs ) . In Section 5 , these results are put in perspective with related work . Our main contributions are the following : • A framework to analyze and compare order-1 optimization schemes of radially-invariant models ; • The first explicit expression of the effective learning rate for Adam ; • The demonstration that , in the presence of BN layers , standard SGD has an adaptive behaviour ; • The identification and study of geometrical phenomena that occur with Adam and impact significantly the training of CNNs with BN . 2 SPHERICAL FRAMEWORK AND EFFECTIVE LEARNING RATE . In this section , we provide background on radial invariance and introduce a generic optimization scheme . Projecting the scheme update on the unit hypersphere leads to the formal definitions of effective learning rate and learning direction . This geometric perspective leads to the first explicit expression of the effective learning rate for Adam . The main notations are summarized in Figure 1 . 2.1 RADIAL INVARIANCE . We consider a family of parametric functions φx : Rin → Rout parameterized by a group of radiallyinvariant parameters x ∈ Rdr { 0 } , i.e. , ∀ρ > 0 , φρx =φx ( possible other parameters of φx are omitted for clarity ) , a dataset D ⊂ Rin ×Rout , a loss function ` : Rout ×Rout → R and a training loss function L : Rd → R defined as : L ( x ) def= 1 |D| ∑ ( s , t ) ∈D ` ( φx ( s ) , t ) . ( 1 ) It verifies : ∀ρ > 0 , L ( ρx ) = L ( x ) . In the context of NNs , the group of radially-invariant parameters x can be the parameters of a single neuron in a linear layer or the parameters of a whole filter in a convolutional layer , followed by BN ( see Appendix A for details , and Appendix B for the application to other normalization schemes such as InstanceNorm ( Ulyanov et al. , 2016 ) , LayerNorm ( Ba et al. , 2016 ) or GroupNorm ( Wu & He , 2018 ) ) . The quotient of the parameter space by the equivalence relation associated to radial invariance is topologically equivalent to a sphere . We consider here the L2 sphere Sd−1 = { u ∈ Rd/‖u‖2 = 1 } whose canonical metric corresponds to angles : dS ( u1 , u2 ) = arccos ( 〈u1 , u2〉 ) . This choice of metric is relevant to study NNs since filters in CNNs or neurons in MLPs are applied through scalar product to input data . Besides , normalization in BN layers is also performed using the L2 norm . Our framework relies on the decomposition of vectors into radial and tangential components . During optimization , we write the radially-invariant parameters at step k ≥ 0 as xk = rkuk where rk = ‖xk‖ and uk = xk/‖xk‖ . For any quantity qk ∈ Rd at step k , we write q⊥k = qk−〈qk , uk〉uk its tangential component relatively to the current direction uk . The following lemma states that the gradient of a radially-invariant loss function is tangential and −1 homogeneous : Lemma 1 ( Gradient of a function with radial invariance ) . If L : Rd → R is radially invariant and almost everywhere differentiable , then , for all ρ > 0 and all x ∈ Rd where L is differentiable : 〈∇L ( x ) , x〉 = 0 and ∇L ( x ) = ρ∇L ( ρx ) . ( 2 ) 2.2 GENERIC OPTIMIZATION SCHEME . There is a large body of literature on optimization schemes ( Sutskever et al. , 2013 ; Duchi et al. , 2011 ; Tieleman & Hinton , 2012 ; Kingma & Ba , 2015 ; Loshchilov & Hutter , 2019 ) . We focus here on two of the most popular ones , namely SGD and Adam ( Kingma & Ba , 2015 ) . Yet , to establish general results that may apply to a variety of other schemes , we introduce here a generic optimization update : xk+1 = xk − ηkak bk , ( 3 ) ak = βak−1 +∇L ( xk ) + λxk , ( 4 ) where xk ∈ Rd is the group of radially-invariant parameters at iteration k , L is the group ’ s loss estimated on a batch of input data , ak ∈ Rd is a momentum , bk ∈ Rd is a division vector that can depend on the trajectory ( xi , ∇L ( xi ) ) i∈J0 , kK , ηk ∈ R is the scheduled trajectory-independent learning rate , denotes the Hadamard element-wise division , β is the momentum parameter , and λ is the L2-regularization parameter . We show how it encompasses several known optimization schemes . Stochastic gradient descent ( SGD ) has proven to be an effective optimization method in deep learning . It can include L2 regularization ( also called weight decay ) and momentum . Its updates are : xk+1 = xk − ηkmk , ( 5 ) mk = βmk−1 +∇L ( xk ) + λxk , ( 6 ) where mk is the momentum , β is the momentum parameter , and λ is the L2-regularization parameter . It corresponds to our generic scheme ( Eqs . 3-4 ) with ak = mk and bk = [ 1 · · · 1 ] > . Adam is likely the most common adaptive scheme for NNs . Its updates are : xk+1 = xk − ηk mk 1− βk+11 √ vk 1− βk+12 + , ( 7 ) mk = β1mk−1+ ( 1− β1 ) ( ∇L ( xk ) + λxk ) , vk = β2vk−1 + ( 1− β2 ) ( ∇L ( xk ) + λxk ) 2 , ( 8 ) where mk is the momentum with parameter β1 , vk is the second-order moment with parameter β2 , and prevents division by zero . ( Here and in the following , the square and the square root of a vector are to be understood as element-wise . ) It corresponds to our generic scheme ( Eqs . 3-4 ) with β=β1 and : ak = mk 1− β1 , bk = 1− βk+11 1− β1 √ vk 1− βk+12 + . ( 9 ) 2.3 IMAGE OPTIMIZATION ON THE HYPERSPHERE . The radial invariance implies that the radial part of the parameter update x does not change the function φx encoded by the model , nor does it change the loss L ( x ) . The goal of training is to find the best possible function encodable by the network . Due to radial invariance , the parameter space projected on the unit hypersphere is topologically closer to the functional space of the network than the full parameter space . It hints that looking at optimization behaviour on the unit hypersphere might be interesting . Thus , we need to separate the quantities that can ( tangential part ) and can not ( radial part ) change the model function . Theorem 2 formulates the spherical decomposition ( Eqs . 3-4 ) in simple terms . It relates the update of radially-invariant parameters in the parameter space Rd and their update on Sd−1 through an exponential map . Theorem 2 ( Image step on Sd−1 ) . The update of a group of radially-invariant parameters xk at step k corresponds to an update of its projection uk on Sd−1 through an exponential map at uk with velocity ηekc ⊥ k , at order 3 : uk+1 = Expuk ( − [ 1 +O ( ( ηek‖c⊥k ‖ ) 2 ) ] ηekc ⊥ k ) , ( 10 ) where Expuk is the exponential map on Sd−1 , and with ck def = rkak bk d−1/2‖bk‖ , ηek def = ηk r2kd −1/2‖bk‖ ( 1− ηk〈ck , uk〉 r2kd −1/2‖bk‖ ) −1 . ( 11 ) More precisely : uk+1 = uk − ηekc⊥k√ 1 + ( ηek‖c⊥k ‖ ) 2 . ( 12 ) The proof is given in Appendix C.1.1 and the theorem is illustrated in the case of SGD in Figure 1 . Note that with typical values in CNN training we have 1− ηk〈ck , uk〉 r2kd −1/2‖bk‖ > 0 , which is a property needed for the proof . Another hypothesis is that steps on the hypersphere are shorter than π . These hypotheses are discussed and empirically verified in Appendix C.1.2 .

Summary: This work studies optimization dynamics for neural network models that are scaling invariant with respect to parameters. A general formulation of optimization algorithms is considered, covering many widely used algorithms like SGD and Adam. The projected dynamics (to the unit sphere) is studied, and the effective learning rate and update direction on the unit sphere are derived. Focusing on the projected dynamics, the equivalence is built between SGD and a type of "Adam". Then, different factors in the Adam dynamics that can potentially influence the optimization performance are identified, and empirically studied.