Paper: Towards Training Without Depth Limits: Batch Normalization Without Gradient Explosion (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
> 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. 

Let us elaborate why this question is important and interesting: 

*Why is it important?* BN and depth can be beneficial in many learning tasks, as they improve signal propagation and expand the functional class and expressivity of the model. However, if one opts to use batch normalization, they cannot train very deep networks due to gradient explosion. Our study shows that by simply changing the weight distribution from Gaussian to the orthogonal ensemble, training very deep networks becomes feasible. Furthermore, we extend this result to non-linear activations by a specific activation shaping mechanism. 

*Why is it interesting?* [1] proves that gradients in MLPs with BN and Gaussian weights explode regardless of the type of activation (including linear). This seems to suggest that gradient explosion with BN is *"unavoidable".* Our study proves that that is not the case.  

> The authors show that orthogonal random weights avoid these issues, but such initialization schemes are not used in practice despite being well known. 

Initializing weight matrices with elements from the orthogonal ensembles has been used in practice, in some cases to great practical benefits as discussed in the "Initialization with orthogonal matrices" paragraph in our paper.

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

 Our paper does not intend to prove benefits of BN, but instead to avoid its issue of gradient explosion. According to [1], the standard initialization scheme of Gaussian weights can not avoid gradient explosion. Under specific hyperparameter choices (batch size of BN) and weight initialization, we prove the explosion is avoidable.

Reviewer comment (Theory):
There are still some simplifying assumptions in the theoretical parts – there are no nonlinearities and the mean subtraction of BN is omitted.

Author response:
> There are still some simplifying assumptions in the theoretical parts – there are no nonlinearities and the mean subtraction of BN is omitted.

While our theoretical results are limited to linear activations and a simplified BN operator, our theory inspires an activation shaping mechanism that allows us to avoid gradient explosion in networks with activations (Section 5). We theoretically study networks with linear activations due to their analytical simplicity and due to the fact that they also suffer from gradient explosion [1].

Reviewer comment (Evaluation):
It is not clear how relevant the empirical results are, though the authors might be able to clarify this for me.

Author response:
> It is not clear how relevant the empirical results are, though the authors might be able to clarify this for me.

Our empirical results elaborated in Sections 4 and 5 show that these techniques enable us to train up to $1000$ layer deep networks with batch normalization. In contrast, using Gaussian weights, training networks with more than $200$ layers is not possible due to gradient explosion.

Reviewer comment (Novelty):
Why is the main question (“Is there any network with batch normalization without gradient explosion and rank collapse issues?”) interesting to practitioners? Given that BN works well without orthogonal initialization.

Author response:
> 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. 

Let us elaborate why this question is important and interesting: 

*Why is it important?* BN and depth can be beneficial in many learning tasks, as they improve signal propagation and expand the functional class and expressivity of the model. However, if one opts to use batch normalization, they cannot train very deep networks due to gradient explosion. Our study shows that by simply changing the weight distribution from Gaussian to the orthogonal ensemble, training very deep networks becomes feasible. Furthermore, we extend this result to non-linear activations by a specific activation shaping mechanism. 

*Why is it interesting?* [1] proves that gradients in MLPs with BN and Gaussian weights explode regardless of the type of activation (including linear). This seems to suggest that gradient explosion with BN is *"unavoidable".* Our study proves that that is not the case.

Reviewer comment (Theory):
How is the activation shaping motivated by theory?

Author response:
> How is the activation shaping motivated by theory?

The high-level intuition that we draw from our theory is that we can ensure that that the log operator norm of the Jacobian of the activation layers, is upper-bounded by a value that forms a Harmonic series wrt the layer depth. The main connection between our theory and our activation shaping is that if we assume that Thm. 1 holds for an MLP with shaped activation, we can deduce the rest of our results. For an elaborate explanation of the connection between our theory and the activation shaping, we refer the reviewer to appendix Section E.

Reviewer comment (Novelty):
What is the practical utility of the activation shaping you propose?

Author response:
> What is the practical utility of the activation shaping you propose?

Our activation shaping scheme enables training very deep networks with BN and non-linear activations. Without orthogonal weights or the shaping parameter $\alpha,$ training networks with BN and activation layers is not feasible for very deep (>200 layers) networks, due to gradient explosion. Since deeper networks may have desirable properties (see Supplementary Figure G), this scheme opens up new avenues for neural architecture design.

Reviewer comment (Theory):
While the non-asymptotic and finite width results may be largely impactful as noted in Strengths, the importance may be diminished by applicability on (1) $\textit{expected}$ grad norm, (2) grad norm bounds of Theorem 5 may be still prohibitively large $\textit{e^{d^5}}$, and (3) "practicality" may be hindered by "$\textit{linear}$ activation" required (this is somewhat alleviated by the last section by introducing more hyper-parameters to bring $sin$ and $tanh$ activations to linear regime). The authors should extend the Discussion or include a section to elaborate more on these points, and possibly extend a rather limited set of experiments to corroborate their arguments in the case of a lack of theory.

Author response:
We thank the reviewer for raising this important point.  Fortunately, based on our our non-asymptotic results with exponential behavior, we can use Markov inequality to convert our results in expectation to with high probability results. Assuming same conditions as Thm. 5 and non-degenerate inputs $\phi(X^0)= O(1),$ we can prove that log-gradients are bounded by $O(d^{6})$ with high probabiltiy. Formally, we prove that there is $C$ such that for any $\delta\in(0,1)$
$$
P\Big(\log\|\nabla_{W_\ell}\mathcal{L}\|_{op}\ge  C \ln(1/\delta)  d^{6}(\phi(X^0)^3+1)\Big) \le 2\delta 
$$
Note that this bound has an additional factor of $d^{1},$ due to the fact that we have to ensure that failure events are uniformly bounded across a possibly infinite number of layers.  

Please see Appendix H in the updated manuscript for the proof and a detailed discussion of the high probability bound.


> - grand norm bounds of Theorem 5 may be still prohibitively large $e^{d^5}$ and 

Our main focus was on depth and not network width, denoted by $d$. The main message our paper conveys is the fact that the upper bound in Thm. 5 holds for an *arbitrarily deep* network. Thus, the main take away message from Thm. 5 is that a BN net with random orthognal weights has bounded gradients. Remarkably, [1] establish that gradient norms grow at exponential rate with depth. Our results shows that this exponential depedency is replaced with constant if we use orthogonal random weights. Yet, we agree with the reviewer that finding the tighter bound for width could be interesting for follow-up studies. 


> - "practicality" may be hindered by "activation" required (this is somewhat alleviated by last section by introducing more hyper-parameters to bring *sin* and *tanh* activations to linear regime).

One of the main practical outcomes of our work is the activation shaping mechanism, showing that gradient explosion in a network with BN can be avoided using orthogonal weights rather than Gaussian. Note that both practically and theoretically, gradient explosion is "unavoidable" using Gaussian weights [1]. To our knowledge, the proposed approach is the only existing practical solution to avoid gradient explosion in a deep network with BN.

 
> - I 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.
>
We have added Section H in the Appendix. Furthermore, experiments presented in Sections 4, 5 and 6 support our theoretical results: going beyond the initialization (Section 4, 6), using non-linear activations (Sections 4, 5, 6), and using standard BN (Sections 4, 5, 6). Furthermore, we explore various depths and datasets for training, showing that the empirical observations are not ad-hoc to a particular data distribution or network configuration. Please let us know whether the above response addresses your comments. We can gladly provide further details.

Reviewer comment (Theory):
Technical: Dimensionality discrepancies throughout Section 3, page 3 and 4. For instance, in the definition of $\varPhi(X)$ on page 4, the $\varPhi(X)$ is not a function of $R^{d \times d}$, and $X$ should be a $d \times n$ matrix instead of previously defined $n \times d$, or $XX^T$ should be used instead of $X^TX$.

Author response:
While the $n=d$ assumption in our analysis makes these notations equivalent, we have nevertheless corrected the notations to avoid any confusions (Section 3.1, paragraph "Isometry gap").

Reviewer comment (Novelty):
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?

Author response:
Our analysis in Theorem 1 proves that deep networks with batch normalization and linear activations are not linear. In particular, these networks can makes their output orthogonal as they grow with depth. This is due to the non-linearity introduced by the batch normalization operator. 

Networks with linear activations and BN are a simplified model used to study the problem of gradient explosion, as these networks also suffer from gradient explosion (Yang & Pennigton et al. 2019). First, we theoretically study these simple models. Based on insights from our theoretical analysis, we propose practical recipes: orthogonal initialization and activation function shaping. Specifically, we use three main insights from networks with linear activations: initialization scheme, shaping mechanism and the choice of batch size for BN.

Reviewer comment (Theory):
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 a technique, tuning pre-activation gain $\alpha$ towards zero, effectively linearises 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 these benefits are maintained?

Author response:
Plese note that Figures G1-G2 show that the shaped $\sin$ and $\tanh$ activations have better accuracy than the linear case, and hence substantiate that the non-linearity in the activations is "effectively" used.

Reviewer comment (Experiments):
In later sections the paper presents the intriguing idea of “implicit bias towards orthogonal matrices”, for instance in Fig 6. But one can read experimental results in Fig 6 another 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 early stopping or otherwise corroborate "implicit bias" claims more extensively?

Author response:
Let us stress that all fully connected weights are initialized to *perfect orthogonality* i.e. $WW^T = W^TW = I$, implying that at initialization the weights have zero isometry gap $\phi(W_\ell)=0$ for all $\ell.$ This would correspond to a flat $\phi(X_\ell)=0$ curve for the $0$th epoch, which is not shown in Figure 6. Furthermore, Figure 6 shows that the middle layers throughout training have lower isometry gap, as evidenced by the U-shaped curves in all epochs. This U-shape indicates that the amount of change in orthogonality, as measured by the isometry gap, is smaller in the middle layers, in contrast with the first and last layers. Thus, if we track the change in orthogonality between epoch $0$ up to epochs $50,100,200,300,$ the middle layers have much smaller change in isometry gap compared to first and last layers. 

>  Could authors show further evolution of training continued beyond the early stopping or otherwise corroborate "implicit bias" claims more extensively?

Our results are not achieved for early stopping as we have trained the networks for 300 epochs. To supplement our observations, we provided Figures G4 where we observe 300 epochs suffice to acheive high training and test accuracies.

Reviewer comment (Novelty):
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]).

