# Batch Normalization Orthogonalizes Representations in Deep Random Networks

Hadi Daneshmand and Amir Joudaki and Francis Bach

###### Abstract.

This paper underlines a subtle property of batch-normalization (BN):
Successive batch normalizations with random linear transformations make hidden
representations increasingly orthogonal across layers of a deep neural
network. We establish a non-asymptotic characterization of the interplay
between depth, width, and the orthogonality of deep representations. More
precisely, under a mild assumption, we prove that the deviation of the
representations from orthogonality rapidly decays with depth up to a term
inversely proportional to the network width. This result has two main
implications: 1) Theoretically, as the depth grows, the distribution of the
representation –after the linear layers– contracts to a Wasserstein-2 ball
around an isotropic Gaussian distribution. Furthermore, the radius of this
Wasserstein ball shrinks with the width of the network. 2) In practice, the
orthogonality of the representations directly influences the performance of
stochastic gradient descent (SGD). When representations are initially aligned,
we observe SGD wastes many iterations to orthogonalize representations before
the classification. Nevertheless, we experimentally show that starting
optimization from orthogonal representations is sufficient to accelerate SGD,
with no need for BN.

INRIA Paris, ETH Zurich, INRIA-ENS-PSL Paris.

##  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 benefits of BN (Santurkar et al.,,
2018; Karakida et al.,, 2019; Arora et al., 2019b, ; Bjorck et al.,, 2018). In
particular, it shown that 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 learning rate 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.,
(2020) 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 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

$\displaystyle{\mathbb{E}}\Bigg{[}\text{orthogonality
gap}\Bigg{]}=\mathcal{O}\left(\left(1-\alpha\right)^{\text{depth}}+\frac{\text{batch
size}}{\alpha\sqrt{\text{width}}}\right)$

holds for $\alpha>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:

$\displaystyle\text{Wasserstein}_{2}(\text{representations},\text{Gaussian})^{2}=\mathcal{O}\left(\left(1-\alpha\right)^{\text{depth}}\left(\text{batch
size}\right)+\frac{\text{batch size}}{\alpha\sqrt{\text{width}}}\right).$

The above equation shows how depth, width, and batch size, interact with the
law 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 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 process 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.,, 2020;
Schoenholz et al.,, 2017; Saxe et al.,, 2014; Bahri et al.,, 2020). In
particular, it is shown that hidden representations in random networks
_without BN_ become correlated as the network grows in depth, thereby
drastically slows 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 ensures the orthogonality of hidden representations;
then, we show that 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 in deep neural networks.

Theoretically, we made the following contributions:

  1. (1)

