# Towards Training Without Depth Limits:
Batch Normalization Without Gradient Explosion

Alexandru Meterez*
ETH Zürich
<EMAIL_ADDRESS>
&Amir Joudaki*
ETH Zürich
<EMAIL_ADDRESS>
&Francesco Orabona
CEMSE, KAUST
<EMAIL_ADDRESS>
Alexander Immer
ETH Zürich
<EMAIL_ADDRESS>
&Gunnar Rätsch
ETH Zürich
<EMAIL_ADDRESS>
&Hadi Daneshmand
MIT LIDS
<EMAIL_ADDRESS>
###### Abstract

Normalization layers are one of the key building blocks for deep neural
networks. Several theoretical studies have shown that batch normalization
improves the signal propagation, by avoiding the representations from becoming
collinear across the layers. However, results on mean-field theory of batch
normalization also conclude that this benefit comes at the expense of
exploding gradients in depth. Motivated by these two aspects of batch
normalization, in this study we pose the following question: “Can a batch-
normalized network keep the optimal signal propagation properties, but _avoid_
exploding gradients?” We answer this question in the affirmative by giving a
particular construction of an _Multi-Layer Perceptron (MLP) with linear
activations_ and batch-normalization that provably has _bounded gradients_ at
any depth. Based on Weingarten calculus, we develop a rigorous and non-
asymptotic theory for this constructed MLP that gives a precise
characterization of forward signal propagation, while proving that gradients
remain bounded for linearly independent input samples, which holds in most
practical settings. Inspired by our theory, we also design an activation
shaping scheme that empirically achieves the same properties for certain non-
linear activations.

††*: Equal contribution††Code is available at:
github.com/alexandrumeterez/bngrad

##  1 Introduction

What if we could train even deeper neural networks? Increasing depth empowers
neural networks, by turning them into powerful data processing machines. For
example, increasing depth allows large language models (LLMs) to capture
longer structural dependencies (Devlin et al., 2018; Liu et al., 2019; Brown
et al., 2020; Goyal et al., 2021; Raffel et al., 2020). Also, sufficiently
deep convolutional networks can outperform humans in image classification (Liu
et al., 2022; Woo et al., 2023; Wu et al., 2021). Nevertheless, increasing
depth imposes an inevitable computational challenge: deeper networks are
harder to optimize. In fact, standard optimization methods exhibit a slower
convergence when training deep neural networks. Hence, computation has become
a barrier in deep learning, demanding extensive research and engineering.

A critical problem is that deeper networks suffer from the omnipresent issue
of rank collapse at initialization: the outputs become collinear for different
inputs as the network grows in depth (Saxe et al., 2013a). Rank collapse is
not only present in MLPs and convolutional networks (Feng et al., 2022; Saxe
et al., 2013a; Daneshmand et al., 2020), but also in transformer architectures
(Dong et al., 2021; Noci et al., 2022). This issue significantly contributes
to the training slowdown of deep neural networks. Hence, it has become the
focus of theoretical and experimental studies (Saxe et al., 2013a; Feng et
al., 2022; Daneshmand et al., 2021; Noci et al., 2022).

One of the most successful methods to avoid rank collapse is Batch
Normalization (BN) (Ioffe and Szegedy, 2015), as proven in a number of
theoretical studies (Yang et al., 2019; Daneshmand et al., 2020, 2021).
Normalization imposes a particular bias across the layers of neural networks
(Joudaki et al., 2023a). More precisely, the representations of a batch of
inputs become more orthogonal after each normalization (Joudaki et al.,
2023a). This orthogonalization effect precisely avoids the rank collapse of
deep neural networks at initialization (Yang et al., 2019; Joudaki et al.,
2023a; Daneshmand et al., 2021; Joudaki et al., 2023b).

While batch normalization effectively avoids rank collapse, it causes
numerical issues. The existing literature proves that batch normalization
layers cause exploding gradients in MLPs in an activation-independent manner
(Yang et al., 2019). Gradient explosion limits increasing depth by causing
numerical issues during backpropagation. For networks without batch
normalization, there are effective approaches to avoid gradient explosion and
vanishing, such as tuning the variance of the random weights based on the
activation and the network width (He et al., 2015; Glorot and Bengio, 2010).
However, such methods cannot avoid gradient explosion in the presence of batch
normalization (Yang et al., 2019). Thus, the following important question
remains unanswered:

_Is there any network with batch normalization without gradient explosion and
rank collapse issues?_

##### Contributions.

We answer the above question affirmatively by giving a specific MLP
construction initialized with orthogonal random weight matrices, rather than
Gaussian. To show that the MLP still has optimal signal propagation, we prove
that the MLP output embeddings become isometric (equation 5), implying the
output representations becomes more orthogonal with depth. For a batch of
linearly independent inputs, we prove

$\displaystyle\mathbb{E}\big{[}\text{isometry
gap}\big{]}=\mathcal{O}\left(e^{-\text{depth}/C}\right),$ (1)

where $C$ is a constant depending only on the network width and input and the
expectation is taken over the random weight matrices. Thus, for sufficiently
deep networks, the representations rapidly approach an orthogonal matrix.
While Daneshmand et al. (2021) prove that the outputs converge to within an
$\mathcal{O}(\text{width}^{-1/2})$-ball close to orthogonality, we prove that
the output representations become perfectly orthogonal in the infinite depth
limit. This perfect orthogonalization turns out to be key in proving our
result about avoiding gradient explosion. In fact, for MLPs initialized with
Gaussian weights and BN, Yang et al. (2019, Theorem 3.9) prove that the
gradients explode at an exponential rate in depth. In a striking contrast, we
prove that gradients of an MLP with BN and orthogonal weights remain bounded
as

$\displaystyle\mathbb{E}\big{[}\log\left(\text{gradient norm for each
layer}\right)\big{]}=\mathcal{O}\left(\text{width}^{5}\right).$ (2)

Thus, the gradient is bounded by a constant that only depends on the network
width where the expectation is taken over the random weight matrices. It is
worth noting that both isometry and log-norm gradient bounds are derived _non-
asymptotically_. Thus, in contrast to the previously studied mean-field or
infinite width regime, our theoretical results hold in practical settings
where the width is finite.

The limitation of our theory is that it holds for a simplification in the BN
module and linear activations. However, our results provide guidelines to
avoid gradient explosion in MLPs with non-linear activations. We
experimentally show that it is possible to avoid gradient explosion for
certain non-linear activations with orthogonal random weights together with
“activation shaping” (Martens et al., 2021). Finally, we experimentally
demonstrate that avoiding gradient explosion stabilizes the training of deep
MLPs with BN.

##  2 Related work

##### The challenge of depth in learning.

Large depth poses challenges for the optimization of neural networks, which
becomes slower by increasing the number of layers. This depth related slowdown
is mainly attributed to: (i) gradient vanishing/explosion, and (ii) the rank
collapse of hidden representations. (i) Gradient vanishing and explosion is a
classic problem in neural networks (Hochreiter, 1998). For some neural
architectures, this issue can be effectively solved. For example, He et al.
(2015) propose a particular initialization scheme that avoids gradient
vanishing/explosion for neural networks with rectifier non-linearities while
Glorot and Bengio (2010) study the effect of initialization on sigmoidal
activations. However, such initializations cannot avoid gradient explosion for
networks with batch normalization (Yang et al., 2019; Lubana et al., 2021).
(ii) Saxe et al. (2013a) demonstrate that outputs become independent from
inputs with growing depth, which is called the rank collapse issue (Daneshmand
et al., 2020; Dong et al., 2021). Various techniques have been developed to
avoid rank collapse such as batch normalization (Ioffe and Szegedy, 2015),
residual connections (He et al., 2016), and self-normalizing activations
(Klambauer et al., 2017). Here, we focus on batch normalization since our
primary goal is to avoid the systemic issue of gradient explosion for batch
normalization.

##### Initialization with orthogonal matrices.

Saxe et al. (2013a) propose initializing the weights with random orthogonal
matrices for linear networks without normalization layers. Orthogonal matrices
avoid the rank collapse issue in linear networks, thereby enabling a depth-
independent training convergence. Pennington et al. (2017) show that MLPs with
sigmoidal activations achieve dynamical isometry when initialized with
orthogonal weights. Similar benefits have been achieved by initializing CNNs
with orthogonal or almost orthogonal kernels (Xiao et al., 2018; Mishkin and
Matas, 2015), and by initializing RNN transition matrices with elements from
the orthogonal and unitary ensembles (Arjovsky et al., 2016; Le et al., 2015;
Henaff et al., 2016). Similarly, we use orthogonal random matrices to avoid
gradient explosion. What sets our study apart from this literature is that our
focus is on batch normalization and the issue of gradient explosion.

##### Networks with linear activation functions.

Due to its analytical simplicity, the identity function has been widely used
in theoretical studies for neural networks. Studies on identity activations
date back to at least two decades. Fukumizu (1998) studies batch gradient
descent in linear neural networks and its effect on overfitting and
generalization. Baldi and Hornik (1995) provide an overview over various
theoretical manuscripts studying linear neural networks. Despite linearity, as
Saxe et al. (2013a, b) observe, the gradient dynamics in a linear MLP are
highly nonlinear. In a line of work, Saxe et al. (2013b, a, 2019) study the
training dynamics of deep neural networks with identity activations and
introduce the notion of dynamical isometry. Baldi and Hornik (1989) and Yun et
al. (2017) study the mean squared error optimization landscape in linear MLPs.
More recently, the optimum convergence rate of gradient descent in deep linear
neural networks has been studied by Arora et al. (2018) and Shamir (2019). Du
and Hu (2019) prove that under certain conditions on the model width and input
degeneracy, linear MLPs with Xavier initialized weights (Glorot and Bengio,
2010) converge linearly to the global optimum. Akin to these studies, we also
analyze networks with linear activations. However, batch normalization is a
non-linear function, hence the network we study in this paper is a highly non-
linear function of its inputs.

##### Mean field theory for random neural networks.

The existing analyses for random networks often rely on mean-field regimes
where the network width tends to infinity (Pennington et al., 2017; Yang et
al., 2019; Li et al., 2022; Pennington and Worah, 2017). However, there is a
discrepancy between mean-field regimes and the practical regime of finite
width. While some analyses attempt to bridge this gap (Joudaki et al., 2023b;
Daneshmand et al., 2021), their results rely on technical assumptions that are
hard to validate. In contrast, our non-asymptotic results hold for standard
neural networks used in practice. Namely, our main assumption for avoiding
rank collapse and gradient explosion is that samples in the input batch are
not linearly dependent, which we will show is necessary. To go beyond mean-
field regimes, we leverage recent theoretical advancements in Weingarten
calculus (Weingarten, 1978; Collins, 2003; Banica et al., 2011; Collins and
Śniady, 2006; Collins et al., 2022).

##  3 Main results

We will develop our theory by constructing networks that do not suffer from
gradient explosion (Sec. 3.3) and still orthogonalize (Sec. 3.1). The
construction is similar to the network studied by Daneshmand et al. (2021): an
MLP with batch normalization and linear activations. Formally, let
$X^{\ell}\in\mathbb{R}^{d\times n}$ denote the representation of $n$ samples
in $\mathbb{R}^{d}$ at layer $\ell$, then

$\displaystyle X_{\ell+1}=\text{\sc{Bn}}(W_{\ell}X_{\ell}),$
$\displaystyle\ell=0,\dots,L,$ (3)

where $W_{\ell}\in\mathbb{R}^{d\times d}$ are random weights. Analogous to
recent theoretical studies of batch normalization (Daneshmand et al., 2021,
2020), we define the BN operator $\text{\sc{Bn}}:\mathbb{R}^{d\times
n}\to\mathbb{R}^{d\times n}$ as

$\displaystyle\text{\sc{Bn}}(X)$
$\displaystyle=\text{diag}(XX^{\top})^{-\frac{1}{2}}X,\quad\text{\sc{Bn}}(X)_{ij}=\frac{X_{ij}}{\sqrt{\sum_{k=1}^{d}X_{ik}^{2}}}\leavevmode\nobreak\
.$ (4)

Note that compared to the standard BN operator, mean reduction in equation 4
is omitted. Our motivation for this modification, similar to Daneshmand et al.
(2021), is purely technical and to streamline our theory. We will
experimentally show that using standard BN modules instead does not influence
our results on gradient explosion and signal propagation (for more details see
Figure G6). A second minor difference is that in the denominator, we have
omitted a $\frac{1}{n}$ factor. However, this only amounts to a constant
scaling of the representations and does not affect our results.

Compared to Daneshmand et al. (2021), we need two main modifications to avoid
gradient explosion: (i) $n=d$, and (ii) $W_{\ell}$ are random _orthogonal_
matrices. More precisely, we assume the distribution of $W_{\ell}$ is the Haar
measure over the orthogonal group denoted by $\mathbb{O}_{d}$ (Collins and
Śniady, 2006). Such an initialization scheme is widely used in deep neural
networks without batch normalization (Saxe et al., 2013a; Xiao et al., 2018;
Pennington et al., 2017). For MLP networks with BN, we prove such
initialization avoids the issue of gradient explosion, while simultaneously
orthogonalizing the inputs.

###  3.1 Tracking signal propagation via orthogonality

As discussed, batch normalization has an important orthogonalization bias that
influences training. Without normalization layers, representations in many
architectures face the issue of rank-collapse, which happens when network
outputs become collinear for arbitrary inputs, hence their directions become
insensitive to the changes in the input. In contrast, the outputs in networks
with batch normalization become increasingly orthogonal through the layers,
thereby enhancing the signal propagation in depth (Daneshmand et al., 2021).
Thus, it is important to check whether the constructed network maintains the
important property of orthogonalization.

##### Isometry gap.

Our analysis relies on the notion of _isometry gap_ ,
$\phi:\mathbb{R}^{d\times d}\to\mathbb{R}$, introduced by Joudaki et al.
(2023a). Isometry gap is defined as

$\displaystyle\phi(X)=-\log\left(\frac{\text{det}(X^{\top}X)^{\frac{1}{d}}}{\frac{1}{d}\text{Tr}(X^{\top}X)}\right).$
(5)

One can readily check that $\phi(X)\geq 0$ and it is zero when $X$ is an
orthogonal matrix, i.e., $XX^{\top}=I_{d}$. The _isometry_ denoted by
$\mathcal{I}:\mathbb{R}^{d\times d}\to\mathbb{R}$ is defined as
$\mathcal{I}(X)=\exp(-\phi(X))$.

##### Geometric interpretation of isometry.

While the formula for isometry gap may seem enigmatic at first, it has a
simple geometric interpretation that makes it intuitively understandable. The
determinant $\text{det}(X^{\top}X)=\text{det}(X)^{2}$ is the squared volume of
the parallelepiped spanned by the columns of $X$, while $\text{Tr}(X^{\top}X)$
is the sum squared-norms of the columns of $X$. Thus, the ratio between the
two provides a scale-invariant notion of volume and isometry. On the one hand,
if there is any collinearity between the columns, the volume will vanish and
the isometry gap will be infinity, $\phi(X)=\infty$. On the other hand,
$\phi(X)=0$ implies $X^{\top}X$ is a scaled identity matrix. We will prove
$\phi$ serves as a Lyapunov function for the chain of hidden representations
$\\{X_{\ell}\\}_{\ell=0}^{\infty}$.