Author response:
We highly appreciate the reviewer for pointing us to further related works on dynamical isometry. We have included and discussed these valuable references in the Related Work part of our paper, paragraph "The challenge of depth in learning".

Reviewer comment (Theory):
The proof relies on the fact that the network 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?

Author response:
> 2. The proof relies on the fact that the network is has linear activations and the weights are orthogonal during training. ...

Having orthogonal weights is not only an assumption for theoretical analysis, but it is necessary to avoid gradient explosion at initialization. For example, [2] proves that the intialization with Gaussian weights causes gradient explosion in depth. Avoiding gradient explosion for batch normalization has been an open problem since 2015 and the literature concluded that the gradient explosion is "unavoidable" with batch normalization [2]. Our study starts with a theoretical setting where the activation function is linear. Taking inspiration from our analysis, we propose an activation shaping mechanism and initializing the weights with orthogonal matrices in order to avoid the gradient explosion for networks with non-linear activations. 

Please note that the established analysis is at *initialization,* when the weights are random. There is no assumption on the training dynamics in the established analysis. However, we observe the weights remain almost orthogonal during training in practice (see Section 6).


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

 
We only considered sine and hyperbolic tangent activations, as proof-of-concept for an activation shaping scheme inspired by our theoretical analysis. For ReLU networks, one can modify the shaping mechanism by introducing two parameters, the gain $\alpha$ and the bias $\beta$ as $\sigma(\alpha (x - \beta) )$, towards avoiding gradient explosion in depth.

