On the Convergence and Robustness of Batch Normalization

1 INTRODUCTION . Batch normalization ( Ioffe & Szegedy , 2015 ) ( BN ) is one of the most important techniques for training deep neural networks and has proven extremely effective in avoiding gradient blowups during back-propagation and speeding up convergence . In its original introduction ( Ioffe & Szegedy , 2015 ) , the desirable effects of BN are attributed to the so-called “ reduction of covariate shift ” . However , it is unclear what this statement means in precise mathematical terms . To date , there lacks a comprehensive theoretical analysis of the effect of batch normalization . In this paper , we study the convergence and stability of gradient descent with batch normalization ( BNGD ) via a modeling approach . More concretely , we consider a simplified supervised learning problem : ordinary least squares regression , and analyze precisely the effect of BNGD when applied to this problem . Much akin to the mathematical modeling of physical processes , the least-squares problem serves as an idealized “ model ” of the effect of BN for general supervised learning tasks . A key reason for this choice is that the dynamics of GD without BN ( hereafter called GD for simplicity ) in least-squares regression is completely understood , thus allowing us to isolate and contrast the additional effects of batch normalization . The modeling approach proceeds in the following steps . First , we derive precise mathematical results on the convergence and stability of BNGD applied to the least-squares problem . In particular , we show that BNGD converges for any constant learning rate ε ∈ ( 0 , 1 ] , regardless of the conditioning of the regression problem . This is in stark contrast with GD , where the condition number of the problem adversely affect stability and convergence . Many insights can be distilled from the analysis of the OLS model . For instance , we may attribute the stability of BNGD to an interesting scaling law governing ε and the initial condition ; This scaling law is not present in GD . The preceding analysis also implies that if we are allowed to use different learning rates for the BN rescaling variables ( εa ) and the remaining trainable variables ( ε ) , we may conclude that BNGD on our model converges for any ε > 0 as long as εa ∈ ( 0 , 1 ] . Furthermore , we discover an acceleration effect of BNGD and moreover , there exist regions of ε such that the performance of BNGD is insensitive to changes in ε , which help to explain the robustness of BNGD to the choice of learning rates . We reiterate that contrary to many previous works , all the preceding statements are precise mathematical results that we derive for our simplified model . The last step in our modeling approach is also the most important : we need to demonstrate that these insights are not specific features of our idealized model . Indeed , they should be true characteristics , at least in an approximate sense , of BNGD for general supervised learning problems . We do this by numerically investigating the convergence , stability and scaling behaviors of BNGD on various datasets and model architectures . We find that the key insights derived from our idealized analysis indeed correspond to practical scenarios . 1.1 RELATED WORK . Batch normalization was originally introduced in ( Ioffe & Szegedy , 2015 ) and subsequently studied in further detail in ( Ioffe , 2017 ) . Since its introduction , it has become an important practical tool to improve stability and efficiency of training deep neural networks ( He et al. , 2016 ; Bottou et al. , 2018 ) . Initial heuristic arguments attribute the desirable features of BN to concepts such as “ covariate shift ” , which lacks a concrete mathematical interpretation and alternative explanations have been given ( Santurkar et al. , 2018 ) . Recent theoretical studies of BN includes ( Ma & Klabjan , 2017 ) , where the authors proposed a variant of BN , the diminishing batch normalization ( DBN ) algorithm and analyzed the convergence of the DBN algorithm , showing that it converges to a stationary point of the loss function . More recently , ( Bjorck et al. , 2018 ) demonstrated that the higher learning rates of batch normalization induce a regularizing effect . Another related work is ( Kohler et al. , 2018 ) , where the authors also considered the convergence properties of BNGD on linear networks ( similar to the least-squares problem ) , as well as other special problems , such as learning halfspaces and extensions . In the OLS case , the authors showed that for a particularly adaptive choice of dynamic learning rate schedule , which can be seen as a fixed effective step size in our terminology ( see equation ( 11 ) and the discussion that immediately follows ) , BNGD converges linearly if λmax is known . Moreover , the analysis also requires setting the rescaling parameter a every step to satisfy a stationarity condition , instead of simply performing gradient descent on a , as is done in the original BNGD . The present research differs from these previous analysis in an important way - we study the BNGD algorithm itself , and not a special variant . More specifically , we consider constant learning rates ( without knowledge of properties of the OLS loss function ) and we perform gradient descent on rescaling parameters . We prove that the convergence occurs for even in this case ( and in fact , for arbitrarily large learning rates for ε , as long as 0 < εa ≤ 1 ) . This poses more challenges in the analysis and contrasts our work with previous analysis on modified versions of BNGD . This is an important distinction ; While a decaying or dynamic learning rate is sometimes used in practice , in the case of BN it is critical to analyze the non-asymptotic , constant learning rate case , precisely because one of the key practical advantages of BN is that a bigger learning rate can be used than that in GD . Hence , it is desirable , as in the results presented in this work , to perform our analysis in this regime . Finally , through the lens of the least-squares example , BN can be viewed as a type of overparameterization , where additional parameters , which do not increase model expressivity , are introduced to improve algorithm convergence and stability . In this sense , this is related in effect to the recent analysis of the implicit acceleration effects of over-parameterization on gradient descent ( Arora et al. , 2018 ) . 1.2 ORGANIZATION . Our paper is organized as follows . In Section 2 , we outline the ordinary least squares ( OLS ) problem and present GD and BNGD as alternative means to solve this problem . In Section 3 , we demonstrate and analyze the convergence of the BNGD for the OLS model , and in particular contrast the results with the behavior of GD , which is completely known for this model . We also discuss the important insights to BNGD that these results provide us with . We then validate these findings on more general supervised learning problems in Section 4 . Finally , we conclude in Section 5 . 2 BACKGROUND . Consider the simple linear regression model where x ∈ Rd is a random input column vector and y is the corresponding output variable . Since batch normalization is applied for each feature separately , in order to gain key insights it is sufficient to the case of y ∈ R. A noisy linear relationship is assumed between the dependent variable y and the independent variables x , i.e . y = xTw + noise where w ∈ Rd is the parameters . Denote the following moments : H : = E [ xxT ] , g : = E [ xy ] , c : = E [ y2 ] . ( 1 ) To simplify the analysis , we assume the covariance matrix H of x is positive definite and the mean E [ x ] of x is zero . The eigenvalues of H are denoted as λi ( H ) , i = 1 , 2 , ... d , . Particularly , the maximum and minimum eigenvalue of H is denoted by λmax and λmin respectively . The condition number of H is defined as κ : = λmaxλmin . Note that the positive definiteness of H allows us to define the vector norms ‖.‖H and ‖.‖H−1 by ‖x‖2H = xTHx and ‖x‖2H−1 = x TH−1x respectively . 2.1 ORDINARY LEAST SQUARES . The ordinary least squares ( OLS ) method for estimating the unknown parameters w leads to the following optimization problem min w∈Rd J0 ( w ) : = 1 2Ex , y [ ( y − x Tw ) 2 ] = c2 − g Tw + 12w THw . ( 2 ) The gradient of J0 with respect to w is ∇wJ0 ( w ) = Hw − g , and the unique minimizer is w = u : = H−1g . The gradient descent ( GD ) method ( with step size or learning rate ε ) for solving the optimization problem ( 2 ) is given by the iterating sequence , wk+1 = wk − ε∇wJ0 ( wk ) = ( I − εH ) wk + εg , ( 3 ) which converges if 0 < ε < 2λmax = : εmax , and the convergence rate is determined by the spectral radius ρε : = ρ ( I − εH ) = maxi { |1− ελi ( H ) | } with ‖u− wk+1‖ ≤ ρε‖u− wk‖ . ( 4 ) It is well known ( for example see Chapter 4 of ( Saad , 2003 ) ) that the optimal learning rate is εopt = 2 λmax+λmin , where the convergence estimate is related to the condition number κ ( H ) : ‖u− wk+1‖ ≤ κ−1κ+1‖u− wk‖ . ( 5 ) 2.2 BATCH NORMALIZATION . Batch normalization is a feature-wise normalization procedure typically applied to the output , which in this case is simply z = xTw . The normalization transform is defined as follows : NBN ( z ) : = z−E [ z ] √ Var [ z ] = x Tw σ , ( 6 ) where σ : = √ wTHw . After this rescaling , NBN ( z ) will be order 1 , and hence in order to reintroduce the scale ( Ioffe & Szegedy , 2015 ) , we multiply NBN ( z ) with a rescaling parameter a ( Note that the shift parameter can be set zero since E [ wTx|w ] = 0 ) . Hence , we get the BN version of the OLS problem ( 2 ) : min w∈Rd , a∈R J ( a , w ) : = 12Ex , y [ ( y − aNBN ( xTw ) ) 2 ] = c2 − wT g σ a+ 1 2a 2 . ( 7 ) The objective function J ( a , w ) is no longer convex . In fact , it has trivial critical points , { ( a∗ , w∗ ) |a∗ = 0 , w∗T g = 0 } , which are saddle points of J ( a , w ) . We are interested in the nontrivial critical points which satisfy the relations , a∗ = sign ( s ) √ uTHu , w∗ = su , for some s ∈ R \ { 0 } . ( 8 ) It is easy to check that the nontrivial critical points are global minimizers , and the Hessian matrix at each critical point is degenerate . Nevertheless , the saddle points are strict ( Details can be found in Appendix ) , which typically simplifies the analysis of gradient descent on non-convex objectives ( Lee et al. , 2016 ; Panageas & Piliouras , 2017 ) . Consider the gradient descent method to solve the problem ( 7 ) , which we hereafter call batch normalization gradient descent ( BNGD ) . We set the learning rates for a and w to be εa and ε respectively . These may be different , for reasons which will become clear in the subsequent analysis . We thus have the following discrete-time dynamical system ak+1 = ak + εa ( wTk g σk − ak ) , ( 9 ) wk+1 = wk + ε ak σk ( g − w T k g σ2k Hwk ) . ( 10 ) We now begin a concrete mathematical analysis of the above iteration sequence .

Summary: This paper provides a theoretical analysis for batch normalization with gradient descent (GDBN) under a simplified scenario, i.e., solving an ordinary least squares problem. The analysis shows that GDBN converges to a stationary point when the learning rate is less than or equal to 1, regardless of the condition number of the problem. Some practical experiments are carried out to justify their theoretical insights. The paper is in general easy to follow.