For MLPs with batch normalization, linear activation, and Gaussian i.i.d.
weights, we prove that representations across layers become increasingly
orthogonal up to a constant inversely proportional to the network width.

  2. (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 contribution111The codes for the
experiments are available at
https://github.com/hadidaneshmand/batchnorm21.git:

  1. (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 initialized
randomly, linear and non-linear networks share similar properties (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 F.

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\geq n$. Let
$H_{\ell}\in{\mathbb{R}}^{d\times n}$ denote $n$ sample representations in
layer $\ell$, with $H_{0}\in{\mathbb{R}}^{d\times n}$ corresponding to $n$
input samples in the batch with $d$ features. Successive representations are
connected by Gaussian weight matrices $W_{\ell}\sim\mathcal{N}(0,I_{d}/d)$ as

(1) $\displaystyle H_{\ell+1}=\frac{1}{\sqrt{d}}BN(W_{\ell}H_{\ell}),\quad
BN(M)=\text{diag}(MM^{\top})^{-\nicefrac{{1}}{{2}}}M,$

where $\text{diag}(M)$ zeros out off-diagonal elements of its input matrix,
and the scaling factor $1/\sqrt{d}$ ensures that all matrices
$\\{H_{k}\\}_{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) showes this
difference does not change the network properties qualitatively. 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=\Omega(n^{2})$. 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.

######  Assumption $\mathcal{A}_{1}(\alpha,\ell)$.

There exists an absolute positive constant $\alpha$ such that the minimum
singular value of $H_{k}$ is greater than (or equal to) $\alpha$ for all
$k=1,\dots,\ell$.

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 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. We
used highly correlated inputs for the BN network and perfectly orthogonal
inputs for the Vanilla network to contrast their differences. While the hidden
representation of vanilla networks has been theoretically studied (Daneshmand
et al.,, 2020; Bjorck et al.,, 2018; Saxe et al.,, 2014), up to the best of
our knowledge, there is no theoretical analysis of BN networks. In the
following section, we formalize and prove this orthogonalizing property for BN
networks.

---
Figure 1. Orthogonality: BN vs. vanilla networks. The horizontal axis shows
the number of layers, and the vertical axis shows the absolute value of cosine
similarity between two samples across the layers ($d=32$). Mean and 95%
confidence intervals of 20 independent runs.

###  3.2. Theoretical analysis

The notion of orthogonality gap plays a central role in our analysis. Given
the hidden representation $H\in{\mathbb{R}}^{d\times n}$, matrix $H^{\top}H$
constitutes inner products between representations of different samples. Note
that $H^{\top}H\in{\mathbb{R}}^{n\times n}$ is different form the covariance
matrix $HH^{\top}/n\in{\mathbb{R}}^{d\times d}$. The orthogonality gap of $H$
is defined as the deviations of $H^{\top}H$ from identity matrix, after proper
scaling. More precisely, define $V:{\mathbb{R}}^{d\times
n}\setminus\mathbf{0}\to{\mathbb{R}}_{+}$ as

(2) $\displaystyle
V(H):=\Big{\|}\left(\frac{1}{\|H\|_{F}^{2}}\right)H^{\top}H-\left(\frac{1}{\|I_{n}\|_{F}^{2}}\right)I_{n}\Big{\|}_{F}.$

The following theorem establishes a bound on the orthogonality of
representation in layer $\ell$.

######  Theorem 1.

Under Assumption Assumption $\mathcal{A}_{1}(\alpha,\ell)$$(\alpha,\ell)$, the
following holds:

(3) $\displaystyle{\mathbb{E}}\left[V(H_{\ell+1})\right]\leq
2\left(1-\frac{2}{3}\alpha\right)^{\ell}+\frac{3n}{\alpha\sqrt{d}}.$

Assumption Assumption $\mathcal{A}_{1}(\alpha,\ell)$ is studied by Daneshmand
et al., (2020) who note that Assumption
$\mathcal{A}_{1}(\alpha,\ell)$$(\alpha,\infty)$ holds as long as
$d=\Omega(n^{2})$. If Assumption $\mathcal{A}_{1}(\alpha,\ell)$ 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
$\text{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\times 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/\sqrt{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. Since these theoretical
properties may not be of interest to the machine learning community, 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 readily check that the orthogonality gap of the hidden
representations in vanilla networks rapidly increases since the
representations converge to a rank one matrix after a proper scaling
(Daneshmand et al.,, 2020).

###  3.3. Experimental validations

Our experiments presented in Fig. 2(a) validate the exponential decay rate of
$V$ with depth. In this plot, we see that $\log(V_{\ell})$ linearly decreases
for $\ell=1,\dots,20$, then it wiggles around a small constant. Our
experiments in Fig. 2(b) suggest that the $\mathcal{O}(1/\sqrt{d})$ dependency
on width is almost tight. Since $V(H_{\ell})$ rapidly converges to a ball, the
average of $V(H_{\ell})$ over layers estimates the radius of this ball. This
plot shows how the average of $V(H_{\ell})$, over 500 layers, changes with the
network width. Since the scale is logarithmic in this plot, slope is close to
$-1/2$ in logarithmic scale, validating the $\mathcal{O}(1/\sqrt{d})$
dependency implied by Theorem 1.

(a) Orthogonality vs. depth

(b) Orthogonality vs. width

Figure 2. Orthogonality gap vs. depth and width. Left: $\log(V(H_{\ell}))$
vertically versus $\ell$ horizontally . Right:
$\log(\frac{1}{500}\sum_{\ell=1}^{500}V(H_{\ell}))$ vertically versus $d$
horizontally. The slop of the right plot is $-0.46$. Mean and 95% confidence
interval of 20 independent runs.

##  4\. Gaussian approximation

###  4.1. Orthogonality yields Gaussian approximation

The established result on the orthogonality gap allows us to show that the
representations after linear layers are approximately Gaussian. If $H_{\ell}$
is orthogonal, the random matrix $W_{\ell}H_{\ell}$ is equal in distribution
to a standard Gaussian matrix due to the invariance of standard Gaussian
distribution under linear orthogonal transformations. We can formalize this
notion of Gaussianity by bounding the Wasserstein-2 distance, denoted by
${\mathcal{W}}_{2}$, between the distribution of $W_{\ell}H_{\ell}$ and
standard Gaussian distribution. The next lemma formally establishes the link
between the orthogonality and the distribution of the representations.

######  Lemma 2.

Given $G\in{\mathbb{R}}^{d\times n}$ with i.i.d. zero-mean $1/d$-variance
Gaussian elements, the following Gaussian approximation holds:

(4)
$\displaystyle{\mathcal{W}}_{2}\left(\mathrm{law}(W_{\ell}H_{\ell}),\mathrm{law}(G/\sqrt{n})\right)^{2}\leq
2n{\mathbb{E}}\left[V(H_{\ell})\right].$

Combining the above result with Theorem 1 yields the result presented in the
next corollary.

######  Corollary 3.

For $G\in{\mathbb{R}}^{d\times n}$ with i.i.d. zero-mean $1/d$-variance
Gaussian elements,

(5)
$\displaystyle{\mathcal{W}}_{2}\left(\mathrm{law}(W_{\ell}H_{\ell}),\mathrm{law}(G/\sqrt{n})\right)^{2}\leq
4n\left(1-\frac{2}{3}\alpha\right)^{\ell}+\frac{6n^{2}}{\alpha\sqrt{d}}$

holds under Assumption Assumption
$\mathcal{A}_{1}(\alpha,\ell)$$(\alpha,\ell)$.

In other words, the distribution of the representations contracts to a
Wasserstein 2 ball around an isotropic Gaussian distribution as the depth
grows. The radius of the Wasserstein 2 ball is at most
$\mathcal{O}(1/\sqrt{\text{width}})$. As noted in the last section, Assumption
$\mathcal{A}_{1}(\alpha,\ell)$ is extensively studied by Daneshmand et al.,
(2020) where it is shown that Assumption
$\mathcal{A}_{1}(\alpha,\ell)$$(\alpha>0,\infty)$ holds as long as
$d=\Omega(n^{2})$.

###  4.2. Deep neural networks as Gaussian processes

Leveraging BN, Corollary 3 establishes the first non-asymptotic Gaussian
approximation for deep random neural networks. For vanilla networks, the
Gaussianity is guaranteed only in the asymptotic regime of infinite width
(Garriga-Alonso et al.,, 2019; Neal,, 1996; Lee et al.,, 2019; Neal,, 1996;
Hazan and Jaakkola,, 2015). Particularly, Matthews et al., (2018) links
vanilla networks to Gaussian processes when their width is infinite and grows
in successive layers, while our Gaussian approximation holds for networks with
finite width across layers. Table 1 briefly compares Gaussian approximation
for vanilla and BN networks.

Network | Width | Depth | Distribution of Outputs
---|---|---|---
Vanilla MLP | infinite | finite | Converges to Gaussian as width $\to\infty$
Vanilla Convnet | infinite | finite | Converges to Gaussian as width $\to\infty$
BN MLP (Cor. 3) | (in)finite | (in)finite | In a $\mathcal{O}(\text{width}^{-\nicefrac{{1}}{{4}}})$-${\mathcal{W}}_{2}$ ball around Gaussian
Table 1. Distribution of representations in random vanilla and BN networks.
For the convolutional network, the width refers to the number of channels.
Results for Vanilla MLPs and Vanilla convolution networks are establish by
Matthews et al., (2018), and Garriga-Alonso et al., (2019), respectively.
Remarkably, Corollary 3 holds for MLP with linear activations.

The link between Gaussian processes and infinite-width neural networks has
inspired several studies to rely on Gaussian representations in deep random
networks (Klambauer et al.,, 2017; De Palma et al.,, 2019; Pennington et al.,,
2018; Schoenholz et al.,, 2017; Yang et al.,, 2019). Assuming the
representations are Gaussian, Klambauer et al., (2017) designed novel
activation functions that improve the optimization performance, De Palma et
al., (2019) studies the sensitivity of random networks, Pennington et al.,
(2018) highlights the spectral universality in random networks, Schoenholz et
al., (2017) studies information propagating through the network layers, and
Yang et al., (2019) studies gradients propagation through the depth. Indeed,
our analysis implies that including BN imposes the Gaussian representations
required for these analyses. Although our result is established for linear
activations, we conjecture a similar result holds for non-linear MLPs (see
Appendix F).

##  5\. The orthogonality and optimization

In the preceding sections, we elaborated on the theoretical properties of BN
networks in controlled settings. In this section, we demonstrate the practical
applications of our findings. In the first part, we focus on the relationship
between depth and orthogonality. Increasing depth drastically slows the
training of neural networks with BN. Furthermore, we observe that as depth
grows, the training slowdown highly correlates with the orthogonality gap.
This observation suggests that SGD needs to orthogonalize deep representations
in order to start classification. This intuition leads us to the following
question: If orthogonalization is a prerequisite for training, can we save
optimization time by starting from orthogonal representations? To test this
experimentally, we devised a weight initialization that guarantees
orthogonality of representations. Surprisingly, even in a network without BN,
our experiments showed that this initialization avoids the training slow down,
affirmatively answering the question.

Throughout the experiments, we use vanilla MLP (without BN) with a width of
800 across all hidden layers, ReLU activation, and used Xavier’s method for
weights intialization (Glorot and Bengio,, 2010). We use SGD with stepsize
$0.01$ and batch size $500$ and for training. The learning task is
classification with cross entropy loss for CIFAR10 dataset (Krizhevsky et
al.,, 2009, MIT license). We use PyTorch (Paszke et al.,, 2019, BSD license)
and Google Colaboratory platform with a single Tesla-P100 GPU with 16GB memory
in all the experiments. The reported orthogonality gap is the average of the
orthogonality gap of representation in the last layer.

###  5.1. Orthogonality correlates with optimization performance

In the first experiment, we show that the orthogonality of representations at
the initialization correlates with optimization speed. For networks with 15,
30, 45, 60, and 75 widths, we register training loss after 30 epochs and
compare it with the initial orthogonality gap. Figure 3(a) shows the training
loss (blue) and the initial orthogonality gap (red) as a function of depth. We
observe that representations are more entangled, i.e., orthogonal, when we
increase depth, coinciding with the training slowdown. Intuitively, the
slowdown is due to the additional time SGD must spend to orthogonalize the
representations before classification. In the second experiment, we validate
this intuitive argument by tracking the orthogonality gap during training.
Figure 3(b) plots the orthogonality gap of output and training loss for a
network with 20 layers. We observe that SGD updates are iteratively
orthogonalizing representations, marked by the reduction in the orthogonality
gap.

(a) gap and loss vs. depth

(b) gap and loss during training

Figure 3. Orthogonality and Optimization Left: the orthogonality gap at
initialization (red, left axis) and the training loss after 30 epochs (blue,
right axis) with depth. Right: the orthogonality gap (red, left axis) and the
training loss in each epoch (blue, right axis). Mean and 95% confidence
interval of 4 independent runs.

###  5.2. Learning with initial orthogonal representations

We have seen that the slowdown in SGD for deeper networks correlates with the
orthogonality gap before training. Here we show that by preemptively
orthogonalizing representations, we avoid the slowdown with depth. While in
MPL with linear activations, hidden representations remain orthogonal simply
by taking orthogonal matrices as weights (Pennington et al.,, 2018; Saxe et
al.,, 2014), the same does not hold for networks with non-linear activations,
such as ReLU. To enforce the orthogonality in the absence of BN, we introduce
a dependency between weights of successive layers that ensures deep
representations remain orthogonal. More specifically, we incorporate the SVD
decomposition of the hidden representation of each layer into the
initialization of the subsequent layer. To emphasis this dependency between
layers and to distinguish it from purely orthogonal weight initialization, we
refer to this as _iterative orthogonalization_.

We take a large batch of samples $n\geq d$, as the input batch for
initialization. Let us assume that weights are initialized up to layer
$W_{\ell-1}$. To initialize $W_{\ell}$, we compute SVD decomposition of the
representations $H_{\ell}=U_{\ell}\Sigma_{\ell}V^{\top}_{\ell}$ where matrices
$U_{\ell}\in{\mathbb{R}}^{d\times d}$ and $V_{\ell}\in{\mathbb{R}}^{n\times
d}$ are orthogonal. Given this decomposition, we initialize $W_{\ell}$ by

(6) $\displaystyle
W_{\ell}=\frac{1}{\|\Sigma_{\ell}^{\nicefrac{{1}}{{2}}}\|_{F}}V_{\ell}^{\prime}\Sigma_{\ell}^{-\nicefrac{{1}}{{2}}}U_{\ell}^{\top},$

where $V^{\prime}_{\ell}\in{\mathbb{R}}^{d\times d}$ is an orthogonal matrix
obtained by slicing $V_{\ell}\in{\mathbb{R}}^{n\times d}$. Notably, the
inverse in the above formula exists when $n$ is sufficiently larger than $d$
222We may inductively assume that $H_{\ell}$ is almost orthogonal by the
choice of $W_{1},\dots,W_{\ell-1}$. Thus, $\Sigma_{\ell}$ is invertible. . It
is easy to check that $V(W_{\ell}H_{\ell})<V(H_{\ell})$ holds for the above
initialization (see Appendix E), similar to BN. By enforcing the
orthogonality, this initialization significantly alleviates the slow down of
training with depth (see Fig. 4), with no need for BN. This initialization is
not limited to MLPs. In Appendix G, we propose a similar SVD-based
initialization for convolutional networks that effective accelerates training
of deep convolutional networks.

Figure 4. Iterative orthogonalization. Horizontal axis: depth. Vertical axis:
the training loss after 30 epochs for Xavier’s initialization (blue), our
initialization (red). Mean and 95% confidence interval of 4 independent runs.

##  6\. Discussion

To recap, we proved the recurrence of random linear transformations and BN
orthogonalizes samples. Our experiments underline practical applications of
this theoretical finding: starting from orthogonal representations effectively
avoids the training slowdown with depth for MLPs. Based on our experimental
observations in Appendix G, we believe that this result extends to standard
convolution networks used in practice. In other words, a proper initialization
ensuring the orthogonality of hidden representations may replace BN in neural
architectures. This future research direction has the potentials to boost the
training of deep neural networks and change benchmarks in deep learning.

Although our theoretical bounds hold for MLPs with linear activations, our
experiments confirm similar orthogonal stability for various neural
architectures. Appendix Figure 5 demonstrates this stability for MLPs with
ReLU activations and a convolutional networks. This plot compares the
evolution of the orthogonality gap, through layers, for BN networks with
vanilla networks. For vanilla networks, we observe that the gap increases with
depth. On the contrary, the gap decreases by adding BN layers and stabilizes
at a term that is constant with regard to depth. Based on our observations, we
conjecture that hidden representations of modern BN networks obey similar
stability. A more formal statement of the conjecture is presented in Appendix
F.

## Acknowledgements

We thank Gideon Dresdner, Vincent Fortuin, and Ragnar Groot Koerkamp for their
helpful comments and discussions. This work was funded in part by the French
government under management of Agence Nationale de la Recherche as part of the
“Investissements d’avenir” program, reference ANR-19-P3IA-0001(PRAIRIE 3IA
Institute). We also acknowledge support from the European Research Council
(grant SEQUOIA 724063).

## References

  * (1) Arora, S., Cohen, N., Golowich, N., and Hu, W. (2019a).  A convergence analysis of gradient descent for deep linear neural networks.  International Conference on Learning Representations.
  * (2) Arora, S., Li, Z., and Lyu, K. (2019b).  Theoretical analysis of auto rate-tuning by batch normalization.  International Conference on Learning Representations.
  * Bahri et al., (2020) Bahri, Y., Kadmon, J., Pennington, J., Schoenholz, S. S., Sohl-Dickstein, J., and Ganguli, S. (2020).  Statistical mechanics of deep learning.  Annual Review of Condensed Matter Physics.
  * Bartlett et al., (2019) Bartlett, P. L., Helmbold, D. P., and Long, P. M. (2019).  Gradient descent with identity initialization efficiently learns positive-definite linear transformations by deep residual networks.  Neural computation.
  * Bietti and Mairal, (2019) Bietti, A. and Mairal, J. (2019).  On the inductive bias of neural tangent kernels.  Advances in Neural Information Processing Systems.
  * Bjorck et al., (2018) Bjorck, J., Gomes, C., Selman, B., and Weinberger, K. Q. (2018).  Understanding batch normalization.  Advances in Neural Information Processing Systems.
  * Bougerol et al., (2012) Bougerol, P. et al. (2012).  Products of random matrices with applications to Schrödinger operators.  Springer Science & Business Media.
  * Daneshmand et al., (2020) Daneshmand, H., Kohler, J., Bach, F., Hofmann, T., and Lucchi, A. (2020).  Batch normalization provably avoids rank collapse for randomly initialised deep networks.  Advances in Neural Information Processing Systems.
  * De Palma et al., (2019) De Palma, G., Kiani, B. T., and Lloyd, S. (2019).  Random deep neural networks are biased towards simple functions.  Advances in Neural Information Processing Systems.
  * Eberle, (2009) Eberle, A. (2009).  Markov processes.  Lecture Notes at University of Bonn.
  * Frankle et al., (2020) Frankle, J., Schwab, D. J., and Morcos, A. S. (2020).  Training batchnorm and only batchnorm: On the expressive power of random features in cnns.  arXiv preprint arXiv:2003.00152.
  * Garriga-Alonso et al., (2019) Garriga-Alonso, A., Rasmussen, C. E., and Aitchison, L. (2019).  Deep convolutional networks as shallow gaussian processes.  International Conference on Learning Representations.
  * Glorot and Bengio, (2010) Glorot, X. and Bengio, Y. (2010).  Understanding the difficulty of training deep feedforward neural networks.  In International Conference on Artificial Intelligence and Statistics.
  * Hazan and Jaakkola, (2015) Hazan, T. and Jaakkola, T. (2015).  Steps toward deep kernel methods from infinite neural networks.  arXiv preprint arXiv:1508.05133.
  * He et al., (2016) He, K., Zhang, X., Ren, S., and Sun, J. (2016).  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition.
  * Huang et al., (2017) Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. (2017).  Densely connected convolutional networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition.
  * Huang et al., (2014) Huang, P.-S., Avron, H., Sainath, T. N., Sindhwani, V., and Ramabhadran, B. (2014).  Kernel methods match deep neural networks on timit.  In ICASSP.
  * Ioffe and Szegedy, (2015) Ioffe, S. and Szegedy, C. (2015).  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In ICML.
  * Jacot et al., (2018) Jacot, A., Gabriel, F., and Hongler, C. (2018).  Neural tangent kernel: Convergence and generalization in neural networks.  Advances in Neural Information Processing Systems.
  * Karakida et al., (2019) Karakida, R., Akaho, S., and Amari, S.-i. (2019).  The normalization method for alleviating pathological sharpness in wide neural networks.  In Advances in Neural Information Processing Systems.
  * Kemeny and Snell, (1976) Kemeny, J. G. and Snell, J. L. (1976).  Markov chains.  Springer-Verlag, New York.
  * Khasminskii, (2011) Khasminskii, R. (2011).  Stochastic stability of differential equations, volume 66.  Springer Science & Business Media.
  * Klambauer et al., (2017) Klambauer, G., Unterthiner, T., Mayr, A., and Hochreiter, S. (2017).  Self-normalizing neural networks.  Advances in Neural Information Processing Systems.
  * Kohler et al., (2018) Kohler, J., Daneshmand, H., Lucchi, A., Zhou, M., Neymeyr, K., and Hofmann, T. (2018).  Exponential convergence rates for batch normalization: The power of length-direction decoupling in non-convex optimization.  arXiv preprint arXiv:1805.10694.
  * Krizhevsky et al., (2009) Krizhevsky, A., Hinton, G., et al. (2009).  Learning multiple layers of features from tiny images.
  * Kushner and Yin, (2003) Kushner, H. and Yin, G. G. (2003).  Stochastic approximation and recursive algorithms and applications.  Springer Science & Business Media.
  * Kushner, (1967) Kushner, H. J. (1967).  Stochastic stability and control.  Technical report.
  * Lee et al., (2019) Lee, J., Bahri, Y., Novak, R., Schoenholz, S. S., Pennington, J., and Sohl-Dickstein, J. (2019).  Deep neural networks as gaussian processes.  International Conference on Learning Representations.
  * Matthews et al., (2018) Matthews, A. G. d. G., Rowland, M., Hron, J., Turner, R. E., and Ghahramani, Z. (2018).  Gaussian process behaviour in wide deep neural networks.  arXiv preprint arXiv:1804.11271.
  * Neal, (1996) Neal, R. M. (1996).  Bayesian learning for neural networks.  Springer Science & Business Media.
  * Paszke et al., (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. (2019).  Pytorch: An imperative style, high-performance deep learning library.  In Advances in Neural Information Processing Systems.
  * Pennington et al., (2018) Pennington, J., Schoenholz, S., and Ganguli, S. (2018).  The emergence of spectral universality in deep networks.  In International Conference on Artificial Intelligence and Statistics.
  * Santurkar et al., (2018) Santurkar, S., Tsipras, D., Ilyas, A., and Madry, A. (2018).  How does batch normalization help optimization?(no, it is not about internal covariate shift).  Advances in Neural Information Processing Systems.
  * Sawa, (1972) Sawa, T. (1972).  Finite-sample properties of the k-class estimators.  Econometrica: Journal of the Econometric Society, pages 653–680.
  * Saxe et al., (2014) Saxe, A. M., McClelland, J. L., and Ganguli, S. (2014).  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  International Conference on Learning Representations.
  * Schoenholz et al., (2017) Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, J. (2017).  Deep information propagation.  International Conference on Learning Representations.
  * Silver et al., (2017) Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., et al. (2017).  Mastering the game of go without human knowledge.  nature.
  * Srivastava et al., (2014) Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. (2014).  Dropout: a simple way to prevent neural networks from overfitting.  The Journal of Machine Learning Research.
  * Yang et al., (2019) Yang, G., Pennington, J., Rao, V., Sohl-Dickstein, J., and Schoenholz, S. S. (2019).  A mean field theory of batch normalization.  International Conference on Learning Representations.

Appendix

##  Appendix A Preliminaries

###  A.1. Notations

Let $v,w\in{\mathbb{R}}^{k}$ then $v\odot w\in{\mathbb{R}}^{n}$ with
coordinates

(7) $\displaystyle[v\odot w]_{i}=v_{i}w_{i}$

Furthermore $v^{\otimes 2}\in{\mathbb{R}}^{k\times k}$ with entities

(8) $\displaystyle[v^{\otimes 2}]_{ij}=v_{i}v_{j}.$

In Table 2, we summarize notations introduced previously.

Notation | Type | Definition
---|---|---
$\ell$ | integer | number of layers
$n$ | integer | batch size
$d$ | integer | width of the network
$k$ | integer | output dimension
$X$ | ${\mathbb{R}}^{d\times n}$ | input matrix
$H_{\ell}$ | ${\mathbb{R}}^{d\times n}$ | hidden representations at $\ell$ (obeying Eq. (4))
$BN$ | ${\mathbb{R}}^{d\times n}\to{\mathbb{R}}^{d\times n}$ | batch normalization layer (defined in Eq.(4))
$\text{Law}(X)$ |  | the law of random matrix $X$
$\sigma_{i}(M)$ | ${\mathbb{R}}^{k_{1}\times k_{2}}\to{\mathbb{R}}_{+}$ | the $i$th largest singular value of matrix $M$
$I_{k}$ | ${\mathbb{R}}^{k\times k}$ | Identity matrix of size $k$
${\bf 1}_{k}$ | ${\mathbb{R}}^{n}$ | all-ones vector
Table 2. Notations

###  A.2. The Markov chain of hidden representations

Recall the chain of the hidden representation, denoted by
$\\{H_{\ell}\in{\mathbb{R}}^{d\times n}\\}$, obeys the following recurrence:

(9) $\displaystyle H_{\ell+1}=\frac{1}{\sqrt{d}}BN(W_{\ell}H_{\ell}),\quad
BN(M)=\left(\text{diag}(MM^{\top})\right)^{-\nicefrac{{1}}{{2}}}M,$

where $W_{\ell}\in{\mathbb{R}}^{d\times n}$ are random weight matrices with
i.i.d. zero-mean Gaussian elements. It is easy to check that the Frobenius
norm of $H_{\ell}$ is one due to the row-wise normalization:

(10) $\displaystyle\text{Tr}\left(BN(H)BN(H)^{\top}\right)$
$\displaystyle=\text{Tr}\left(\text{diag}(HH^{\top})^{-\nicefrac{{1}}{{2}}}HH^{\top}\text{diag}(HH^{\top})^{-\nicefrac{{1}}{{2}}}\right)$
$\displaystyle=\text{Tr}\left(\text{diag}(HH^{\top})^{-1}HH^{\top}\right)$
$\displaystyle=d.$

###  A.3. A Lyapunov function for the orthogonal stability

We establish stochastic stability of the chain of the hidden representations.
Our analysis is in terms of Lyapunov function
$\widehat{V}:{\mathbb{R}}^{d\times n}$ defined as

(11) $\displaystyle\widehat{V}(H)=\frac{1}{n}-(\sigma_{n}(H))^{2},$

where $\sigma_{n}(H)$ is the minimum singular value of matrix $H$. Next lemma
proves $\widehat{V}(H_{\ell})$ bounds the orthogonality gap.

######  Lemma 4.

For all hidden representations $H_{\ell}$, the following holds:

$\displaystyle V(H_{\ell})\leq 2n\widehat{V}(H_{\ell}).$

###### Proof.

Let $\sigma_{1},\dots,\sigma_{n}$ be the singular values of $H_{\ell}$. Given
these singular values, one can compute $V(H_{\ell})$ as

$\displaystyle(V(H_{\ell}))^{2}=\sum_{i=1}^{n}\left(\sigma_{i}^{2}-\frac{1}{n}\right)^{2}.$

According to Eq. (10), $\sum_{i=1}^{2}\sigma_{i}^{2}=1$ holds. The proof is an
immediate consequence of this propery.

$\displaystyle V^{2}(H_{\ell})$
$\displaystyle=\sum_{i=1}^{n}\sigma_{i}^{4}-2\underbrace{\left(\sum_{i=1}^{n}\sigma_{i}^{2}\right)}_{=1}\frac{1}{n}+\frac{1}{n}$
$\displaystyle=\sum_{i=1}^{n}\sigma_{i}^{4}-\frac{1}{n}.$

Fixing $\sigma_{n}$, the maximum of $\sum_{i=1}^{n}\sigma_{i}^{4}-\frac{1}{n}$
subject to $\sum_{i=1}^{n}\sigma_{i}^{2}=1$ is met when
$\sigma_{1}^{2}=1-(n-1)\sigma_{n}^{2}$ and $\sigma_{2}=\dots=\sigma_{n}$.
Therefore,

$\displaystyle V^{2}(H_{\ell})\leq
2(n-1)^{2}\underbrace{\left(\frac{1}{n}-\sigma_{n}^{2}\right)^{2}}_{=\widehat{V}^{2}(H_{\ell}(X))}$

holds true. Taking the square root of both sides concludes the proof. ∎

##  Appendix B Proof of Theorem 1

The proof of Theorem 1 relies on the following Theorem that characterizes the
change of $\widehat{V}$ in consecutive layers.

######  Theorem 5.

Sequence $\\{H_{\ell}\\}$ obeys

$\displaystyle{\mathbb{E}}\left[\widehat{V}(H_{\ell+1})|H_{\ell}\right]\leq\left(1-\frac{2}{3}\left(\frac{1}{n}-\widehat{V}(H_{\ell})\right)\right)\widehat{V}(H_{\ell})+\frac{1}{\sqrt{d}}.$

Notably, the above result does not rely on Assumption Assumption
$\mathcal{A}_{1}(\alpha,\ell)$. Assuming that Assumption Assumption
$\mathcal{A}_{1}(\alpha,\ell)$ holds, we complete the proof of Theorem 1.
Combining the last Theorem by this assumption, we get

$\displaystyle{\mathbb{E}}\left[\widehat{V}(H_{\ell+1})\right]\leq\left(1-\frac{2}{3}\left(\alpha\right)\right){\mathbb{E}}\left[\widehat{V}(H_{\ell})\right]+\frac{1}{\sqrt{d}}$

Induction over $\ell$ yields

$\displaystyle{\mathbb{E}}\left[\widehat{V}(H_{\ell+1}(X))\right]$
$\displaystyle\leq\left(1-\frac{2}{3}\alpha\right)^{\ell}{\mathbb{E}}\left[\widehat{V}(H_{1})\right]+\left(\sum_{k=1}^{\ell}(1-\frac{2}{3}\alpha)^{k}\right)\frac{1}{\sqrt{d}}$
$\displaystyle\leq\left(1-\frac{2}{3}\alpha\right)^{\ell}{\mathbb{E}}\left[\widehat{V}(H_{1})\right]+\frac{3}{2\alpha\sqrt{d}}$

An application of Lemma 4 completes the proof:

$\displaystyle{\mathbb{E}}\left[V(H_{\ell+1})\right]$ $\displaystyle\leq
2n{\mathbb{E}}\left[\widehat{V}(H_{\ell+1})\right]$ $\displaystyle\leq
2\left(1-\frac{2}{3}\alpha\right)^{\ell}+\frac{3n}{2\alpha\sqrt{d}}$

###  B.1. Stability analysis without Assumption $\mathcal{A}_{1}(\alpha,\ell)$

Using Theorem 5, we can prove stability of the chain $\\{H_{\ell}\\}$ without
Assumption Assumption $\mathcal{A}_{1}(\alpha,\ell)$. After rearrangement of
terms in Theorem 5, we get

$\displaystyle{\mathbb{E}}\left[\widehat{V}(H_{\ell+1})|H_{\ell}\right]-\widehat{V}(H_{\ell})\leq-\frac{2}{3}\widehat{V}(H_{\ell})\left(\frac{1}{n}-\widehat{V}(H_{\ell})\right)+\frac{}{\sqrt{d}}$

Taking the expectation over $H_{\ell}$ and average over $\ell$ yields

$\displaystyle{\mathbb{E}}\left[\frac{1}{\ell}\sum_{k=1}^{\ell}\widehat{V}(H_{k})\left(\frac{1}{n}-\widehat{V}(H_{k})\right)\right]\leq\left(\frac{3{\mathbb{E}}\left[\widehat{V}(H_{0})\right]}{2\ell}\right)+\frac{3}{2\sqrt{d}}$

##  Appendix C Proof of Theorem 5

###  C.1. Spectral decomposition.

Consider the SVD decomposition of $H_{\ell}$ as
$H_{\ell}=U\text{diag}(\sigma)V^{\top}$ where $U$ and $V$ are orthogonal
matrices. Given this decomposition, we get

(12) $\displaystyle
W_{\ell}H_{\ell}=\underbrace{W_{\ell}U}_{W}\text{diag}(\sigma)V^{\top}$

Since $W_{\ell}$ is Gaussian and $U$ is orthogonal, entities of matrix $W$ are
also i.i.d. standard normal. This decomposition will be repeatedly used.

###  C.2. Concentration

Consider matrix $C_{\ell+1}:=H_{\ell+1}^{\top}H_{\ell+1}$ whose eigenvalues
are $\sigma_{1}^{2},\dots,\sigma_{n}^{2}$. The SVD decomposition of $H_{\ell}$
allows writing $C_{\ell+1}$ as the average of vectors

$\displaystyle
C_{\ell+1}=\frac{1}{d}\sum_{i=1}^{d}\left(\frac{w_{i}\odot\sigma}{\|w_{i}\odot\sigma\|_{2}}\right)^{\otimes
2}$

where $w_{i}\in{\mathbb{R}}^{n}$ are rows of matrix $W$ in Eq.(12) and
$\sigma\in{\mathbb{R}}^{n}$ is the vector of singular values of $H_{\ell}$.
Thus, conditioned on $\sigma$, $C_{\ell+1}$ is an empirical average of i.i.d.
random vectors. This allows us to prove that this empirical average is
concentrated around its expectation. The next lemma states this concentration.

######  Lemma 6.

The following concentration always holds

$\displaystyle{\mathbb{E}}_{W_{\ell}}\|C_{\ell+1}-{\mathbb{E}}_{W_{\ell}}\left[C_{\ell+1}\right]\|^{2}\leq
1/d$

where

(13)
$\displaystyle{\mathbb{E}}_{W_{\ell}}\left[C_{\ell+1}\right]=\text{diag}(p_{1}(\sigma),\dots,p_{n}(\sigma)),\quad
p_{i}(\sigma):={\mathbb{E}}\left[\frac{\sigma_{n}^{2}w_{n}^{2}}{\sum_{k=1}^{n}\sigma_{k}^{2}w_{k}^{2}}\right],\quad
w_{i}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\mathcal{N}(0,1)$

The concentration of $C_{\ell+1}$ allows us to prove that the Lyapunov
function $\widehat{V}(H_{\ell+1})$ is concentrated around $1/n-p_{n}(\sigma)$.

######  Lemma 7.

The following holds

$\displaystyle{\mathbb{E}}_{W_{\ell}}\left[\left(\widehat{V}(H_{\ell+1})-(1/n-p_{n}(\sigma))\right)^{2}\right]\leq\frac{1}{d}$

.

The last lemma allows us to predict the value of random variable
$\widehat{V}(H_{\ell+1})$ by deterministic term $1/n-p_{n}(\sigma)$.

###  C.3. Contraction.

The decay in $\widehat{V}(H_{\ell+1})$ with $\ell$ is due to term
$1/n-p_{n}(\sigma)$ in the last lemma. This term is less than (or equal to)
$V(H_{\ell})$.

######  Lemma 8.

For $p_{n}(\sigma)$ defined in Eq. (13), the following holds:

$\displaystyle\left(\frac{1}{n}-p_{n}(\sigma)\right)\leq\left(1-\frac{2}{3}\left(\frac{1}{n}-\widehat{V}(H_{\ell})\right)\right)\widehat{V}(H_{\ell}).$

Combining the last lemma by Lemma 7 concludes the proof of Theorem 5:

(14)
$\displaystyle{\mathbb{E}}_{W_{\ell}}\left[\widehat{V}(H_{\ell+1})\right]\leq\left(1-\frac{2}{3}\left(\frac{1}{n}-V(H_{\ell})\right)\right)\widehat{V}(H_{\ell})+\frac{1}{\sqrt{d}}.$

To complete the proof, we prove Lemmas 6, 7, and 8.

###  C.4. Proof of Lemma 6

Given the spectral decomposition of $H_{\ell}$ in Eq. (12), we compute element
$ij$ of $C_{\ell+1}$, which is denoted by $[C_{\ell+1}]_{ij}$:

$\displaystyle[C_{\ell+1}]_{ij}=[A^{\top}A]_{ij}=\frac{1}{d}\sum_{k=1}^{d}A_{ki}A_{kj},\quad
A_{ki}=W_{ki}\sigma_{i}/\sqrt{v_{k}}$

where $v_{k}=\sum_{m=1}^{n}W_{km}^{2}\sigma_{m}^{2}$. Since $W_{km}$ are zero
mean and unit variace, we get

$\displaystyle{\mathbb{E}}[C_{\ell+1}]_{ij}$ $\displaystyle=0$
$\displaystyle{\mathbb{E}}[C_{\ell+1}]_{ij}^{2}$
$\displaystyle=\frac{1}{d^{2}}\sum_{k=1}^{d}A_{ki}^{2}A_{kj}^{2}+\frac{1}{d^{2}}\sum_{k,k^{\prime}}^{d}\underbrace{{\mathbb{E}}\left[A_{ki}^{2}A_{kj}^{2}A_{k^{\prime}i}^{2}A_{k^{\prime}j}^{2}\right]}_{=0}$
$\displaystyle=\frac{1}{d^{2}}\sum_{k=1}^{d}A_{ki}^{2}A_{kj}^{2}$

holds for $i\neq j$. By summing up over $i\neq j$, we get

$\displaystyle\sum_{i\neq
j}{\mathbb{E}}[C_{\ell+1}]_{ij}=\frac{1}{d^{2}}\left(\sum_{k}\left(\sum_{i}A_{ki}^{2}\right)^{2}-\sum_{ik}A_{ki}^{4}\right)$

For the diagonal elements, we get

$\displaystyle{\mathbb{E}}\left[C_{\ell+1}\right]_{ii}$
$\displaystyle=\frac{1}{d}\sum_{k=1}^{d}{\mathbb{E}}A_{ki}^{2}$
$\displaystyle=\frac{1}{d}\sum_{k=1}^{d}{\mathbb{E}}A_{ki}^{2}$
$\displaystyle=\underbrace{\frac{1}{d}\sum_{k=1}^{d}{\mathbb{E}}\left[\frac{W_{ki}^{2}\sigma_{i}^{2}}{\sum_{j=1}^{n}W_{kj}^{2}\sigma_{j}^{2}}\right]}_{p_{i}(\sigma)}$

The variance of $[C_{\ell+1}]_{ii}$ is bounded as

$\displaystyle\text{var}([C_{\ell+1}]_{ii})$
$\displaystyle={\mathbb{E}}\left(\frac{1}{d}\sum_{k=1}^{d}(A_{ki}^{2}-p_{i}(\sigma))\right)^{2}$
$\displaystyle=\frac{1}{d^{2}}\sum_{k=1}^{d}(A_{ki}^{2}-p_{i}(\sigma))^{2}$
$\displaystyle\leq\frac{1}{d^{2}}\sum_{k=1}^{d}A_{ki}^{4}$

Combining results for diagonal and off-diagonal elements yields

$\displaystyle{\mathbb{E}}\|C_{\ell+1}-{\mathbb{E}}_{W_{\ell}}\left[C_{\ell+1}\right]\|_{F}^{2}$
$\displaystyle=\sum_{ij}\text{var}([C_{\ell+1}]_{ij})$
$\displaystyle\leq\frac{1}{d^{2}}\left(\sum_{i}A_{ki}^{2}\right)^{2}=1/d$

###  C.5. Proof of Lemma 7

Notably, the eigenvalues of $C_{\ell+1}$ are squared singular values of
$H_{\ell+1}$. Let $\lambda_{n}(C)$ denote the $n$th largest eigenvalue of
matrix $C$.

$\displaystyle{\mathbb{E}}_{W_{\ell}}\left[\left(V(C_{\ell+1})-\left(\frac{1}{n}-p_{n}(\sigma)\right)\right)^{2}\right]$
$\displaystyle={\mathbb{E}}_{W_{\ell}}\left[\left(\lambda_{n}\left(C_{\ell+1}\right)-\lambda_{n}\left({\mathbb{E}}_{W_{\ell}}\left[C_{\ell+1}\right]\right)\right)^{2}\right]$
$\displaystyle\leq{\mathbb{E}}_{W_{\ell}}\left[\|C_{\ell+1}-{\mathbb{E}}_{W_{\ell}}\left[C_{\ell+1}\right]\|_{F}^{2}\right]$
$\displaystyle\leq\frac{1}{d}$

where the last inequality relies on Lemma 6.

###  C.6. Proof of Lemma 8

The proof is based an application of moment generating function that allows us
to compute expectations of ratios of random variables.

######  Lemma 9 (Lemma 1 in (Sawa,, 1972)).

Let $X_{1}$ be a random variable that is positive with probability one and
$X_{2}$ be an arbitrary random variable. Suppose that there exists a joint
moment generating function of $X_{1}$ and $X_{2}$:

$\displaystyle\phi(\theta_{1},\theta_{2})={\mathbb{E}}\left[\exp(\theta_{1}X_{1}+\theta_{2}X_{2})\right]$

for $\theta_{1}\leq\epsilon$ and $|\theta_{2}|<\epsilon$ where $\epsilon$ is
some positive constant. Then

$\displaystyle{\mathbb{E}}\left[\frac{X_{2}}{X_{1}}\right]=\int_{-\infty}^{0}\left[\frac{\partial\phi(\theta_{1},\theta_{2})}{\partial\theta_{2}}\right]_{\theta_{2}=0}d\theta_{1}$

To estimate $p_{n}(\sigma)$, we set $X_{2}:=\sigma_{i}^{2}w_{i}^{2}$ and
$X_{1}=\sum_{j}\sigma_{j}^{2}w_{j}^{2}$, which obtains

$\displaystyle\phi(\theta_{1},\theta_{2})$
$\displaystyle={\mathbb{E}}\left[\exp(\theta_{1}X_{1}+\theta_{2}X_{2})\right]$
$\displaystyle=(2\pi)^{-n/2}\int_{-\infty}^{\infty}\exp((\theta_{1}+\theta_{2})\sigma_{i}^{2}w_{i}^{2}+\sum_{j\neq
i}\theta_{1}\sigma_{j}^{2}w_{j}^{2})\exp(-\sum_{k}w_{k}^{2}/2)dw$
$\displaystyle=(2\pi)^{-n/2}\int_{-\infty}^{\infty}\exp(\left(-0.5+(\theta_{1}+\theta_{2})\sigma_{i}^{2}\right)w_{i}^{2})dw_{i}\left(\prod_{j\neq
i}\int_{-\infty}^{\infty}\exp(\left(-0.5+\theta_{1}\sigma_{j}^{2}\right)w_{j}^{2})dw_{j}\right)$
$\displaystyle=\frac{1}{\sqrt{1-2(\theta_{1}+\theta_{2})\sigma_{i}^{2}}}\left(\prod_{j\neq
i}\frac{1}{\sqrt{1-2\theta_{1}\sigma_{j}^{2}}}\right).$

Taking derivative with respect to $\theta_{2}$ yields

$\displaystyle\frac{\partial\phi}{\theta_{2}}(\theta_{1},0)=\frac{\sigma_{i}^{2}}{(1-2\theta_{1}\sigma_{i}^{2})^{3/2}}\left(\prod_{j\neq
i}\frac{1}{\sqrt{1-2\theta_{1}\sigma_{j}^{2}}}\right)$

Using the result of the last lemma, we get

$\displaystyle
p_{i}(\sigma)=\int_{-\infty}^{0}\frac{\sigma_{i}^{2}}{(1-2\theta\sigma_{i}^{2})}\left(\prod_{j}\frac{1}{\sqrt{1-2\theta\sigma_{j}^{2}}}\right)d\theta$

Therefore,

(15) $\displaystyle p_{n}(\sigma)=\sigma_{n}^{2}f_{n}(\sigma),\quad
f_{n}(\sigma):=\int_{-\infty}^{0}\frac{1}{(1-2\theta\sigma_{n}^{2})^{\nicefrac{{3}}{{2}}}\prod_{j\neq
n}(1-2\theta\sigma_{j}^{2})^{\nicefrac{{1}}{{2}}}}$

Since $\sum_{i=1}^{n}\sigma_{i}^{2}=1$ (see Eq. (10)),$f_{n}(\sigma)$ in the
above formulation is minimized when the $\sigma_{j}^{2}$s are all equal for
all $j\neq n$. Let $\sigma_{n}^{2}:=1/n-\delta$ and
$\sigma_{j}^{2}:=1/n+\delta/(n-1)$ for all $j\neq i$. This allows us to
establish a lowerbound on $f_{n}(\sigma)$ as

(16) $\displaystyle f_{n}(\sigma)$
$\displaystyle\geq\int_{0}^{\infty}\underbrace{\left(1+2\theta(\frac{1}{n}-\delta)\right)^{-\frac{3}{2}}\left(1+2\theta(\frac{1}{n}+\frac{\delta}{n-1})\right)^{-\frac{n-1}{2}}}_{g(\delta):=}$

Next lemma proves that $g(\delta)$ is a convex function for
$\delta\in[0,1/n]$.

######  Lemma 10.

The function $g(\delta)$, which is defined in Eq. (16), is a convex function
on domain $\delta\in[0,1/n]$.

The convexity of $g(\delta)$ yields

$\displaystyle g^{\prime\prime}(\delta)\geq 0\;\forall\delta$
$\displaystyle\implies g(\delta)\geq g(0)+\delta g^{\prime}(0)$

The above bound allow us to bound $f_{n}(\sigma)$ as

$\displaystyle f_{n}(\sigma)$
$\displaystyle\geq\int_{0}^{\infty}g(0)+\delta\int_{0}^{\infty}2\theta(1+\frac{2\theta}{n})^{-\frac{n}{2}-2}$
$\displaystyle=1+\delta\frac{2n}{n+2}$

Recall $\delta=\frac{1}{n}-\sigma_{n}^{2}$. Combining the above inequality by
Eq. (15) concludes the proof of the Lemma 8:

$\displaystyle\frac{1}{n}-p_{n}(\sigma)\leq\left(1-\frac{2n}{(n+2)}\sigma_{n}^{2}\right)\left(\frac{1}{n}-\sigma_{n}^{2}\right)$

###  C.7. Proof of Lemma 10

We can show convexity of $g(\delta)$ by showing $g^{\prime\prime}(\delta)\geq
0$ for all $\delta\in(0,1)$. To this end, we define the helper function
$h_{1}(\delta)$ (for the compactness of notations) as

$\displaystyle
h_{1}(\delta):=(1+2\theta(\frac{1}{n}-\delta))^{-5/2}(1+2\theta(\frac{1}{n}+\frac{\delta}{n-1}))^{-\frac{n-1}{2}-1}.$

Given $h_{1}$, the derivative of $g$ reads as

$\displaystyle g^{\prime}(\delta)$
$\displaystyle=h_{1}(\delta)\left((-\frac{3}{2})(-2\theta)\left(1+2\theta(\frac{1}{n}+\frac{\delta}{n-1})\right)-\frac{n-1}{2}\frac{2\theta}{n-1}\left(1+2\theta(\frac{1}{n}-\delta)\right)\right)$
$\displaystyle=\theta
h_{1}(\delta)\left(3+\theta\frac{6}{n}+\theta\delta\frac{6}{n-1}-1-\theta\frac{2}{n}+2\theta\delta\right)$
$\displaystyle=\theta
h_{1}(\delta)\underbrace{\left(2+\theta\frac{4}{n}+\theta\delta\frac{2n+4}{n-1}\right)}_{h_{2}(\delta):=}$
$\displaystyle=\theta h_{1}(\delta)h_{2}(\delta).$

One can readily check that $h_{2}^{\prime}(\delta)\geq 0$. Hence,
$h_{1}^{\prime}(\delta)\geq 0$ ensures the convexity of $g(\delta)$. Consider
the following helper

$h_{3}(\delta):=\left(1+2\theta(\frac{1}{n}-\delta)\right)^{-7/2}\left(1+2\theta(\frac{1}{n}+\frac{\delta}{n-1})\right)^{-\frac{n-1}{2}-2}.$

Given $h_{3}$, we compute $h^{\prime}_{1}$ as

$\displaystyle h_{1}^{\prime}(\delta)$
$\displaystyle=h_{3}(\delta)\left((-\frac{5}{2})(-2\theta)(1+2\theta(\frac{1}{n}+\frac{\delta}{n-1}))+(-\frac{n-3}{2})\frac{2\theta}{n-1}(1+2\theta(\frac{1}{n}-\delta))\right)$
$\displaystyle=h_{3}(\delta)\theta\left(5(1+2\theta/n+\frac{2\theta\delta}{n-1})-\frac{n-3}{n-1}(1+2\theta/n-2\theta\delta)\right)$
$\displaystyle=h_{3}(\delta)\theta\underbrace{\left(\frac{4n-2}{n-1}+2\theta\frac{4n-2}{n-1}+2\theta\delta\frac{n+2}{n-1}\right)}_{h_{4}(\delta):=}$
$\displaystyle=\theta h_{3}(\delta)h_{4}(\delta).$

Clearly we have $h_{3}(\delta),h_{4}(\delta)\geq 0$ for all valid choices of
$\delta$ which finishes the proof.

##  Appendix D Proof of Lemma 2

The main idea is based on a particular coupling of random matrices
$W_{\ell}H_{\ell}$ and $G$. Consider the truncated SVD decomposition of
$H_{\ell}$ as $H_{\ell}=U\text{diag}(\sigma)V^{\top}$ where
$U\in{\mathbb{R}}^{d\times n}$ and $V\in{\mathbb{R}}^{n\times n}$ are
orthogonal matrices. Due to the orthogonality, the law of $W_{\ell}U$ is the
same as those of $GV$. By coupling $W_{\ell}U=GV$, we get

$\displaystyle\left({\mathcal{W}}_{2}(\mu(W_{\ell}H_{\ell}),\mu(G/\sqrt{n}))\right)^{2}$
$\displaystyle=\inf_{\text{all the
couplings}}{\mathbb{E}}\|W_{\ell}U\text{diag}(\sigma)V^{\top}-GVV^{\top}/\sqrt{n}\|_{F}^{2}$
$\displaystyle\leq{\mathbb{E}}\|GV\left(\text{diag}(\sigma)-I/\sqrt{n}\right)V^{\top}\|_{F}^{2}$
$\displaystyle={\mathbb{E}}\text{Tr}(GV\left(\text{diag}(\sigma)-I/\sqrt{n}\right)V^{\top}V\left(\text{diag}(\sigma)-I/\sqrt{n}\right)V^{\top}G^{\top})$
$\displaystyle=\text{Tr}(V\left(\text{diag}(\sigma)-I/\sqrt{n}\right)\left(\text{diag}(\sigma)-I/\sqrt{n}\right)V^{\top}{\mathbb{E}}\left[G^{\top}G\right])$
$\displaystyle={\mathbb{E}}\text{Tr}(\left(\text{diag}(\sigma)-I/\sqrt{n}\right)\left(\text{diag}(\sigma)-I/\sqrt{n}\right)V^{\top}V)$
$\displaystyle={\mathbb{E}}\|\left(\text{diag}(\sigma)-I/\sqrt{n}\right)\|_{F}^{2}$
$\displaystyle={\mathbb{E}}\left[\sum_{i=1}^{n}(\sigma_{i}-1/\sqrt{n})^{2}\right]$
$\displaystyle={\mathbb{E}}\left[\sum_{i=1}^{n}\left(\sigma_{i}^{2}-1/n\right)^{2}/\left(\sigma_{i}+1/\sqrt{n}\right)^{2}\right]$
$\displaystyle\leq
n{\mathbb{E}}\left[\sum_{i=1}^{n}\left(\sigma_{i}^{2}-1/n\right)^{2}\right]$
$\displaystyle\leq n{\mathbb{E}}\left[V^{2}(H_{\ell})\right]$
$\displaystyle\leq 2n{\mathbb{E}}\left[V(H_{\ell})\right].$

##  Appendix E Orthogonality gap for the iterative initialization

Recall the proposed initialization for weights based on SVD decomposition
$H_{\ell}=U_{\ell}\Sigma_{\ell}V_{\ell}^{\top}$.

$\displaystyle
W_{\ell}=\frac{1}{\|\Sigma_{\ell}^{\nicefrac{{1}}{{2}}}\|_{F}}V_{\ell}^{\prime}\Sigma_{\ell}^{-\nicefrac{{1}}{{2}}}U_{\ell}^{\top}.$

Here, we show that

(17) $\displaystyle V(H_{\ell})>V(W_{\ell}H_{\ell})$

holds as long as $V(H_{\ell})\neq 0$ and Assumption
$\mathcal{A}_{1}(\alpha,\ell)$$(\alpha,\ell)$ holds. Given singular values
$H_{\ell}$, we get

(18) $\displaystyle V(H_{\ell})$
$\displaystyle=\sum_{i=1}^{n}\left(\sigma_{i}^{2}-\frac{1}{n}\right)^{2}$
$\displaystyle=\sum_{i}\sigma_{i}^{4}-1/n,$

where we used $\sum_{i=1}^{n}\sigma_{i}^{2}=1$ (see Eq. (2)). Now, we compute
$V(H_{\ell})$ using the singular values.

$\displaystyle
W_{\ell}H_{\ell}=\frac{1}{\|\Sigma_{\ell}^{\nicefrac{{1}}{{2}}}\|_{F}}V_{\ell}^{\prime}\Sigma^{\nicefrac{{1}}{{2}}}_{\ell}V_{\ell}$

Hence, the following holds

$\displaystyle V(W_{\ell}H_{\ell})$
$\displaystyle=\sum_{i=1}^{n}\left(\frac{\sigma_{i}}{\sum_{j}\sigma_{j}}-\frac{1}{n}\right)^{2}$
$\displaystyle=\frac{1}{(\sum_{i=1}^{n}\sigma_{i})^{2}}-1/n$

Combining with Eq. (18), we get

$\displaystyle
V(H_{\ell})-V(W_{\ell}H_{\ell})=\sum_{i=1}^{n}\sigma_{i}^{4}-\frac{1}{(\sum_{i=1}^{n}\sigma_{i})^{2}}.$

To show that the right side of the above equation is positive, we need to
prove

$\displaystyle\sum_{i}\sigma_{i}^{4}\left(\sum_{i=1}^{n}\sigma_{i}\right)^{2}>1=\left(\sum_{i}\sigma_{i}^{2}\right)^{4}$

holds. Using Cauchy-Schwarz inequality, we get

$\displaystyle\left(\sum_{i}\sigma_{i}^{2}\right)^{4}$
$\displaystyle=\left(\sum_{i}\sigma_{i}^{\nicefrac{{3}}{{2}}}\sigma_{i}^{\nicefrac{{1}}{{2}}}\right)^{4}$
$\displaystyle\leq\left(\sum_{i}\sigma_{i}\right)^{2}\left(\sum_{i}\sigma_{i}^{2}\sigma_{i}\right)^{2}$
$\displaystyle\leq\left(\sum_{i}\sigma_{i}\right)^{2}\left(\sum_{i}\sigma_{i}^{4}\right)^{2},$

where the equality in the above inequality is met only for
$\sigma_{i}^{2}=\frac{1}{n}$ (so $V(H_{\ell})=0$) under Assumption
$\mathcal{A}_{1}(\alpha,\ell)$.

##  Appendix F Activation functions and the orthogonality

We have shown hidden representations become increasingly orthogonal through
layers of BN MLPs with linear activations. We observed a similar results for
MLP with ReLU activations and also a simple convolution network (see Figure
5). Here, we elaborate more on the role activation functions in our analysis.

(a) MLP

(b) ConvNet

Figure 5. Orthogonality gap of layers for BatchNorm vs vanilla network with
ReLU, for MLP and a basic convolutional network on CIFAR10. Batch size was set
to $8$, the number of hidden layers for both networks was set to $30$, with
$100$ width for MLP and $100$ channels for ConvNet in all hidden layers. For
the ConvNet, the kernel size was set to $3$, across all hidden layers, with
zero-padding to make convolutional feature maps equal sized. The standard
deviations were computed for $50$ different randomly sampled batches.

For an arbitrary activation function $F$, the hidden representations make the
following Markov chain:

$\displaystyle Q_{\ell+1}=BN(W_{\ell}F(Q_{\ell})).$

###  F.1. Conjecture

Suppose that $\\{Q_{\ell}\\}$ is ergodict (Eberle,, 2009), and admits a unique
invariant distribution denoted by $\nu$. We define function
$L:{\mathbb{R}}^{d\times n}\to{\mathbb{R}}_{+}$ as

(19) $\displaystyle
L(Q_{\ell})=\left\|Q_{\ell}^{\top}Q_{\ell}-{\mathbb{E}}_{Q\sim\nu}\left[Q^{\top}Q\right]\right\|_{F}.$

We conjecture that there exits an integer $k$ such that for $\ell>k$ the
following holds:

(20)
$\displaystyle{\mathbb{E}}\left[L(Q_{\ell})\right]=\mathcal{O}\left(\frac{1}{\sqrt{d}}\right).$

Figure 6 experimentally validates that the above result holds for standard
activations such as ReLU, sigmoid, and tanh.

Figure 6. Validating the conjecture. Horizontal axis: $\log(d)$; vertical
axis: $\log(\frac{1}{1000}\sum_{k=1}^{1000}L(Q_{k}))$. The difference between
plots for the activation functions is negligible, so not visible.

###  F.2. Odd activation

We observe that Theorem 1 holds for odd activations such as $\sin$ and
$\tanh$. For these activation result of Theorem 1.

$\displaystyle{\mathbb{E}}\left[V(H_{\ell+1})\right]=\mathcal{O}\left((1-\alpha)^{\ell}+\frac{1}{\sqrt{d}}\right)$

holds under Assumption $\mathcal{A}_{1}(\alpha,\ell)$$(\alpha,\ell)$. This
also aligns with our conjecture for all activations as we observed that $V=Q$
for these activations.

##  Appendix G Experiments for convolutional networks

The introduced initialization scheme in Section 5 extends to convnets. In
convolutional networks, hidden representation are in the tensor form
$H_{\ell}\in{\mathbb{R}}^{d\times m\times m\times n}$ where $d$ is the number
of filters and $k$ is the image dimensions. Let matrix $W_{\ell}^{d\times
k^{2}}$ are weights of convolutions layer $\ell$ with kernel size $k$ and $d$
filters. A 2D convolution is a matrix multiplication combined by projections
as:

(21)
$\displaystyle\text{conv2d}(W_{\ell},H_{\ell})=\text{fold}\left(W_{\ell}\underbrace{\text{unfold}(H_{\ell})}_{H^{\prime}_{\ell}}\right)$

where the unfolding extract batches of $k\times k$ from images in tensor
$H_{\ell}$, folding operation combines the computed convolution for the
extracted batches. We use the SVD decomposition
$H^{\prime}_{\ell}=U_{\ell}\Sigma_{\ell}V_{\ell}$ to initialize weights
$W_{\ell}$, exactly the same as Eq. (6):

$\displaystyle
W_{\ell}=\frac{1}{\|\Sigma_{\ell}^{\nicefrac{{1}}{{2}}}\|_{F}}V_{\ell}^{\prime}\Sigma_{\ell}^{-\nicefrac{{1}}{{2}}}U_{\ell}^{\top},$

where $V_{\ell}^{\prime}$ is a slice of $V_{\ell}$. We experimentally validate
the performance of this initialization for two convolution networks with ReLU
activations consist of 20 and 80 convolutions. Table 3 outlines the details of
these neural architectures.

The network with 20 layers | The network with 80 layers
---|---
Conv2d(3, 64)+ MaxPool2d(2) + RELU | Conv2d(3, 64)+ MaxPool2d(2) + RELU
Conv2d(64, 192) + MaxPool2d(2) + ReLU | Conv2d(64, 192) + MaxPool2d(2) + ReLU
Conv2d(192, 256) \+ Conv2d(256, 256) +ReLU | Conv2d(192, 256) Conv2d(256, 256) + ReLU
$\\{$Conv2d(256, 256)+ReLU$\\}\times 16$ |  $\\{$Conv2d(256, 256)+ReLU$\\}\times 76$
MaxPool2d(2) + ReLU | MaxPool2d(2) + ReLU
Dropout(0.5) + Linear(1024,4096) + ReLU | Dropout(0.5) + Linear(1024,4096) + ReLU
Dropout(0.5) + Linear(4096,4096) + ReLU | Dropout(0.5) + Linear(4096,4096) + ReLU
Linear(4096,10) | Linear(4096,10)
Table 3. Convolutional networks used in the experiment. These neural
architectures are obtained from adding layers to AlexNet (Srivastava et al.,,
2014). The kernel size is 3 for all the convolutional layers.

Figure 7 shows the performance of SGD (with stepsize $0.001$ and batchsize 30)
for the proposed initialization. This initialization slows SGD for the shallow
network with 20 compared to standard Xavier’s initialization Glorot and
Bengio, (2010). However, this initialization outperforms Xavier’s
initialization when the depth is significantly large. This result
substantiates the role of the orthogonality in training and shows that the
experimental result of Figure 4 extends to convolutional neural networks.

20 layers | 80 layers
---|---
|
Figure 7. Training convolutional networks. Convergence of SGD for Xavier
initialization (Glorot and Bengio,, 2010) (in blue) and also the proposed
initialization method that ensures the orthogonality of hidden representations
(in red). Mean and 95% confidence interval of 4 independent runs.