##### Theory for orthogonalization.

The following theorem establishes the link between orthogonality of
representations and depth.

######  Theorem 1.

There is an absolute constant $C$ such that for any layer $\ell\leq L$ we have

$\displaystyle\mathbb{E}\phi(X_{\ell+1})\leq\phi(X_{0})e^{-\ell/k},$ where
$\displaystyle k:=Cd^{2}(1+d\phi(X_{0}))\leavevmode\nobreak\ .$ (6)

Theorem 1 states that if the samples in the input batch are not linearly
dependent, representations approach orthogonality at an exponential rate in
depth. The orthogonalization in depth ensures the avoidance of the rank
collapse of representations, which is a known barrier to training deep neural
networks (Daneshmand et al., 2020; Saxe et al., 2013a; Bjorck et al., 2018).

Figure 1 compares the established theoretical decay rate of $\phi$ with the
practical rate. Interestingly, the plot confirms that the rate depends on
width in practice, akin to the theoretical rate in Theorem 1.

Figure 1: Isometry gap (y-axis, log-scale) in depth for an MLP with orthogonal
weights, over randomly generated data. As predicted by Theorem 1, isometry gap
of representations vanishes at an exponential rate. The solid traces are
averaged over $10$ independent runs, and the dashed traces show the
theoretical prediction from Theorem 1.

It is worth mentioning that the condition on the input samples to not be
linearly dependent is necessary to establish this result. One can readily
check that starting from a rank-deficient input, neither matrix products, nor
batch-normalization operations can increase the rank of the representations.
Since this assumption is quantitative, we can numerically verify it by
randomly drawing many input mini-batches and check if they are linearly
dependent. For CIFAR10, CIFAR100, MNIST and FashionMNIST, we empirically
tested that most batches across various batch sizes are full-rank (see Section
D for details on the average rank of a batch in these datasets).

Theorem 1 distinguishes itself from the existing orthogonalization results in
the literature (Yang et al., 2019; Joudaki et al., 2023a) as it is non-
asymptotic and holds for networks with finite width. Since practical networks
have finite width and depth, non-asymptotic results are crucial for their
applicability to real-world settings. While Daneshmand et al. (2021) provide a
non-asymptotic bound for orthogonalization, the main result relies on an
assumption that is hard to verify.

Proof idea of Theorem 1. We leverage a recent result established by Joudaki et
al. (2023a), proving that the isometry gap does not decrease with BN layers.
For all non-degenerate matrices $X\in\mathbb{R}^{d\times d}$, the following
holds

$\displaystyle\mathcal{I}(\text{\sc{Bn}}(X))\geq\left(1+\frac{\mathrm{variance}\\{\|X_{j\cdot}\|\\}_{j=1}^{d}}{(\mathrm{mean}\\{\|X_{j\cdot}\|\\}_{j=1}^{d})^{2}}\right)\mathcal{I}(X)\leavevmode\nobreak\
.$

Using the above result, we can prove that matrix multiplication with
orthogonal weights also does not decrease isometry as stated in the next
lemma.

######  Lemma 2 (Isometry after rotation).

Let $X\in\mathbb{R}^{d\times d}$ and $W\in\mathbb{R}^{d\times d}$ be an
orthogonal matrix and $X^{\prime}=WX$; then,

$\displaystyle\mathcal{I}(\text{\sc{Bn}}(X^{\prime}))\geq\left(1+\frac{\mathrm{variance}\\{\|X^{\prime}_{j\cdot}\|\\}_{j=1}^{d}}{(\mathrm{mean}\\{\|X^{\prime}_{j\cdot}\|\\}_{j=1}^{d})^{2}}\right)\mathcal{I}(X)\leavevmode\nobreak\
.$ (7)

It is straightforward to check that there exists at least an orthogonal matrix
$W$ for which $\mathcal{I}(\text{\sc{Bn}}(WX))=1$ (see Corollary A.3). Thus,
$\mathcal{I}(\cdot)$ strictly increases for some weight matrices, as long as
$X$ is not orthogonal. When the distribution of $W$ is the Haar measure over
the orthogonal group, we can leverage recent developments in Weingarten
calculus (Weingarten, 1978; Banica et al., 2011; Collins and Śniady, 2006;
Collins et al., 2022) to calculate a rate for the isometry increase in
expectation:

######  Theorem 3.

Suppose $W\sim\mathbb{O}_{d}$ is a matrix drawn from $\mathbb{O}_{d}$ such
that the distribution of $W$ and $UW$ are the same for all orthogonal matrices
$U$. Let $\\{\lambda_{i}\\}_{i=1}^{d}$ be the eigenvalues of $XX^{\top}$.
Then,

$\displaystyle\mathbb{E}_{W}\left[\mathcal{I}(\text{\sc{Bn}}(WX))\right]\geq\left(1-\frac{\sum_{k=1}^{d}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}\right)^{-1}\mathcal{I}(X)$
(8)

holds for all $X=\text{\sc{Bn}}(\cdot)$, with equality for orthogonal
matrices.

The structure in $X$ induced by Bn ensures its eigenvalues lie in the interval
$(0,1]$, in that the multiplicative factor in the above inequality is always
greater than one. In other words, $\mathcal{I}(\cdot)$ increases by a constant
factor in expectation that depends on how close $X$ is to an orthogonal
matrix.

The connection between Theorem 3 and the main isometry gap bound stated in
Theorem 1 is established in the following Corollary (recall
$\phi=-\log\mathcal{I}$).

######  Corollary 4 (Isometry gap bound).

Suppose the same setup as in Theorem 3, where $X^{\prime}=WX$. Then, we have:

$\displaystyle\mathbb{E}_{W}[\phi(X^{\prime})|X]\leq\phi(X)+\log\left(1-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}\right).$
(9)

Notice that the term
$\frac{\sum_{k=1}^{d}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}=\mathcal{O}(\frac{1}{d})$,
yielding
$\log\left[1-\frac{\sum_{k=1}^{d}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}\right]\leq
0$. The rest of proof is based on an induction over the layers, presented in
Appendices A and B.

###  3.2 Orthogonalization and gradient explosion

There is a subtle connection between orthogonalization and gradient explosion.
Suppose the input batch is rank-deficient, i.e., degenerate. As elaborated
above, since all operations in our MLP can be formulated as matrix products,
they cannot recover the rank of the representations, which thus remain
degenerate. By perturbing the input such that it becomes full-rank, the output
matrix becomes orthogonal, hence non-degenerate at an exponential rate in
depth as proven in Theorem 1.

Figure 2: Logarithmic plot for the gradient norm of the first layer for
networks with different number of layers evaluated on degenerate (orange) and
non-degenerate (blue) inputs. The degenerate inputs contain repeated samples
from CIFAR10 in the batch, measured at initialization for MLPs of various
depths. While gradients explode for degenerate inputs, there is no explosion
for non-degenerate inputs. Traces are averaged over $10$ independent runs.

Thus, a slight change in inputs leads to a significant change in outputs from
degeneracy to orthogonality. Considering that the gradient measures changes in
the loss for infinitesimal inputs changes, the large changes in outputs
potentially lead to gradient explosion. While this is only an intuitive
argument, we observe that in practice the gradient does explode for degenerate
inputs, as shown in Figure 2.

Nonetheless, in Figure 2 we observe that for non-degenerate inputs the
gradient norm does not explode. In fact, we observe that inputs are often non-
degenerate in practice (see Table D1 for details). Thus, an important question
is whether the gradient norm remains bounded for non-degenerate input batches.
Remarkably, we can not empirically verify that for _all degenerate inputs_ the
gradient norm remains bounded. Therefore, a theoretical guarantee is necessary
to ensure avoiding gradient explosion.

###  3.3 Avoiding gradient explosion in depth

So far, we have proven that the constructed network maintains the
orthogonalization property of BN. Now, we turn our focus to the gradient
analysis. The next theorem proves that the constructed network does not suffer
from gradient explosion in depth for non-degenerate input matrices.

######  Theorem 5.

Let the loss function $\mathcal{L}:\mathbb{R}^{d\times d}\to\mathbb{R}$ be
$\mathcal{O}(1)$-Lipschitz, and input batch $X_{0}$ be non-degenerate. Then,
there exists an absolute constant $C$ such that for all $\ell\leq L$ it holds

$\displaystyle\mathbb{E}\left[\log\|\nabla_{W_{\ell}}\mathcal{L}(X_{\ell})\|\right]\leq
Cd^{5}(\phi(X_{0})^{3}+1)$ (10)

where the expectation is over the random orthogonal weight matrices.

######  Remark 1.

For degenerate inputs $\phi(X_{0})=\infty$ holds, in that the bound becomes
vacuous.

######  Remark 2.

The $\mathcal{O}(1)$-Lipschitz condition holds in many practical settings. For
example, in a classification setting, MSE and cross entropy losses obey the
$\mathcal{O}(1)$-Lipschitz condition (see Lemma C.2).

Figure 3: Logarithmic plot for the gradient norm of the first layer for
networks with different number of layers evaluated on CIFAR10. For Gaussian
weights (orange) the gradient-norm grows at an exponential rate, as predicted
by Yang et al. (2019, Theorem 3.9), while for orthogonal weights (blue)
gradients remain bounded by a constant, validating Theorem 5. Traces are
averaged over $10$ runs and shaded regions denote the $95\%$ confidence
intervals.

Note that the bound is stated for the expected value of log-norm of the
gradients, which can be interpreted as bits of precision needed to store the
gradient matrices. Thus, the fact that depth does not appear in any form in
the upper bound of Theorem 5 points out that training arbitrarily deep MLPs
with orthogonal weights will not face numerical issues that arise with
Gaussian weights (Yang et al., 2019) as long as the inputs are non-degenerate.
Such guarantees are necessary to ensure backpropagation will not face
numerical issues.

Theorem 5 states that as long as the input samples are not linearly dependent,
the gradients remain bounded for any arbitrary depth $L$. As discussed in the
previous section and evidenced in Figure 2, this is necessary to avoid
gradient explosion. Therefore, the upper bound provided in Theorem 5 is tight
in terms of inputs constraints. Furthermore, as mentioned before, random
batches sampled from commonly used benchmarks, such as CIFAR10, CIFAR100,
MNIST, and FashionMNIST, are non-degenerate in most practical cases (see
Section D for more details). Thus, the assumptions and thereby assertions of
the theorem are valid for all practical purposes.

To the best of our knowledge, Theorem 5 is the first non-asymptotic gradient
analysis that holds for networks with batch normalization and finite width.
Previous results heavily rely on mean field analyses in asymptotic regimes,
where the network width tends to infinity (Yang et al., 2019). While mean-
field analyses have brought many insights about the rate of gradient
explosion, they are often specific to Gaussian weights. Here, we show that
non-Gaussian weights can avoid gradient explosion, which has previously been
considered “unavoidable” (Yang et al., 2019). Figure 3 illustrates this
pronounced discrepancy.

Proof idea of Theorem 5. The first important observation is that, due to the
chain rule, we can bound the log-norm of the gradient of a composition of
functions, by bounding the summation of the log-norms of the input-output
Jacobian of each layer, plus two additional terms corresponding to the loss
term and the gradient of the first layer in the chain. If we discount the
effect of the first and last terms, the bulk of the analysis is dedicated to
bounding the total sum of log-norms of per layer input-output Jacobian, i.e.,
the fully connected and batch normalization layers. The second observation is
that because the weights are only rotations, their Jacobian has eigenvalues
equal to $1.$ Thus, the log-norm of gradients corresponding to fully connected
layers vanish. What remains is to show that for any arbitrary depth $\ell,$
the log-norm of gradients of batch normalization layers also remains bounded.
The main technical novelty for proving this step is showing that the log-norm
of the gradient of Bn layers is upper bounded by the isometry gap of pre-
normalization matrices. Thus, we can invoke the exponential decay in isometry
gap stated in Theorem 1 to establish a bound on the log-norm of the gradient
of these layers. Finally, since the decay in isometry gap is exponentially
fast, the bound on the total sum of log-norm of the gradients amounts to a
geometric sum that remains bounded for any arbitrary depth $\ell.$

##  4 Implications on training

In this section, we experimentally validate the benefits of avoiding gradient
explosion and rank collapse for training. Thus far, we have proved that our
constructed neural network with BN does not suffer from gradient explosion in
Theorem 5, and does not have the rank collapse issue in depth via the
orthogonalization property established in Theorem 1. We find that the
constructed MLP is therefore less prone to numerical issues that arise when
training deep networks.

Figure 4: Contrasting the training accuracy of MLPs with BN and shaped sin,
shaped tanh and identity activations, on the CIFAR10 dataset. The identity
activation performs much worse than the nonlinearities, confirming that the
$\sin$ and $\tanh$ networks are not operating in the linear regime. The
networks are trained with vanilla SGD and the hyperparameters are width $100$,
batch size $100$, learning rate $0.001$.

By avoiding gradient explosion and rank collapse in depth, we observe that the
optimization with vanilla minibatch Stochastic Gradient Descent (SGD) exhibits
an almost depth-independent convergence rate for linear MLPs. In other words,
the number of iterations to get to a certain accuracy does not vary widely
between networks with different depths. Figure 4 (c) shows the convergence of
SGD for CIFAR10, with learning rate $0.001$, for MLPs with width $d=100$ and
batch size $n=100.$ While the SGD trajectory strongly diverges from the
initial conditions that we analyze theoretically, Figure 4 shows that the
gradients remain stable during training, as well as the fact that different
depths exhibit largely similar accuracy curves.

While the empirical evidence for our MLP with linear activation is
encouraging, non-linear activations are essential parts of feature learning
(Nair and Hinton, 2010; Klambauer et al., 2017; Hendrycks and Gimpel, 2016;
Maas et al., 2013). However, introducing non-linearity violates one of the key
parts of our theory, in that it prevents representations from reaching perfect
isometry (see Figure G5 in the Appendix for details on the connection between
non-linearities and gradient explosion in depth). Intuitively, this is due to
the fact that non-linear layers, as opposed to rotations and batch
normalization, perturb the isometry of representations and prevent them from
reaching zero isometry gap in depth. This problem turns out to be not just a
theoretical nuisance, but to play a direct role in the gradient explosion
behavior. While the situation may seem futile at first, it turns out that
activation shaping (Li et al., 2022; Zhang et al., 2022; Martens et al., 2021;
He et al., 2023; Noci et al., 2023) can alleviate this problem, which is
discussed next. For the remainder of this section, we focus on the training of
MLPs with non-linear activations, as well as standard batch normalization and
fully connected layers.

##  5 Activation shaping based on the theoretical analysis

In recent years, several works have attempted to overcome the challenges of
training very deep networks by parameterizing activation functions. In a
seminal work, Martens et al. (2021) propose deep kernel shaping, which is
aimed at facilitating the training of deep networks without relying on skip
connections or normalization layers, and was later extended to LeakyReLU in
tailored activation transformations (Zhang et al., 2022). In a similar
direction, Li et al. (2022) propose activation shaping in order to avoid a
degenerate output covariance. Li et al. (2022) propose shaping the negative
slope of LeakyReLU towards identity to ensure that the output covariance
matrix remains non-degenerate when the networks becomes very deep.

