From: ICLR.cc/2023/Conference/Program_Chairs
title: Paper Decision
decision: Reject
metareview:_summary,_strengths_and_weaknesses: This paper studies the partial (layer-to-layer) Jacobians of wide neural networks and proposes a simple numerical test to select initialization hyperparameters based on the concept of criticality. The analysis marginally extends prior work to include architectures with LayerNorm and skip connections, and finds that certain architectures such as MLP-mixers can exhibit "everywhere criticality." 

The reviewers agreed that the empirical estimate is interesting and could be valuable. They also appreciate the identification of the special behavior of LayerNorm + skip connections. Some reviewers appreciated the clarity of exposition in general, but there was significant confusion regarding the extent of novelty in the theoretical contributions of this work. In particular, there were concerns that proper credit was not given to the prior works that have derived very similar results, making it unclear what precisely the novel developments are in this paper.

All told, this paper presents some developments to the theory of initialization for wide and deep neural network architectures, but the extent of those developments is not clearly characterized. Moreover, there are some concerns that the novel contributions are quite narrow, and that the practical impact of those insights remains speculative. For these reasons, I do not recommend publication of this manuscript at ICLR.
justification_for_why_not_higher_score: In order to merit acceptance, the paper needs to more precisely frame its novel contributions in relation to prior work. The relatively narrow scope of these contributions would likely then require further experimentation to demonstrate the practical relevance of the insights.
justification_for_why_not_lower_score: N/A
summary_of_AC-reviewer_meeting: From the AC-reviewer meeting:

Positive points:
* Relatively clear exposition
* Interesting empirical measure
* Some good validation of the theory
* Could be used for architecture design, etc.

Negative points:
* Overclaiming of results
* Presentation doesn't highlight main contributions
* No demonstration of practical value

While the discussion phase did give the authors the opportunity to address these negative points, their answers did not convince all the reviewers.

Following the AC-reviewer meeting, reviewer 6A97 noted that a more appropriate score would be closer to 6.5-7 and sMx1 acknowledged more concerns than originally noted. Both reviewers converged to more lukewarm perspectives, characterizing the paper as "decent" and "not bad", while mentioning that the above negative points kept it from being "great". These changes in opinion aren't reflected in the reviews but do fold into the final decision.

From: ICLR.cc/2023/Conference/Paper3195/Authors
title: Response to $\sigma_b^2=0$ Concern
comment: To avoid any possible confusion, We clarify that "disabling" biases is different from setting $\sigma_b=0$. The difference being that biases set to 0 still have non-zero gradients and receive non-zero gradient updates, while disabled/turned-off biases do not. This distinction is important since APJN deals with gradients.

SOTA architectures that use LayerNorm, e.g. Transformers, MLP-Mixer and ConvNeXt[Liu, et. al.. A ConvNet for the 2020s], do not disable their biases. They do, however, set $\sigma_b=0$. Importantly, these architectures utilize the combination of LayerNorm *and* skip connections, rendering them in the everywhere-critical regime. Consequently, in these cases, the choice of $\sigma_b$ and $\sigma_w$ is not crucial for trainability. However, in architectures that use *only* LayerNorm, the choice of $\sigma_b$ plays a crucial role.

We suspect that when the reviewer wrote "you would typically disable the bias parameters of the preceding layer", they had ResNet architecture in mind. It is indeed common to turn off biases of the convolutional layers preceding *BatchNorm* in ResNet. To understand why this is okay with BatchNorm and *not* LayerNorm, it is important to look at the dimensions along which these normalization layers act. Consider preactivations with dimensions $(B, C, H ,W)$ where $B$ is the batch dimension, $C$ is the channel dimension, $H$ is height, and $W$ is width. BatchNorm acts along the batch dimension B. Since, in convolutional layers, biases are shared (identical) across the batch, the following BatchNorm layer will completely nullify the effect of biases. However, this is not the case with LayerNorm. In the aforementioned architectures, LayerNorm acts along either $C$ or $(H, W)$. These dimensions do not share biases, and hence, LayerNorm does not nullify their effect.

From: ICLR.cc/2023/Conference/Paper3195/Reviewer_PgVN
title: In practice, $\sigma_b^2 = 0$
comment: I thank the authors for highlighting the actual effects due to the components. I would argue that it would have made things much easier (for me) if the paper was constructed around these entities. 

When chaining layers with layernorm, you would typically "disable" the bias parameters of the preceding layer.
Wouldn't this imply that $\sigma_b^2 = 0$ can/should be assumed in practice?

From: ICLR.cc/2023/Conference/Paper3195/Authors
title: Follow Up Responses
comment: 
1. We would like to clarify that skip connections with $\mu=1$ *do not* cancel the effect of LayerNorm. The combination of the two achieves everywhere-criticality, which neither of the techniques can achieve on its own. To drive this point home, we offer the following concrete example. Consider a ReLU network:
- With *only* LayerNorm ($\mu=0$) (Eq. (96), general proof in Appendix B.3):
    $$\chi^{\star}_{\mathcal J} = \frac{\sigma_w^2/2}{\sigma_w^2/2 + \sigma_b^2} \\,. $$
    In this case, the only critical points lie on the $\sigma_b=0$ line.
- With *only* skip-connections (Eq. (107), general proof in Appendinx D): 
    $$\chi^{\star}_{\mathcal J} = \sigma_w^2/2 + \mu^2 \\,.$$
    Clearly, when $\mu=1$, the only critical points lie on the $\sigma_w=0$ line.
- With the combination of LN and skip-connections (Eq. (109), general proof in Appendix E): 
    $$\chi^{\star}_{\mathcal J} = \frac{\sigma_w^2/2}{(\sigma_w^2/2 + \sigma_b^2)(1-\mu^2)^{-1}} + \mu^2 \\, ,$$
    where the extra denominator comes from the contribution of LayerNorm. Indeed, we see that when $\mu=1$, $\chi^{\star}\_{\mathcal J}$ is identically 1, rendering the network *everywhere*-critical.

