Towards Training Without Depth Limits: Batch Normalization Without Gradient Explosion
Authors: Alexandru Meterez
Amir Joudaki
Francesco Orabona
Alexander Immer
Gunnar Ratsch
Hadi Daneshmand
Venue: ICLR 2024 poster (2024)
Keywords: mlp
batch-normalization
optimization
depth
calculus
theory
deep-learning
non-asymptotic
Primary area: optimization
TL;DR: We show that MLPs with orthogonal weights and batch normalization orthogonalize the inputs, without gradient explosion at arbitrary depths.

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 *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 non-linear activations.

Published as a conference paper at ICLR 2024
TOWARDS TRAINING WITHOUT DEPTH LIMITS: BATCH NORMALIZATION WITHOUT GRADIENT EXPLOSION
Alexandru Meterez" D-INFK, ETH Zürich ameterez@ethz.ch
Amir Joudaki" D-INFK, Zürich a 4joudaki@ethz.co
Francesco Orabona
CEMSE, KAUST
Francesco@orabonaa
Alexander Immer D-INFK, ETH Zürich simmer@ethz.co
Gunnar Rätsch D-INFK, ETH Zürich raetschEinf aarfffethz.c..
Hadi Daneshmand
MIT LIDS/Boston University
hdanesh@mit.
ABSTRACT
Normalization layers are one of the key building blocks for deep neural networks. Several theoretical studies have shown that batch normalization improves the sig- nal propagation, by avoiding the representations from becoming collinear across the layers. However, results conclude that this benefit comes at the expense of exploding gradients in depth. Motivated by these two aspects of batch normalization, in this study pose the following question: "Can a normmlizddd 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 natch-normalizztion that provably has bounded gradients at any depth. Based on Weingarten calculus, we develop a rigorous and asymmtotiiicc theory for this constructed MLP that gives a precise characteriza- tion of forward signal propagation, while proving that gradients remain bounded for linearly independent input samples, which holds most practical settings. In- spired by theory, we also design an activation shaping scheme that empirically achieves the same properties for certain non-linead activations.
rrean-field theory of batch normalization also
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) capture longer structural dependencies (Devlin al., 2018: Liu al., 2019 Brown et al., 2020; Goyal et al., 2021; Raffel et al., 2020). Also, sufficiently deep convolu- tional networks can outperform humans 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. convergence when training deep neural networks. Hence, computation has become barrier in deep learning, demanding extensive research and engineering.
fact, standard optimization methods exhibit a slower
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 al., 2022: Saxe et al., 2013a; Daneshmand et al., 2020), but also in transformer architectures (Dong al., 2021; Noci al., 2022). This issue significantly contributes the training slowdown of deep neural networks. Hence, has become the focus theoretical and experimental studies (Saxe et al., 201 3a; Feng al., 2022; Daneshmand et al., 2021: Noci et al., 2022).
One of the most successful methods to avoid rank collapse is Batch Normalization (BN) (loffe & Szegedy, 2015), as proven in a number of theoretical studies (Yang et al., 2019; Daneshmand al.,
Equal contribution
Code available httpps:/ggithhe http:///ggthhhhs.ccmm////aeneeee hxw.a//eennnuueeeeeee//nnraaaaa

---

Published as a conference paper at ICLR 2024
2020: 2021). Normalization imposes particular bias across the layers of neural networks (Joudaki et al., 2023b). More precisely, the representations of a batch of inputs become more orthogonal after each normalization (Joudaki et al., 2023b). This orthogonalization effect precisely avoids the rank collapse of deep neural networks at initialization (Yang et al., 2019; Joudaki et al., 2023b; Daneshmand et al., 2021; Joudaki et al., 2023a).
While batch normalization effectively avoids rank collapse, it causes numerical issues. The ex- isting literature proves that batch normalization layers cause exploding gradients in MLPs in an seiivaiionindeeeenden 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 Bengio, 2010). However, such methods cannot avoid gradient explosion in the presence of batch normaliza- tion (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 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
E iisometry gap) = O (e (e-depth/C)
(1)
where 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 O ((width" -bbl) close to orthogonality, we prove that the output representations become perfectly orthogonal in the infinite depth limit. This perfect orthogonalizatton turns out 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
B[ log (gradient norm for each layer)] ] = (width²).
Thus, the gradient is bounded by 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 the previously studied mean-field or infinite width regime, our theoretical results hold in practical settings where the width finite.
The limitation of our theory is that it holds for simplification in the BN module and linear ac- tivations. 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 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 MI.ss with BN.
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 repre- sentations. (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 partieular initialization scheme that avoids gradient vanishing/explosion for neural networks with rectifier non- con-lincassties while Glorot & Bengio (2010) study the effect of initializa- tion sigmoidal activations. However, such initializations cannot avoid gradient explosion for
2

---

Published as a conference paper at ICLR 2024
networks with batch normalization (Yang et al., 2019; Lubana et al., 2021). (ii) Saxe 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 (loffe & Szegedy, 2015), residual connections (He et al., 2016), and seff-normalizing activations (Klambauer et al., 2017). A related line of work has shown how signal propagation can be achieved without batch normalization in feed- forward networks (Burkholz & Dubatovka, 2019) and ResNets (Brock et al., 2021). Other works on CNNs (Blumenfeld et al., 2020) have shown that symmetry breaking is vital element for achieving signal propagation with stable gradients in deep models. 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 depth-independen training con- vergence. Pennington et al. (2017) show that MLPs with sigmoidal activations achieve dynamical isometry when initialized with orthogonal weights. Similar benefits have been achieved by initial- izing CNNs with orthogonal or almost orthogonal kernels (Xiao et al., 2018; Mishkin & 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 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 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 & Hornik (1995) provide an overview over various theoretical manuscripts studying linear neural networks. Despite linearity, as Saxe al. (2013arb) observe, the gradient dynamics in 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 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 & Hu (2019) prove that under certain conditions on the model width and input degeneracy. linear MLPs with Xavier initialized weights (Glorot & Bengio, 2010) converge linearly to the global optimum. Akin these studies, also analyze networks with linear activations. However, batch normalization is a non-linear function, hence the network study in this paper is highly non-linear function of its inputs.
3 MAIN RESULTS
We will develop our theory by constructing networks that do not suffer from gradient explo- sion (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 E dd denote the representation of samples in Rd layer E, then
X441 (W,
t= L,
where W2 E Rdxd are random weights, n is the mini-batch size and d' is the feature dimension. Analogous to recent theoretical studies batch normalization (Daneshmand et al., 2021; 2020), we define the BN operator BN Rdxn + Bddx as
X
BN(X) diag(XX x BN(X),
(4)
E X2
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 H6).
3

---

Published as a conference paper at ICLR 2024
A second minor difference is that in the denominator, we have omitted factor. However, this only amounts to a constant scaling of the representations and does not affect our results.
Compared to Daneshmand et al. (2021), need two main modifications avoid gradient explo- sion: (i) n = d, and (ii) We are random orthogonal matrices. More precisely, we assume the distribution of We is the Haar measure over the orthogonal group denoted by Ca (Collins & Sniady, 2006). Such initialization scheme is widely used in deep neural networks without batch normal- ization (Saxe et al., 201 3a; Xiao et al., 2018; Pennington et al., 2017). For MLP networks with BN, prove such initialization avoids the issue of gradient explosion, while simultaneously orthogo- nalizing 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 nor- malization become increasingly orthogonal through the layers, thereby enhancing the signal propa- gation 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 the notion of isometry gap, Rdxn_ R. introduced by Joudaki et al. (2023b). Isometry gap is defined as
det(xXT)) TT((XX))
o(X) llog
(5)
One can readily check that 6(X) > 0 and it is zero when X is orthogonal matrix, i.e., XXT Id- The isometry denoted by I Raxn - IR is defined as I(X) exp(- ((X))
Geometric interpretation of isometry. While the formula for isometry gap may seem enigmatic at first, it has simple geometric interpretation that makes it in- tuitively understandable. The determinant det(X" xx) det(X)² is the squared volume of the parallelepiped spanned by the columns of X, while Tr( X) is the squared-norms of the columns X. Thus, the ratio between the provides scale-invariant notion of volume and isometry. On the one hand, if there any collinearity between the columns, the volume will vanish and the isometry gap will be infinity, o(X) oo. On the other hand, ((X) = 0 implies T is scaled identity matrix. We will prove serves as Lyapunov function for the chain of hidden representations {X
10-4
Width
e(Xx))
50
100
10-3
50 (theory)
100 (theory)
10
50
Figure I: Isometry gap (y-axis, log-scale) in depth for an MLP with orthogonal weights, over ran- domly generated data. As predicted by Theo- rem I, 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
Theory for orthogonalization. of representations and depth. The following theorem establishes the link between orthogonality
Theorem 1. There is an absolute constant such that for any layer < we have
Eø(X++1) ((No)e-
where
k := Cal de(
Theorem states that if the samples in the input batch are not linearly dependent, representations exponential rate in depth. The orthogonalization in depth ensures the avoidance of the rank collapse of representations, which a known barrier to training deep neural networks (Daneshmand et al., 2020; Saxe et al., 2013a; Bjorck et al., 2018).
approach orthogonality
4

---

Published as a conference paper at ICLR 2024
Figure compares the established theoretical decay rate of ø with the practical rate. Interestingly, the plot confirms that the rate depends on width in practice, akin to the theoretical rate in Theorem It is worth mentioning that the condition 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, batch-normalizzttion 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 CIFARIO, CIFARI00, MNIST and MNShi we empirically tested that most batches across various batch sizes are full-rank (see Section E for details on the average rank of batch in these datasets).
Theorem 1 distinguishes itself from the existing orthogonalization results in the literature (Yang et al., 2019; Joudaki et al., 2023b) 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 ap- plicability real-world settings. While Daneshmand et al. (2021) provide a non-asymptotic bound for orthogonalization, the main result relies
assumption that is hard to verify.
Proof idea of Theorem We leverage a recent result established by Joudaki et al. (2023b), proving that the isometry gap does not decrease with BN layers. For all non-degenerate matrices X E Rdxd the following holds
I(BN(X)) 2 1+ variance ( [| 1-
I(X).
(mean ( 1
Using the above result, 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 Igddd and W B.dxd be an orthogonal matrix and X' = WX; then,
variance (1X35 ( I(X).
I(BN(X')) 2 (1+ (mean (mean[|1)].1135 [13.1115=1)
It is straightforward to check that there exists at least orthogonal matrix W for which I(BN(W = 1 (see Corollary B.3). Thus, I(.) strictly increases for some weight matrices, as long as X is not orthogonal. When the distribution of W is the Haar measure over the orthogo- nal group, we can leverage recent developments in Weingarten calculus (Weingarten, 1978; Banica et al., 2011; Collins & Sniady, 2006; Collins et al., 2022) to calculate a rate for the isometry increase in expectation:
Theorem 3. Suppose W ~ Gg is a matrix drawn from Og such that the distribution of and UW are the same for all orthogonal matrices U. Let {}i} be the eigenvalues of XX... Then,
Ew [(B((WX)]2 (1
D+ 1 1
-1
I(X)
(8)
(d
holds for all X = BN((-), with equality for orthogonal matrices.
The structure in X induced by BN ensures its eigenvalues lie in the interval (0, 1), in that the multi- plicative factor in the above inequality is always greater than one. In other words, I(-) 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 estab- lished in the following Corollary (recall & log I).
Corollary 4 (Isometry gap bound). Suppose the same setup in Theorem 3, where X W. Then, we have:
Ew (wlo(X')xx X ((X)= + log (1- (1- Ex(ds 1)2 -1)2
(9)
2d 2)
Notice that the term (4 2x = O yielding log
(4m 5. 2d'
The rest of proof is based induction over the layers, presented in Appendices B and C.
5

---

Published as a conference paper at ICLR 2024
3.2
ORTHOGONALIZATION AND GRADIENT EXPLOSION
There is subtle connection between orthogonalization and gradient explosion. Suppose the input batch is rank-deficient, i.e., degenerate. elaborated above, since all operations in our MLP can be formulated as matrix products, they cannot recover the rank of the representations, which thus re- main degenerate. By perturbing the input such that becomes full-rank, the output matrix becomes orthogonal, hence non-degenerate at exponential rate in depth as proven in Theorem 1.
Thus, a slight change in inputs leads to a signif- icant change in outputs from degeneracy to or- thogonality. Considering that the gradient mea-
30
log Vawll All
20
Inputs Non-degenerate
sures changes in the loss for infinitesimal in-
puts changes, the large changes in outputs po-
Degenerate
tentially lead gradient explosion. While this
only an intuitive argument, we observe that in practice the gradient does explode for degen- erate inputs, as shown in Figure 2.
Depth
Figure 2: Logarithmic plot for the gradient norm of the first layer for networks with different num- ber of layers evaluated on degenerate (orange) and non-degenerate (blue) inputs. The degenerate in- puts contain repeated samples from CIFARIO in the batch, measured at initialization for MIPs of various depths. While gradients explode for de- generate inputs, there is no explosion for non- degenerate inputs. Traces are averaged over 10 independent runs.
Nonetheless, in Figure 2 we observe that for non non-degenerate inputs the gradient norm does explode. In fact, we observe that inputs are often non-degenerate in practice (see Table E1 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 degener- ate inputs the gradient norm remains bounded. Therefore, theoretical guarantee is necessary to ensure avoiding gradient explosion.
3.3 AVOIDING GRADIENT EXPLOSION 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 L. did > IR be O(1) Lipshiiis and input batch Xo be non- degenerate. Then, there exists an absolute constant C. such that for all S. L it holds
E (V s Cd +1)
(10)
where the expectation over the random orthogonal weight matrices. Remark 1. For degenerate inputs ((Xo) = oc holds. in that the bound becomes vacuous.
Remark 2. The Lipschitt condition holds in many practical settings. For example, in clas- sification setting, MSE and cross entropy losses obey the O11sshipsh condition (see Lemma D.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 al., 2019) as long as the inputs are non-degenerate. Such guarantees are necessary ensure backpropagation will not face numerical issues.
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 is tight in terms of inputs constraints. Furthermore, mentioned before, random batches sampled from commonly used benchmarks, such as CIFAR10, CIFAR 100, MNIST, and FashionMNIST,, are non non-degenerate in most practical cases (see Section E for more details). Thus, the assumptions and thereby assertions of the theorem are valid for all practical purposes.
Theorem 5 states that long the input samples

---

Published as a conference paper at ICLR 2024
To the best of our knowledge, The- orem is the first non-asymptotic gradient analysis that holds for net- works with batch normalization and finite width. Previous results heavily rely on mean field analyses in asymp- totic 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 explo- sion, which has previously been con- sidered "unavoidable" (Yang et al., 2019). Figure illustrates this pro- nounced discrepancy.
10.0
log |/ Vw, Cll
7.5
Initialization
Orthogonal
5.0
Normal
2.5
0.0
250
500
1000
Depth
Figure 3: Logarithmie plot for the gradient norm of the first layer for networks with different number of layers evaluated on CIFARIO. 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 constant, validating Theorem 5. Traces are averaged over 10 runs and shaded regions denote the 95% confidence intervals.
Proof idea of Theorem 5. The first important observation is that, due to the chain rule, we can bound the log-norm the gradient of 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 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 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 f, the log- fog-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 nnemaaiaaliationn matrices. Thus, exponential decay in isometry gap stated in Theorem 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 L.
invoke the
4 IMPLICATIONS ON TRAINING
In this section, 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 Thcorem 5, and does not have the rank collapse issue in depth via the orthogonalization property established in Theorem We find that the constructed MLP is therefore less prone to numerical issues that arise when training deep networks.
(a) sin
(b) tanh
(c) identity
1.00
1.00
Train accuracy
0.75
0.75
0.4
Depth
100
0.50
0.50
0.2
1000
0.25
0.25
10a
10a
10
10a
10a
Epoch
Epoch
Epoch
Figure 4: Contrasting the training accuracy of MLPs with BN and shaped sin, shaped tanh and identity activations, on the CIFARI0 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.
7

---

Published as a conference paper at ICLR 2024
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-independeen conver- gence rate for linear MLPs. In other words, the number of iterations to get to certain accuracy does not vary widely between networks with different depths. Figure 4 (c) shows the convergence of SGD for CIFARIO, 10, with learning rate 0.001, for MLPs with width d = 100 and batch size 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, 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 non-linear activa- tions are essential parts of feature learning (Nair & Hinton, 2010; Klambauer et al., 2017: Hendryeks & 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 H5 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 normal- ization, perturb the isometry of representations and prevent them from reaching zero isometry gap in depth. This problem turns out to be not just theoretical nuisance, but to play 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 al., 2022; Martens et al., 2021; He et al., 2023; Noci et al., 2023) can alleviate this problem, which 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 net- works by parameterizing activation functions. In 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 LeakyReLU in tailored activa- tion transformations (Zhang et al., 2022). In a similar direction, Li et al. (2022) propose activation shaping order to avoid degenerate output covariance. While the mechanism proposed by Li al. (2022) covers both smooth and non-smooth activations, we focus on their result for LeakyReLU, which consists of shaping the negative slope of the activation 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 con- junction with normalization layers. In fact, in networks with batch normalization. even linear ac- tivations 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 normal- ization layers may seem fruitless. Remarkably, we empirically demonstrate that we can leverage activation shaping to avoid gradient explosion in depth by using pre-activation gain at each layer.
Without gain
With gain
40
log Cll, Cl
log Viw, Cll
Activation
identity
tanh
250
1000
250
1000
Depth
Depth
Figure 5: Logarithmic plot contrasting the effect of gain 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.
8

---

Published as a conference paper at ICLR 2024
Our activation shaping consists of tuning the gain of the activation, ie., tuning for ((zz). We consider non-linear activations a E (tanh, sin}.
The special property that both tanh and sin activations have in common is that they are centered, (0) = 0, are differentiable around the origin o' (0) = 1, and have bounded gradients o'x)) I 1, Vx. Therefore, by tuning the per-layer pre-activation gain og towards 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 function of the gain parameter ag. Formally, we consider an MLP with shaped activations using gain Oy for the 9th layer, that has the update rule
Xe+1 (a,,
(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 model of depth I. and gain ae at layer We show in Figure 5 that by tuning the gain properly, we are able 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 to have a bounded total sum of rates across layers, by ensuring faster decay than harmonic series (see App. F for more details activation shaping). Figure illustrates that this activation shaping strategy effectively avoids gradient explosion while maintaining the signal propagation and orthogonality of the outputs in depth. Furthermore, Figure shows that the training accuracy remains largely depth-independent. For further experiments using activation shaping, see Appendix H.
slope of the log norm of the gradients R(e, oz))
6 DISCUSSION
Isometry gap
(b) Gradient norm
10a
Epoch
V ClF
x(W))
10 2
50
109
100
200
10-1
300
10
500
1000
500
1000
Figure 6: Implicit orthogonality bias of SGD. Training an MIP with width d = 100, batch size
100, and depth T, 1000, activation tanh, using SGD with Ir 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.
Implicit bias of SGD towards orthogonality in optimization. Optimization over orthogonal ma- trices 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 & Martinez-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 constraint based on Householder matrices.
While these studies enforce orthogonality constraints, one of our most striking empirical observa- tions 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 or- thogonal during training. One could hypothesize that this is due small gradients in these layers. In Figure 6, we observe that the gradients of these middle layers are not negligible. Thus, 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 ample direction for future research.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS & CONTRIBUTIONS
AM: proof' of Section B, driving experiments, and writing the paper. AJ: proof's of Section C and Section D, designing the activation shaping scheme, and writing the paper. FO: proposing the idea of using orthogonal weights to achieve perfect isometry, reading the proofs, help with writing. Al: reading the proofs, help with experiments and paper writing. GR: help with experimental designs for activation shaping and paper writing. HD: proposing the idea and leading the proofs, proposed using orthogonal weights to avoid gradients explosion, writing the initial draft.
AJ is funded through Swiss National Science Foundation Project Grant #200550 Andre Kahles. HD acknowledges support from the NSF TRIPODS program (award DMS-2022448), Al acknowl- edges support from the Max Planck ETH Center for Learning Systems. AJ and Al were partially funded by ETH Core funding (to G.R.).
REFERENCES
Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International conference on machine learning, pp. 1120-1128, PMLR, 2016.
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.
Pierre Baldi and Kurt Horiik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1))53-58, 1989.
Pierre F Baldi and Kurt Hornik. Learning in linear neural networks: A survey. IEEE Transactions on neural networks, 6(4):837-858, 1995.
Teodor Banica, Benoit Collins, and Jean-Mare Schlenker. On polynomial integrals over the orthog- onal group. Journal of Combinatorial Theory, Series A, 118(3):778795,, 2011.
Nils Bjorck, Carla P Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normal- ization. Advances in neural information processing systems, 31, 2018.
Yaniv Blumenfeld, Dar Gilboa, and Daniel Soudry, Beyond signal propagation: is feature diver- sity necessary in deep neural network initialization? In International Conference on Machine Learning. pp. 960-969. PMLR, 2020.
Andrew Brock, Soham De, and Samuel L Smith. Characterizing signal propagation to close the performance gap in unnormalized resnets. arXiv preprint arXiv:2101.08692 2021.
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: 33:1877-1901, 2020.
Rebekka Burkholz and Alina Dubatovka. Initialization of relus for dynamical isometry. Advances in Neural Information Processing Systems, 32, 2019.
Benoît Collins. Moments and cumulants of polynomial random variables on unitary groups, the Itzykson-Zude Zuber integral, and free probability. International Mathematics Research Notices, 2003 (17):953-982, 2003.
Benoît Collins and Sho Matsumoto. On some properties of orthogonal Weingarten functions. Jour- nal of Mathematical Physics, 30(11), 2009.
Benoît Collins and Piotr Sniady. Integration with respect to the haar measure on unitary, orthogonal and symplectic group. Communications in Mathematical Physics, 264(3):773-795, 2006.
10

---

Published as a conference paper at ICLR 2024
Benoit Collins, Sho Matsumoto, and Jonathan Novak. The Weingarten calculus. arXiv preprint arXiv:2109.14800 2022.
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:18887-18398. 2020.
Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes repre- sentations in deep random networks. Advances in Neural Information Processing Systems, 34: 4896-4906, 2021.
Jacob Devlin, Ming-Wa Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Yihe Dong, Jean Jean- Sann-Bapistte Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pp. 2793-2803 PMLR, 2021.
Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In International Conference on Machine Learning, pp. 1655-1664. PMLR, 2019.
Ruili Feng, Kecheng Zheng, Yukun Huang, Deli Zhao, Michael Jordan, and Zheng-Jun Jun Zha. Rank diminishing in deep neural networks. Advances Neural Information Processing Systems, 35: 33054-33065, 2022.
Kenji Fukumizu. Effect of batch learning in multilayer neural networks. Gen, 1(04):1E-03, 1998.
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, pp. 249-2566 256. JMLR Workshop and Conference Proceedings, 2010.
Naman Goyal, Jingfei Du, Myle Ott, Giri Anantharaman, and Alexis Conncau. Larger-scale trans- formers for multilingual masked language modeling. arXiv preprint arXiv:2105.00577 2021.
Bobby He, James Martens, Guodong Zhang. Aleksandar Botev, Andrew Brock, Samuel L. Smith, and Yee Whye Teh. Deep transformers without shorteuts: Modifying self-attention for faithful signal propagation. arXiv preprint 2301 2023.
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, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Mikael Henaff, Arthur Szlam, and Yann LeCun. Recurrent orthogonal networks and aag-memory tasks. In International Conference Machine Learning, pp. 2034-2042. PMLR, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXiv preprint arXiv:1606. 08415, 2016.
Sepp Hochreiter. The vanishing gradient problem during learning recurrent neural nets and problem solutions. International Journal of Uncertainty, Fuzziness and Knowledge-Baud Based Systems, 6(02): 107-116, 1998.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456 pmlr, 2015.
Cijo Jose, Moustapha Cissé, and Francois Fleuret. Kronecker recurrent units. In International Conference Machine Learning, pp. 2389. PMLR, 2018

---

Published as a conference paper at ICLR 2024
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 Confer- on Machine Learning, 2023a.
Amir Joudaki, Hadi Daneshmand, and Francis Bach. On the impact of activation and normalization in obtaining isometric embeddings at initialization. arXiv preprint 2tth:2335518899, 2023b.
Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-normalizing neural networks. Advances in neural information processing systems, 30, 2017.
Quoc Le, Navdeep Jaitly, and Geoffrey Hinton. A simple way to initialize recurrent networks of rectified linear units. arXiv preprini arXiv: /504 150550044000001 2015.
Mario Lezcano- Casado and David Martunez-Rubio Cheap orthogonal constraints in neural net- works: A simple parametrization of the orthogonal and unitary group. In International Confer- ence on Machine Learning, pp. 3794-3803. PMLR, 2019.
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-10088, 2022.
Yinhan Liu, Myle Out, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pre- training approach. arXiv preprint arXiv:1907.11692 2019.
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, pp. 11976-1198,, 2022.
Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka. Beyond BatchNorm: towards unified understanding of normalization in deep learning. Advances in Neural Information Processing Systems, 34:4778-4799 2021.
Andrew Maas, Awni Y Hannun, and Andrew Y Ng. Rectifier nonlinearities improve neural net- work acoustic models. In Proc. ICML, volume 30, pp. 3. Atlanta, GA, 2013.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, 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:2110001763 2021.
Zakaria Mhammedi, Andrew Hellicar, Ashfaqur Rahman, and James Bailey. Efficient orthogo- nal parametrisation of recurrent neural networks using Householder reflections. In International Conference on Machine Learning, pp. 2401-2409. PMLR, 2017.
Dmytro Mishkin and Jiri Matas. All you need is good init. arXiv preprint arXiv:1511.06422 2015.
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), pp. 807-814, 2010.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aure- lien Lucchi. Signal propagation in Transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 15:27198-27211 2022.
Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, and Daniel Roy. The shaped transformer: Attention models in the infinite depth-and- limit. arXiv preprint arXiv: 17059 2023.
Jeffrey Pennington and Pratik Worah. Nonlinear random matrix theory for deep learning. Advances in neural information processing systems, 30, 2017.
12

---

Published as a conference paper at ICLR 2024
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learn- ing through dynamical isometry: theory and practice. Advances neural information processing systems. 30, 2017.
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-555, 2020.
Andrew M Saxe, James McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics of learning in deep linear neural networks. arXiv preprint arXiv: 2013a.
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.
Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116 (23):11537-11566, 2019.
Ohad Shamir. Exponential convergence time of gradient descent for one-dimensional deep linear neural networks. In Conference on Learning Theory, pp. 2691-2713, PMLR, 2019.
Eugene Vorontsov, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal. On orthogonality and learn- ing recurrent networks with long term dependencies. In International Conference on Machine Learning, pp. 3570-3578. PMLR, 2017.
Don Weingarten. Asymptotic behavior of group integrals in the limit of infinite rank. Journal of Mathematical Physics, 19(5):99--1011, 1978.
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, pp- 2023.
Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. Cyt: Introducing convolutions vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 22-31, 2021.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and mean field theory of CNNs: How to train 10,000-layer vanilla con- volutional neural networks. In International Conference on Machine Learning, pp. 5393-5402. PMLR, 2018.
Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel Schoenholz. A mean field theory of batch normalization. ICLR, 2019.
Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Global optimality conditions for deep neural networks. arXiv preprint arXiv: / 707.02444.4444 2017.
Guodong Zhang, Aleksandar Botev, and James Martens. Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. arXiv preprint uxXiv:2233.88120, 2022.
13

---

Published as a conference paper at ICLR 2024
TABLE OF CONTENTS
The appendix is structured in the following sections:
Appendix A: supplemental related work
Appendix B: contains the proof based on Weingarten calculus for the conditional isometry gap decay rate, conditioned on the previous layer in the network.
Appendix C: contains the proof for the isometry gap exponential rate decay in depth, based on the conditional rate established in Appendix B.
Appendix D: contains the proof for the log- nog-norm gradient bound at arbitrary depths, based on the orthogonality result established in Appendix C.
Appendix E: contains the numerical verification of the average rank of the input batch in common datasets.
Appendix F: contains the full derivation and additional explanations of the activation shap- ing heuristic.
Appendix G: contains further results on CIFARIO, MNIST, FashionMNIST and CI- FARI00, showcasing the implicit orthogonality bias of SGD during training.
Appendix H: contains supplemental experiments, including train and test results MÑIST, FashionMNNIST, CIFAR100, as well as the effect of mean reduction and non- linearities without shaping the gradient log-norms.
A SUPPLEMENTAL RELATED WORK
Due to space constraints, we expand upon the related works in this section.
Mean field theory for random neural networks. The existing analyses for random networks often rely on mean-fiell regimes where the network width tends to infinity (Pennington et al., 2017; Yang et al., 2019; Li et al., 2022: Pennington & Worah, 2017). However, there is discrepancy between mean-field regimes and the practical regime of finite width. While some analyses attempt to bridge this gap (Joudaki et al., 2023a; Daneshmand 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 necessary. To go beyond mean mean-field regimes, we leverage recent theoretical advancements in Weingarten calculus (Weingarten, 1978; Collins, 2003; Banica al., 2011; Collins Sniady, 2006; Collins et al., 2022).
B CONDITIONAL ORTHOGONALIZZTION
Isometry after rotation. Our analysis is based on (Joudaki et al., 2023b, Corollary 3), which restate in the following Lemma:
Lemma B.1. For all non-degenerate matrices E Badd we have:
variance (I 1155
I(BN(X)) 2 1+
I(X).
(12)
(mean [|| 1155
Lemma B.1 proves isometry bias of BN. The next lemma proves that isometry does not change under rotation
Lemma B.2 (Isometry after rotation). Let X E gdxd and W ~ Og be a random orthogonal matrix and X' = W. Then,
I(BN(X')) 2 1+ variance (11,,nce)1xx111113 11]
I(X).
(13)
(mean ( )
14

---

Published as a conference paper at ICLR 2024
Proof. Using properties of the determinant, we have
det(X'XX") det(W))det(XX"") det(XXX)
(14)
where the last equation holds since W an orthogonal matrix. Furthermore,
Tr(X'x Xx W
(15)
=Tr(XX W W
(16)
=I
-T(XXx)).
(17)
Combining the last two equations with Lemma B.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 B.3 (Increasing isometry). Let X I and denote its singular value decomposition X = Udiag ((o,)?=)) V where Ú and V are orthogonal matrices. Then, we have:
I(BN(U H)) =1.
(18)
Proof. Let diag ((o:)=)) be the diagonal matrix containing the singular values of X.
Then,
we have:
BN( x)=
AN(UTUVV
(19)
BN(SVT)
(20)
diag(SV VTSIISVVVSSSSS
SV
(21)
ST
(22)
-V
(23)
which has maximum isometry 1 since it's 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 stated in the last corollary). Hence, it is possible to increase isometry with random orthogonal matrices.
Theorem B.4. Suppose W ~ Og is a matrix drawn from Oa such that W d WU for all orthogonal matrices U. Let {Ai be the eigenvalues of X.X Then,
Ew [[(B((XX)] X
I(X)
(24)
(A. 1)2
1
(d+2)
holds for all X = BN((), with equality for orthogonal matrices.
Remark 3. Note that the assumption on X = BN((), can be viewed as an induction hypothesis, in that we can recursively apply this theorem to arrive at a quantitative rate at depth.
Notably, ΣL = d x = BN(-). Hence, one can expect that Σis X3 < dl for all full random matrices in form of X BN(-).
Proof. We need to compute the variance/mean ratio in Lemma 2. Let X E dd have SVD de- composition X Udiag(a,) where U and V are orthogonal matrix and of Aς. Since
15

---

X.W _s Published as a conference paper at ICLR 2024
the distribution of W is invariant transformations with orthogonal matrices, the distribution of W equates those of X' = WWdiaggo.V'. It easy to check that
d
| ||
0
W jt
(25)
i=1
Thus,
d
Σ
X
12
Σ A
=
(26)
1=1
where the last equality holds due to the batch normalization. Thus, we have
d 2 -
d
d²
E
E
(27)
C 1 1x
|X 113
2
E
||
We need to estimate
1
E [| = E
d W
W.2
(28)
ik
k=1
Since square root function is concave, we have vx S 1 + {x - 1). Thus
d Σ (A+ 1)
d
E ||
s E
1
0.5
[+0.5 0.5 Σ
(Ax 1)333
(29)
k=1
1
1
+
(Ax 1)(4 - 1) [w W
(30)
4
k.g
1
+
1
Σ (Ax 1)(2 - 1) W kt9
+
(As (1<A -
[
W
4
E
E2
(31)
where in the first equality we have used the fact that the cross terms reduce, where the expectations are applications of Weingarten calculus:
0.5 Σ (AR 1)W3 + Σ (AR
1)W?
E
0.5
A W
0.5
A
W
jk
jk
(32)
d
d
= 0.5
A
[W2]] ik
+ 0.5
A
[
jkc
A=1
k=1
(33)
d
0.5 d
0.5
A
+
1
(34)
d
d
k=1
k=1
=0
(35)
The main quantity we must compute is an expectation of polynomials taken over the Haar measure of the Orthogonal group (n)) To carry out the computation, we make use of Weingarten calcu- lus (Banica et al., 2011; Collins & Sniady, 2006; Weingarten, 1978). More specifically, we make use of the of the Weingarten formula, studied by (Collins & Šniady, 2006; Collins et al., 2022):
Jonn Tini Tinja Tinair dp)O(n))= = Σ Σ (1) (o -7),, (36)
MEMid OEM2a
16

---

Published as a conference paper at ICLR 2024
where (n(( is the Haar measure of Orthogonal group. For in depth explanation of each quantity in the Weingarten formula, we refer the reader Collins et al. (2022, Section 5.2).
The quantity focus is Ew [Wa W.. W Wy sg]. We will do the computation multiple cases, based on the equalities of k, q. Notice that i * j in all cases. suffices we focus on the two (A + q) and Eg Ew [ w jk (A: q). Jq.
distinct cases: E1 Ew [ W
We first compute Ej.
Following the procedure from Collins et al. (2022, Section 5.2), we take the index sequences to be (i,) and j (k. k. g. g). Similarly, we get Ao( A-(J) = 1 only if = (3,44) [3,4]) and {(1,2): [3, 4)).
Consdering as permutations, we get:
2
3
(37)
1
2
3
4
2
3
4 4
7
1
2
3
(38)
2
3
4
(39)
2
3
4
where o -T has coset-type [1, 1]. Finally, we plug the results back into the formula and we obtain:
d + E = Ew [W3 w3. ([1, 1)) d(d + (d 1)
where the last equality is based on the results in Section 7 of (Collins & Matsumoto, 2009).
We compute Eg.
Similar to the previous expression, we take the index sequences to be to be = (t, j.j) and
j (k,, k k k). Thus, obtain Ag (i) (j( 1 only if f(1,2), (3, (3,4)) 4)} and π {{1, [3. }. T2 {f1, 3), 4]..
= {[1,4), (2, 3)}. Similarly, we compute o -ー for i E 3.3). Notice that is the identity permutation, thus yielding - Tu with the coset-
types [], [2] respectively, for each E [1,2,3}.
Plugging back into the original equation, we obtain:
E2 Ew [W.. Wg (1, 2Wg ((2))
(40)
d+1 d(d () 11) d d-1
+2 -1 d(d + 2) 1)
(41)
(42)
d(d 2(d=2)(d-1)
Thus, plugging back into the original inequality, we obtain:
d+1 m0e - 1)(4g - 1) d(d + 2)(d 1))
d-1
E X|
11+ 4
+ -(he - -1)(x 1)
d(d 2) 1)
1gg
kmg
(43)
1
(Ae 1)(2, 1) E(E (Ax 1)(Ax - 1)
(44)
4d(d + 2) - 1)
kzgg
k=g
Exe 1)²
= 2d(d + 2d - 11)
(45)
17

---

Published as a conference paper at ICLR 2024
where we have used that Sy + S= (x.((x - 1)(2, 1) (Ei=i(A 1))² equality.
0 in the last
Thus, we obtain:
27
E
Σ
11X;,1
E Σ 12 + 2 Σ |
(46)
dd+2 2 ix
(47)
sd+ (d. - d) (1-
Ex (dx 1)2
(48)
2d(d+ + 2) - -1)
d Σ Ex(A 1)²
(49)
2(d +
Corollary B.5 (Isometry gap bound). Suppose the same setup as in Theorem 3. Then, we have:
Ew Ew(o(X') xX)) + log (1 (1- Ee(hs -1)2)
(50)
2d (
Remark 4. Notice that the term Σk- Ch O(2), yielding log 2dP(d+2)
01e 11)
I 0.
(d+(2)
Proof. From Lemma 2, we know that:
dc
log Iogg((B(X))) s -logI((X) log
(51)
EΣj=1 |13..1)
2
I(gN(XX)) < - log I(X) log de + +oo Σ ||
(52)
d Ewl- log-log2ne I(BN(X')))A < log log de + Ew log Σ
|1x3.11
X
(53)
d Σ
log - log d? + log Ew
|x
X
(54)
j=1
S - logI(X) log d + + d2 Σt(Ae - -1)2
(55)
2(d (dd+2 2)
log log (1- 11-E(Am Ex 1)2
(56)
2d 2)
where in inequality 54 we have used the fact that Ellog X < log E[X] and in inequality 55 we have used the bound obtained in proof Theorem 3, equation 49. Thus, we obtain:
En(Av - 1)2
[x(((') S ((X) + loog (1
(57)
2d 2d +
18

---

Published as a conference paper at ICLR 2024
C ISOMETRY GAP DECAY RATE
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 C.1 (Isometry gap and orthogonality). Ifo(X) s 16a then eigenvalues x X are within
[1 c, cl.
in the following Note that, in order to simplify the calculations, we use the fact that a[d42] de proofs.
1
Based on the conditional expectation in Corollary B.5, we have:
E [[(x S e(x)) xx
(58)
2d²
Now, we prove a lemma that is conditioned on the previous layer isometry gap being smaller or larger than
Lemma C.2 (Isometry gap conditional bound), For X being the representations of an MLP under our selting, we have:
E
( Xe, xx,6 s
1
S (Xx)
(59)
16d
2d²
(x
Xe, xxx)(x)>
>
1
S ((x))
1
(60)
16d
32d3
Proof of Lemma C.2. Let A 1 ck. and assume without loss of generality that Σ Ek=1 6k Then, using the numerical inequality log(1 x)) 2 x², when
0.
| S
we have:
(x)
(61)
d
d Σ Am
d
max leal S 2
6(X) = d
log(1 + 1
Σ
(cx + = ox(X).
(62)
d
km
Altogether, have
max |el $
o(X) S xx(X).
(63)
Now, we can restate the condition terms of an inequality on the isometry gap. Thus, we can write:
3e
3e
dø(X) XX)
Σ log A
Σ log(1 + Es) 1=1
Σ 6;
(64)
6
4c,
+ 46;
i=1
where we used the fact that E; As d implying E, cs = 0 and also used the inequality log(1+z) <
t+t when x 2 -1 for E;'s. Note that because |e,l < 3. we get 6 + 4e, 2 4, all terms on the x 6-4z right-hand side are positive, implying that each term is bounded by the upper bound:
3e?
I do(X)
Vi (1,2....,d}
(65)
de;
363 1 Since 6 + 4c; 2 2, we can multiply both sides 6-
By construction, we have 6, 2 - and
<
16
by 4c; and conclude 3c? 6+4c.
O.
We can now solve the quadratic equation and obtain
1-V74 < 6 1+y74 which numerically becomes -0.35 < S 0.4, implying [csl < 0.5.
By solving the quadratic equation above we can guarantee that
((X) <
max leal S 5
(66)
16d
19

---

Published as a conference paper at ICLR 2024
Furthermore, we can restate the condition maximum using:
max
[ea] max c S Σ f (d)
(67)
k
and conclude that
xx(X) S
max lel S = o(X) = ox(X).
(68)
4d
Using this statement, we have
xx S
((X) s xx(x = x(x) <
(69)
16d
16d
If we negate and flip the two sides we arrive at
1
((X)
xx(X)= >
(70)
16d
16d
Thus, we can simplify the recurrence
(x(*)
Bio x
S
(71)
2d
as follows
1
(x
x2 x(x))
5
I ((x*)
1
(72)
16d
2d
E
((x+t))
x2,
((x))
1
S $(x))
1
(73)
16d
32d3
where we used equation 66 the first one and equation 70 in the second one.
Proof of Theorem 1. From Lemma B.I, we know that s(x)) 2 o(X²) 2 ... 6(84) 2 0, for
any layer 0 5 l < L. Thus, we get using equation 60:
E (oxxtt))x'' ,(x))
1
1
S
6(x))
(74)
16d
32d
6(X)
(75)
32d S (XT)
o(X*),
(76)
320
where in the last step have used the fact that ((x) < o(x)).
20

---

Published as a conference paper at ICLR 2024
Thus, combine equation 59 and equation 60 and obtain:
Blo(x(t)
= E o(x e+
X x(x)
1
1
6009 Nे.
(77)
16d
+ E ((xx+))
x x(xx) > 16d
1
)((x)>>
(78)
Ina
$ max
B )
x x(x*)
1
E
(o(xett)
x²
x(x))
1
16d
16d
(79)
1
< max
1
6(x²)
(80)
2dd²
32d (xx)
min
1
6(x9)
(81)
2d2
32d 6(X))
6(xx)
(82)
max 2d2 32d 6(x")
S exp
6(x))
(83)
max 32d di X))
exp o(X*)
(84)
By iterated expectations over X we get:
Blo(x*+1) S exp (-1) [[((X*) S exp (- o(xx).
(85)
Note that since max "02,,22d(('))) S 2dL (1+ + 1)) we can conclude the proof.
21

---

Published as a conference paper at ICLR 2024
D GRADIENT NORM BOUND
In the following section, we denote by H = w the pre-normalization values. Moreover, define as FL Rddd - RRxd where C is the number of output classes, the functional composition an layers MLP, following the update rule defined equation 3, i.e.:
Fr.(xt) BN(W"Jr.
(86)
Let us restate the theorem for completeness of the appendix: Theorem D.1 (Restated Thm. 5). For any -iisschtt loss function L and non-degenerate input Xo, we have:
dc
S d (o(Xo)* +1- e
log
(87)
ow
holds for all & S L, where possibly L + oc.
In partieular, the following lemma guarantees that the Lipschitz conditions are met for practical loss functions:
Lemma D.2. a classification setting, cross entropy and mean squared error losses are ((1)-
Lipschitz.
The main idea for the feasibility of this theorem the presence of perfectly isometric weight ma- trices that 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- 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 establish that the log-norm sum is bounded.
Proof Thm. D.I. Now, considering an layer deep model, where can possibly be + oo, we can finalize the proof of Theorem 5. Consider an MLP model as defined in equation 3. Let H_ W X be the logits of the model, where H E RCxd WL E RCxd W an orthogonal matrix and C' is the number of output classes. Denote C(H+,)) .y) the loss of model for an input matrix, with ground truth y. Then, applying the chain rule, we have:
OC OC OHL 0XL ox ex++1 OH
(88)
aw
OH ax L OXL-1 ax++1 OH aw*
By taking the logarithm of the norm of each factor and applying Lemma D.12, we get:
dC aw?
dC
HHL
k+1+1
0x ++1
01
log
<
log
+ log
+ Σ
log
+ log
+ log
TH0
T.
0x*
OH
OW?
k=t+1
||
1| [H
(89)
0C OHL
0 ow
< log
log (BN((H+)) (H+*)|
log
(90)
DHL
where log
llog W
0, since the orthogonal matrix WL has operator norm 1.
710
|
aH' Since बमा the main part is to bound the Jacobian log norms, which is provided by the following lemma: Lemma D.3. We have
x and x = 1N is batch normalized, this means that
aH
I d. Thus,
SWT
log (Bx(HH*) S do((()) +1 - 3264
(91)
22

---

Published as a conference paper at ICLR 2024
Finally, can plug the bound from Lemma D.3 in equation 90 and obtain the conclusion:
0L
ac
+ log + d (D(X)) + -e $2d*
log
log
(92)
ow?
Note that, for L - 0o, get:
0C ow
0C
log
S log
+ log d - d +
(93)
0H
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 D.2.
Proof of Lemma D.3. The proof of this lemma chiefly relying on the following bound on the Jacobian of batch normalization layers, which we will state and prove beforehand.
Lemma D.4 (Log-norm bound). If X E Indxd is the input to BN layer, its Jacobian operator norm is bounded by
log [JBx(X))lop I dØ(X) +1.
(94)
Furthermore, if (X)) S I. then we have
T6a
log [JJxxxX)lloo S <2/do(X).
(95)
Based on the lemma above, we shall define as the hitting time, corresponding to the first layer in our case, that the isometry gap drops below the critical value of 1
16d
{ e: ((x)))
1
s ==min
(96)
16d
So, we first bound the total log-grad norm for layers 1 up to and subsequently $ + up to L:
log || (X)
log JAx JJxx(X (Xo)
(97)
$
(do(x)) + +1
Σ dd((x))
(98)
S=1+1
(do(X))
+1) +2
Σ t=$+1
V do(x))
(99)
S(do(x)) +1) +2
V do(x)),
(100)
S=
where we have used that o(X) as an upper bound on ((X*).
Thus, taking expectation we get:
E log |1 (X) S (do(x)) + 1)B[S] + Σ Ed/o(xx) E=S+1
(101)
Note that is a random variable, which why the expectation over the number of layers appears at the last line. Thus, we can bound the log-norm by bounding [E[S] and the summation separately.
Lemma D.5 (stopping time bound). We have E[S] S ( if ((X)) > Tea and E[S] =
if ((x( < Tod
Lemma D.6 (second phase bound), We have
Σ
Eddo(x)) < 32d 320)55((x)00 (1- 1--esaee 3264
S=1+1
23

---

Published as a conference paper at ICLR 2024
Thus, have the following cases, based whether ((X)) is below over the 16d threshold. If we plug the bounds in equation 101 we get the following.
If ((X)) I
16a then:
E log [J,, I 32d x(x))
0.5 (1 e 3244
(102)
S ((X)) 0.5 (1 1 326
(103)
and if ((x)
16d then:
E log [1) 3 (1122 do(X))) + 32d (1 (1-e) 3207
(104)
S dd + d (1-6 3zid4
(105)
dd (do(xo)² +1 - e $2ú0
(106)
In fact, the maximum of the two bounds is
E log || (X)llp S 5 (o(X0) +1 - e." 322
(107)
Proof of Lemma D.6. By the bound in Lemma C.2, have
E 6(x ((xxt))((xx)
1
s o(x²)
(1- 2d
(108)
16d
Since assumed 2 S, the conditional inequality always holds and thus have the Markov bound
E $2S- - q = Pr Prf fo(x4t1)
211 o(x))
1
1
(109)
4d22
4d
4d
define as failure the event = (6(x++1) 2 (1-++)) )x with probability q. and conversely as success the event A with probability decrease by least factor of 1 पह is bounded by the failure probability T..
In other words, the probability that o(x*+) does not
Since o(x+1)) < ((x), then under the assumption that 2 S we can upper bound v Vo(X+T)
with ((X)) in case of failure with probability and with v 1d8 )x(X)) in case of success
24

---

Published as a conference paper at ICLR 2024
with probability 1 q.
V Vo(x+1)16(2))
E
(110)
-E (
(x 1) (1 (1-g) 9)
A 9
(111)
(1
Vox + / (Xx))
(1 - ②)
(112)
4d2
=oox
a+
1
9)
(113)
4d22
-Vo(x)
1
1
4d2
-
4d2
q
monotonic in q
(114)
s SVO(X))
1
1
1-
1
plug q S 1
4d
4d
4d
4d
(115)
=o((x)
1
1
1
rearranging terms (116)
4d
4d
4d2
Vo(x*)
1
+
1
vI x 1
T
for 2
4d
8dl2
4d²
2
(117)
=o((X)
(118)
32d
Thus, for E 2 S, we have
= Vo(x+1110((3) s By (xx( (1-
(119)
32d1
The summation starts from below / dol Xs); and will decay by rate 1 bounded by the geometric series:
which is upper
32d
& S 3d)
L+1
vd((Xs)
Σ
1 32d
1
1 -
(120)
32d
k=0
< 32.
3200455(X0)
(1
3264
(121)
Proof of Lemma D.5. By Lemma C.2 we have
Pr Pr n 2
1
exp (- ('max(2dd) 122 32dd((X)) 16d)(x))
(122)
16d
Thus, have
i ('max(2d' .2dd ((x))( 16do(x)).
Pr(S 2 0 S exp
(123)
25

---

Published as a conference paper at ICLR 2024
Since S is non-negative integer valued random variable, thus bound E[S] as:
E|S|
PIS 2
(124)
t=1
< 16dgp x
Σ
exp
(125)
k
8=1
1
(126)
exp( xt) -1
16do(X")
(127)
16d0
32x
(128)
S 512d"
(129)
PROOF LEMMA D.4: BOUNDING BN GRAD-NORM WITH ISOMETRY GAP
The proof of the Lemma relies two main observations that are crystallized in the following lem- mas that first establish bound on Jacobian operator norm based on the inverse of smallest eigen- value, and then establish lower bound for the smallest eigenvalue using the isometry gap.
Lemma D.7. Let X Raxd and let {}:]a= be the eigenvalues of X XXT. Then, we have that:
1 [x) S
(130)
Xd
where JBN is the Jacobian of the BN(() operator.
Using the above lemma we have log [N (X) S log Xa- The following lemma upper bounds this quantity using isometry gap:
Lemma D.8. The minimum eigenvalue of Gram matrix that is Irace-normalized is lower-bounded
by the isometry gap as log Xa I dé(X) + 1. Furthermore, if 6(X) I 1 then log Ad S
2. dx().
Plugging these two values have the bounds
log [l]JJB(X))looo S do(() +1,
(131)
(XX) < 16d
log (X S /d(x
(132)
Now we can turn our attention to the proof of the Lemmas used in the proof. The proof of re- lationship between minimum eigenvalue and isometry gap is obtained by merely a few numerical inequalities:
Proof of Lemma D.8. Let {A:]d= be the eigenvalues of XTX. Since the matrix is trace- normalized, we have Dd d = d.
The arithmetic mean of the top d values can be written as
(133)
d-1
A=1
Thus, have that their geometric mean is bounded by the same value. Therefore, have the following bound:
\ s As 1+ (1+1-1 1-
Ad 1-11 -
II k=1
(134)
clogg(X) < log(>d) (d 1) log 1+ 1- Xd
(135)
26

---

Published as a conference paper at ICLR 2024
where in the second inequality have taken logarithm of both sides. Now, apply the
numerical inequality log(1 + z) I x to conclude:
ddo(X) & log Xa + Ad
(136)
Since re is non-negan non-negative, this clearly implies the first inequality: - log Ad S do(X) + 1.
For second inequality first we use the numerical inequality +1 S (x-1)² Vo E [0
to conclude
(1 - Xa)2 dd(()
(137)
Aa 2 1 2dd((x)
(138)
= log Aa S -log(1 - v2dd(())
(139)
We can now use the inequality - log(1 x) 1 22x for 0 S x to conclude that
log Aa - 2 (x
(140)
when 2 < 2/4 which is equivalent to ((X) v 1
l6d
For proving Lemma D.4, we first analyze Bn operator on a row, and then invoke this bound and the special structure JBN to derive the main proof.
Lemma D.9. Let f Rd Id defined as f(x) =
be the elementwise normalization of the x. |all
x
Then:
Jz(x)
|0
10
||0|
6 x,
(141)
where & the outer product.
Proof. To begin, notice that for x E Rd we have Then the Jacobian entries become:
B|
m x
Denote by y: = [(xx)] =
m
Ox
[
Oy:
|
|=
(142)
0r;
|
|
|
Oy:
(143)
0z,
|
|
Assembling the equations into matrix form, we obtain:
3z(x) Le
Ex.
(144)
Corollary D.10. The (x) has the eigenvalue
with multiplicity de and with multiplicity T
Lemma D.11. Let XXT Ei=, Adnuu, where X UAUT is the eigendecomposition of XXT. Then, we have that miny X ming Au-
Proof.
1x (xx D33
Σ Auuz
min Av.
Σ Wife
min Ae:
(145)
2m
where in the last equality we have used the fact that orthogonal. Since this true for all rows j. get min, |X,-1|2 ming At.
27

---

Published as a conference paper at ICLR 2024
Having established the above properties, we equipped prove that the Jacobian of a BN layer bounded by the inverse of the minimum eigenvalue its input Gram matrix.
Proof of Lemma D.7. To begin, notice that since BN : I df + dd implies that (Bx(X) E xx Denote X' Bx(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.:
0x O, vi j. Vk,
(146)
This creates block-diagonal structure in JBN(X), with d blocks on the main diagonal, where each block has size x d and equal to J (X,,), where f as defined Lemma D.9. Therefore, due to the block-diagonal structure, we know that:
X JBN (X)
U (
(147)
(=1
where A[-] denotes the eigenvalue spectrum. From Corollary D.10, we know that
A [X))
U
(148)
||x
<=1
with their respective multiplicites. Finally, using Lemma D.I.I 11, this implies:
|| = max X
(149)
min; X
ming AR
Proof of Lemma D.2. Assuming the the logits are passed through softmax layer, we analyze the case of Cross Entropy Loss for one sample i in C-classes classification problem. Denoting
H.L. we have:
C
C(zi, yi) =
Σ
yi logPs
(150)
1=1
where p;
is the probability vector for sample i after passing through the softmax func- FC E5
tion.
Computing the partial derivatives, we obtain:
Öp: Pi(1 - pk), i == k
(151)
0z:
0p:
Ozk PiPk: i k.
(152)
Finally, we can compute the gradient of the loss with respect the logits:
a log(p:)
C
V/
(153)
Ozy
1 Op:
1/s pt 0zz
(154)
Σ (BEPk) (- Pk))
(155)
Pb
Σ V: - Vv - -Pk)
(156)
IV L1
Pk
Uk
+ +1 -Px)ll 2.
(157)
28

---

Published as a conference paper at ICLR 2024
Since the gradient of the loss with regard each sample is bounded, conclude that the operator norm of the Jacobian of the loss with regards to the logits matrix H also bounded.
In similar analysis, now shift our attention towards the Mean Squared Error (MSE) loss:
1 C
((z,, m/i
Σ y1 pi)²
We want to compute the gradient of the loss with respect each logit zy:
-(pp:)
V
Σ
2(y
pi)
(158)
C
8zk
6 -
2
Σ
(yi - pi)
Op:
- C
2
Σ
yi
psl
0p: Zz
S 2.
(159)
2/
Ozk
By substituting these derivatives into the gradient equation, we can derive the gradient for each logit with respect to the MSE loss.
Lemma D.12. Let X Rexd be the hidden representations layer >> O as defined in equation 3. Then, we have that:
0x
log
S log JBN
(160)
0x
Proof of Lemma D.12. By definition, we have that:
H' = x
(161)
x = BN(H").
(162)
Therefore, applying the chain rule, we get:
0x*+1 ax OH
TH ox = JBx
(163)
OX
Taking the logarithm of the norm of this quantity, we reach the conclusion:
0x
log
S log JJN(HH)
||
log
| log [Bx ||
(164)
0x
where have used the fact that the spectrum of the orthogonal matrix e contains only the singular value 1 with multiplicity d.
29

---

Published as a conference paper at ICLR 2024
E 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 CIFARIO, MNIST, FashionMNIST and CIFARI00. We test this assumption by randomly sampling 100 input batches of sizes 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.ram approach the dimensionality of some datasets, i.e. FashionMNIST, MNIST. We show in Table E1 the average rank with the standard deviation for each n. over 100 randomly sampled batches.
function provided in Py PyTorc. We stop at size 512 since we
Table E1: Average rank of Gram matrix 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.
Datasct
16
32
64
128
256
91 512
CIFARIO
16.0 + 0.0
32.0 = 0.0
64.0 - 0.0
127.99 + 0.09
221.06 4: 2.89
203.70 = 3.58
MNIST
0.0
32.0
64.0
128.00
250.48
1.53
2.82
MNShI
16.0
0.0
32.0 0.0
64.0 0.0
128.00
0.00
238.19
1. 3.27
275.37
4.20
CIFARIOU
16.0 + 0.0
32.0
0.0
61.0 0.0
127.92
+ 0.27
218.11 + 3.69
201.51
3.95
We would like to remark that these datasets are fairly simple in terms of dimensionality and seman- tics, 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 be- tween samples in 512 size batch is 0.81, 0.40, 0.58, 0.81 for CIFARIO, MNIST, FashionMNIST and CIFARIO6 100 respectively.
30

---

Published as a conference paper at ICLR 2024
F 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 ag be the pre-activation gain.
More formally, since the gradient norm has exponential growth in depth, shown in Figure H5, 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 layer as the slope of the log norm of the gradients:
R(E, oz) log V
CII
log V VWe-1of)
(165)
10
Since the rate function is not perfectly linear and has noisy peaks, measure the slope with 10 layer gap in order to capture the true behaviour instead of the noise.
Our goal is choose such that the sum of the rates across the layers in depth is bounded by constant that does not depend on the depth the model, i.e. R(E, ar) S B, where / is independent of L. One choice to achieve this is to pick gain such that the sum of the rates behaves like decaying harmonic sum in depth.
To this end, measure the rate of explosion multiple layers in 1000 layer deep model, for various gains a which are constant across the layers in Figure F1 and notice that it behaves as R(C) = cja.. order have the sum rates across layers behave like bounded harmonic series in depth, we must choose the gain such that it decays roughly as a E- where & > 1 results in convergence. Therefore, can obtain heuristic for picking gain such that the gradients remain bounded in depth as or k/ where refer c/c2 the gain exponent.
sin
tanh
0.10
100
Rate
0.05
Rabe
0.05
200
0
0.00
0.00
0.25
0.50
0.75
0.25
0.50
0.75
Figure Fl: 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 & as function of gain a. The black trace shows the fitted function c105- Traccs are averaged over independent runs, with the shades showing the 95% confidence interval.
This reduces the problem picking the exponent such that the sum stays bounded. We show how the behaviour the explosion rate at the carly layers, for various models, is impacted by the exponent in Figure F2. Note that for several exponent values, able reduce the exponential explosion rate and obtain trainable models, which we show in Appendix H.
31

---

Published as a conference paper at ICLR 2024
sin
Depth
100
Rate
Rtte
200
0.02
0.02
500
1000
-0.8
0.6
0.2
0.8
0.4
Exponent
Exponent
Figure F2: 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 function of the gain exponent. Traces are averaged over 10 independent runs, where the shade shows the 95% confidence interval. Rate measured at 10 avoid the any transient effects of the function
G IMPLICIT ORTHOGONALITY DURING TRAINING
In this section, we provide empirical evidence that our architecture during training maintains orthog- onality across depths, while maintaining bounded gradients. Figure G1 shows the evolution of the isometry gap of the weight matrices We 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 the loss gradients with regards to matrices W; in Figure G2.
These experiments are performed on MLP with orthogonal weight matrices and batch normaliza- tion, with sin and tanh activations. The width set to 100, batch size 100 and learning rate 0.001. The gain exponent is set to fixed value for all experiments. The measurements are performed single batch of size 100 from CIFAR 10, after each epoch of training on the same dataset.
32

---

Published as a conference paper at ICLR 2024
sin
tanh
10-1
10-1
Epoch
((W))
50
10-2
10-2
100
10-3
10-3
300
50
50
8
8
sin
tanh
10-1
10
Epoch
8(W2)
50
100
10-3
10-3
200
300
0
200
400
0
200
400
8
tanh
Epoch
10-2
10
50
((W)
200
10-4
10
300
500
1000
500
1000
8
Figure G1: 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 small isometry gap. During training, the isometry gap also remains low, suggesting the matrices remain close to being orthogonal.
33

---

Published as a conference paper at ICLR 2024
sin
tanh
10a
10
Epoch
V C1l
50
100
100
109
300
10
50
sin
tanh
10a
Epoch
kG All
50
10a
109
100
200
300
10-1
10
0
200
400
0
200
400
tanh
10
10a
Epoch
[Vw,C1 C|
50
10"
10a
200
300
10-1
10-1
500
1000
500
1000
Figure G2: 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 Figure GI not due to the weights not being updated during SGD.
34

---

Published as a conference paper at ICLR 2024
H OTHER EXPERIMENTS
In this section we provide the train and test accuracies of deep MLPs on 4 popular image datasets, namely MNIIS, FaShionMNIST, CIFARIO, CIFARI00. Hyperparameters and measurements pro- cedure are described in Section 4.
SUPPLEMENTARY TRAIN AND TEST RESULTS ON MNIST, FASHIONMNIST, CIFARIO. CIFARI00
sin
tanh
identity
1.0
1.00
Train accuracy
0.75
0.75
Depth
100
0.5
0.50
0.50
500
0.25
0.25
10a
10
10a
10
109
Epoch
Epoch
Epoch
sin
tanh
identity
1.00
1.00
Test accuracy
0.75
0.75
Depth
100
0.50
0.50
0.50
500
1000
0.25
0.25
0.25
10t
10a
10a
10a
10a
10a
Epoch
Epoch
Epoch
Figure H1: Contrasting the train and test accuracy of MLPs with gained sin, tanh and identity acti- vations 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.
SUPPLEMENTAL FIGURES
We present empirical results in Figure 2 showing that degenerate input batches are hard constraint for orthogonalization without gradient explosion. For MLPs with different depths, we show that by repeating samples batch of size 10 we get an exponential gradient explosion, which is unavoid- able theoretically.
Furthermore, we show how non non-linearities affect the gradient explosion rate Figure H5. Using standard batch normalization and fully connected layers from PyTorch we show that non- lineaiinessssess maintain 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 prove that the gradients remain bounded depth.
INFLUENCE OF MEAN REDUCTION ON THE GRADIENT BOUND
In this section, compare whether adding reduction and the additional factor of in the denominator of the batch normalization module influences our gradient bound. As expected, we show in Figure H6 that in both cases, for the identity activation, the result remains similar, with the gradients remaining bounded in depth.
35

---

Published as a conference paper at ICLR 2024
sin
tanh
identity
1.00
1.00
Train accuracy
0.75
0.75
Depth
0.50
0.50
0.50
0.25
1000
0.25
0.25
0.00
10
10
10'
10
10'
10a
Epoch
Epoch
Epoch
5in
tanh
identity
Test accuracy
0.75
0.75
0.75
Depth
0.50
100
0.50
0.50
500
0.25
1000
0.25
0.25
0.00
10a
10a
10a
10a
10a
10a
Epoch
Epoch
Epoch
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 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 H2:
sin
tanh
identity
Train accuracy
0.6
Depth
0.4
0.4
100
0.2
0.2
0.2
1000
0.0
10
10a
10
10a
10
10a
Epoch
Epoch
Epoch
tanh
identity
0.2
0.2
Teat accuracy
Depth
0.10
100
0.1
0.1
0.05
1000
0.0
0.00
10a
10>
10
10a
10
10a
Epoch
Epoch
Epoch
Figure H3: Contrasting the train and test accuracy of MLPs with gained sin, tanh and identity activations on CIFARI00. 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.
36

---

Published as a conference paper at ICLR 2024
sin
tanh
identity
Test accuracy
0.4
0.4
Depth
0.3
0.3
0.3
0.2
0.2
1000
0.1
0.1
10'
10a
10'
10a
10'
10
Epoch
Epoch
Epoch
Figure H4: Contrasting the train test accuracy of MLPs with gained sin, tanh and identity activations on CIFARIO. 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.
40
10
Activation
10
identity
sin
10
tanh
0
0
500
1000
0
500
1000
Depth
Depth
Figure H5: Left: average log-norm of gradients (log-scale y-axis) at the first layer for networks with different depths, evaluted on 10. Right: Isometry gap (log-scale y-axis) at the last layer for networks with different depths, evaluted on CIFAR1O. The MIPP is initialized with orthogonal weights and batch normalization, with standard modules, with sin, tanh, identity non-linearities. After stabilizing the isometry gap, the non- linearitsss have an exponential gradient explosion.
log [ Vw,, C
Mean reduction
With
Without
0
250
500
750
1000
Depth
Figure H6: 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 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.
I HIGH PROBABILITY LOG-GRADIENT BOUND
Theorem L.I. In the same conditions as Theorem D.I, it there is constant C such that:
Pr (1og ow. 2 Cloo(115)) ( + +1) s 25
for any 8 E (0,
For example for 5 1/2, have VVV Cllop S (6( +1) with at most 1/2 probability.
37

---

Published as a conference paper at ICLR 2024
Proof. We follow the same proof Theorem D.1 by decomposing the log gradient to sum of log gradients of each layer. Note that Lemma D.4 holds deterministically, as it is valid by construction. Thus, what remains is to derive probabilistic version of Lemma D.3. We start from equation 100 which holds deterministically, to arrive at a probabilistic bound:
Σ log JB s ((do(X)) +2
Vdo(x9)
8: min (e: -
1 16d2
E=
(166)
We derive a probabilistic bound for each term.
First term.
Based on Lemma D.5 we have E[S] - d (Xx)). Thus, there is Co such that
Co d' (x))²
ECu : E|S] S 2
Thus, by Markov inequality we have
B
P(S
B Cod ((x)2
First, we discretize the layers into blocks of size B, where B is defined above. With slight abuse of notation, define B; as the end of the ith block of size B, and let E, be the event E, = (o(XB:) > Ina J. which is the event of failure for & to drop below the threshold of 16 after the last layer from block i.
By the inequality established above, we know that in each block, & can either fall below the thresh- old, or stay above, with probability at most 1/2. Moreover, knowing that son-increasing, we know that E(E))E<)) S 1/2.
Thus, by the non non- -on-increasing property of Φ, the probability of failure after blocks of size B is at most the probability that á did not fall below the threshold in any of the & successive blocks:
P(Es) S P(EIA ... EG-1) En) ... Fx Fa)
(167)
(F
(168)
1
(169)
2
Thus, obtain the probability:
P(S BB) P(S kCod" 2
(170)
Thus, connecting back to gradients, we obtain:
Pr DIIBA(HH'll 2 oC 5
(171)
E=0
Second term Starting from equation 109, we have
l 2 = Pr Pr (((((x(t) 2 (1 (1- (01-1/44)d(X")) < 1
4d
Let E, denote the event that o(X++)) does not decrease by - 1/4d compared to its previous
layer: E, 1 ((x(==1)) 2 (1- (1- Due to the non-increasing property of o, we know that Pr(Ex) = Pr(Ed E. Repeating this for steps, we obtain:
8+s-1
82S - Prfo(xtton 2 (1-1/40d)6((X))
II ==
Pr(E:+1 <
1
(172)
4d
38

---

Published as a conference paper at ICLR 2024
Inspired by this, consider the sequence to. l,... defined
E/ = S lx = x--1 + 4ckd
Thus, get:
Pr ( 2 (1
1
)0(( 5(1
1
jackd
Se
1d2
Thus, have the union bound
Pr { 0C
2
-
Φ ((x(+t))) 2 21 (1-
c
ck
4d
c
1 1 -
Pr ((t)) (1
1
) 2
1 - 2e e
4dp
c
1 c
Q:
Note that in the event that Q holds, we can the isometry gaps in each [&x-1, 6x] interval as:
1 jdا ) for all 8 (4x-1,
Q o(X,) S 16d 4d2
We can upper bound 28-s+1 v d by using the numbers of items and upper bound on each
block. Thus, assuming for all & have ((x(+1)) I (1 40z (x(())), ), we can derive
1
Q
Σ Vdo(Xi) s
E(ex
(bx 1
(1
4d2 )
16d
S=++1
k=1
cd22
k<1
jk/2
4d²
< cd2
Σ
k:1
using (1 x1) < x/2
8.12
k=1
8/d2
= cd2
using Σ ka = a/(1 a)²
(1/8d2y2
I 64cd
Thus we have:
Pr
/do(Xe) s 64ed
2 Pr (2) 2 1 - c
2
S+1+1
which yields
Pr
log ||N (H) 2 64cd
< Pr(Q)) = e
--
S=1+1
Combining first and second term bounds for the Jacobian log-norms we have
Pr
log MBN (1I))))lllpp > &Cod ((x))) 5-
(173)
-
P
log [JBx 2 64ed
5
(174)
1 1- c
And thus we have
Pr Σ log JBN JJB((HH)))))) llop 2 (CCdd((x0) + 64cd)) 52 +
1 c-
Thus, we can find C such that
Pr [oog| 2 kCd + +1)) --+++
We can finish the proof by defining 8 2 and change of variables A (1/)).
39

---

Published as a conference paper at ICLR 2024
GRADIENTS IN RESIDUAL NETWORKS J
log [[ F
Activation
identity
0
tanh
sin
500
1000
1500
2000
Depth
Figure J1: Logarithmic plot for the gradient norm of the first layer for residual networks with batch normalization initialized with Gaussian weights at different depths, evaluated CIFARIO. The traces show that the gradients do not remain bounded in depth for different activations.
40


Reviews:
[{"review_id": "QZqrbUSwE7", "reviewer": "Reviewer_692A", "summary": "The paper shows the existence of a batch-normalized network that avoids rank collapse and does not suffer from exploding gradients as depth increases. To construct a network that satisfies the above properties, the paper uses an initialization scheme where the weights are random orthogonal matrices.", "questions": "Do you obtain figures similar to F1 when using ReLU activations?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The isometry and log gradient norm bounds are non-asymptotic bounds. Furthermore, the usage of Weingarten calculus to obtain rates for isometry increase is interesting.", "weaknesses": "The paper uses a modification of batch normalization (in particular the mean centering operation is removed) when proving bounds on the isometry gap and log gradient norms (however, the experiments seem to suggest that this is not an issue).\n\nThe paper also considers the setting where the number of training examples is equal to the dimension of the problem..", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BFGc9svZnQ", "reviewer": "Reviewer_eDWh", "summary": "The authors address the problem of gradient explosion with depth in networks with Batch Normalization (BN) outlined by Yang et al. [1]. They show that infinitely deep linear networks can be trained with BN while avoiding the problem of gradient explosion as long the weights are orthogonal.\n\n[1] Yang, Greg, et al. A Mean Field Theory of Batch Normalization. International Conference on Learning Representations. 2018.", "questions": "See above section for questions.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors present a proof for avoiding gradient explosion in infinitely deep linear networks with BN as long as the inputs are full rank and the weights of the network are orthogonal.\n\n2. In linear MLPs with upto 1000 layers, they show that orthogonality is approximately maintained in the middle layers as required by the proposed theorem.\n\n3. The presented proof and argument is easy to follow.", "weaknesses": "1. The authors have not included literature on recent work on inducing dynamical isometry in feedforward networks which has shown to also avoid the need for BN. The discussion seems relevant for this paper (see [1], [2], [3]).\n\n2. The proof relies on the fact that the network is has linear activations and the weights are orthogonal during training. Although, this seems to somewhat hold empirically for linear networks, how reasonable is it to expect such a condition to hold for nonlinear activations. The weights will likely not stay orthogonal during training for nonlinear activations. Further, can the analysis be extended to nonlinear ReLU - like activations, by potentially using a first order approximation of the ReLU [4] or other approximations?\n\n3. Since the proof requires for the network weights to be orthogonal, the network is only restricted to learn a rotation of the inputs. This seems restrictive from an expressiveness point of view unless I misunderstand something.\n\n[1] Rebekka Burkholz and Alina Dubatovka. Initialization of relus for dynamical isometry. Advances in Neural Information Processing Systems, 2019.\n\n[2] Yaniv Blumenfeld, Dar Gilboa, and Daniel Soudry. Beyond signal propagation: is feature diver-\nsity necessary in deep neural network initialization? In International Conference on Machine\nLearning, 2020\n\n[3] Andrew Brock, Soham De, and Samuel L. Smith. Characterizing signal \npropagation to close the performance gap in unnormalized ResNets. International Conference on Learning Representations. 2020.\n\n[4] Rebekka Burkholz. Most activation functions can win the lottery without excessive depth. Advances in Neural Information Processing Systems 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oqQ1qMo0dQ", "reviewer": "Reviewer_GLi5", "summary": "The paper addresses limits of gradient descent loss minization of (very) deep networks using batch normalization (BN). In particular, it mends explosion of gradients during backdrop training in such networks. After a neat introduction a novel theoretical results are laid out in an accessible manner and main idea of proofs is presented (linked to full proofs in Appendix). Using recent results on random matrix theory, namely that BN does not decrease isometry gap (a measure of deviance of sample covariance $XX^T$ from orthogonality), paper proofs its main Theorem 1. It claims that expected isometry gap has an upper bound whose rate of decrease is exponential in depth under assumed linear independence of inputs. It is followed by Theorem 5, which proves that, under mild smoothness and loss assumptions, the $\\textit{expected}$ log norm of gradient of the constructed MLP network, i.e. deep network with linear activations and simplified BN, has finite upper bound $~d^5$.\n\nFollowing section supports the results by limited experiments, showing that rank collapse does not occur when MLP using (simplified) BN as proposed is used.\n\nPaper follows to discuss possible extensions to treat non-linear networks and proposes activation function shaping through additional pre-activation scale hyperparameter, that is tuned towards zero, effectively linearising explored models using $tanh(\\cdot)$ and $sin(\\cdot)$ activations.", "questions": "Q1: Related to Weaknesses, ad 3.) \"practicality\", Could authors present their view on what benefits are provided by deep linear networks compared to shallower linear nets? \n\nQ2: To address aforementioned limitations of theoretical results (see Weaknesses #1, 1-3),, e.g., to support non-linear activation, authors propose activation shaping combined with use of BN. From their argument it seems that such technique, tuning pre-activation gain $\\alpha$ towards zero, effectively linearise the network. Section 5 claims that this technique still maintains the benefits of non-linear activations demonstrated in Fig.5. Could authors elaborate on how exactly are benefits maintained? \n\nQ3: In later sections paper presents intriguing idea of “implicit bias towards orthogonal matrices”, for instance in Fig 6. But one can read experimental results in Fig 6. other way, namely that it suggests that while middle layers have been initialised already “almost orthogonal” their orthogonality gap increased the most during training compared to other layers! Could authors show further evolution of training continued beyond the early stopping or otherwise corroborate \"implicit bias\" claims more extensively?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. To the best of reviewers knowledge presented are novel results (Theorem 1, 5) that are non-asymptotic and hold for networks with finite width, as opposed to previous works. This is very promising and of large interest and value to ML community increasing potential impact of the manuscript. \n\n2. Very well and accessibly written exposition of rather technical methods (random matrix theory, Weingarten calculus) used in proofs.\n\n3. After theoretical results paper (creatively) suggests activation shaping technique in combination with BN increasing and demonstrating practical utility of results.", "weaknesses": "#1: While the non-asymptotic and finite width results may be largely impactful as noted in Strengths, the importance may be diminished by applicability on \n1. $\\textit{expected}$ grad norm\n2. grand norm bounds of Theorem 5 may be still prohibitively large $\\textit{e^{d^5}}$ and \n3. \"practicality\" may be hindered by \"$\\textit{linear}$ activation\" required (this is somewhat alleviated by last section by introducing more hyper-parameters to bring $sin$ and $tanh$ activations to linear regime).\n\nI believe authors should extend the Discussion or include section to elaborate more on these points. Possibly extend a rather limited experiments to corroborate their arguments in case of a lack of theory.\n\n#2: Technical: Dimensionality discrepancies throughout Section 3, page 3 and 4. For instance:\n - in the definition of $\\varPhi(X)$ on page 4, the $\\varPhi(X)$ is not function of $R^{d \\times d}$.\n - $X$ should be $d \\times n$ matrix instead of previously defined $n \\times d$ or $XX^T$ should be used instead of $X^TX$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xAn4Ah06bW", "reviewer": "Reviewer_uof1", "summary": "The paper is a theoretical study on batch normalization. A simplified setting is considered where BN does not use mean subtraction and where there are no other non-lineary functions. The paper presents two theorems that hold when the weight matrices are random orthogonal matrices. Theorem 1 states that the representations become increasingly orthogonal with the depth of the network. Theorem 5 states that the gradients do not explode with depth. All results hold in expectation over the random weights. The authors also provide a method for modifying non-linear activation functions by changing the gain.", "questions": "Why is the main question (“Is there any network with batch normalization without gradient explosion and rank collapse issues?”) interesting to practioners? Given that BN works well without orthogonal initialization.\n\nHow is the activation shaping motivated by theory?\n\nWhat is the practical utility of the activation shaping you propose?\n\nDoes your result hold for networks with mean reduction?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper presents two theorems which are relevant and non-asymptotic.\n\nNormalization is used everywhere in ML, so studying it has a high impact.", "weaknesses": "The main question of the paper is “Is there any network with batch normalization without gradient explosion and rank collapse issues?”. It is not clear why this question is interesting. The authors show that orthogonal random weights avoid these issues, but such initialization schemes are not used in practice despite being well known. Thus, in practice, BN works for reasons that are unrelated to the specific weight initialization scheme that is proposed. So the paper is not really “explaining” why BN works.\n\nThere are still some simplifying assumptions in the theoretical parts – there are no nonlinearities and the mean subtraction of BN is omitted. \n\nIt is not clear how relevant the empirical results are, though the authors might be able to clarify this for me.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)