Demystifying Batch Normalization in ReLU Networks: Equivalent Convex Optimization Models and Implicit Regularization

1 INTRODUCTION . Deep neural networks have achieved dramatic progress in the past decade . This dramatic progress largely hinged on improvements in terms of optimization techniques . One of the most prominent recent optimization techniques is Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) . BN is an operation that can be introduced in between layers to normalize the output of each layer and it has been shown to be extremely effective in stabilizing and accelerating training of deep neural networks . Hence , it became standard in numerous state-of-the-art architectures , e.g. , ResNets ( He et al. , 2016 ) . Despite its empirical success , it is still theoretically elusive why BN is extremely effective for training deep neural networks . Therefore , we investigate the mechanisms behind the success of BN through convex duality . 1.1 RELATED WORK . Batch Normalization : One line of research has focused particularly on alternatives to BN , such as Layer Normalization ( Ba et al. , 2016 ) , Instance Normalization ( Ulyanov et al. , 2016 ) , Weight Normalization ( Salimans & Kingma , 2016 ) , and Group Normalization ( Wu & He , 2018 ) . Although these techniques achieved performance competitive with BN , they do not provide any theoretical insight about its empirical success . Another line of research studied the effects of BN on neural network training and identified several benefits . For example , Im et al . ( 2016 ) showed that training deep networks with BN reduces dependence on the parameter initialization . Wei et al . ( 2019 ) analyzed BN via mean-field theory to quantify its impact on the geometry of the optimization landscape . They reported that BN flattens the optimization landscape so that it enables the use of larger learning rates . In addition , Bjorck et al . ( 2018 ) ; Santurkar et al . ( 2018 ) ; Arora et al . ( 2018 ) showed that networks trained with BN achieve faster convergence and generalize better . Furthermore , Daneshmand et al . ( 2020 ) proved that BN avoids rank collapse so that gradient-based algorithms , e.g. , Stochastic Gradient Descent ( SGD ) , are able to effectively train deep networks . Even though these studies are important to understand the benefits of BN , they fail to provide a theoretically complete characterization for training deep networks with BN . Convex Neural Networks : Recently , a series of papers ( Pilanci & Ergen , 2020 ; Ergen & Pilanci , 2021 ; Sahiner et al. , 2021a ; b ) studied ReLU networks through the lens of convex optimization theory . Particularly , Pilanci & Ergen ( 2020 ) introduced exact convex representations for two-layer ReLU networks , which can be trained in polynomial-time via standard convex solvers . However , this work is restricted to two-layer fully connected networks with scalar outputs . Later on , Ergen & Pilanci ( 2021 ) first extended this approach to two-layer scalar output Convolutional Neural Networks ( CNNs ) with average and max pooling and provided further improvements on the training complexity . These results were extended to two-layer fully convolutional networks and two-layer networks with vector outputs ( Sahiner et al. , 2021a ; b ) . However , these convex approaches are restricted to two-layer ReLU networks without BN , thus , do not reflect the exact training framework in practice , i.e. , regularized deep ReLU networks with BN . Unlike these prior works , we not only analyze training problems with batch normalization but also extend the analysis to deeper architectures in Section 4 . Furthermore , we prove that the optimal layer weight can be obtained as closed-form solutions in the high-dimensional and/or overparameterized regimes ( see Theorem 2.1 , 2.4 , 4.1 , 4.2 , G.2 ) . 1.2 OUR CONTRIBUTIONS . • We introduce an exact convex framework to explicitly characterize optimal solutions to ReLU network training problems with weight-decay regularization and BN . Thus , we obtain closed-form solutions for the optimal layer weights in the high-dimensional and overparameterized regime . • We prove that regularized ReLU network training problems with BN can be equivalently stated as a finite-dimensional convex problem . As a corollary , we also show that the equivalent convex problems involve whitened data matrices unlike the original non-convex training problem . Hence , using convex optimization , we reveal an implicit whitening effect introduced by BN . • We demonstrate that GD applied to BN networks provides an implicit regularization effect by learning high singular value directions of the training data more aggressively , whereas this regularization is absent for GD applied to the equivalent whitened data formulation . We propose techniques to explicitly regularize BN networks to capture this implicit regularization effect . • Unlike previous studies , our derivations extend to deep ReLU networks with BN , CNNs , BN after ReLU1 , vector output networks , and arbitrary convex loss functions . 1.3 PRELIMINARIES . Notation : We denote matrices and vectors as uppercase and lowercase bold letters , respectively , where a subscript indicates a certain element or column . We use 0 ( or 1 ) to denote a vector or matrix of zeros ( or ones ) , where the sizes are appropriately chosen depending on the context . We also use In to denote the identity matrix of size n. To represent Euclidean , Frobenius , and nuclear norms , we use ‖ · ‖2 , ‖ · ‖F , and ‖ · ‖∗ , respectively . Lastly , we denote the elementwise 0-1 valued indicator function and ReLU activation as 1 [ x ≥ 0 ] and ( x ) + = max { x , 0 } , respectively . We2 consider an L-layer ReLU network with layer weights W ( l ) ∈ Rml−1×ml , where m0 = d and mL = C are the input and output dimensions , respectively . Given a training data X ∈ Rn×d and a label matrix Y ∈ Rn×C , we particularly focus on the following regularized training problem min θ∈Θ L ( fθ , L ( X ) , Y ) + βR ( θ ) , ( 1 ) where we compactly represent the parameters as θ : = { W ( l ) , γl , αl } Ll=1 and the corresponding parameter space as Θ : = { { W ( l ) , γ ( l ) , α ( l ) } Ll=1 : W ( l ) ∈ Rml−1×ml , γl ∈ Rml , αl ∈ Rml , ∀l ∈ [ L ] } . We note that γl and αl are the parameters of the BN operator , for which we discuss the details below . In addition , L ( · , · ) is an arbitrary convex loss function , including squared , hinge , and cross entropy loss , andR ( · ) is the regularization function for the layer weights with the tuning parameter β > 0 . We also compactly define the network output as fθ , L ( X ) : = A ( L−1 ) W ( L ) , where A ( l ) : = ( BNγ , α ( A ( l−1 ) W ( l ) ) ) + we denote the lth layer activations as A ( l ) ∈ Rn×ml , and A ( 0 ) = X . Here , BNγ , α ( · ) represents the BN operation introduced in Ioffe & Szegedy ( 2015 ) and applies matrices columnwise . 1Presented in Appendix G. 2All the proofs and some extensions are presented in Appendix . Remark 1.1 . Note that above we use BN before ReLU activations , which is common practice and consistent with the way introduced in Ioffe & Szegedy ( 2015 ) . However , BN can be placed after ReLU as well , e.g. , Chen et al . ( 2019 ) , and thus in Section G of Appendix , we will also consider architectures where BN layers are placed after ReLU . For a layer with weight matrix W ( l ) ∈ Rml−1×ml and arbitrary batch of activations denoted as A ( l−1 ) b ∈ Rs×ml−1 , BN applies to each column j independently as follows BNγ , α ( A ( l−1 ) b w ( l ) j ) : = ( Is − 1s11 T ) A ( l−1 ) b w ( l ) j ‖ ( Is − 1s11T ) A ( l−1 ) b w ( l ) j ‖2 γ ( l ) j + α ( l ) j 1√ n , ( 2 ) where γ ( l ) and α ( l ) are trainable parameters that scale and shift the normalized value . In this work , we focus on the full-batch case , i.e , A ( l−1 ) b = A ( l−1 ) ∈ Rn×ml . This corresponds to training the network with GD as opposed to mini-batch SGD . We note that our empirical findings with GD indicate identical if not better performance compared to the mini-batch case , which is also consistent with the previous studies Lian & Liu ( 2019 ) ; Summers & Dinneen ( 2020 ) . Throughout the paper , we consider a regression framework with squared loss and standard weightdecay regularization . Extensions to general convex loss functions are presented in Appendix B.1 . Moreover , below , we first focus on scalar outputs i.e . C = 1 , and then extend it to vector outputs . 1.4 OVERVIEW OF OUR RESULTS . Here , we provide an overview of our main results . To simplify the notation , we consider L-layer ReLU networks with scalar outputs , i.e. , mL = C = 1 thus the label vector is y ∈ Rn , and extend the analysis to vector output networks with the label matrix Y ∈ Rn×C in the next sections . The regularized training problems for an L-layer network with scalar output and BN is given by p∗L : = min θ∈Θ 1 2 ‖fθ , L ( X ) − y‖22 + β 2 L∑ l=1 ( ∥∥∥γ ( l ) ∥∥∥2 2 + ∥∥∥α ( l ) ∥∥∥2 2 + ∥∥∥W ( l ) ∥∥∥2 F ) , ( 3 ) where we use γ ( L ) = α ( L ) = 0 as dummy variables for notational simplicity . Lemma 1.1 . The problem in ( 3 ) is equivalent to the following optimization problem min θ∈Θs 1 2 ‖fθ , L ( X ) − y‖22 + β ∥∥∥w ( L ) ∥∥∥ 1 , ( 4 ) where Θs : = { θ ∈ Θ : γ ( L−1 ) j 2 + α ( L−1 ) j 2 = 1 , ∀j ∈ [ mL−1 ] } . Using the equivalence in Lemma 1.1 , we now take dual of ( 4 ) with respect to the output layer weights w ( L ) to obtain 3 p∗L ≥ d∗L : = max v −1 2 ‖v − y‖22 + 1 2 ‖y‖22 s.t . max θ∈Θs ∣∣∣∣v > ( BNγ , α ( A ( L−2 ) w ( L−1 ) ) ) + ∣∣∣∣ ≤ β . ( 5 ) Since the original formulation in ( 3 ) is a non-convex optimization problem , any solution v in the dual domain yields a lower bound for the primal problem , i.e. , p∗L ≥ d∗L . In this paper , we first show that strong duality holds in this case , i.e. , p∗L = d ∗ L , and then derive an exact equivalent convex formulation for the non-convex problem ( 3 ) . Furthermore , we even obtain closed-form solutions for the layer weights in some cases so that there is no need to train a network in an end-to-end manner .

Summary: The paper studies batch normalization in deep neural networks. For a two-layer network with scalar output and batch normalization, a dual of the problem is derived. It is then shown that in the high-dimensional regime, the dual can be further simplified so that an optimal solution can be computed in closed form. In the general case, the concept of hyperplane arrangements can be used to formulate an equivalent finite-dimensional convex program which can be solved in polynomial time.   The analysis is then extended for vector output networks and different architectures including L-layer neural networks and CNNs. It is then shown that the convex problems tend to fit low singular value directions which will lead to poorer generalization. As a remedy, the authors propose a truncated variant of the problem by obtaining a low-rank approximation of the data.  In the experimental section, it is shown on the CIFAR-100 dataset that the closed form solution leads to superior results compared to the solution obtained by gradient descent. Moreover it is shown that the truncated variant of the approach is needed to achieve a good generalization performance.