To summarize, LayerNorm doesn't get canceled, it modifies the formula.

Theorem 1.3 is a summary of the results proved in Appendix E. For comparison with the case with only skip connections, we also point the reviewer to Appendix D.

2. We would like to emphasise the following points:
- The crucial difference between our calculations and the reviewer's example is the inclusion of separate variance for the bias : $\sigma_b^2$. By everywhere-criticality we mean arbitrary choice of $\sigma_w$ *and* $\sigma_b$. Upon using only LayerNorm, the criticality is always confined to a line on this plane. We discuss the example given by the reviewer next.
    
- The model proposed by the reviewer with ReLU activations has the following APJN: $$\chi_{\mathcal J}^{\star} = \frac{\sigma_w^2 \gamma_{ReLU}^2}{\sigma_w^2 \gamma_{ReLU}^2 + 0} = 1 \\,,$$ where the $0$ term is from the fact that the model has no bias. There is indeed the freedom to arbitrarily scale $\sigma_w$. This is a result of (i) The scale-invariant property of LayerNorm, which allows one to scale the $\sigma_w$ and $\sigma_b$ along a critical line; (ii) the positive homogeneity of ReLU i.e. $\mathrm{ReLU}(a \cdot x) = a\cdot \mathrm{ReLU}(x)$ for $a>0$, which makes $\sigma_b=0$ (x-axis) the critical line.
    More importantly, there is no freedom to choose $\sigma_b$ in this case -- it has to be set to 0 for criticality to hold. Thus, this case does not give everywhere-criticality.
- Similar arguments apply to other activation functions. The scale invariant property of LayerNorm still allows one to scale $\sigma_w$ and $\sigma_b$ along a critical line. Although, the lack of positive homogeneity means that this line is not the x-axis. In other words, one has the freedom to scale $\sigma_w$ and $\sigma_b$ while maintaining a fixed ratio between them. We stress that this still a line and not the entire plane.
- The results and calculations for ReLU, erf and GELU activation functions can be found in Appendices G, H and I respectively.

### Other Comments
- We stated a brief motivation in the abstract: "We derive and implement a simple and cheap numerical test that allows one to select optimal initialization for a broad class of deep neural networks; including fully
    connected, convolutional and attention layers.". We agree that it does not fully elaborate the reasons for choosing APJN. We will incorporate the detailed motivation in section 1.1 of the main text as well as the appendix, in the camera-ready version. 
 - Eq. (17) in Proposition 2.6 includes $L \longrightarrow \infty$ limit. We have also said the same in words, below Eq. (17). We think this will make it  clear to the reader that Proposition 2.6 applies to deep networks.
 - The key limitation of the approach in [Schoenholz et al., 2016] is as follows: With unbounded activation functions like ReLU, when $\sigma_w > 2$, the variance of preactivations is no longer finite as depth goes to infinity. As a consequence, the quantity $\chi_1$ in [Schoenholz et al., 2016], though calculable, is no longer meaningful. If one were to repeat their calculation for ReLU, one would find that $\chi_1=1$ for all $\sigma_w \geq 2$. Hence, $\chi_1$ cannot give us the phase boundary in this case. This issue has been first addressed in [Roberts et al, 2022] using $\chi_{\perp}$, as we mention above Section 2.1. 
- We will make it a point to include [1] and [2] in Section 5.
- We thank the reviewer for pointing out the typo in Theorem 2.4.

