Batch Normalization Orthogonalizes Representations in Deep Random Networks

1 Introduction . Batch Normalization ( BN ) ( Ioffe and Szegedy , 2015 ) enhances training across a wide range of deep network architectures and experimental setups ( He et al. , 2016 ; Huang et al. , 2017 ; Silver et al. , 2017 ) . The practical success of BN has inspired research into the underlying mechanism of BN ( Santurkar et al. , 2018 ; Karakida et al. , 2019 ; Arora et al. , 2019b ; Bjorck et al. , 2018 ) . BN influences first-order optimization methods by avoiding the rank collapse in deep representation ( Daneshmand et al. , 2020 ) , direction-length decoupling of optimization ( Kohler et al. , 2018 ) , influencing the convergence of the steepest descent ( Arora et al. , 2019b ; Bjorck et al. , 2018 ) , and smoothing the optimization objective function ( Santurkar et al. , 2018 ; Karakida et al. , 2019 ) . However , the benefits of BN go beyond its critical role in optimization . For example , Frankle et al . ( 2021 ) shows that BN networks with random weights also achieve surprisingly high performance after only minor adjustments of their weights . This striking result motivates us to study the representational power of random networks with BN . We study hidden representations across layers of a laboratory random BN with linear activations . Consider a batch of samples passing through consecutive BN and linear layers with Gaussian weights . The representations of these samples are perturbed by each random linear transformation , followed by a non-linear BN . At first glance , the deep representations appear unpredictable 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . after many stochastic and non-linear transformations . Yet , we show that these transformations orthogonalize the representations . To prove this statement , we introduce the notion of “ orthogonality gap ” , defined in Section 3 , to quantify the deviation of representations from a perfectly orthogonal representation . Then , we prove that the orthogonality gap decays exponentially with the network depth and stabilizes around a term inversely related to the network width . More precisely , we prove E [ orthogonality gap ] = O ( ( 1− α ) depth + batchsize α √ width ) holds for α > 0 that is an absolute constant under a mild assumption . In probability theoretic terms , we prove stochastic stability of the Markov chain of hidden representations ( Kushner , 1967 ; Kushner and Yin , 2003 ; Khasminskii , 2011 ) . The orthogonality of deep representations allows us to prove that the distribution of the representations after linear layers contracts to a Wasserstein-2 ball around isotropic Gaussian distribution as the network depth grows . Moreover , the radius of the ball is inversely proportional to the network width . Omitting details , we prove the following bound holds : Wasserstein2 ( representations , Gaussian ) 2 = O ( ( 1− α ) depth ( batchsize ) + ( batchsize ) 2 α √ width ) . The above equation shows how depth , width , and batch size , interact with the Gaussian approximation of the representations . Since the established rate is exponential with depth , the distribution of the representations stays in a Wasserstein ball around isotropic Gaussian distribution after a few layers . Thus , BN not only stabilizes the distribution of the representations , which is its main promise ( Ioffe and Szegedy , 2015 ) , but also enforces Gaussian isotropic distribution in deep layers . There is growing interest in bridging the gap between neural networks , as the most successful parametric methods for learning , and Gaussian processes and kernel methods , as well-understood classical models for learning ( Jacot et al. , 2018 ; Matthews et al. , 2018 ; Lee et al. , 2019 ; Bietti and Mairal , 2019 ; Huang et al. , 2014 ) . This link is inspired by studying random neural networks in the asymptotic regime of infinite width . The seminal work by Neal ( 1996 ) sparks that a single-layer network resembles a Gaussian process as its width goes to infinity . However , increasing the depth may significantly shift the distribution of the representations away from Gaussian ( Ioffe and Szegedy , 2015 ) . This distributional shift breaks the link between Gaussian processes and deep neural networks . To ensure Gaussian representations , Matthews et al . ( 2018 ) suggests increasing the width of the network proportional to the network depth . Here , we show that BN ensures Gaussian representations even for deep networks with finite width . This result bridges the link between deep neural networks and Gaussian processes in the regime of finite width . Many studies rely on deep Gaussian representations in an infinite width setting ( Yang et al. , 2019 ; Schoenholz et al. , 2017 ; Pennington et al. , 2018 ; Klambauer et al. , 2017 ; De Palma et al. , 2019 ) . Our non-asymptotic Gaussian approximation can be incorporated into their analysis to extend these results to the regime of finite width . Since training starts from random networks , representations in these networks directly influence training . Hence , recent theoretical studies has investigated the interplay between initial hidden representations and training ( Daneshmand et al. , 2020 ; Bjorck et al. , 2018 ; Frankle et al. , 2021 ; Schoenholz et al. , 2017 ; Saxe et al. , 2014 ; Bahri et al. , 2020 ) . In particular , it is known that hidden representations in random networks without BN become correlated as the network grows in depth , thereby drastically slowing training ( Daneshmand et al. , 2020 ; He et al. , 2016 ; Bjorck et al. , 2018 ; Saxe et al. , 2014 ) . On the contrary , we prove that deep representations in networks with BN are almost orthogonal . We experimentally validate that initial orthogonal representations can save training time that would otherwise be needed to orthogonalize them . By proposing a novel initialization scheme , we ensure the orthogonality of hidden representations . Such an initialization effectively avoids the training slowdown with depth for vanilla networks , with no need for BN . This observation further motivates studying the inner workings of BN to replace or improve it for deep learning . Theoretically , we made the following contributions : 1 . For MLPs with batch normalization , linear activation , and Gaussian weights , we prove that representations across layers become increasingly orthogonal up to a constant inversely proportional to the network width . 2 . Leveraging the orthogonality , we prove that the distribution of the representations contracts to a Wasserstein ball around a Gaussian distribution as the depth grows . Up to the best of our knowledge , this is the first non-asymptotic Gaussian approximation for deep neural networks with finite width . Experimentally , we made the following contribution1 : 3 . Inspired by our theoretical understanding , we propose a novel weight initialization for standard neural networks that ensure orthogonal representations without BN . Experimentally , we show that this initialization effectively avoids training slowdown with depth in the absence of BN . 2 Preliminaries . 2.1 Notations . Akin to Daneshmand et al . ( 2020 ) , we focus on a Multi-Layer Perceptron ( MLP ) with batch normalization and linear activation . Theoretical studies of linear networks is a growing research area ( Saxe et al. , 2014 ; Daneshmand et al. , 2020 ; Bartlett et al. , 2019 ; Arora et al. , 2019a ) . When weights are initialized randomly , linear and non-linear networks share similar properties such as the rank collapse issue studied in ( Daneshmand et al. , 2020 ) . For ease of analysis , we assume activations are linear . Yet , we will argue that our findings extend to non-linear in Appendix G. We use n to denote batch size , and d to denote the width across all layers , which we further assume is larger than the batch size d ≥ n. Let H ` ∈ Rd×n denote representations for n samples in layer ` , with H0 ∈ Rd×n corresponding to n input samples in the batch with d features . Successive representations are connected by Gaussian weight matrices W ` ∼ N ( 0 , Id/d ) as H ` +1 = 1√ d BN ( W ` H ` ) , BN ( M ) = diag ( MM > ) − 1/2M , ( 1 ) where diag ( . ) zeros out off-diagonal elements of its input matrix , and the scaling factor 1/ √ d ensures that all matrices { Hk } k have unit Frobenius norm ( see Appendix 2 ) . The BN function in Eq . ( 1 ) differs slightly from the commonly used definition for BN as the mean correction is omitted . However , Daneshmand et al . ( 2020 ) shows this difference does not change the network properties qualitatively . Appendix G further discusses about the role of the mean correction term in our analysis . Readers can find all the notations in Appendix Table 2 . 2.2 The linear independence of hidden representations . Daneshmand et al . ( 2020 ) observe that if inputs are linearly independent , then their hidden representations remain linearly independent in all layers as long as d = Ω ( n2 ) . Under technical assumptions , Daneshmand et al . ( 2020 ) establishes a lower-bound on the average of the rank of hidden representations over infinite layers . Based on this study , we assume that the linear independence holds and build our analysis upon that . This avoids restating the technical assumptions of Daneshmand et al . ( 2020 ) and also further technical refinements of their theorems . The next assumption presents the formal statement of the linear independence property . AssumptionA1 ( α , ` ) . There exists an absolute positive constant α such that the minimum singular value of Hk is greater than ( or equal to ) α for all k = 1 , . . . , ` . The linear independence of the representations is a shared property across all layers . However , the representations constantly change when passing through random layers . In this paper , we mathematically characterize the dynamics of the representations across layers . 3 Orthogonality of deep representations . 3.1 A warm-up observation . To illustrate the difference between BN and vanilla networks , we compare hidden representations of two input samples across the layers of these networks . Figure 1 plots the absolute value of 1Implementations are available at https : //github.com/hadidaneshmand/batchnorm21.git cosine similarity of these samples across layers . This plot shows a stark contrast between vanilla and BN networks : While representations become increasingly orthogonal across layers of a BN network , they become increasingly aligned in a vanilla network . More specifically , we observe that BN is able to orthognalized almost aligned representations ; while , the vanilla network provides almost align representation of two orthogonal samples in deep layers . While this behaviour has been theoretically studied for vanilla networks ( Daneshmand et al. , 2020 ; Bjorck et al. , 2018 ; Saxe et al. , 2014 ) , up to the best of our knowledge , it is not theoretically analyzed for BN networks for networks with finite width regime . In the following section , we formalize and prove this orthogonalizing property for BN networks with finite widths . 3.2 Theoretical analysis . The notion of orthogonality gap plays a central role in our analysis . Given the hidden representation H ∈ Rd×n , matrix H > H constitutes inner products between representations of different samples . Note that H > H ∈ Rn×n is different form the covariance matrix HH > /n ∈ Rd×d . The orthogonality gap of H is defined as the deviations of H > H from identity matrix , after proper scaling . More precisely , define V : Rd×n \ 0→ R+ as V ( H ) : = ∥∥∥ ( 1‖H‖2F ) H > H − ( 1 ‖In‖2F ) In ∥∥∥ F . ( 2 ) The following theorem establishes a bound on the orthogonality of representation in layer ` . Theorem 1 . Under Assumption A1 ( α , ` ) , the following holds : E [ V ( H ` +1 ) ] ≤ 2 ( 1− 2 3 α ) ` + 3n α √ d . ( 3 ) Assumption A1 is studied by Daneshmand et al . ( 2020 ) who note that A1 ( α , ∞ ) holds as long as d = Ω ( n2 ) . If A1 does not hold , one can still prove that there is a function of representations that decays with depth up to a constant ( see Appendix C ) . The above result implies that BN is an approximation algorithm for orthogonalizing the hidden representations . If we replace diag ( M ) −1 by ( M ) −1 in BN formula , in Eq . ( 1 ) , then all the hidden representation will become exactly orthogonal . However , computing the inverse of a non-diagonal d × d matrix is computationally expensive , which must repeat for all layers throughout training , and differentiation must propagate back through this inversion . The diagonal approximation in BN significantly reduces the computational complexity of the matrix inversion . Since the orthogonality gap decays at an exponential rate with depth , the approximate orthogonality is met after only a few layers . Interestingly , this yields a desirable cost-accuracy trade-off : For a larger width , the orthogonality is more accurate , due to term 1/ √ d in the orthogonality gap , and also the computational gain is more significant . From a different angle , Theorem 1 proves the stochastic stability of the Markov chain of hidden representations . In expectation , stochastic processes may obey an inherent Lyapunov-type of stability ( Kushner , 1967 ; Kushner and Yin , 2003 ; Khasminskii , 2011 ) . One can analyze the mixing and hitting times of Markov chains based on the stochastic stability of Markov chains , ( Kemeny and Snell , 1976 ; Eberle , 2009 ) . In our analysis , the orthogonality gap is a Lyapunov function characterizing the stability of the chain of hidden representations . This stability opens the door to more theoretical analysis of this chain , such as studying mixing and hitting times . For , example the stability can be used for mixing analysis of the chain { H1 , . . . , H ` , . . . } . Let π denote the stationary distribution of this chain . Under a particular stochastic stability condition called geometric drift condition , |E [ ϕ ( H ` ) ] − EH∼π [ ϕ ( H ) ] | ≤ α ` |E [ ϕ ( H0 ) ] − EH∼π [ ϕ ( H ) ] | holds for α ∈ ( 0 , 1 ) and measurable function ϕ : Rd×n → R ( see Theorem 3.6 of ( Hairer , 2010 ) ) . The drift condition holds if there exists a Lyapunov function L : Rd×n → [ 0 , 1 ] and constant K ≥ 0 such that the following holds : E [ L ( H ` +1 ) |H ` ] ≤ γL ( H ` ) +K . We prove that the above condition holds in Appendix Eq ( 12 ) . However , such theoretical analyses may not be of interest to the machine learning community , hence we focus on the implications of these results for understanding the underpinnings of neural networks . It is helpful to compare the orthogonality gap in BN networks to studies on vanilla networks ( Bjorck et al. , 2018 ; Daneshmand et al. , 2020 ; Saxe et al. , 2014 ) . The function implemented by a vanilla linear network is a linear transformation as the product of random weight matrices . The spectral properties of the product of i.i.d . random matrices are the subject of extensive studies in probability theory ( Bougerol et al. , 2012 ) . Invoking these results , one can check that the orthogonality gap of the hidden representations in vanilla networks rapidly increases since the rank of representations converges to a one matrix as the depth grows ( Daneshmand et al. , 2020 ) . We refer readers to Appendix D for more details about vanilla networks . Remarkably , Agarwal et al . ( 2021 ) proves if activation functions obey a self-normalization property , then a specific kernel of hidden representations becomes well-condition as the network depth grows . However , it is not clear how to impose the self-normalization property . Here , we establish the whitening for an explicitly defined normalization used in practice .

Summary: This paper studies the effect of batch normalization on representations at random initialization.  They formally show that representations become increasingly orthogonal with layer depth up to a term inversely proportional to layer width. Using this, they prove that the distribution of the output contracts to a Wassertein ball around a Gaussian distribution and demonstrate empirically that orthogonal outputs save training time.  Lastly, they propose an initialization that ensures orthogonality even without batch normalization, which they empirically test.