Since kernel and activation shaping aim to replace normalization, they have
not been used in conjunction with normalization layers. In fact, in networks
with batch normalization, even linear activations have non-degenerate outputs
(Daneshmand et al., 2021; Yang et al., 2019) and exploding gradients (Yang et
al., 2019). Thus, shaping activations towards identity in the presence of
normalization layers may seem fruitless. Remarkably, we empirically
demonstrate that we can leverage activation shaping to avoid gradient
explosion in depth by using a pre-activation gain at each layer.

Figure 5: Logarithmic plot contrasting the effect of gain on the gradient at
initialization of the first layer, for networks with different number of
layers initialized with orthogonal weights, BN and different activations,
evaluated on CIFAR10. The networks have hyperparameters width $100$, batch
size $100$. Traces are averaged over $10$ independent runs, with the shades
showing the $95\%$ confidence interval.

Inspired by our theory, we develop a novel activation shaping scheme for
networks with BN. The main strategy consists of shaping the activation
function towards a linear function across the layers. Our activation shaping
consists of tuning the gain of the activation, i.e., tuning $\alpha$ for
$\sigma(\alpha x).$ We consider non-linear activations
$\sigma\in\\{\tanh,\sin\\}.$

The special property that both $\tanh$ and $\sin$ activations have in common
is that they are centered, $\sigma(0)=0,$ are differentiable around the origin
$\sigma^{\prime}(0)=1,$ and have bounded gradients $\sigma^{\prime}(x)\leq
1,\forall x$. Therefore, by tuning the per-layer pre-activation gain
$\alpha_{\ell}$ towards $0$, the non-linearities behave akin to the identity
function. This observation inspires us to study the relationship between the
rate of gradient explosion for each layer as a function of the gain parameter
$\alpha_{\ell}$. Formally, we consider an MLP with shaped activations using
gain $\alpha_{\ell}$ for the $\ell$th layer, that has the update rule

$\displaystyle
X_{\ell+1}=\sigma(\alpha_{\ell}\text{\sc{Bn}}(W_{\ell}X_{\ell}))\leavevmode\nobreak\
.$ (11)

Since the gradient norm has an exponential growth in depth, as shown in Figure
5, we can compute the slope of the linear growth rate of log-norm of gradients
in depth. We define the rate of explosion for a model of depth $L$ and gain
$\alpha_{\ell}$ at layer $\ell$ as the slope of the log norm of the gradients
$R(\ell,\alpha_{\ell})$. We show in Figure 5 that by tuning the gain properly,
we are able to reduce the exponential rate of the log-norm of the gradients by
diminishing the slope of the rate curve and achieve networks trainable at
arbitrary depths, while still maintaining the benefits of the non-linear
activation. The main idea for our activation shaping strategy is to have a
bounded total sum of rates across layers, by ensuring faster decay than a
harmonic series (see App. E for more details on activation shaping). Figure 5
illustrates that this activation shaping strategy effectively avoids gradient
explosion while maintaining the signal propagation and orthogonality of the
outputs in depth. Furthermore, Figure 4 shows that the training accuracy
remains largely depth-independent. For further experiments using activation
shaping, see Appendix G.

##  6 Discussion

##### Implicit bias of SGD towards orthogonality in optimization.

Optimization over orthogonal matrices has been an effective approach for
training deep neural networks. Enforcing orthogonality during training ensures
that the spectrum of the weight matrices remains bounded, which prevents
gradient vanishing and explosion in depth. Vorontsov et al. (2017) study how
different orthogonality constraints affect training performance in RNNs. For
example Lezcano-Casado and Martınez-Rubio (2019) leverage the exponential map
on the orthogonal group, Jose et al. (2018) decompose RNN transition matrices
in Kronecker factors and impose soft constraints on each factor and Mhammedi
et al. (2017) introduce a constraint based on Householder matrices.

Figure 6: Implicit orthogonality bias of SGD. Training an MLP with width
$d=100,$ batch size $n=100,$ and depth $L=1000,$ activation $\tanh$, using SGD
with $\text{lr}=0.001$ (a) Isometry gap (y-axis; log-scale) of weight matrices
across all layers throughout training. (b) Gradient norms at each layer during
training.

While these studies _enforce_ orthogonality constraints, one of our most
striking empirical observations is that when our MLP grows very deep, the
middle layers remain almost orthogonal even after many steps of SGD. As shown
in Figure 6, for $1000$ layer networks, the middle layers remain orthogonal
during training. One could hypothesize that this is due to small gradients in
these layers. In Figure 6, we observe that the gradients of these middle
layers are not negligible. Thus, in our MLP construction, both with linear
activation and with activation shaping, the gradient dynamics have an
_implicit bias_ to optimize over the space of orthogonal matrices. The
mechanisms underlying this implicit orthogonality bias will be an ample
direction for future research.

#### Author Contributions

Alexandru Meterez: proofs of Section A, driving experiments, and writing the
paper. Amir Joudaki: proofs of Section B and Section C, designing the
activation shaping scheme, and writing the paper. Francesco Orabona: proposing
the idea of using orthogonal weights to achieve prefect isometry, reading the
proofs, help with writing. Alexander Immer: reading the proofs, help with
experiments and paper writing. Gunnar Rätsch: help with experimental designs
for activation shaping and paper writing. Hadi Daneshmand: proposed using
orthogonal weights to avoid gradients explosion, leading the proofs help with
paper writing.

#### Acknowledgments

Amir Joudaki is funded through Swiss National Science Foundation Project Grant
#200550 to Andre Kahles. We acknowledge support from the NSF TRIPODS program
(award DMS-2022448). Alexander Immer acknowledges support from the Max Planck
ETH Center for Learning Systems. Amir Joudaki and Alexander Immer were
partially funded by ETH Core funding (to G.R.).

## References

  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.  RoBERTa: A robustly optimized BERT pretraining approach.  _arXiv preprint arXiv:1907.11692_ , 2019.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Goyal et al. (2021) Naman Goyal, Jingfei Du, Myle Ott, Giri Anantharaman, and Alexis Conneau.  Larger-scale transformers for multilingual masked language modeling.  _arXiv preprint arXiv:2105.00572_ , 2021.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551, 2020.
  * Liu et al. (2022) Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie.  A ConvNet for the 2020s.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 11976–11986, 2022.
  * Woo et al. (2023) Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon, and Saining Xie.  ConvNeXt V2: Co-designing and scaling convnets with masked autoencoders.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 16133–16142, 2023.
  * Wu et al. (2021) Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang.  Cvt: Introducing convolutions to vision transformers.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 22–31, 2021.
  * Saxe et al. (2013a) Andrew M Saxe, James L McClelland, and Surya Ganguli.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  _arXiv preprint arXiv:1312.6120_ , 2013a.
  * Feng et al. (2022) Ruili Feng, Kecheng Zheng, Yukun Huang, Deli Zhao, Michael Jordan, and Zheng-Jun Zha.  Rank diminishing in deep neural networks.  _Advances in Neural Information Processing Systems_ , 35:33054–33065, 2022.
  * Daneshmand et al. (2020) Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, and Aurelien Lucchi.  Batch normalization provably avoids ranks collapse for randomly initialised deep networks.  _Advances in Neural Information Processing Systems_ , 33:18387–18398, 2020.
  * Dong et al. (2021) Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas.  Attention is not all you need: Pure attention loses rank doubly exponentially with depth.  In _International Conference on Machine Learning_ , pages 2793–2803. PMLR, 2021.
  * Noci et al. (2022) Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi.  Signal propagation in Transformers: Theoretical perspectives and the role of rank collapse.  _Advances in Neural Information Processing Systems_ , 35:27198–27211, 2022.
  * Daneshmand et al. (2021) Hadi Daneshmand, Amir Joudaki, and Francis Bach.  Batch normalization orthogonalizes representations in deep random networks.  _Advances in Neural Information Processing Systems_ , 34:4896–4906, 2021.
  * Ioffe and Szegedy (2015) Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International conference on machine learning_ , pages 448–456. pmlr, 2015.
  * Yang et al. (2019) Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S Schoenholz.  A mean field theory of batch normalization.  _ICLR_ , 2019.
  * Joudaki et al. (2023a) Amir Joudaki, Hadi Daneshmand, and Francis Bach.  On the impact of activation and normalization in obtaining isometric embeddings at initialization.  _arXiv preprint arXiv:2305.18399_ , 2023a.
  * Joudaki et al. (2023b) Amir Joudaki, Hadi Daneshmand, and Francis Bach.  On bridging the gap between mean field and finite width deep random multilayer perceptron with batch normalization.  _International Conference on Machine Learning_ , 2023b.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification.  In _Proceedings of the IEEE international conference on computer vision_ , pages 1026–1034, 2015.
  * Glorot and Bengio (2010) Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pages 249–256. JMLR Workshop and Conference Proceedings, 2010.
  * Martens et al. (2021) James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszcz, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S Schoenholz.  Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping.  _arXiv preprint arXiv:2110.01765_ , 2021.
  * Hochreiter (1998) Sepp Hochreiter.  The vanishing gradient problem during learning recurrent neural nets and problem solutions.  _International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems_ , 6(02):107–116, 1998.
  * Lubana et al. (2021) Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka.  Beyond BatchNorm: towards a unified understanding of normalization in deep learning.  _Advances in Neural Information Processing Systems_ , 34:4778–4791, 2021.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Klambauer et al. (2017) Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter.  Self-normalizing neural networks.  _Advances in neural information processing systems_ , 30, 2017.
  * Pennington et al. (2017) Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli.  Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice.  _Advances in neural information processing systems_ , 30, 2017.
  * Xiao et al. (2018) Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington.  Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks.  In _International Conference on Machine Learning_ , pages 5393–5402. PMLR, 2018.
  * Mishkin and Matas (2015) Dmytro Mishkin and Jiri Matas.  All you need is a good init.  _arXiv preprint arXiv:1511.06422_ , 2015.
  * Arjovsky et al. (2016) Martin Arjovsky, Amar Shah, and Yoshua Bengio.  Unitary evolution recurrent neural networks.  In _International conference on machine learning_ , pages 1120–1128. PMLR, 2016.
  * Le et al. (2015) Quoc V Le, Navdeep Jaitly, and Geoffrey E Hinton.  A simple way to initialize recurrent networks of rectified linear units.  _arXiv preprint arXiv:1504.00941_ , 2015.
  * Henaff et al. (2016) Mikael Henaff, Arthur Szlam, and Yann LeCun.  Recurrent orthogonal networks and long-memory tasks.  In _International Conference on Machine Learning_ , pages 2034–2042. PMLR, 2016.
  * Fukumizu (1998) Kenji Fukumizu.  Effect of batch learning in multilayer neural networks.  _Gen_ , 1(04):1E–03, 1998.
  * Baldi and Hornik (1995) Pierre F Baldi and Kurt Hornik.  Learning in linear neural networks: A survey.  _IEEE Transactions on neural networks_ , 6(4):837–858, 1995.
  * Saxe et al. (2013b) Andrew M Saxe, James L McClellans, and Surya Ganguli.  Learning hierarchical categories in deep neural networks.  In _Proceedings of the Annual Meeting of the Cognitive Science Society_ , volume 35, 2013b.
  * Saxe et al. (2019) Andrew M Saxe, James L McClelland, and Surya Ganguli.  A mathematical theory of semantic development in deep neural networks.  _Proceedings of the National Academy of Sciences_ , 116(23):11537–11546, 2019.
  * Baldi and Hornik (1989) Pierre Baldi and Kurt Hornik.  Neural networks and principal component analysis: Learning from examples without local minima.  _Neural networks_ , 2(1):53–58, 1989.
  * Yun et al. (2017) Chulhee Yun, Suvrit Sra, and Ali Jadbabaie.  Global optimality conditions for deep neural networks.  _arXiv preprint arXiv:1707.02444_ , 2017.
  * Arora et al. (2018) Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu.  A convergence analysis of gradient descent for deep linear neural networks.  _arXiv preprint arXiv:1810.02281_ , 2018.
  * Shamir (2019) Ohad Shamir.  Exponential convergence time of gradient descent for one-dimensional deep linear neural networks.  In _Conference on Learning Theory_ , pages 2691–2713. PMLR, 2019\.
  * Du and Hu (2019) Simon Du and Wei Hu.  Width provably matters in optimization for deep linear neural networks.  In _International Conference on Machine Learning_ , pages 1655–1664. PMLR, 2019.
  * Li et al. (2022) Mufan Li, Mihai Nica, and Dan Roy.  The neural covariance SDE: Shaped infinite depth-and-width networks at initialization.  _Advances in Neural Information Processing Systems_ , 35:10795–10808, 2022.
  * Pennington and Worah (2017) Jeffrey Pennington and Pratik Worah.  Nonlinear random matrix theory for deep learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Weingarten (1978) Don Weingarten.  Asymptotic behavior of group integrals in the limit of infinite rank.  _Journal of Mathematical Physics_ , 19(5):999–1001, 1978.
  * Collins (2003) Benoît Collins.  Moments and cumulants of polynomial random variables on unitary groups, the Itzykson-Zuber integral, and free probability.  _International Mathematics Research Notices_ , 2003(17):953–982, 2003.
  * Banica et al. (2011) Teodor Banica, Benoit Collins, and Jean-Marc Schlenker.  On polynomial integrals over the orthogonal group.  _Journal of Combinatorial Theory, Series A_ , 118(3):778–795, 2011.
  * Collins and Śniady (2006) Benoît Collins and Piotr Śniady.  Integration with respect to the haar measure on unitary, orthogonal and symplectic group.  _Communications in Mathematical Physics_ , 264(3):773–795, 2006.
  * Collins et al. (2022) Benoit Collins, Sho Matsumoto, and Jonathan Novak.  The Weingarten calculus.  _arXiv preprint arXiv:2109.14890_ , 2022.
  * Bjorck et al. (2018) Nils Bjorck, Carla P Gomes, Bart Selman, and Kilian Q Weinberger.  Understanding batch normalization.  _Advances in neural information processing systems_ , 31, 2018.
  * Nair and Hinton (2010) Vinod Nair and Geoffrey E Hinton.  Rectified linear units improve restricted boltzmann machines.  In _Proceedings of the 27th international conference on machine learning (ICML-10)_ , pages 807–814, 2010.
  * Hendrycks and Gimpel (2016) Dan Hendrycks and Kevin Gimpel.  Gaussian error linear units (GELUs).  _arXiv preprint arXiv:1606.08415_ , 2016.
  * Maas et al. (2013) Andrew L Maas, Awni Y Hannun, and Andrew Y Ng.  Rectifier nonlinearities improve neural network acoustic models.  In _Proc. ICML_ , volume 30, page 3. Atlanta, GA, 2013.
  * Zhang et al. (2022) Guodong Zhang, Aleksandar Botev, and James Martens.  Deep learning without shortcuts: Shaping the kernel with tailored rectifiers.  _arXiv preprint arXiv:2203.08120_ , 2022.
  * He et al. (2023) Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andrew Brock, Samuel L Smith, and Yee Whye Teh.  Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation.  _arXiv preprint arXiv:2302.10322_ , 2023.
  * Noci et al. (2023) Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, and Daniel M Roy.  The shaped transformer: Attention models in the infinite depth-and-width limit.  _arXiv preprint arXiv:2306.17759_ , 2023.
  * Vorontsov et al. (2017) Eugene Vorontsov, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal.  On orthogonality and learning recurrent networks with long term dependencies.  In _International Conference on Machine Learning_ , pages 3570–3578. PMLR, 2017.
  * Lezcano-Casado and Martınez-Rubio (2019) Mario Lezcano-Casado and David Martınez-Rubio.  Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group.  In _International Conference on Machine Learning_ , pages 3794–3803. PMLR, 2019.
  * Jose et al. (2018) Cijo Jose, Moustapha Cissé, and Francois Fleuret.  Kronecker recurrent units.  In _International Conference on Machine Learning_ , pages 2380–2389. PMLR, 2018.
  * Mhammedi et al. (2017) Zakaria Mhammedi, Andrew Hellicar, Ashfaqur Rahman, and James Bailey.  Efficient orthogonal parametrisation of recurrent neural networks using Householder reflections.  In _International Conference on Machine Learning_ , pages 2401–2409. PMLR, 2017.
  * Collins and Matsumoto (2009) Benoît Collins and Sho Matsumoto.  On some properties of orthogonal Weingarten functions.  _Journal of Mathematical Physics_ , 50(11), 2009.

