Towards Training Without Depth Limits: Batch Normalization Without Gradient Explosion

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?

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

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

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.

Related work

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.

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.

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.

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).

Main results

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.

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.

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

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.

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.

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

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:

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}$).

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

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.

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.

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 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.

where the expectation is over the random orthogonal weight matrices.

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

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).

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.

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.

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.

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.

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\}.$

Discussion

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.

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

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

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

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

Using properties of the determinant, we have

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

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

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.

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

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.

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.

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

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)$.

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

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

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

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]:

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].

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:

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

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

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 $,,$ respectively, for each $i\in\{1,2,3\}$.

Plugging back into the original equation, we obtain:

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

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.

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

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$.

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:

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:

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

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:

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

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:

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

Furthermore, we can restate the condition on maximum using:

If we negate and flip the two sides we arrive at

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

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

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

Let us restate the theorem for completeness of the appendix:

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

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

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.

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

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$.

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

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. ∎

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.

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

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}$:

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

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.

In fact, the maximum of the two bounds is

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:

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.

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:

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

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:

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

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

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:

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$ to conclude

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

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.

Assembling the equations into matrix form, we obtain:

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

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}$.

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.

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:

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

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

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:

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:

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

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:

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

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

Therefore, applying the chain rule, we get:

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

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.

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

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.

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.

Appendix F Implicit orthogonality during training

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.

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.

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.

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.