[1] Frankle, Jonathan, David J. Schwab, and Ari S. Morcos. "Training batchnorm and only batchnorm: On the expressive power of random features in cnns." arXiv preprint arXiv:2003.00152 (2020).

[2] Yang, Greg, et al. "A mean field theory of batch normalization." arXiv preprint arXiv:1902.08129 (2019).

Reviewer comment (Theory):
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.

Author response:
> 3. 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 ...

 Theorem 1 proves that the network does more than a rotation: the network makes outputs orthogonal for each non-degenerate input matrix. While the rotation operation is angle preserving, orthogonalization increases the angle between every pair of vectors. Indeed, batch normalization, as a non-linear function, makes the input-output function siginficantly non-linear as the network grows in depth. Empirical observations show that networks with batch normalization and random weights become more expressive by increasing depth [1].

[1] Frankle, Jonathan, David J. Schwab, and Ari S. Morcos. "Training batchnorm and only batchnorm: On the expressive power of random features in cnns." arXiv preprint arXiv:2003.00152 (2020).

Reviewer comment (Experiments):
The paper also considers the setting where the number of training examples is equal to the dimension of the problem.

Author response:
We thank the reviewer for pointing out this potentially confusing issue. We assume the *mini-batch* size is equal to the width. However, the number of samples in the training dataset can be arbitrarily large. Notably, batch normalization layers use the statistics of intermediate representations for input batches. The batch size is a hyperparameter that can be adjusted to achieve a better performance. Here, we suggest a precise value for this hyperparameter to avoid the issue of gradient explosion. 

We have made the distinction between number of training samples and mini-batch size clear in the revised text, in Section 3, Equation 3.

Reviewer comment (Experiments):
Do you obtain figures similar to F1 when using ReLU activations?

Author response:
The reviewer is raising an interesting and important question. The proposed activation shaping scheme, which allows us to avoid gradient explosion for sine and hyperbolic tangent, relies on the fact that they are both differentiable at $x=0.$ Intuitively, when $\sigma'(0)$ exists, we can linearize $\sigma(\alpha x )$ more accurately when $\alpha$ is smaller. This property breaks down for ReLU, which is not differentiable at the origin. It is certainly worth exploring in a follow-up study how to extend the proposed activation shaping scheme for non-differentiable functions like ReLU.