## Table of contents

The appendix is structured in the following sections:

  * •

Appendix A: contains the proof based on Weingarten calculus for the
conditional isometry gap decay rate, conditioned on the previous layer in the
network.

  * •

Appendix B: contains the proof for the isometry gap exponential rate decay in
depth, based on the conditional rate established in Appendix A.

  * •

Appendix C: contains the proof for the log-norm gradient bound at arbitrary
depths, based on the orthogonality result established in Appendix B.

  * •

Appendix D: contains the numerical verification of the average rank of the
input batch in common datasets.

  * •

Appendix E: contains the full derivation and additional explanations of the
activation shaping heuristic.

  * •

Appendix F: contains further results on CIFAR10, MNIST, FashionMNIST and
CIFAR100, showcasing the implicit orthogonality bias of SGD during training.

  * •

Appendix G: contains supplemental experiments, including train and test
results on MNIST, FashionMNIST, CIFAR100, as well as the effect of mean
reduction and non-linearities without shaping on the gradient log-norms.

##  Appendix A Conditional orthogonalization

##### Isometry after rotation.

Our analysis is based on [Joudaki et al., 2023a, Corollary 3], which we
restate in the following Lemma:

######  Lemma A.1.

For all non-degenerate matrices $X\in\mathbb{R}^{d\times d}$, we have:

$\displaystyle\mathcal{I}(\text{\sc{Bn}}(X))\geq\left(1+\frac{\text{variance}\\{\|X_{j\cdot}\|\\}_{j=1}^{d}}{(\text{mean}\\{\|X_{j\cdot}\|\\}_{j=1}^{d})^{2}}\right)\mathcal{I}(X)\leavevmode\nobreak\
.$ (12)

Lemma A.1 proves isometry bias of Bn. The next lemma proves that isometry does
not change under rotation

######  Lemma A.2 (Isometry after rotation).

Let $X\in\mathbb{R}^{d\times d}$ and $W\sim\mathbb{O}_{d}$ be a random
orthogonal matrix and $X^{\prime}=WX$. Then,

$\displaystyle\mathcal{I}(\text{\sc{Bn}}(X^{\prime}))\geq\left(1+\frac{\text{variance}\\{\|X^{\prime}_{j\cdot}\|\\}_{j=1}^{d}}{(\text{mean}\\{\|X^{\prime}_{j\cdot}\|\\}_{j=1}^{d})^{2}}\right)\mathcal{I}(X)\leavevmode\nobreak\
.$ (13)

###### Proof.

Using properties of the determinant, we have

$\displaystyle\text{det}(X^{\prime}X^{\prime\top})=\text{det}(W)^{2}\text{det}(XX^{\top})=\text{det}(XX^{\top}),$
(14)

where the last equation holds since $W$ is an orthogonal matrix. Furthermore,

$\displaystyle\text{Tr}(X^{\prime}X^{\prime\top})$
$\displaystyle=\text{Tr}(WXX^{\top}W^{\top})$ (15)
$\displaystyle=\text{Tr}(XX^{\top}\underbrace{W^{\top}W}_{=I})$ (16)
$\displaystyle=\text{Tr}(XX^{\top})\leavevmode\nobreak\ .$ (17)

Combining the last two equations with Lemma A.1 concludes the proof. ∎

##### Increasing isometry with rotations and Bn.

The last lemma proves the isometry bias does not decrease with rotation and
Bn. However, this does not prove a strict decrease in isometry with Bn and
rotation. The next lemma proves there exists an orthogonal matrix for which
the isometry is strictly increasing.

######  Corollary A.3 (Increasing isometry).

Let $X\in\mathbb{R}^{d\times d}$ and denote its singular value decomposition
$X=U\text{diag}\left(\\{\sigma_{i}\\}_{i=1}^{d}\right)V^{\top}$, where $U$ and
$V$ are orthogonal matrices. Then, we have:

$\displaystyle\mathcal{I}(\text{\sc{Bn}}(U^{\top}H))=1\leavevmode\nobreak\ .$
(18)

###### Proof.

Let $S=\text{diag}\left(\\{\sigma_{i}\\}_{i=1}^{d}\right)$ be the diagonal
matrix containing the singular values of $X$. Then, we have:

$\displaystyle\text{\sc{Bn}}(U^{\top}X)$
$\displaystyle=\text{\sc{Bn}}(U^{\top}USV^{\top})$ (19)
$\displaystyle=\text{\sc{Bn}}(SV^{\top})$ (20)
$\displaystyle=\text{diag}(SV^{\top}VS)^{-\frac{1}{2}}SV^{\top}$ (21)
$\displaystyle=S^{-1}SV^{\top}$ (22) $\displaystyle=V^{\top},$ (23)

which has maximum isometry 1 since it’s an orthonormal matrix. ∎

Thus, there exists a rotation that increases isometry with Bn for each non-
orthogonal matrix. The proof of the last corollary is based on a
straightforward application of Lemma 2.

##### Orthogonalization with randomness.

The isometric is non-decreasing in Lemma 2 and provably increases for a
certain rotation matrix (as stated in the last corollary). Hence, it is
possible to increase isometry with random orthogonal matrices.

######  Theorem A.4.

Suppose $W\sim\mathbb{O}_{d}$ is a matrix drawn from $\mathbb{O}_{d}$ such
that $W\stackrel{{\scriptstyle d}}{{=}}WU$ for all orthogonal matrices $U$.
Let $\\{\lambda_{i}\\}_{i=1}^{d}$ be the eigenvalues of $XX^{\top}$. Then,

$\displaystyle\mathbb{E}_{W}\left[\mathcal{I}(\text{\sc{Bn}}(WX))\right|X]\geq\left(\frac{1}{1-\frac{\sum_{k=1}^{d}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}}\right)\mathcal{I}(X)$
(24)

holds for all $X=\text{\sc{Bn}}(\cdot)$, with equality for orthogonal
matrices.

######  Remark 3.

Note that the assumption on $X=\text{\sc{Bn}}(\cdot),$ can be viewed as an
induction hypothesis, in that we can recursively apply this theorem to arrive
at a quantitative rate at depth.

Notably, $\sum_{i=1}^{d}\lambda_{i}=d$ if $X=\text{\sc{Bn}}(\cdot)$. Hence,
one can expect that $\sum_{i=1}^{d}\lambda_{i}^{2}<d^{2}$ for all full random
matrices $X$ in form of $X=\text{\sc{Bn}}(\cdot)$.

###### Proof.

We need to compute the variance/mean ratio in Lemma 2. Let
$X\in\mathbb{R}^{d\times d}$ have SVD decomposition
$X=U\text{diag}\\{\sigma_{i}\\}_{i=1}^{d}V^{\top}$ where $U$ and $V$ are
orthogonal matrix and $\sigma_{i}^{2}=\lambda_{i}$. Since the distribution of
$W$ is invariant to transformations with orthogonal matrices, the distribution
of $W$ equates those of $X^{\prime}=W\text{diag}\\{\sigma_{i}\\}V^{\top}$. It
easy to check that

$\displaystyle\|X_{j\cdot}^{\prime}\|=\sqrt{\sum_{i=1}^{d}\sigma_{i}^{2}W_{ji}^{2}}=\sqrt{\sum_{i=1}^{d}\lambda_{i}W_{ji}^{2}}\leavevmode\nobreak\
.$ (25)

Thus,

$\displaystyle\sum_{j=1}^{d}\|X_{j\cdot}^{\prime}\|^{2}=\sum_{i=1}^{d}\lambda_{i}=d,$
(26)

where the last equality holds due to the batch normalization. Thus, we have

$\mathbb{E}\left[\frac{d\sum_{j=1}^{d}\|X^{\prime}_{j\cdot}\|^{2}}{(\sum_{j=1}^{d}\|X^{\prime}_{j\cdot}\|)^{2}}\right]=\mathbb{E}\left[\frac{d^{2}}{(\sum_{j}\|X^{\prime}_{j\cdot}\|)^{2}}\right]\geq\frac{d^{2}}{\mathbb{E}\left[\left(\sum_{j=1}^{d}\|X^{\prime}_{j\cdot}\|\right)^{2}\right]}\leavevmode\nobreak\
.$ (27)

We need to estimate

$\displaystyle\mathbb{E}\left[\|X_{i}^{\prime}\|\|X_{j}^{\prime}\|\right]=\mathbb{E}\left[\left(\sum_{k=1}^{d}\lambda_{k}W_{ik}^{2}\right)^{\frac{1}{2}}\left(\sum_{k=1}^{d}\lambda_{k}W_{jk}^{2}\right)^{\frac{1}{2}}\right]\leavevmode\nobreak\
.$ (28)

Since square root function is concave, we have $\sqrt{x}\leq
1+\frac{1}{2}(x-1)$. Thus

$\displaystyle\mathbb{E}\left[\|X^{\prime}_{i}\|\|X^{\prime}_{j}\|\right]$
$\displaystyle\leq\mathbb{E}\left[\left(1+0.5\sum_{k=1}^{d}(\lambda_{k}-1)W_{ik}^{2}\right)\left(1+0.5\sum_{k=1}^{d}(\lambda_{k}-1)W_{jk}^{2}\right)\right]$
(29)
$\displaystyle=1+\frac{1}{4}\sum_{k,q}(\lambda_{k}-1)(\lambda_{q}-1)\mathbb{E}\left[W_{ik}^{2}W_{jq}^{2}\right]$
(30) $\displaystyle=1+\frac{1}{4}\left[\sum_{k\neq
q}(\lambda_{k}-1)(\lambda_{q}-1)\underbrace{\mathbb{E}\left[W_{ik}^{2}W_{jq}^{2}\right]}_{E_{1}}+\sum_{k=q}(\lambda_{k}-1)(\lambda_{k}-1)\underbrace{\mathbb{E}\left[W_{ik}^{2}W_{jk}^{2}\right]}_{E_{2}}\right],$
(31)

where in the first equality we have used the fact that the cross terms reduce,
where the expectations are applications of Weingarten calculus:

$\displaystyle\mathbb{E}\left[0.5\sum_{k=1}^{d}(\lambda_{k}-1)W_{ik}^{2}+0.5\sum_{k=1}^{d}(\lambda_{k}-1)W_{jk}^{2}\right]$
$\displaystyle=\mathbb{E}\left[0.5\sum_{k=1}^{d}\lambda_{k}W_{ik}^{2}+0.5\sum_{k=1}^{d}\lambda_{k}W_{jk}^{2}-1\right]$
(32)
$\displaystyle=0.5\sum_{k=1}^{d}\lambda_{k}\mathbb{E}\left[W_{ik}^{2}\right]+0.5\sum_{k=1}^{d}\lambda_{k}\mathbb{E}\left[W_{jk}^{2}\right]-1$
(33)
$\displaystyle=\frac{0.5}{d}\sum_{k=1}^{d}\lambda_{k}+\frac{0.5}{d}\sum_{k=1}^{d}\lambda_{k}-1$
(34) $\displaystyle=0\leavevmode\nobreak\ .$ (35)

The main quantity we must compute is an expectation of polynomials taken over
the Haar measure of the Orthogonal group $O(n)$. To carry out the computation,
we make use of Weingarten calculus [Banica et al., 2011, Collins and Śniady,
2006, Weingarten, 1978]. More specifically, we make use of the of the
Weingarten formula, studied by [Collins and Śniady, 2006, Collins et al.,
2022]:

$\displaystyle\int_{O(n)}r_{i_{1}j_{1}}r_{i_{2}j_{2}}\dots
r_{i_{2d}j_{2d}}d\mu(O(n))=\sum_{\sigma\in\mathcal{M}_{2d}}\sum_{\sigma\in\mathcal{M}_{2n}}\Delta_{\sigma}(\bm{i})\Delta_{\sigma}(\bm{j})\text{Wg}^{O}(\sigma^{-1}\tau),$
(36)

where $\mu(O(n))$ is the Haar measure of Orthogonal group. For an in depth
explanation of each quantity in the Weingarten formula, we refer the reader to
Collins et al. [2022, Section 5.2].

The quantity we focus on is
$\mathbb{E}_{W}\left[W_{ik}W_{ik}W_{jq}W_{jq}\right]$. We will do the
computation on multiple cases, based on the equalities of $k,q$. Notice that
$i\neq j$ in all cases. It suffices if we focus on the two distinct cases:
$E_{1}=\mathbb{E}_{W}\left[W_{ik}^{2}W_{jq}^{2}\right]$ $(k\neq q)$ and
$E_{2}=\mathbb{E}_{W}\left[W_{ik}^{2}W_{jk}^{2}\right]$ $(k=q)$.

We first compute $E_{1}$.

Following the procedure from Collins et al. [2022, Section 5.2], we take the
index sequences to be $\bm{i}=(i,i,j,j)$ and $\bm{j}=(k,k,q,q)$. Similarly, we
get $\Delta_{\sigma}(\bm{i})=\Delta_{\tau}(\bm{j})=1$ only if
$\sigma=\\{\\{1,2\\},\\{3,4\\}\\}$ and $\tau=\\{\\{1,2\\},\\{3,4\\}\\}$.

Consdering $\sigma,\tau$ as permutations, we get:

$\displaystyle\sigma$ $\displaystyle=\begin{pmatrix}1&2&3&4\\\
1&2&3&4\end{pmatrix},$ (37) $\displaystyle\tau$
$\displaystyle=\begin{pmatrix}1&2&3&4\\\ 1&2&3&4\end{pmatrix},$ (38)
$\displaystyle\sigma^{-1}\tau$ $\displaystyle=\begin{pmatrix}1&2&3&4\\\
1&2&3&4\end{pmatrix},$ (39)