From: ICLR.cc/2023/Conference/Paper3195/Reviewer_PgVN
title: One more question and a few comments
comment: ### Why is the combination of layer norm and skip-connections so important?

   1. On one side, theorem&nbsp;1.3 seems to imply that the effects due to layer norm are cancelled for $\mu = 1$, which is exactly the setting for which everywhere-criticality is claimed. However, this would mean that the skip connections are sufficient to achieve everywhere-criticality (which does not seem to be the case from figure&nbsp;1). It might also be useful to point to the proof so that I can understand what is happening.
   2. On the other side, everywhere-criticality could be achieved in an "easier" way by using a formula of the form $h_i^{l+1}(x) = \sum_j w_{ij}^{l+1} \phi(\gamma_\phi \tilde{h}_j^l(x)),$ where $\gamma_\phi$ is the critical scaling factor (e.g. $\gamma_\mathrm{ReLU} = \sqrt{2}$ if I don't mistake) for a the activation function. This simply makes use of the fact that layer-norm is scale-invariant.

---

### Other Comments
 - I don't understand why the motivation for APJN has not been included in the revised paper. Even if you do not have space for it in the main paper, I think it would be useful to include it at least in the appendix.
 - With "not properly tested", I actually meant that I did not find any experiments on the effect of depth. I admit that this was formulated too vaguely in my original view. However, the reference to the  (previously hidden) experiment in the appendix is what I was looking for. With this information, it would be advisable to warn the reader (in the main text) that proposition&nbsp;2.6 only holds for deep networks.
 - I think the revision was successful in improving readability.
 - The requirement for the function to be bounded in Schoenholz et al. (2016) is only required for fixed points to exist. The recurrence relation should still hold (under mild assumptions, e.g. bounded variance). I admit that the result is new in the sense that it is applied to the APJN, but I do not think it is appropriate to put so much emphasis on the novelty of theorem 2.5.
 - It would be good to include references to [1] and [2] in section 5 for the camera-ready version.

PS: concerning theorem&nbsp;2.4, shouldn't it be "...kernel is deterministic, and can _be_ determined ..."?

From: ICLR.cc/2023/Conference/Paper3195/Authors
title: Response to Reviewer PgVN (2/2)
comment: 
## Quality
Clarified in the **Response** section in part 1/2.

## Clarity
  - We have modified our paper to make the reading flow more steady.
  - We thank the reviewer for pointing out a subtlety in (9). We checked numerically that it holds for activation functions commonly used in practice.
  - We have rephrased proposition 2.6 to make it easier to understand.


## Originality
  - We have mentioned the key relation between our work and existing work above Section 2.1. Namely, to the work on signal propagation and the book by Roberts et. al.[3].
  - As we mentioned in Section 2.1, even though the work (Schoenholz et al., 2016) looks very similar to ours, it is different. They defined $\chi_1$, which differs from our $\chi_{\mathcal J}$ by a depth and input dependent factor. This factor disappears when the activation function is bounded. However, for unbounded activation functions like ReLU and GELU, one would find that $\chi_1=1$ holds regardless of the choice of initialization. Consequently, one cannot find the phase boundary for these activation functions using the $\chi_1=1$ condition.
  - The results in eq. (23) are necessary to demonstrate the subsequent results. We have cited Yang \& Schoenholz, 2017 properly above Theorem 4.4.
  - We thank the reviewer for attracting our attention to the fact that Ba et al. (2016), indeed, apply LayerNorm to preactivations. We have fixed the footnote accordingly.
    
  - In the presence of LayerNorm on preactivations, as we define it, the critical line is $\sigma_b^2 = c \sigma_w^2$. For ReLU we have $c=0$. Thus, it is a coincidence that Kaiming's initialization still holds for ReLU. Note further that $\sigma_w$ does not have to equal to $\sqrt{2}$ and can take *any* value. The network will remain critical due to the effect of LayerNorm. Our study of signal propagation with LayerNorm is, to the best of our knowledge, new.  
    
  - As addressed in the Response section, many SOTA architectures for vision[1] and language[2] tasks utilize the combination of LayerNorm and residual connections, motivating its careful study. BatchNorm introduces correlations between batches, and hence, requires a slightly modified definition of the APJN involving an average over batches. BatchNorm, in the absence of skip connections, makes the network *chaotic* for any init as was first shown in [4]. When skip connections are introduced, BatchNorm network also becomes everywhere-critical. Moreover, we have mentioned the related works (eg. Yang et al., 2019) in Remark 5.4.


## Others

 - We have fixed the typos pointed out by the reviewer.
 - We could not locate the typo in Theorem 2.4.
 - Eq. (10) for $\mathcal K^{l+1}(x, x)$ is recursive with the help of Lemma 2.2. We now write this explicitly in Theorem 2.4.



[1] Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, 2016.

[2] Chowdhery et al., PaLM: Scaling Language Modeling with Pathways, 2022.

[3] Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory: An Effective Theory Approach to Understanding Neural Networks. Cambridge University Press, 2022.

[4] Yang, G., Pennington, J., Rao, V., Sohl-Dickstein, J., \& Schoenholz, S. S. (2019). A Mean Field Theory of Batch Normalization. ICLR 7.

From: ICLR.cc/2023/Conference/Paper3195/Authors
title: Response to Reviewer PgVN (1/2)
comment: We thank the reviewer for the constructive feedback.

# Response
One of our main contributions, as mentioned in the abstract, is that the combination of LayerNorm on preactivations and residual connections makes networks everywhere critical. Interestingly, this way of combining LayerNorm and residual connections is used in SOTA architectures for both vision[1] and language[2] tasks. We also want to clarify that we are not proposing using LayerNorm on preactivations as something new -- such architectures have been used before. Moreover, the question of ordering between LayerNorm and skip connection was investigated experimentally in [Xiong et, al. *On layer normalization in the transformer architecture*. PMLR, 2020]. In our work, we provide a theoretical justification for why one ordering is likely to be preferable over the other. We do so by showing that the interaction between the LayerNorm and skip connections makes the network insensitive to initialization by rendering the network critical at any init. Of course, different initializations would require us to re-tune other hyperparameters.

In Fig. 1, black lines indicate the critical lines obtained by first taking the limit of infinite width and then of infinite depth. The phase diagram (colors) was plotted numerically using our method and used to compare to theoretical values. Next we address the observations made by the reviewer.

- In the "GELU $\mu=0$" panel, it is known that the $(\sigma_w, \sigma_b)=(2, 0)$ fixed point is unstable, and the finite $\sigma_b$ fixed point is half-stable [3]. Consequently, the difference between the theoretical result and numerical result is to be expected. The training accuracy plotted in Fig.2 suggests that network is trainable around the critical region depicted in Fig.1, which indicates the utility of numerical evaluation of APJN for finite networks.

- In the Erf "$\mu=0.9$" panel, there was a bug in the theoretical derivation of the critical line (we could not find results for arbitrary $\mu$ in the literature and had to derive it ourselves) -- the factor of $1/(1-\mu^2)$ was missing while plotting the line. The numerical plot is accurate. We have corrected Fig. 1 and Fig. 2 for this bug and updated the supplementary material. We are very grateful to the reviewer for this very keen observation.

- We would like to emphasize that we did not introduce the APJN to study the infinite width limit. We introduced APJN to figure out how to initialize real finite-width and finite-depth networks. We used large width limit to benchmark how well APJN does in a theoretically-controlled setting. In practice, one would discard the network initialized with large or small APJN and train the network with APJN $\approx 1$. In the "GELU $\mu=0$" case it is sufficient to initialize the network in the white region of the Fig.1. The `experimental reality' is the color plot, while the black line is an idealized model. For the purpose of training a finite network we are interested in the former.

The motivation for APJN can be summarized as follows: We want a simple, cheap, and flexible tool to analyze gradient exploding and vanishing problems for most modern architectures. Our choice of APJN satisfies these requirements since (1) Jacobian is the dominating factor while considering the order of gradients; (2) We compute Frobenius norm with Hutchinson's trace estimator in practice, which is cheaper than spectral norm calculation; and we introduced $1/N_l$ factor to make sure the quantity will be finite in the infinite width limit for the purposes of the aforementioned benchmark; (3) The freedom of choosing $(l_0, l)$ pairs allows one to analyze inhomogeneous architectures with non-repeating building blocks. In these cases, the gradient may fluctuate at the intermediate blocks (or even within the blocks). Sweeping over various values of $l_0$ allows to control these intermediate gradient fluctuations. Finally, it is possible to enforce $\mathcal J^{l, l+1}=1$ for all $l$ via optimization, leading to an automatic initialization scheme, discussed elsewhere.

We do disagree with the reviewer's Summary. Namely, with the two claims
- "The estimator is not properly tested." In our opinion, combination of Fig.1 and Fig.2 is proper testing. We further tested it with a modern architecture in Fig. 4. In all cases APJN $\approx 1$ correlates with trainability. The detailed response to the reviewer's criticism on this point is provided in part 2/2 of our response.

- Reviewer's statement: ``Most of the theoretical analyses turn out to be known results from a body of work on mean field theory.'' 
We believe that our observation that LayerNorm + skip connections lead to an architecture that is not sensitive to initialization due to criticality is both (i) not previously known, and two (ii) sufficiently important to be published, along with the *new and practical* method of mapping out the phase diagram.

From: ICLR.cc/2023/Conference/Paper3195/Authors
title: Response to Reviewer 7FZf
comment: We thank the reviewer for the constructive feedback.

# Response
The results in Section 2 are not all new, however we believe that the section is essential for the flow of the paper. Moreover, as we mentioned before Section 2.1, our test of criticality is easier to implement. It is particularly useful when the network architecture cannot be analyzed analytically, *e.g.* MLP-Mixer. We would also like to point out that our work can be used for both NTK initialization and standard initialization ($1/fan\\_ in$ is included in the distribution of weight matrix $w_{ij}$) without any change, whereas previous works based on NTK initialization cannot be directly applied to standard initialization and need to be tweaked [1].

We believe that there must be a misunderstanding on the reviewer's side (or our failure to communicate). The exponential law (see Eq.(18)) is, indeed, the `trivial' generic behaviour. However, Theorem 2.7 is non-trivial, precisely because it states that at criticality the generic exponential behavior is replaced by a *power law*. To re-iterate, Theorem 2.7 says that if a network is at criticality then its APJN scales like $l^{-\zeta}$, which is a power law, not an exponential. If the network is off criticality, APJN scales like $(\chi_{\mathcal J}^{\star})^l$, which is exponential and leads to the exploding/vanishing gradients problem.

## Other Comments
We have clarified that *at initialization* the Jacobian is a random matrix.

Eq(10) for $\mathcal K^{l+1}(x, x)$ is recursive with the help of Lemma 2.2. We now write this explicitly in Theorem 2.4.

[1] Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory: An Effective Theory Approach to Understanding Neural Networks. Cambridge University Press, 2022.

From: ICLR.cc/2023/Conference/Paper3195/Authors
title: Response to Reviewer sMx1
comment: We thank the reviewer for the constructive feedback.

# Response
 We introduced the $(l_0, l)$ pair with a more general architecture in mind. As we mentioned in the first paragraph of the conclusions, the APJN would be useful for analyzing criticality when the network is inhomogenous. Especially for the case where there are no layer-independent $(\sigma_w, \sigma_b)$ pairs that can make the network critical. Moreover, there could be cases where gradients can drastically fluctuate across different layers, yet the network satisfies $\mathcal J^{1, L}=1$ or $\mathcal J^{L-2, L-1}=1$. To analyze those complicated cases, one has to scan over different $l_0$ choices and require all $\mathcal J^{l_0, l}=1$, which will let one find a proper initialization.


## Clarity
1. We have changed the equations as the reviewer suggested. 

2. The concept of critical exponent $\zeta$ is common in statistical mechanics, where it characterizes universal scaling (i.e. independent of microscopic details) near a phase transition. There are works that found similar scaling behaviour, at criticality, for other quantities [1,2]. To the best of our knowledge, the scaling of the Jacobian was not studied. The interesting new element of Theorem 1.3 is that the critical exponent can change continuously with $(\sigma_b, \sigma_w)$, while the network remains critical for all values of these hyperparameters.

3. We do not include a plot for $\tanh$ and we have modified the last paragraph of Section 4 to avoid confusion. Generally, the results for $\mathrm{erf}$ and $\tanh$ are extremely close to each other. The advantage of $\mathrm{erf}$ is purely theoretical: the critical line can be calculated in a closed form. 

4. By trainability, we meant training accuracy. We have modified the caption of Fig 2.

5. We only use $\chi_{\mathcal J}$. We have corrected those typos.

[1] Lechao Xiao, Jeffrey Pennington, Samuel S. Schoenholz, Disentangling Trainability and Generalization in Deep Neural Networks, 2019.

[2] Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory: An Effective Theory Approach to Understanding Neural Networks. Cambridge University Press, 2022.

From: ICLR.cc/2023/Conference/Paper3195/Authors
title: Response to Reviewer 6A97
comment: We thank the reviewer for the constructive feedback.

# Response
We have modified the last paragraph before Section 1.1 to incorporate the reviewer's suggestion. We mention the connection between our work and the work done by others before section 2.1 -- we would offer a more detailed comparison between our work and existing works if we were less constrained by space. We have modified the sentence "Very deep networks will not train unless initialized critically." in the main text, for clarity. Fig. 2 illustrates this idea, where we plot training accuracy (to which we referred to as trainability) for fairly deep neural networks. 

We have revised the Proposition 2.6. The justification for it is as follows: $\mathcal J^{L-2, L-1}$ is entirely determined by the distribution of preactivations, i.e., NNGP kernel. NNGP kernel, in turn, can be determined recursively via (10) and Lemma 2.2. Since the NNGP kernel recurrence relation has a fixed point as $L \rightarrow \infty$, so does $\mathcal J^{L-2, L-1}$. For additional clarity, we added a sentence that refers to Fig. 4. The non-trivial part of the Proposition 2.6 is that $\mathcal J^{L-2, L-1}$ at large $L$ provides an accurate estimate of $\chi^{l}_{\mathcal J} = \mathbb E_\theta \left[ \frac{\sigma_w^2}{N_l}  \sum\_{i=1}^{N\_{l}} \left(\phi^\prime(h_i^{l}(x)) \right)^2 \right]$.

We thank the reviewer for bringing the work [Arora et al., 2018] to our attention. We have added a reference in the main text.

We introduce the $(l_0, l)$ pair with more general architectures in mind. As we mention in the first paragraph of the conclusions, the APJN would be useful for analyzing criticality even when the network is inhomogenous. It is especially useful in the case where there are *no* layer-independent $(\sigma_w, \sigma_b)$ pairs that can make the network critical. Moreover, there could be cases where gradients can drastically fluctuate across different layers, yet the network satisfies $\mathcal J^{1, L}=1$ or $\mathcal J^{L-2, L-1}=1$. To analyze those complicated cases, one has to scan over different $l_0$ choices and require all $\mathcal J^{l_0, l}\approx 1$, which will lead one to a proper initialization.

## Clarity, Quality, Novelty And Reproducibility
We have fixed typos and confusing symbols pointed out by the reviewer. Some important clarifications are:
 - By algebraic dependence we mean power law.
 - Lemma 2.2, we used $\mathcal O$ which is different from big $O$ we used in this paper. $\mathcal O(h^l(x))$ represents a Lipschitz continuous function with the variable $h^l(x)$. The $\mathcal O$ was meant to stand for an ``$\mathcal{O}$bservable''.
 - Figure 2 is training accuracy, we have improved the caption.
 - We removed the confusing $O$ in Remark B.7.

From: ICLR.cc/2023/Conference/Paper3195/Reviewer_6A97
confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
summary_of_the_paper: This paper discusses a notion for assessing the trainability (meaning, do the gradients avoid vanishing or exploding during training) of deep neural networks of various architectures -- the Frobenius norm, squared, of the averaged partial Jacobian (APJN) between the pre-activations at layers $\ell_0 \leq \ell$. The paper applies the APJN as an empirical diagnostic of the criticality (is the network trainable, in the above sense) and also to demonstrate that for deep MLP networks with residual connections and with Lipschitz activations with LayerNorm applied to pre-activations, the initialization scheme does not impact the criticality of the network. This result formalizes a view on deep network architecture design -- either fix a network for which the criticality (and thus, loosely, trainability) is sensitive to initialization, and find the right initialization, or define a network using methods like LayerNorm and residual connections so that the network is trainable regardless of initialization. Additionally, an upper bound on the rate of divergence of the gradients in non-critical settings depending on residual connection strength $\mu$ and the activation functions is given. Several experiments are performed to validate the theoretical contribution as well.
strength_and_weaknesses: * Strengths: The paper provides a nice characterization of their proposed measure, the APJN, and suggests that it is a useful diagnostic to assess the trainability of deep neural networks empirically. The paper also more formally makes some connections present in the existing literature, and presents a nice takeaway that more precisely characterizes the importance of LayerNorm and residual connections from the perspective of initialization. In particular I appreciated the development of the correlation length concept, it was clearly explained and provided nice intuition. Finally, several experiments on smaller datasets are conducted to validate the theory and the gap between the infinite width setting and realistic settings.

* Weaknesses: The primary weaknesses in my view are as follows: 
     * I believe that the notion of trainability/criticality in this paper has not been directly shown to imply generalization in the literature, and it is more of a heuristic notion. However, many papers study this notion and so it is not a big issue. It would however be nice to clarify with more detail the relation between criticality/trainability and actual model performance in a more formal sense (or at least highlight that this is unknown). It would also be useful to provide a more formal definition of criticality from the previous works that studied it, and to make the connection between the criticality conditions that APJN satisfies and the existing notions of criticality clearer. It would be good to clarify the meaning (and justification) of "Very deep networks will not train unless initialized critically." at the end of the first section.
     * I don't see a justification of Proposition $2.6$ anywhere -- is this just a (reasonable) assumption for the empirical results? It would be nice to highlight / include a pointer to Figures 3 and 4 which essentially are justifying this assumption, if I understood correctly.
     * At least one other paper (https://arxiv.org/abs/1802.05296) studies the notion of partial Jacobians and provides conditions that depend on the norms of these partial Jacobians (see Definition 5, the interlayer cushion, in the mentioned paper) in the setting of giving generalization bounds. It probably makes sense to mention this work and any others, as while I do not think I have seen the APJN specifically in previous works, there are probably quite a few more than just the three papers that were cited in the related work that study it. This issue is not a big deal as long as the APJN has not been introduced before or analyzed, but it does impact a little bit the claims of introducing this notion in my opinion.
    * I don't understand why one would ever care about $\ell_0$ other than $\ell_0 = 1$ or $\ell_0 = L - 1$ (the latter being a computationally tractable, and apparently accurate, approximation) -- is there any reason to care about other $\ell_0$ as a diagnostic? If not, then perhaps this could be clarified when the paper mentions that the growth of the APJN only depends on $\ell_0$ via some constant.

clarity,_quality,_novelty_and_reproducibility: * Clarity: The paper was clear overall, with a few suggestions highlighted in the weaknesses mentioned above. I have listed some minor points below: 
    * 2nd sentence of abstract: add a comma after "infinity"
    * 7th line of intro: "Furthermore, for special choice of parameterization..." choice --> choices
    * Below Definition $1.1$: "sqaured" is misspelled. 
    * I'm not sure why the term "algebraically" is used -- shouldn't it be "polynomially" to contrast with exponential situations? If I read correctly, this term is referring to the polynomial growth of $\ell$ with some degree $-\zeta$.
    * In the statement of Thm $1.3$, the forward reference to Eq. 24 is a bit jarring. It would be nice to include the definition in the statement.
    * In Lemma $2.2$, the big-O notation used in Eq. 8 is unclear to me - could it be clarified?
    * Figure 2 was a bit unclear to me in terms of exactly what is being plotted - is Figure 2 plotting accuracy after training? More details are needed.
    * Remark $5.2$ -- "the architecture design" -- don't need "the".
    * Lemma B.1 -- need to complete the sentence with punctuation. 
    * Before Definition B.2 -- need to complete the sentence with punctuation.
    * Remark B.7 -- $\mathbb{E}_{\theta}\left[O\right]$ is not defined?
    * Before Remark B.8 -- Equation reference missing
    * Figure $7$ -- formatting on "log-log" should be fixed.


* Quality: The paper was interesting and the ideas introduced useful and insightful. 


* Novelty: The paper was sufficiently novel for publication.


* Reproducibility: The experiments seem easy to reproduce, though I have not done so myself.
summary_of_the_review: Overall, the paper introduced a seemingly novel diagnostic for evaluating the criticality/trainability of deep neural networks and provided some thorough characterizations of the trainability in a variety of situations, in the infinite-width limit. The paper backed up its results with experiments, and used the results to highlight and justify the ways methods like LayerNorm and residual connections can be viewed as a method of making the initialization problem easier to solve. The settings studied are limited, and it is not directly clear how these issues (trainability, etc) impact generalization performance, but it seems to be a useful approximation to understand better and to more thoroughly study. The ideas in the paper may also be helpful to further develop both the practice and theory of deep networks.

==== POST-REBUTTAL ====

After reading the other reviews and rebuttals, I remain convinced the paper should be accepted.

==== POST-DISCUSSION =====

 After the discussion, overall, I think the paper proposes an interesting measure to investigate when deciding how to initialize deep networks, and provides some justification via a combination of existing theory and limited experiments that it should be interesting. My feeling was still that the paper was relatively clear, though perhaps there is some over-claiming in the theorem statements and attribution should be a bit clearer/cleaned up. I think it is reasonable to accept (I would vote something closer to 7 if that were an option), but I do agree that the paper would be significantly improved by a) cleaning up the theorems / equations (probably don't need to write down each case the way they did), b) tweaking the language to avoid some over-claiming, and c) additional experiments validating the measure in practice. However, I view these changes along the axis from being a decent paper to a great paper and still lean towards acceptance.
correctness: 4: All of the claims and statements are well-supported and correct.
technical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
empirical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
flag_for_ethics_review: NO.
recommendation: 8: accept, good paper

From: ICLR.cc/2023/Conference/Paper3195/Reviewer_PgVN
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary_of_the_paper: This work introduces the concept of Averaged Partial Jacobian Norm (APJN).
The APJN can be used to derive criticality measures &mdash;i.e. how trainable a network is&mdash; for arbitrary parts of a neural network.
From the APJN, an empirical estimator can be derived to empirically compute how "trainable" a given neural network is.
Furthermore, the criticality of layer normalisation, skip connections and the MLP-Mixer architecture is analysed.
strength_and_weaknesses:  - [+] An empirical estimate of criticality should be useful to easily tune certain hyper-parameters of neural networks (e.g. initialisation).
 - [-] The quality of the proposed empirical estimate is not properly tested.
 - [-] There are a lot of forward and backward references, making the paper hard to read.
 - [-] Assuming that the new architecture refers to applying Layer Normalisation before the activation function, there is nothing *new* about it (Ba et al., 2016).
 - [-] Most of the results from the criticality analysis are already known.
clarity,_quality,_novelty_and_reproducibility: ### Quality

 - There seems to be no motivation for why the APJN should be an interesting value to consider.
   To be more precise, the exact formulation feels arbitrary:
    1. the choice for computing the Jacobian over a subset of layers, instead of just one or all layers,
    2. the use of the Frobenius norm instead of something like the spectral norm,
    3. the multiplication factor $\frac{1}{N_l}$, which seems to imply that this is actually an average over squared $L_2$-norms of gradient vectors of individual output entries.
 - Figure&nbsp;1 seems to be the only experiment to verify that the APJN of the last layer can actually be used to estimate criticality.
   Although these results look promising on first sight, closer inspection reveals a few issues.
   E.g. for "erf, $\mu = 0.9$" the APJN seems to be quite off in general.
   Similarly, for "GELU, $\mu = 0.0$", the empirical and theoretical results do not seem to agree for small $\sigma_b^2$.
   Without a proper analysis of these failures, it would be hard to rely on these estimates.

### Clarity

 - Although I am familiar with most of the related work, I needed multiple read-throughs to distill the actual contributions from this paper.
   This seems to be a strong indicator that this paper would benefit from a thorough rewrite for it to be useful to the broader ML community.
 - The many references to various parts in the paper, typically break the reading flow.
   Since the references often span multiple pages, I found myself spending quite some time going back and forth to put everything together.
 - The derivation in eq.&nbsp;(9) appears to be trivial, but I would argue that it can not be that easy.
   E.g. using the uniform law of large numbers (Jennrich, 1969), I see how $\frac{1}{N_l} \sum_{i=1}^{N_l} \phi(h_i^l)^2 \overline{\to}^{N_l \to \infty} \mathbb{E}_{h_i^l}[\phi(h_i^l)^2]$, from which the result would follow, could hold.
   However, there are a number of assumptions on $f(x) = \phi(x)^2$ that would need to be satisfied for this to work.
 - The statement of proposition&nbsp;2.6. is very confusing.
   The authors mention $l_0$ and $l$ in the text, but the formula only contains $L-2$ and $L-1$, which are both not specified.
   Also, it is unclear what "next to the output" means in this context.
   I would argue it must be the derivative from pre-activations of the last layer w.r.t. the penultimate layer, but I am not entirely sure whether this is correct.

### Originality

 - I would argue that the related work section should include a paragraph to establish the relation of this work to mean field theory.
   After all, the results in eq.&nbsp;(21),&nbsp;(23) and&nbsp;(25) appear to be the kind of results you could obtain from a mean field analysis.
   Most of the relevant papers are cited at some point, but I think it would be useful to explicitly highlight similarities and differences with this line of research.
 - I fail to see what is new about Theorem&nbsp;2.5.
   This seems to be exaclty the same result as eq.&nbsp;(8) in (Schoenholz et al., 2016).
 - The results in eq.&nbsp;(23) seem to be exactly those in Table&nbsp;l ($\underline{\boldsymbol{\chi}}$) from (Yang & Schoenholz, 2017).
 - Applying layer normalisation to pre-activations is exactly what Ba et al. (2016) did,
   unlike what footnote&nbsp;2 suggests.
   In eq.&nbsp;(3) of (Ba et al., 2016) normalisation is defined in terms of $a_i^l$, which should be the pre-activations according to eq.&nbsp;(1).
 - Ba et al. (2016) clearly state (see table&nbsp;1) that layer normalisation makes weights scale- and shift-invariant.
   Therefore, the following excerpt from the manuscript should probably be more explicit about how well this is already known:
   > One fortunate outcome of both theory and experiment is that when LayerNorm is applied to preactivations, ReLU networks can still be initialized using He initialization (He et al., 2015) which, in our convention, is $(\sqrt{2}, 0)$.
 - It would probably be useful to point out how the layer-norm results relate to existing analyses of batch-normalisation (e.g. Yang et al., 2019).
 - Section&nbsp;5 states that residual connections are "commonly" combined with layer normalisation.
   To the best of my knowledge, it is more common to use batch normalisation, because that is also what He et al. (2016) used.
   I think it would be useful to point out a few works that use layer normalisation at this point.
   It also raises the question why the authors chose to study layer normalisation instead of batch normalisation.

### Other, minor comments

 - At some points in the paper, sentences are hard to read due to lack of punctuation.
   (e.g. first sentence of the introduction).
 - The starting point for the recursion in eq.&nbsp;(1) is not defined.
 - The weight distribution for $w_{ij}^{l+1}$ was probably scaled by $\frac{1}{N^{l}}$ instead of $\frac{1}{N^{l-1}}$ (cf. Lecun et al., 1998b and He et al., 2015).
 - Typo in theorem 2.4. 
 - Eq.&nbsp;(10) does not seem to be a recursive formula.
   Also the formulas in eq.&nbsp;(21) are not recurence relations.
   In eq.&nbsp;(23), only the first formula would be a recurrence relation.
 - A lot of formulas would be easier to read (and shorter) if $\phi(x)^2$ (cf. eq.&nbsp;21) were used instead of $\phi(x)\phi(x)$ (e.g. eq.&nbsp;12).
 - Please, directly refer to the corresponding sections in the appendix instead of just "the appendix".

### References

 - Jennrich, R. I. (1969). Asymptotic properties of non-linear least squares estimators. The Annals of Mathematical Statistics, 40(2), 633-643.
 - LeCun, Y., Bottou, L., Orr, G. B., & Müller, K.-R. (1998b). Efficient BackProp. In G. B. Orr & K.-R. Müller (Eds.), Neural Networks: Tricks of the Trade (1st ed., pp. 9-50). Springer.
  - Yang, G., Pennington, J., Rao, V., Sohl-Dickstein, J., & Schoenholz, S. S. (2019). A Mean Field Theory of Batch Normalization. ICLR 7.
summary_of_the_review: I would argue that the most interesting contribution of this paper would be a practical estimator for criticality.
However, the proposed estimator is not properly tested to be able to assess when/how this could be used.
Most of the theoretical analyses turn out to be known results from a body of work on mean field theory.
Therefore, I do not see what added value the APJN has since it basically builds on the same principles.
Finally, I found the paper very hard to read and I believe that a rewrite would be necessary to really benefit the broader ML community.

---

### update

After the discussion with the authors, I see that there were some misconceptions from my side that lead to an overly negative review.
On top of the empirical measure for criticality, the effects due to the combination of layernorm and skip connections could be considered a novel (theoretical) contribution.
However, the bias parameter of a linear layer in front of a normalisation layer is typically initialised to all zeros in practice and thus $\sigma_b^2 = 0$. As a result, everywhere criticality should also be achieved with layernorm only.
Overall, I would argue that the paper in its current form still does not do a good job of clearly communicating the insights (and their limitations).
Therefore, I believe that the ideas in this paper would benefit from a more structural rewrite than the rebutttal period would allow.

Therefore, I decided the correctness score from 2 to 3 and my recommendation from 3 to 5.
correctness: 3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.
technical_novelty_and_significance: 2: The contributions are only marginally significant or novel.
empirical_novelty_and_significance: Not applicable
flag_for_ethics_review: NO.
recommendation: 5: marginally below the acceptance threshold

From: ICLR.cc/2023/Conference/Paper3195/Reviewer_sMx1
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary_of_the_paper: The paper investigates criticality conditions in deep networks, using the NNGP (neural network Gaussian process) framework in the infinite width limit, focusing on architectures with LayerNorm and residual connections.
It introduces an estimator for trainability or training stability, the Averaged partial Jacobian norm (APJN), related to the scale (number of layers) over which gradients explode or vanish, which can help selecting the scale of parameter initialization.
It then shows theoretically and empirically that some architectures are always trainable, as they are critical regardless of the initialization scale, confirming previous experimental findings.
strength_and_weaknesses: Strengths
-------------
1. A formal analysis, building on existing work on criticality, and expanding it to popular architecture choices
2. Theoretical confirmation of practical observations regarding critical behavior of networks with various non-linearities, skip-connections, and LayerNorm
3. Practical confirmation on finite-depth and width MLP Mixer networks
4. The presentation is overall clear and easy to follow, previous work and novel contributions are usually well introduced
5. Source code is included

Weaknesses
------------------
1. Some confusion in notations (see below)
2. Some clarity concerns: the definition of APJN always refer to an $(l_0, l)$ pair of layers, but it's not clear what the role of $l_0$ is. It looks like a free parameter one could chose, but what are the consequences of picking one? I get that setting $l_0 = 0$ would have challenges as $N_0$ cannot go to infinity, but could one simply set $l_0 = 1$ without loss of generality?
clarity,_quality,_novelty_and_reproducibility: Clarity
--------
Good overall, the paper is easy enough to read despite the formalism. It is well structured, with the essential information in the main article and additional content in the appendix.

Minor points:
1. Squared terms explicitly repeated in formulas (e.g., derivatives in Eq. 3) would be easier to read with $(...)^2$, as the reader would not need to try and play differences to make sure both expressions are the same.
2. Is the characterization of the architecture-dependent critical exponent $\zeta$ novel, or should there be a citation before Theorem 1.3?
3. The end of Section 4 mentions the tanh activation function and refers to Fig. 2, but tanh is not in Fig. 2.
4. What is the scale of "trainability" in Fig. 2? What does a trainability of 0.5 mean, for instance?
5. Is there supposed to be a difference between $\chi_{\mathcal J}$ and $\chi_J$?

Quality
---------
Good. The theoretical analysis seems solid. The experiments are well designed to show the practical consequences of initialization in different settings, in terms of gradient scale, and of trainability. The bibliography seems complete.

Novelty
-----------
The theoretical analysis is new, and more general than previous works, which are mentioned in Section 1.2, which makes connections with this work.

Reproducibility
--------------------
Excellent. The description of algorithms is clear and not ambiguous, and source code is provided.
summary_of_the_review: The paper provides a theoretical analysis of criticality in different variants of deep neural architectures, proposes a practical surrogate for trainability, and provides experiments supporting its usefulness.
These are useful, well-supported results, that confirm and justify some previous findings, the presentation is clear enough to understand and reproduce.
I think this is an overall good paper, and should be accepted.

**Update after discussions**

After live discussions with AC and other reviewers, I've come to add additional weight to the following negative points:
* clarity still leaves to be desired, the reader has to work too much to extract the main points
* the insight has limited value given the existing best practices: the proposed indicator is not really tested outside of current usual settings.

According to that discussion, I'm updating my rating. I still think the paper should be accepted, and would have rated it a 7 (rather than 6) if the option was available.
correctness: 3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.
technical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
empirical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
flag_for_ethics_review: NO.
recommendation: 6: marginally above the acceptance threshold

From: ICLR.cc/2023/Conference/Paper3195/Reviewer_7FZf
confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
summary_of_the_paper: This paper studies the weight initialization scales/magnitudes for sufficiently wide and sufficiently deep neural networks. Noticing that an inappropriate initialization magnitude may multiplicatively accumulate via forward and backward propagation through the large number of layers leading to non-trainability of the network due to exponential explosion or vanishing, the authors focus on the critical initialization magnitudes in order to keep the trainability. Particularly, the paper shows that, with Layer Normalization, the criticality of initialization holds for a large range of magnitudes (or even all initialization).

strength_and_weaknesses: The main result of the paper is novel, while the other results and discussions are somewhat incremental. 

> As the gradient explosion or vanishing had been one of the major issues in training deep neural networks, many researchers have realized the importance of the magnitude of initialization, and proposed methods to solve or mitigate this issue. Many of these attempts were indeed about finding the critical initialization, such as Kaiming initialization, NTK initialization etc.. For example, in the NTK initialization [1], there is a scaling factor $1/\sqrt{m_l}$ on each layer so that the output of each neuron is at the order of a constant (neither too large nor too small). In addition, in [2], there is an extra factor $c_\sigma$ to normalize each layer in order to make the initialization critical. Given these prior results, the result and discussion in Section 2 about MLP seem repetitive and not new. 

> Theorem 2.7 is more like a definition of $\zeta$ instead of a theorem. It is not a surprise at all that APJN is exponential to the number of layers, as the network is MLP. 

> The major contribution tells an interesting result. Although it has been wide known that both LayerNorm and skip connect mitigate the gradient explosion and vanishing issues, it is still surprising to see that the combination of the two can totally address this issue, in the sense that any initialization does not lead to gradient explosion or vanishing. 

Other comments:

> Below Eq.(2), I cannot quite agree that the partial Jacobian is a random matrix. Once the network is trained away from its initialization, I don’t see how the partial Jacobian is still random.

> In Theorem 2.4, Eq. (10), the NNGP kernel is not in a recursive format. 

References:
[1] Neural Tangent Kernel: Convergence and Generalization in Neural Networks.
[2] Gradient Descent Finds Global Minima of Deep Neural Networks

clarity,_quality,_novelty_and_reproducibility: see comments in the 'strength and weaknesses" section
summary_of_the_review: Initialization scale has long been an essential topic in training deep neural networks. While many aspects/results of the analysis have been known by the machine learning community, I have never seen the conclusion that, with a proper combination of LayerNorm and skip connection, any initialization scale does not result in gradient explosion and vanishing. 
correctness: 3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.
technical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
empirical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
flag_for_ethics_review: NO.
recommendation: 6: marginally above the acceptance threshold