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 considers deep linear networks with batch normalization. The main contributions are: 1. Under a linear independence condition, it is proved that the orthogonality gap is the sum of a term that decreases exponentially fast with the depth, and a term that is basically number_of_samples / sqrt{width}.  2. Under similar conditions, it is proved that the distribution of the hidden representation is close to the Gaussian distribution, and a Wasserstein distance bound is proved. 3. It is shown empirically that a larger orthogonality gap at initialization makes training slower, and that SGD decreases the orthogonality gap during training. This paper further proposes an "iterative orthogonalization" algorithm, which can enforce orthogonality and accelerate training when the depth is large.