where $\sigma^{-1}\tau$ has coset-type $[1,1]$. Finally, we plug the results
back into the formula and we obtain:

$E_{1}=\mathbb{E}_{W}\left[W_{ik}^{2}W_{jq}^{2}\right]=\text{Wg}^{O}([1,1])=\frac{d+1}{d(d+2)(d-1)},$

where the last equality is based on the results in Section 7 of [Collins and
Matsumoto, 2009].

We compute $E_{2}$.

Similar to the previous expression, we take the index sequences to be to be
$\bm{i}=(i,i,j,j)$ and $\bm{j}=(k,k,k,k)$. Thus, we obtain
$\Delta_{\sigma}(\bm{i})=\Delta_{\tau}(\bm{j})=1$ only if
$\sigma=\\{\\{1,2\\},\\{3,4\\}\\}$ and $\tau_{1}=\\{\\{1,2\\},\\{3,4\\}\\}$,
$\tau_{2}=\\{\\{1,3\\},\\{2,4\\}\\}$, $\tau_{1}=\\{\\{1,4\\},\\{2,3\\}\\}$.
Similarly, we compute $\sigma^{-1}\tau_{i}$ for $i\in\\{1,2,3\\}$. Notice that
$\sigma$ is the identity permutation, thus yielding
$\sigma^{-1}\tau_{i}=\tau_{i}$, with the coset-types $[1,1],[2],[2]$
respectively, for each $i\in\\{1,2,3\\}$.

Plugging back into the original equation, we obtain:

$\displaystyle E_{2}=\mathbb{E}_{W}\left[W_{ik}^{2}W_{jk}^{2}\right]$
$\displaystyle=\text{Wg}^{O}([1,1])+2\text{Wg}^{O}([2])$ (40)
$\displaystyle=\frac{d+1}{d(d+2)(d-1)}+2\frac{-1}{d(d+2)(d-1)}$ (41)
$\displaystyle=\frac{d-1}{d(d+2)(d-1)}\leavevmode\nobreak\ .$ (42)

Thus, plugging back into the original inequality, we obtain:

$\displaystyle\mathbb{E}\|X^{\prime}_{i}\|\|X^{\prime}_{j}\|$
$\displaystyle\leq 1+\frac{1}{4}\left[\sum_{k\neq
q}(\lambda_{k}-1)(\lambda_{q}-1)\frac{d+1}{d(d+2)(d-1)}+\sum_{k=q}(\lambda_{k}-1)(\lambda_{k}-1)\frac{d-1}{d(d+2)(d-1)}\right]$
(43) $\displaystyle=1+\frac{1}{4d(d+2)(d-1)}\left[\underbrace{\sum_{k\neq
q}(\lambda_{k}-1)(\lambda_{q}-1)}_{S_{\neq}}-\underbrace{\sum_{k=q}(\lambda_{k}-1)(\lambda_{k}-1)}_{S_{=}}\right]$
(44) $\displaystyle=1-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2d(d+2)(d-1)},$ (45)

where we have used that
$S_{\neq}+S_{=}=\sum_{k,q}(\lambda_{k}-1)(\lambda_{q}-1)=(\sum_{k=1}^{d}(\lambda_{k}-1))^{2}=0$
in the last equality.

Thus, we obtain:

$\displaystyle\mathbb{E}\left[\left(\sum_{j}\|X^{\prime}_{j\cdot}\|\right)^{2}\right]$
$\displaystyle=\mathbb{E}\left[\sum_{j}\|X_{j\cdot}^{\prime}\|^{2}+2\sum_{i<j}\|X^{\prime}_{i\cdot}\|\|X^{\prime}_{j\cdot}\|\right]$
(46)
$\displaystyle=d+2\sum_{i<j}\mathbb{E}\left[\|X^{\prime}_{i\cdot}\|\|X^{\prime}_{j\cdot}\|\right]$
(47) $\displaystyle\leq
d+(d^{2}-d)\left(1-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2d(d+2)(d-1)}\right)$
(48)
$\displaystyle=d^{2}-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2(d+2)}\leavevmode\nobreak\
.$ (49)

∎

######  Corollary A.5 (Isometry gap bound).

Suppose the same setup as in Theorem 3. Then, we have:

$\displaystyle\mathbb{E}_{W}[\phi(X^{\prime})|X]\leq\phi(X)+\log\left(1-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}\right).$
(50)

######  Remark 4.

Notice that the term
$\frac{\sum_{k=1}^{d}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}=\mathcal{O}(\frac{1}{d})$,
yielding
$\log\left[1-\frac{\sum_{k=1}^{d}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}\right]\leq
0$.

###### Proof.

From Lemma 2, we know that:

$\displaystyle-\log\mathcal{I}(\text{\sc{Bn}}(X^{\prime}))$
$\displaystyle\leq-\log\mathcal{I}(X)-\log\frac{d^{2}}{\left(\sum_{j=1}^{d}\|X_{j\cdot}^{\prime}\|\right)^{2}}$
(51) $\displaystyle\implies-\log\mathcal{I}(\text{\sc{Bn}}(X^{\prime}))$
$\displaystyle\leq-\log\mathcal{I}(X)-\log{d^{2}}+\log\left(\sum_{j=1}^{d}\|X_{j\cdot}^{\prime}\|\right)^{2}$
(52)
$\displaystyle\implies\mathbb{E}_{W}[-\log\mathcal{I}(\text{\sc{Bn}}(X^{\prime}))|X]$
$\displaystyle\leq-\log\mathcal{I}(X)-\log{d^{2}}+\mathbb{E}_{W}\left[\log\left(\sum_{j=1}^{d}\|X_{j\cdot}^{\prime}\|\right)^{2}\Bigg{|}X\right]$
(53)
$\displaystyle\leq-\log\mathcal{I}(X)-\log{d^{2}}+\log\mathbb{E}_{W}\left[\left(\sum_{j=1}^{d}\|X_{j\cdot}^{\prime}\|\right)^{2}\Bigg{|}X\right]$
(54)
$\displaystyle\leq-\log\mathcal{I}(X)-\log{d^{2}}+\log\left(d^{2}-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2(d+2)}\right)$
(55)
$\displaystyle\leq-\log\mathcal{I}(X)+\log\left(1-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}\right),$
(56)

where in inequality 54 we have used the fact that $\mathbb{E}[\log
X]\leq\log\mathbb{E}[X]$ and in inequality 55 we have used the bound obtained
in proof Theorem 3, equation 49. Thus, we obtain:

$\displaystyle\mathbb{E}_{W}[\phi(X^{\prime})|X]\leq\phi(X)+\log\left(1-\frac{\sum_{k}(\lambda_{k}-1)^{2}}{2d^{2}(d+2)}\right)\leavevmode\nobreak\
.$ (57)

∎

##  Appendix B Isometry gap decay rate in depth

Before we start with the main part of our analysis, let us establish a simple
result on the relation between isometry gap and orthogonality:

######  Lemma B.1 (Isometry gap and orthogonality).

If $\phi(X)\leq\frac{c}{16d},$ then eigenvalues of $X^{\top}X$ are within
$[1-c,1+c].$

Note that, in order to simplify the calculations, we use the fact that
$\frac{1}{d(d+2)}\approx\frac{1}{d^{2}}$ in the following proofs.

Based on the conditional expectation in Corollary A.5, we have:

$\displaystyle\mathbb{E}\left[\phi(X^{\ell+1})|X^{\ell}\right]\leq\phi(X^{\ell})-\frac{\sigma_{\lambda}(X^{\ell})}{2d^{2}}\leavevmode\nobreak\
.$ (58)

Now, we prove a lemma that is conditioned on the previous layer isometry gap
being smaller or larger than $\frac{1}{16d}$.

######  Lemma B.2 (Isometry gap conditional bound).

For $X^{\ell}$ being the representations of an MLP under our setting, we have:

$\displaystyle\mathbb{E}\left[\phi(X^{\ell+1})\middle|X_{\ell},\phi(X^{\ell})\leq\frac{1}{16d}\right]$
$\displaystyle\leq\phi(X^{\ell})\left(1-\frac{1}{2d^{2}}\right),$ (59)
$\displaystyle\mathbb{E}\left[\phi(X^{\ell+1})\middle|X_{\ell},\phi(X^{\ell})>\frac{1}{16d}\right]$
$\displaystyle\leq\phi(X^{\ell})-\frac{1}{32d^{3}}\leavevmode\nobreak\ .$ (60)

###### Proof of Lemma B.2.

Let $\lambda_{k}=1+\epsilon_{k}$, and assume without loss of generality that
$\sum_{k=1}^{d}\epsilon_{k}=0$. Then, using the numerical inequality
$\log(1+x)\geq x-x^{2}$, when $|x|\leq\frac{1}{2}$ we have:

$\displaystyle\sigma_{\lambda}(X)=\frac{1}{d}\sum_{k=1}^{d}\epsilon_{k}^{2}$
(61)
$\displaystyle\max_{k}|\epsilon_{k}|\leq\frac{1}{2}\implies\phi(X)=-\frac{1}{d}\sum_{k=1}^{d}\log(1+\epsilon_{k})\leq-\frac{1}{d}\sum_{k=1}^{d}(\epsilon_{k}-\epsilon_{k}^{2})=\sigma_{\lambda}(X)\leavevmode\nobreak\
.$ (62)

Altogether, we have

$\displaystyle\max_{i}|\epsilon_{i}|\leq\frac{1}{2}\implies\phi(X)\leq\sigma_{\lambda}(X)\leavevmode\nobreak\
.$ (63)

Now, we can restate the condition in terms of an inequality on the isometry
gap. Thus, we can write:

$\displaystyle
d\phi(X)=-\sum_{i=1}^{d}\log\lambda_{i}=-\sum_{i=1}^{d}\log(1+\epsilon_{i})\geq-\sum_{i=1}^{d}\left(\epsilon_{i}-\frac{3\epsilon_{i}^{2}}{6+4\epsilon_{i}}\right)=\sum_{i=1}^{d}\frac{3\epsilon_{i}^{2}}{6+4\epsilon_{i}},$
(64)

where we used the fact that $\sum_{i}\lambda_{i}=d$ implying
$\sum_{i}\epsilon_{i}=0$ and also used the inequality $\log(1+x)\leq
x-\frac{6+x}{6+4x}$ when $x\geq-1$ for $\epsilon_{i}$’s. Note that because
$|{\epsilon}_{i}|\leq\frac{1}{2}$, we get $6+4\epsilon_{i}\geq 4,$ all terms
on the right-hand side are positive, implying that each term is bounded by the
upper bound:

$\displaystyle\frac{3\epsilon_{i}^{2}}{6+4\epsilon_{i}}\leq
d\phi(X)\quad\forall i\in\\{1,2,\dots,d\\}\leavevmode\nobreak\ .$ (65)

By construction, we have $\epsilon_{i}\geq-1$ and
$\frac{3\epsilon_{i}^{2}}{6+4\epsilon_{i}}\leq\frac{1}{16}.$ Since
$6+4\epsilon_{i}\geq 2$, we can multiply both sides by $6+4\epsilon_{i}$ and
conclude $3\epsilon_{i}^{2}-\frac{6+4\epsilon_{i}}{16}\leq 0$. We can now
solve the quadratic equation and obtain
$\frac{1-\sqrt{74}}{24}\leq\epsilon_{i}\leq\frac{1+\sqrt{74}}{24}$ which
numerically becomes $-0.35\leq\epsilon_{i}\leq 0.4$, implying
$|\epsilon_{i}|<0.5$.

By solving the quadratic equation above we can guarantee that

$\phi(X)\leq\frac{1}{16d}\implies\max_{i}|\epsilon_{i}|\leq\frac{1}{2}\implies\phi(X)\leq\sigma_{\lambda}(X)\leavevmode\nobreak\
.$ (66)

Furthermore, we can restate the condition on maximum using:

$\displaystyle\max_{k}|\epsilon_{k}|=\sqrt{\max_{k}\epsilon_{k}^{2}}\leq\sqrt{\sum_{k}\epsilon_{k}^{2}}=\sqrt{d\sigma_{\lambda}(X)}$
(67)

and conclude that

$\displaystyle\sigma_{\lambda}(X)\leq\frac{1}{4d}\implies\max_{i}|{\epsilon}_{i}|\leq\frac{1}{2}\implies\phi(X)\leq\sigma_{\lambda}(X)\leavevmode\nobreak\
.$ (68)

Using this statement, we have

$\displaystyle\sigma_{\lambda}(X)\leq\frac{1}{16d}\implies\phi(X)\leq\sigma_{\lambda}(X)\implies\phi(X)\leq\frac{1}{16d}\leavevmode\nobreak\
.$ (69)

If we negate and flip the two sides we arrive at

$\phi(X)>\frac{1}{16d}\implies\sigma_{\lambda}(X)>\frac{1}{16d}\leavevmode\nobreak\
.$ (70)

Thus, we can simplify the recurrence

$\displaystyle\mathbb{E}[\phi(X^{\ell+1})|X^{\ell}]\leq\phi(X^{\ell})-\frac{\sigma_{\lambda}(X^{\ell})}{2d^{2}}$
(71)

as follows

$\displaystyle\mathbb{E}\left[\phi(X^{\ell+1})\middle|X^{\ell},\phi(X^{\ell})\leq\frac{1}{16d}\right]$
$\displaystyle\leq\phi(X^{\ell})\left(1-\frac{1}{2d^{2}}\right),$ (72)
$\displaystyle\mathbb{E}\left[\phi(X^{\ell+1})\middle|X^{\ell},\phi(X^{\ell})>\frac{1}{16d}\right]$
$\displaystyle\leq\phi(X^{\ell})-\frac{1}{32d^{3}},$ (73)

where we used equation 66 in the first one and equation 70 in the second one.
∎

###### Proof of Theorem 1.

From Lemma A.1, we know that
$\phi(X^{0})\geq\phi(X^{1})\geq\dots\geq\phi(X^{L})\geq 0$, for any layer
$0\leq\ell\leq L$. Thus, we get using equation 60:

$\displaystyle\mathbb{E}\left[\phi(X^{\ell+1})\middle|X^{\ell},\phi(X^{\ell})>\frac{1}{16d}\right]$
$\displaystyle\leq\phi(X^{\ell})-\frac{1}{32d^{3}}$ (74)
$\displaystyle=\left(1-\frac{1}{32d^{3}\phi(X^{\ell})}\right)\phi(X^{\ell})$
(75)
$\displaystyle\leq\left(1-\frac{1}{32d^{3}\phi(X^{0})}\right)\phi(X^{\ell}),$
(76)

where in the last step we have used the fact that
$\phi(X^{\ell})\leq\phi(X^{0})$.

Thus, we can combine equation 59 and equation 60 and obtain:

$\displaystyle\mathbb{E}[\phi(X^{\ell+1})|X^{\ell}]$
$\displaystyle=\mathbb{E}\left[\phi(X^{\ell+1})\middle|X^{\ell},\phi(X^{\ell})\leq\frac{1}{16d}\right]\bm{1}_{\phi(X^{\ell})\leq\frac{1}{16d}}$
(77)
$\displaystyle+\mathbb{E}\left[\phi(X^{\ell+1})\middle|X^{\ell},\phi(X^{\ell})>\frac{1}{16d}\right]\bm{1}_{\phi(X^{\ell})>\frac{1}{16d}}$
(78)
$\displaystyle\leq\max\left(\mathbb{E}\left[\phi(X^{\ell+1})\middle|X^{\ell},\phi(X^{\ell})\leq\frac{1}{16d}\right],\mathbb{E}\left[\phi(X^{\ell+1})\middle|X^{\ell},\phi(X^{\ell})>\frac{1}{16d}\right]\right)$
(79)
$\displaystyle\leq\max\left(1-\frac{1}{2d^{2}},1-\frac{1}{32d^{3}\phi(X^{0})}\right)\phi(X^{\ell})$
(80)
$\displaystyle=\left(1-\min\left(\frac{1}{2d^{2}},\frac{1}{32d^{3}\phi(X^{0})}\right)\right)\phi(X^{\ell})$
(81)
$\displaystyle=\left(1-\frac{1}{\max(2d^{2},32d^{3}\phi(X^{0}))}\right)\phi(X^{\ell})$
(82)
$\displaystyle\leq\exp\left[-\underbrace{\frac{1}{\max(2d^{2},32d^{3}\phi(X^{0}))}}_{k}\right]\phi(X^{\ell})$
(83)
$\displaystyle=\exp\left(-\frac{1}{k}\right)\phi(X^{\ell})\leavevmode\nobreak\
.$ (84)

By iterated expectations over $X^{\ell}$ we get:

$\displaystyle\mathbb{E}[\phi(X^{\ell+1})]\leq\exp\left(-\frac{1}{k}\right)\mathbb{E}[\phi(X^{\ell})]\leq\exp\left(-\frac{\ell}{k}\right)\phi(X^{0})\leavevmode\nobreak\
.$ (85)

Note that since $\max(2d^{2},32d^{3}\phi(X^{0}))\leq
2d^{2}(1+16d\phi(X^{0}))$, we can conclude the proof. ∎

##  Appendix C Gradient norm bound

In the following section, we denote by $H^{\ell}=W^{\ell}X^{\ell}$ the pre-
normalization values. Moreover, we define as
$\mathcal{F}_{L}:\mathbb{R}^{d\times d}\to\mathbb{R}^{C\times d}$, where $C$
is the number of output classes, as the functional composition of an $L$
layers MLP, following the update rule defined in equation 3, i.e.:

$\displaystyle\mathcal{F}_{L}(X_{L})=\text{\sc{Bn}}(W^{L}\mathcal{F}_{L-1}(X_{L-1}))\leavevmode\nobreak\
.$ (86)

Let us restate the theorem for completeness of the appendix:

######  Theorem C.1 (Restated Thm. 5).

For any $\mathcal{O}(1)$-Lipschitz loss function $L$ and non-degenerate input
$X_{0}$, we have:

$\displaystyle\log\left\|\frac{\partial\mathcal{L}}{\partial
W^{\ell}}\right\|\lesssim
d^{5}\left(\phi(X_{0})^{3}+1-e^{-\frac{L}{32d^{4}}}\right)$ (87)

holds for all $\ell\leq L$, where possibly $L\to\infty$.

In particular, the following lemma guarantees that the Lipschitz conditions
are met for practical loss functions:

######  Lemma C.2.

In a classification setting, cross entropy and mean squared error losses are
$\mathcal{O}(1)$-Lipschitz.

The main idea for the feasibility of this theorem is the presence of perfectly
isometric weight matrices that are orthonormal, and the linear activation that
does not lead to vanishing or exploding gradients. The only remaining layers
to be analyzed are the batch normalization layers. Thus, our main goal is to
show that the sum of log-norm gradient of Bn layers remains bounded even if
the network has infinite depth. To do so, we shall relate the norm of the
gradient of those layers to the isometry gap of representations, and use the
bounds from the previous section to establish that the log-norm sum is
bounded.

###### Proof of Thm. C.1.

Now, considering an $L$ layer deep model, where $L$ can possibly be
$L\to\infty$, we can finalize the proof of Theorem 5. Consider an MLP model as
defined in equation 3. Let $H^{L}=W^{L}X^{L}$ be the logits of the model,
where $H^{L}\in\mathbb{R}^{C\times d}$, $W^{L}\in\mathbb{R}^{C\times d}$,
$W^{L}$ is an orthogonal matrix and $C$ is the number of output classes.
Denote as $\mathcal{L}(H^{L},y)$ the loss of model for an input matrix, with
ground truth $y$. Then, applying the chain rule, we have:

$\displaystyle\frac{\partial\mathcal{L}}{\partial W^{\ell}}$
$\displaystyle=\frac{\partial\mathcal{L}}{\partial H^{L}}\frac{\partial
H^{L}}{\partial X^{L}}\frac{\partial X^{L}}{\partial
X^{L-1}}\dots\frac{\partial X^{\ell+2}}{\partial X^{\ell+1}}\frac{\partial
X^{\ell+1}}{\partial H^{\ell}}\frac{\partial H^{\ell}}{\partial
W^{\ell}}\leavevmode\nobreak\ .$ (88)

By taking the logarithm of the norm of each factor and applying Lemma C.12, we
get:

$\displaystyle\log\left\|\frac{\partial\mathcal{L}}{\partial
W^{\ell}}\right\|$
$\displaystyle\leq\log\left\|\frac{\partial\mathcal{L}}{\partial
H^{L}}\right\|+\log\underbrace{\left\|\frac{\partial H^{L}}{\partial
X^{L}}\right\|}_{\left\|W^{L}\right\|}+\sum_{k=\ell+1}^{L}\log\left\|\frac{\partial
X^{k+1}}{\partial X^{k}}\right\|+\log\underbrace{\left\|\frac{\partial
X^{\ell+1}}{\partial
H^{\ell}}\right\|}_{\left\|J_{\text{\sc{Bn}}}(H^{\ell})\right\|}+\log\left\|\frac{\partial
H^{\ell}}{\partial W^{\ell}}\right\|$ (89)
$\displaystyle\leq\log\left\|\frac{\partial\mathcal{L}}{\partial
H^{L}}\right\|+\sum_{k=\ell}^{L}\log\left\|J_{\text{\sc{Bn}}}(H^{k})\right\|+\log\left\|\frac{\partial
H^{\ell}}{\partial W^{\ell}}\right\|\leavevmode\nobreak\ .$ (90)

where $\log\underbrace{\left\|\frac{\partial H^{L}}{\partial
X^{L}}\right\|}_{\left\|W^{L}\right\|}=\log{\left\|W^{L}\right\|}=0$, since
the orthogonal matrix $W^{L}$ has operator norm $1$.

Since $\frac{\partial H^{\ell}}{\partial W^{\ell}}=X^{\ell}$ and
$X^{\ell}=\text{\sc{Bn}}(H^{\ell-1})$ is batch normalized, this means that
$\left\|\frac{\partial H^{\ell}}{\partial W^{\ell}}\right\|\leq d$. Thus, the
main part is to bound the Jacobian log-norms, which is provided by the
following lemma:

######  Lemma C.3.

We have

$\displaystyle\sum_{k=1}^{L}\log\left\|J_{\text{\sc{Bn}}}(H^{k})\right\|\lesssim
d^{5}(\phi(X_{0})^{3}+1-e^{-\frac{L}{32d^{4}}})\leavevmode\nobreak\ .$ (91)

Finally, we can plug the bound from Lemma C.3 in equation 90 and obtain the
conclusion:

$\displaystyle\log\left\|\frac{\partial\mathcal{L}}{\partial
W^{\ell}}\right\|\lesssim\log\left\|\frac{\partial\mathcal{L}}{\partial
H^{L}}\right\|+\log
d+d^{5}\left(\phi(X_{0})^{3}+1-e^{-\frac{L}{32d^{4}}}\right)\leavevmode\nobreak\
.$ (92)

Note that, for $L\to\infty$, we get:

$\displaystyle\log\left\|\frac{\partial\mathcal{L}}{\partial
W^{\ell}}\right\|\lesssim\log\left\|\frac{\partial\mathcal{L}}{\partial
H^{L}}\right\|+\log d+d^{5}(\phi(X_{0})^{3}+1)\leavevmode\nobreak\ .$ (93)

In order to conclude the bound, it suffices to show that the norm of the
gradient of the loss with respect to the logits is bounded, which is the
objective of Lemma C.2. ∎

###### Proof of Lemma C.3.

The proof of this lemma is chiefly relying on the following bound on the
Jacobian of batch normalization layers, which we will state and prove
beforehand.

######  Lemma C.4 (Log-norm bound).

If $X\in\mathbb{R}^{d\times d}$ is the input to a Bn layer, its Jacobian
operator norm is bounded by

$\displaystyle\log\|J_{\text{\sc{Bn}}}(X)\|_{op}\leq
d\phi(X)+1\leavevmode\nobreak\ .$ (94)

Furthermore, if $\phi(X)\leq\frac{1}{16d}$, then we have

$\displaystyle\log\|J_{\text{\sc{Bn}}}(X)\|_{op}\leq
2\sqrt{d\phi(X)}\leavevmode\nobreak\ .$ (95)

Based on the lemma above, we shall define $S$ as the hitting time,
corresponding to the first layer in our case, that the isometry gap drops
below the critical value of $\frac{1}{16d}$:

$\displaystyle
S=\min\left\\{\ell:\phi(X^{\ell})\leq\frac{1}{16d}\right\\}\leavevmode\nobreak\
.$ (96)

So, we first bound the total log-grad norm for layers $1$ up to $S$, and
subsequently $S+1$ up to $L$:

$\displaystyle\log\|J_{\mathcal{F}_{L}}(X)\|_{op}$
$\displaystyle\leq\sum_{\ell=1}^{L}\log\|J_{\text{\sc{Bn}}}(X^{\ell})\|_{op}$
(97)
$\displaystyle\leq\sum_{\ell=1}^{S}(d\phi(X^{\ell})+1)+2\sum_{\ell=S+1}^{L}\sqrt{d\phi(X^{\ell})}$
(98)
$\displaystyle\leq\sum_{\ell=1}^{S}(d\phi(X^{0})+1)+2\sum_{\ell=S+1}^{L}\sqrt{d\phi(X^{\ell})}$
(99)
$\displaystyle=S(d\phi(X^{0})+1)+2\sum_{\ell=S+1}^{L}\sqrt{d\phi(X^{\ell})},$
(100)

where we have used that $\phi(X^{0})$ as an upper bound on $\phi(X^{\ell})$.

Thus, taking expectation we get:

$\displaystyle\mathbb{E}\log\|J_{\mathcal{F}_{L}}(X)\|_{op}\leq(d\phi(X^{0})+1)\mathbb{E}[S]+2\sum_{\ell=S+1}^{L}\mathbb{E}\sqrt{d\phi(X^{\ell})}\leavevmode\nobreak\
.$ (101)

Note that $S$ is a random variable, which is why the expectation over the
number of layers appears at the last line. Thus, we can bound the log-norm by
bounding $\mathbb{E}[S]$ and the summation separately.

######  Lemma C.5 (stopping time bound).

We have $\mathbb{E}[S]\lesssim 512d^{4}\phi(X^{0})^{2}$ if
$\phi(X_{0})>\frac{1}{16d}$, and $\mathbb{E}[S]=0$ if
$\phi(X_{0})\leq\frac{1}{16d}$.

######  Lemma C.6 (second phase bound).

We have

$\sum_{\ell=S+1}^{L}\mathbb{E}\sqrt{d\phi(X^{\ell})}\leq
32d^{4.5}\phi(X_{0})^{0.5}\left(1-e^{-\frac{L}{32d^{4}}}\right)\leavevmode\nobreak\
.$

Thus, we have the following 2 cases, based on whether $\phi(X_{0})$ is below
or over the $\frac{1}{16d}$ threshold. If we plug the bounds in equation 101
we get the following.

If $\phi(X_{0})\leq\frac{1}{16d}$, then:

$\displaystyle\mathbb{E}\log\|J_{\mathcal{F}_{L}}(X)\|_{op}$
$\displaystyle\leq
32d^{4.5}\phi(X_{0})^{0.5}\left(1-e^{-\frac{L}{32d^{4}}}\right)$ (102)
$\displaystyle\lesssim
d^{4.5}\phi(X_{0})^{0.5}\left(1-e^{-\frac{L}{32d^{4}}}\right),$ (103)

and if $\phi(X_{0})>\frac{1}{16d}$ then:

$\displaystyle\mathbb{E}\log\|J_{\mathcal{F}_{L}}(X)\|_{op}$
$\displaystyle\leq
512d^{4}\phi(X_{0})^{2}(1+d\phi(X_{0}))+32d^{4}\left(1-e^{-\frac{L}{32d^{4}}}\right)$
(104) $\displaystyle\lesssim
d^{5}\phi(X_{0})^{3}+d^{4}\left(1-e^{-\frac{L}{32d^{4}}}\right)$ (105)
$\displaystyle\lesssim
d^{4}\left(d\phi(X_{0})^{3}+1-e^{-\frac{L}{32d^{4}}}\right)\leavevmode\nobreak\
.$ (106)

In fact, the maximum of the two bounds is

$\displaystyle\mathbb{E}\log\|J_{\mathcal{F}_{L}}(X)\|_{op}\lesssim
d^{5}\left(\phi(X_{0})^{3}+1-e^{-\frac{L}{32d^{4}}}\right)\leavevmode\nobreak\
.$ (107)

∎

###### Proof of Lemma C.6.

By the bound in Lemma B.2, we have

$\displaystyle\mathbb{E}\left[\phi(X^{\ell+1})\middle|\phi(X^{\ell})\leq\frac{1}{16d}\right]\leq\phi(X^{\ell})\left(1-\frac{1}{2d^{2}}\right)\leavevmode\nobreak\
.$ (108)

Since we assumed $\ell\geq S$, the conditional inequality always holds and
thus we have the Markov bound

$\displaystyle\ell\geq S\implies
q:=\Pr\left\\{\phi(X^{\ell+1})\geq\left(1-\frac{1}{4d^{2}}\right)\phi(X^{\ell})\right\\}\leq\frac{1-\frac{1}{2d^{2}}}{1-\frac{1}{4d^{2}}}\leq
1-\frac{1}{4d^{2}}\leavevmode\nobreak\ .$ (109)

We define as failure the event
$\bar{A}=\\{\phi(X^{\ell+1})\geq(1-\frac{1}{4d^{2}})\phi(X^{\ell})\\}$ with
probability $q$, and conversely as success the event $A$ with probability
$1-q$. In other words, the probability that $\phi(X^{\ell+1})$ does not
decrease by at least a factor of $1-\frac{1}{4d^{2}}$ is bounded by the
failure probability $1-\frac{1}{4d^{2}}.$

Since $\phi(X^{\ell+1})\leq\phi(X^{\ell})$, then under the assumption that
$\ell\geq S$ we can upper bound $\sqrt{\phi(X^{\ell+1})}$ with
$\sqrt{\phi(X^{\ell})}$ in case of failure with probability $q$, and with
$\sqrt{(1-\frac{1}{4d^{2}})\phi(X^{\ell})}$ in case of success with
probability $1-q$:

$\displaystyle\mathbb{E}$
$\displaystyle\left[\sqrt{\phi(X^{\ell+1})}|\phi(X^{\ell})\right]$ (110)
$\displaystyle=\mathbb{E}\left[\sqrt{\phi(X^{\ell+1})}|\phi(X^{\ell}),\bar{A}\right]q+\mathbb{E}\left[\sqrt{\phi(X^{\ell+1})}|\phi(X^{\ell}),A\right](1-q)$
(111)
$\displaystyle\leq\sqrt{\phi(X^{\ell})}q+\sqrt{\phi(X^{\ell})\left(1-\frac{1}{4d^{2}}\right)}(1-q)$
(112)
$\displaystyle=\sqrt{\phi(X^{\ell})}\left(q+\sqrt{1-\frac{1}{4d^{2}}}(1-q)\right)$
(113)
$\displaystyle=\sqrt{\phi(X^{\ell})}\left(\sqrt{1-\frac{1}{4d^{2}}}+\left(1-\sqrt{1-\frac{1}{4d^{2}}}\right)q\right)$
monotonic in $q$ (114)
$\displaystyle\leq\sqrt{\phi(X^{\ell})}\left(\sqrt{1-\frac{1}{4d^{2}}}+\left(1-\sqrt{1-\frac{1}{4d^{2}}}\right)\left(1-\frac{1}{4d^{2}}\right)\right)$
plug $q\leq 1-\frac{1}{4d^{2}}$ (115)
$\displaystyle=\sqrt{\phi(X^{\ell})}\left(1-\frac{1}{4d^{2}}+\sqrt{1-\frac{1}{4d^{2}}}\frac{1}{4d^{2}}\right)$
rearranging terms (116)
$\displaystyle\leq\sqrt{\phi(X^{\ell})}\left(1-\frac{1}{4d^{2}}+\left(1-\frac{1}{8d^{2}}\right)\frac{1}{4d^{2}}\right)$
$\sqrt{1-x}\leq 1-\frac{x}{2}$ for $x\geq 0$ (117)
$\displaystyle=\sqrt{\phi(X^{\ell})}\left(1-\frac{1}{32d^{4}}\right)\leavevmode\nobreak\
.$ (118)

Thus, for $\ell\geq S$, we have

$\displaystyle\mathbb{E}\sqrt{\phi(X^{\ell+1})}=\mathbb{E}_{X^{\ell}}\mathbb{E}\left[\sqrt{\phi(X^{\ell+1})}|\phi(X^{\ell})\right]\leq\mathbb{E}\sqrt{\phi(X^{\ell})}\left(1-\frac{1}{32d^{4}}\right)\leavevmode\nobreak\
.$ (119)

The summation starts from below $\sqrt{d\phi(X_{S})},$ and will decay by rate
$1-\frac{1}{32d^{4}},$ which is upper bounded by the geometric series:

$\displaystyle\sqrt{d\phi(X_{S})}\sum_{k=0}^{L}\left(1-\frac{1}{32d^{4}}\right)^{k}$
$\displaystyle\leq\sqrt{d\phi(X_{0})}32d^{4}\left(1-\left(1-\frac{1}{32d^{4}}\right)^{L+1}\right)$
(120) $\displaystyle\leq
32d^{4.5}\phi(X_{0})^{0.5}\left(1-e^{-\frac{L}{32d^{4}}}\right)\leavevmode\nobreak\
.$ (121)

∎

###### Proof of Lemma C.5.

By Lemma B.2 we have

$\displaystyle\Pr\left\\{\phi(X^{\ell})\geq\frac{1}{16d}\right\\}\leq\exp\left(-\frac{\ell}{\max(2d^{2},32d^{3}\phi(X^{0}))}\right)16d\phi(X^{0})\leavevmode\nobreak\
.$ (122)

Thus, we have

$\displaystyle\Pr\\{S\geq\ell\\}\leq\exp\left(-\frac{\ell}{\max(2d^{2},32d^{3}\phi(X^{0}))}\right)16d\phi(X^{0})\leavevmode\nobreak\
.$ (123)

Since $S$ is a non-negative integer valued random variable, we can thus bound
$\mathbb{E}[S]$ as:

$\displaystyle\mathbb{E}[S]$
$\displaystyle=\sum_{\ell=1}^{\infty}P\\{S\geq\ell\\}$ (124)
$\displaystyle\leq
16d\phi(X^{0})\sum_{\ell=1}^{\infty}\exp\left(\frac{-\ell}{k}\right)$ (125)
$\displaystyle=16d\phi(X^{0})\frac{1}{\exp(\frac{1}{k})-1}$ (126)
$\displaystyle\leq 16d\phi(X^{0})k$ (127)
$\displaystyle=16d\phi(X^{0})\max(2d^{2},32d^{3}\phi(X^{0}))$ (128)
$\displaystyle\lesssim 512d^{4}\phi(X^{0})^{2}\leavevmode\nobreak\ .$ (129)

∎

###  1 Proof of Lemma C.4: Bounding BN grad-norm with isometry gap

The proof of the Lemma relies on two main observations that are crystallized
in the following lemmas that first establish a bound on Jacobian operator norm
based on the inverse of smallest eigenvalue, and then establish a lower bound
for the smallest eigenvalue using the isometry gap.

######  Lemma C.7.

Let $X\in\mathbb{R}^{d\times d}$ and let $\\{\lambda_{i}\\}_{i=1}^{d}$ be the
eigenvalues of $XX^{\top}$. Then, we have that:

$\displaystyle\|J_{\text{\sc{Bn}}}(X)\|_{op}^{2}\leq\frac{1}{\lambda_{d}},$
(130)

where $J_{\text{\sc{Bn}}}$ is the Jacobian of the $\text{\sc{Bn}}(\cdot)$
operator.

Using the above lemma we have
$\log\|J_{\text{\sc{Bn}}}(X)\|_{op}\leq-\log\lambda_{d}$. The following lemma
upper bounds this quantity using isometry gap:

######  Lemma C.8.

The minimum eigenvalue of a Gram matrix that is trace-normalized is lower-
bounded by the isometry gap as $-\log\lambda_{d}\leq d\phi(X)+1.$ Furthermore,
if $\phi(X)\leq\frac{1}{16d},$ then $-\log\lambda_{d}\leq 2\sqrt{d\phi(X)}$.

Plugging these two values we have the bounds

$\displaystyle\log\|J_{\text{\sc{Bn}}}(X)\|_{op}$ $\displaystyle\leq
d\phi(X)+1,$ (131)
$\displaystyle\phi(X)\leq\frac{1}{16d}\implies\log\|J_{\text{\sc{Bn}}}(X)\|_{op}$
$\displaystyle\leq 2\sqrt{d\phi(X)}\leavevmode\nobreak\ .$ (132)

Now we can turn our attention to the proof of the Lemmas used in the proof.
The proof of relationship between minimum eigenvalue and isometry gap is
obtained by merely a few numerical inequalities:

###### Proof of Lemma C.8.

Let $\\{\lambda_{i}\\}_{i=1}^{d}$ be the eigenvalues of $X^{\top}X$. Since the
matrix is trace-normalized, we have $\sum_{k=1}^{d}\lambda_{k}=d$.

The arithmetic mean of the top $d-1$ values can be written as

$\displaystyle\frac{1}{d-1}\sum_{k=1}^{d-1}\lambda_{k}=1+\frac{1-\lambda_{d}}{d-1}\leavevmode\nobreak\
.$ (133)

Thus, we have that their geometric mean is bounded by the same value.
Therefore, we have the following bound:

$\displaystyle\prod_{k=1}^{d}\lambda_{k}$
$\displaystyle\leq\lambda_{d}\left(1+\frac{1-\lambda_{d}}{d-1}\right)^{d-1}$
(134) $\displaystyle\implies d\log\mathcal{I}(X)$
$\displaystyle\leq\log(\lambda_{d})+(d-1)\log\left(1+\frac{1-\lambda_{d}}{d-1}\right),$
(135)

where in the second inequality we have taken logarithm of both sides. Now, we
can apply the numerical inequality $\log(1+x)\leq x$ to conclude:

$\displaystyle-d\phi(X)\leq\log\lambda_{d}+1-\lambda_{d}\leavevmode\nobreak\
.$ (136)

Since $\lambda_{d}$ is non-negative, this clearly implies the first
inequality: $-\log\lambda_{d}\leq d\phi(X)+1$.

For the second inequality first we use the numerical inequality
$\log(x)+1-x\leq-\frac{(x-1)^{2}}{2},\forall x\in[0,1]$ to conclude

$\displaystyle\frac{(1-\lambda_{d})^{2}}{2}$ $\displaystyle\leq d\phi(X)$
(137) $\displaystyle\implies\lambda_{d}$ $\displaystyle\geq
1-\sqrt{2d\phi(X)}$ (138) $\displaystyle\implies-\log\lambda_{d}$
$\displaystyle\leq-\log(1-\sqrt{2d\phi(X)})\leavevmode\nobreak\ .$ (139)

We can now use the inequality $-\log(1-x)\leq\sqrt{2}x$ for $0\leq
x\leq\frac{1}{2}$ to conclude that

$\displaystyle-\log\lambda_{d}\leq 2\sqrt{d\phi(X)}$ (140)

when $2\sqrt{d\phi(X)}\leq\frac{1}{2}$, which is equivalent to
$\phi(X)\leq\frac{1}{16d}$.

∎

For proving Lemma C.4, we first analyze Bn operator on a row, and then invoke
this bound and the special structure $J_{\text{\sc{Bn}}}$ to derive the main
proof.

######  Lemma C.9.

Let $f:\mathbb{R}^{d}\to\mathbb{R}^{d}$ defined as $f(x)=\frac{x}{\|x\|}$ be
the elementwise normalization of the $x$. Then:

$\displaystyle J_{f}(x)=\frac{1}{\|x\|}I_{d^{2}}-\frac{1}{\|x\|^{3}}x\otimes
x,$ (141)

where $\otimes$ is the outer product.

###### Proof.

To begin, notice that for $x\in\mathbb{R}^{d}$ we have
$\frac{\partial\|x\|}{\partial x}=\frac{x}{\|x\|}$. Denote by
$y_{i}:=\left[f(x)\right]_{i}=\frac{x_{i}}{\|x\|}$. Then the Jacobian entries
become:

$\displaystyle\frac{\partial y_{i}}{\partial x_{i}}$
$\displaystyle=\frac{\|x\|-\frac{x_{i}}{\|x\|}x_{i}}{\|x\|^{2}}=\frac{1}{\|x\|}-\frac{1}{\|x\|^{3}}x_{i}x_{i},$
(142) $\displaystyle\frac{\partial y_{i}}{\partial x_{j}}$
$\displaystyle=\frac{-\frac{x_{j}}{\|x\|}x_{i}}{\|x\|^{2}}=-\frac{1}{\|x\|^{3}}x_{i}x_{j}\leavevmode\nobreak\
.$ (143)

Assembling the equations into matrix form, we obtain:

$\displaystyle J_{f}(x)=\frac{1}{\|x\|}I_{d^{2}}-\frac{1}{\|x\|^{3}}x\otimes
x\leavevmode\nobreak\ .$ (144)

∎

######  Corollary C.10.

The $J_{f}(x)$ has the eigenvalue $\frac{1}{\|x\|}$ with multiplicity
$d^{2}-1$ and $0$ with multiplicity $1$.

######  Lemma C.11.

Let $XX^{\top}=\sum_{i=1}^{d}\lambda_{i}u_{i}u_{i}^{\top}$, where
$XX^{\top}=U\Lambda U^{\top}$ is the eigendecomposition of $XX^{\top}$. Then,
we have that $\min_{j}\|X_{j\cdot}\|^{2}\geq\min_{k}\lambda_{k}$.

###### Proof.

$\displaystyle\|X_{j\cdot}\|^{2}=(XX^{\top})_{jj}=\sum_{i=1}^{d}\lambda_{i}u_{ji}^{2}\geq\min_{k}\lambda_{k}\sum_{i=1}^{d}u_{ji}^{2}=\min_{k}\lambda_{k},$
(145)

where in the last equality we have used the fact that $U$ orthogonal. Since
this is true for all rows $j$, we get
$\min_{j}\|X_{j\cdot}\|^{2}\geq\min_{k}\lambda_{k}$. ∎

Having established the above properties, we are equipped to prove that the
Jacobian of a BN layer is bounded by the inverse of the minimum eigenvalue of
its input Gram matrix.

###### Proof of Lemma C.7.

To begin, notice that since
$\text{\sc{Bn}}:\mathbb{R}^{d^{2}}\to\mathbb{R}^{d^{2}}$, implies that
$J_{\text{\sc{Bn}}}(X)\in\mathbb{R}^{d^{2}\times d^{2}}$. Denote
$X^{\prime}=\text{\sc{Bn}}(X)$. Since the normalization happens on each row
independently of the other rows, the only non-zero derivatives in the Jacobian
correspond to changes in output row $i$ with regards to the same input row
$i$, i.e.:

$\displaystyle\frac{\partial X^{\prime}_{ik}}{\partial X_{jl}}=0,\quad\forall
i\neq j,\forall k,l\leavevmode\nobreak\ .$ (146)

This creates a block-diagonal structure in $J_{\text{\sc{Bn}}}(X)$, with $d$
blocks on the main diagonal, where each block has size $d\times d$ and is
equal to $J_{f}(X_{i\cdot})$, where $f$ is as defined in Lemma C.9. Therefore,
due to the block-diagonal structure, we know that:

$\displaystyle\lambda\left[J_{\text{\sc{Bn}}}(X)\right]=\bigcup_{i=1}^{d}\lambda\left[J_{f}(X_{i\cdot})\right],$
(147)

where $\lambda[\cdot]$ denotes the eigenvalue spectrum. From Corollary C.10,
we know that

$\displaystyle\lambda\left[J_{\text{\sc{Bn}}}(X)\right]=\left\\{\frac{1}{\|X_{i\cdot}\|}\right\\}_{i=1}^{d}\cup\\{0\\}$
(148)

with their respective multiplicites. Finally, using Lemma C.11, this implies:

$\displaystyle\|J_{\text{\sc{Bn}}}(X)\|_{2}^{2}=\left[\max_{i}\frac{1}{\|X_{i\cdot}\|}\right]^{2}=\left[\frac{1}{\min_{i}{\|X_{i\cdot}\|}}\right]^{2}\leq\frac{1}{\min_{k}\lambda_{k}}\leavevmode\nobreak\
.$ (149)

∎

###### Proof of Lemma C.2.

Assuming the the logits are passed through a softmax layer, we analyze the
case of Cross Entropy Loss for one sample $i$ in a $C$-classes classification
problem. Denoting $z_{i}=H^{L}_{i\cdot}$, we have:

$\displaystyle\mathcal{L}(z_{i},y_{i})=-\sum_{i=1}^{C}y_{i}\log p_{i},$ (150)

where $p_{i}=\frac{e^{z_{i}}}{\sum_{j=1}^{C}e^{z_{j}}}$ is the probability
vector for sample $i$ after passing through the softmax function.

Computing the partial derivatives, we obtain:

$\displaystyle\frac{\partial p_{i}}{\partial z_{i}}$
$\displaystyle=p_{i}(1-p_{k}),\quad i=k$ (151) $\displaystyle\frac{\partial
p_{i}}{\partial z_{k}}$ $\displaystyle=-p_{i}p_{k},\quad i\neq
k\leavevmode\nobreak\ .$ (152)

Finally, we can compute the gradient of the loss with respect to the logits:

$\displaystyle\nabla_{z_{k}}\mathcal{L}$
$\displaystyle=\sum_{i=1}^{C}\left(-y_{i}\frac{\partial\log(p_{i})}{\partial
z_{k}}\right)$ (153)
$\displaystyle=\sum_{i=1}^{C}\left(-y_{i}\frac{1}{p_{i}}\frac{\partial
p_{i}}{\partial z_{k}}\right)$ (154) $\displaystyle=\sum_{i\neq
k}(y_{i}p_{k})+(-y_{k}(1-p_{k}))$ (155) $\displaystyle=p_{k}\sum_{i\neq
k}y_{i}-y_{k}(1-p_{k})$ (156)
$\displaystyle\implies\|\nabla_{z_{k}}\mathcal{L}\|$
$\displaystyle\leq\left\|p_{k}\sum_{i\neq
k}y_{k}\right\|+\|y_{k}(1-p_{k})\|\leq 2\leavevmode\nobreak\ .$ (157)

Since the gradient of the loss with regard to each sample is bounded, we can
conclude that the operator norm of the Jacobian of the loss with regards to
the logits matrix $H^{L}$ is also bounded.

In a similar analysis, we now shift our attention towards the Mean Squared
Error (MSE) loss:

$\mathcal{L}(z_{i},y_{i})=\frac{1}{C}\sum_{i=1}^{C}(y_{i}-p_{i})^{2}\leavevmode\nobreak\
.$

We want to compute the gradient of the loss with respect to each logit
$z_{k}$:

$\displaystyle\|\nabla_{z_{k}}\mathcal{L}\|$
$\displaystyle=\frac{1}{C}\sum_{i=1}^{C}2(y_{i}-p_{i})\frac{\partial(-p_{i})}{\partial
z_{k}}$ (158) $\displaystyle\implies\|\nabla_{z_{k}}\mathcal{L}\|$
$\displaystyle\leq\frac{2}{C}\sum_{i=1}^{C}\left|(y_{i}-p_{i})\frac{\partial
p_{i}}{\partial
z_{k}}\right|\leq\frac{2}{C}\sum_{i=1}^{C}|y_{i}-p_{i}|\cdot\left|\frac{\partial
p_{i}}{\partial z_{k}}\right|\leq 2\leavevmode\nobreak\ .$ (159)

By substituting these derivatives into the gradient equation, we can derive
the gradient for each logit with respect to the MSE loss.

∎

######  Lemma C.12.

Let $X^{\ell}\in\mathbb{R}^{d\times d}$ be the hidden representations of layer
$\ell>0$ as defined in equation 3. Then, we have that:

$\displaystyle\log\left\|\frac{\partial X^{\ell+1}}{\partial
X^{\ell}}\right\|\leq\log\left\|J_{\text{\sc{Bn}}}(H^{\ell})\right\|\leavevmode\nobreak\
.$ (160)

###### Proof of Lemma C.12.

By definition, we have that:

$\displaystyle H^{\ell}$ $\displaystyle=W^{\ell}X^{\ell},$ (161)
$\displaystyle X^{\ell+1}$
$\displaystyle=\text{\sc{Bn}}(H^{\ell})\leavevmode\nobreak\ .$ (162)

Therefore, applying the chain rule, we get:

$\displaystyle\frac{\partial X^{\ell+1}}{\partial X^{\ell}}=\frac{\partial
X^{\ell+1}}{\partial H^{\ell}}\frac{\partial H^{\ell}}{\partial
X^{\ell}}=J_{\text{\sc{Bn}}}(H^{\ell})W^{\ell}\leavevmode\nobreak\ .$ (163)

Taking the logarithm of the norm of this quantity, we reach the conclusion:

$\displaystyle\log\left\|\frac{\partial X^{\ell+1}}{\partial
X^{\ell}}\right\|\leq\log\left\|J_{\text{\sc{Bn}}}(H^{\ell})\right\|+\log\left\|W^{\ell}\right\|=\log\left\|J_{\text{\sc{Bn}}}(H^{\ell})\right\|,$
(164)

where we have used the fact that the spectrum of the orthogonal matrix
$W^{\ell}$ contains only the singular value $1$ with multiplicity $d$. ∎

##  Appendix D Linear independence in common datasets

In this section, we empirically verify the assumption that popular datasets do
not suffer from rank collapse in most practical settings.

We provide empirical evidence for CIFAR10, MNIST, FashionMNIST and CIFAR100.
We test this assumption by randomly sampling $100$ input batches of sizes
$n=16,32,64,128,256,512$ from each of these datasets and then measuring the
rank of the Gram matrix of these randomly sampled batches using the
matrix_rank() function provided in PyTorch. We stop at size $512$ since we
approach the dimensionality of some datasets, i.e. FashionMNIST, MNIST. We
show in Table D1 the average rank with the standard deviation for each $n$,
over $100$ randomly sampled batches.

Table D1: Average rank of Gram matrix of input batches of size $n$ from different datasets. Mean and standard deviation are computed over $100$ randomly selected input batches, where the samples are chosen without replacement. Dataset | $n=16$ | $n=32$ | $n=64$ | $n=128$ | $n=256$ | $n=512$
---|---|---|---|---|---|---
CIFAR10 | $16.0\pm 0.0$ | $32.0\pm 0.0$ | $64.0\pm 0.0$ | $127.99\pm 0.09$ | $221.06\pm 2.89$ | $203.70\pm 3.58$
MNIST | $16.0\pm 0.0$ | $32.0\pm 0.0$ | $64.0\pm 0.0$ | $128.00\pm 0.00$ | $250.48\pm 1.53$ | $318.04\pm 2.82$
FashionMNIST | $16.0\pm 0.0$ | $32.0\pm 0.0$ | $64.0\pm 0.0$ | $128.00\pm 0.00$ | $238.19\pm 3.27$ | $275.37\pm 4.20$
CIFAR100 | $16.0\pm 0.0$ | $32.0\pm 0.0$ | $64.0\pm 0.0$ | $127.92\pm 0.27$ | $218.11\pm 3.69$ | $201.51\pm 3.95$

We would like to remark that these datasets are fairly simple in terms of
dimensionality and semantics, which can lead to correlated samples.
Furthermore, the rank degeneracy can be alleviated even in the larger batch
sizes through various data augmentations techniques. Note that these datasets
have a high degree of correlation between samples. Most notably, the average
cosine similarity between samples in a $512$ size batch is
$0.81,0.40,0.58,0.81$ for CIFAR10, MNIST, FashionMNIST and CIFAR100
respectively.

##  Appendix E Activation shaping

In this section, we explain the full procedure for shaping the activation, as
well as expand on the heuristic we use to choose the pre-activation gain.
Under the functional structure of the MLP in equation 11, let $\alpha_{\ell}$
be the pre-activation gain.

More formally, since the gradient norm has an exponential growth in depth, as
shown in Figure G5, we can compute the linear growth rate of log-norm of
gradients in depth. We define the rate of explosion for a model of depth $L$
and gain $\alpha$ at layer $\ell$ as the slope of the log norm of the
gradients:

$\displaystyle
R(\ell,\alpha_{\ell})=\frac{\log\|\nabla_{W_{\ell}}\mathcal{L}\|-\log\|\nabla_{W_{\ell-10}}\mathcal{L}\|}{10}\leavevmode\nobreak\
.$ (165)

Since the rate function is not perfectly linear and has noisy peaks, we
measure the slope with a $10$ layer gap in order to capture the true behaviour
instead of the noise.

Our goal is to choose $\alpha$ such that the sum of the rates across the
layers in depth is bounded by a constant that does not depend on the depth of
the model, i.e. $R(\ell,\alpha_{\ell})\leq\beta$, where $\beta$ is independent
of $L$. One choice to achieve this is to pick a gain such that the sum of the
rates behaves like a decaying harmonic sum in depth.

To this end, we measure the rate of explosion at multiple layers in a $1000$
layer deep model, for various gains $\alpha$ which are constant across the
layers in Figure E1 and notice that it behaves as
$R(\ell)=c_{1}\alpha^{c}_{2}$. In order to have the sum of rates across layers
behave like a bounded harmonic series in depth, we must choose the gain such
that it decays roughly as $\alpha^{c_{2}}=\ell^{-k}$ where $k>1$ results in
convergence. Therefore, we can obtain a heuristic for picking a gain such that
the gradients remain bounded in depth as $\alpha_{\ell}=\ell^{-k/c_{2}}$,
where we refer to $k/c_{2}$ as the gain exponent.

Figure E1: Explosion rate of the log norm of the gradients at initialization
for an MLP model with orthogonal weights and batch normalization, for sin and
tanh nonlinearities measured for a $1000$ layer deep model at layers $\ell$ as
a function of gain $\alpha$. The black trace shows the fitted function
$c_{1}\alpha^{c}_{2}$. Traces are averaged over $10$ independent runs, with
the shades showing the $95\%$ confidence interval.

This reduces the problem to picking the exponent such that the sum stays
bounded. We show how the behaviour of the explosion rate at the early layers,
for various models, is impacted by the exponent in Figure E2. Note that for
several exponent values, we able to reduce the exponential explosion rate and
obtain trainable models, which we show in Appendix G.

Figure E2: Explosion rate of the log norm of the gradients at initialization
for an MLP model with orthogonal weights and batch normalization, for sin and
tanh nonlinearities at depths $100,200,500,1000$ as a function of the gain
exponent. Traces are averaged over $10$ independent runs, where the shade
shows the $95\%$ confidence interval. Rate is measured at $\ell=10$ to avoid
the any transient effects of the function

##  Appendix F Implicit orthogonality during training

In this section, we provide empirical evidence that our architecture during
training maintains orthogonality across depths, while maintaining bounded
gradients. Figure F1 shows the evolution of the isometry gap of the weight
matrices $W_{\ell}$ during training, for models at different depths and
different nonlinearities. In order to show that these weights are updated
gradient descent, we also show the evolution of the norm of the loss gradients
with regards to matrices $W_{\ell}$ in Figure F2.

These experiments are performed on an MLP with orthogonal weight matrices and
batch normalization, with sin and tanh activations. The width is set to $100$,
batch size $100$ and learning rate $0.001$. The gain exponent is set to a
fixed value for all experiments. The measurements are performed on a single
batch of size $100$ from CIFAR10, after each epoch of training on the same
dataset.

Figure F1: Contrasting the isometry gap of weight matrices during training for
MLPs of depth $100$ (top), $500$ (middle), $1000$ (bottom). The middle layers
become increasingly more orthogonal with depth, while maintaining a small
isometry gap. During training, the isometry gap also remains low, suggesting
the matrices remain close to being orthogonal.

Figure F2: Contrasting the Frobenius norm of the gradients of the loss with
respect to the weights during training for MLPs of depth $100$ (top), $500$
(middle), $1000$ (bottom). The gradients do not vanish during training and
across different depths for all layers, suggesting that the orthogonality
evidenced in Figure F1 is not due to the weights not being updated during SGD.

##  Appendix G Other experiments

In this section we provide the train and test accuracies of deep MLPs on 4
popular image datasets, namely MNIST, FashionMNIST, CIFAR10, CIFAR100.
Hyperparameters and measurements procedure are described in Section 4.

###  1 Supplementary train and test results on MNIST, FashionMNIST, CIFAR10,
CIFAR100

Figure G1: Contrasting the train and test accuracy of MLPs with gained sin,
tanh and identity activations on MNIST. The identity activation performs much
worse than the nonlinearities, indicating the fact that the sin and tanh
networks are not operating in the linear regime. The network is trained with
vanilla SGD and the hyperparameters are width $100$, batch size $100$,
learning rate $0.001$.

Figure G2: Contrasting the train and test accuracy of MLPs with gained sin,
tanh and identity activations on FashionMNIST. The identity activation
performs much worse than the nonlinearities, indicating the fact that the sin
and tanh networks are not operating in the linear regime. The networks are
trained with vanilla SGD and the hyperparameters are width $100$, batch size
$100$, learning rate $0.001$.

Figure G3: Contrasting the train and test accuracy of MLPs with gained sin,
tanh and identity activations on CIFAR100. The identity activation performs
much worse than the nonlinearities, indicating the fact that the sin and tanh
networks are not operating in the linear regime. The networks are trained with
vanilla SGD and the hyperparameters are width $100$, batch size $100$,
learning rate $0.001$. Figure G4: Contrasting the train test accuracy of MLPs
with gained sin, tanh and identity activations on CIFAR10. The identity
activation performs much worse than the nonlinearities, indicating the fact
that the sin and tanh networks are not operating in the linear regime. The
networks are trained with vanilla SGD and the hyperparameters are width $100$,
batch size $100$, learning rate $0.001$.

###  2 Supplemental figures

We present empirical results in Figure 2 showing that degenerate input batches
are a hard constraint for orthogonalization without gradient explosion. For
MLPs with different depths, we show that by repeating samples in a batch of
size $10$ we get an exponential gradient explosion, which is unavoidable
theoretically.

Furthermore, we show how non-linearities affect the gradient explosion rate in
Figure G5. Using standard batch normalization and fully connected layers from
PyTorch we show that non-linearities maintain a large isometry gap. This is a
critical issue for our theoretical framework, since we take advantage of the
fact that the identity activation achieves perfect orthogonality in order to
prove that the gradients remain bounded in depth.

Figure G5: Left: average log-norm of gradients (log-scale y-axis) at the first
layer for networks with different depths, evaluted on CIFAR10. Right: Isometry
gap (log-scale y-axis) at the last layer for networks with different depths,
evaluted on CIFAR10. The MLP is initialized with orthogonal weights and batch
normalization, with standard modules, with $\sin,\tanh,\text{identity}$ non-
linearities. After stabilizing the isometry gap, the non-linearities have an
exponential gradient explosion.

###  3 Influence of mean reduction on the gradient bound

In this section, we compare whether adding mean reduction and the additional
factor of $\frac{1}{n}$ in the denominator of the batch normalization module
influences our gradient bound. As expected, we show in Figure G6 that in both
cases, for the identity activation, the result remains similar, with the
gradients remaining bounded in depth.

Figure G6: Comparing the gradient explosion rate in networks with standard
batch normalization (blue) and networks with the simplified batch
normalization operator from our theoretical framework (orange). Notice that
the 2 traces are similar in terms of gradient explosion. Traces are averaged
over $10$ runs with the shaded regions showing the $95\%$ confidence interval.
Samples are from CIFAR10.