Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper studies the operator norm of the input-output jacobian fully-connected networks with layernorm. Under the assumptions that the width approaches infinity, the authors provide a lower bound for the operator norm. Using this result, they argue that for very depth neural network, either this norm is very large or the network's representation power shrinks. 
main_review: Understanding the representation power, trainability of networks (e.g., vanishing/exploding gradients) is an important research direction in deep learning. Inspired by the "mean-field" approach (e.g., Poole), the authors study the lower bound of the operator norm of the input-output Jacobian map in the large width limit and discover a trade-off between representation power and trainability (exploding gradients). Although I find this line of research of great interest, the current paper has two main drawbacks that I couldn't recommend for acceptance.

1. Similar results have been considered in existing work, e.g. [1]. 

2. The proof of the main theorem is not correct (see details below). 

In addition, it seems most of the observations from the current paper have been found in existing work (e.g., Schoenholz and follow-up work). The key difference is that the current work does not require the number of layers to go to infinity. 

--- 

[Strength]: the paper is well-written, easy to follow, and with a very clear goal.

[Weakness]: 
(1) There is a large gap in the proof of Theorem 1. 
(2) Missing discussion of the line of research using random matrix theory to understand the input-output Jacobian [1], which also consider the operator norm of the input-out Jacobian and draws a very similar conclusion, e.g., the squared operator norm must grow linearly with the number of layers; see eq (17) and follow up discussion in [1].  

---

In what follows, I elaborate (1) and (2) since they are related. 

The biggest issue I see in the proof is the equation above (A.1) on page 11. The authors mixed the calculation of finite width networks (on the left of the equation) and infinite width network calculation together (on the right). More precisely, the authors exchange the order of the two limits $\lim_{width\to\infty} $ and $\limsup_{x_\alpha\to x_\beta}$. The exchangeability of the two limits is questionable to me. In the order: $\lim_{width\to\infty} \limsup_{x_\alpha\to_\beta}$, we need to handle a product of random matrices (if we compute the Jacobian). This is indeed a core contribution of [1], who uses free probability theory to compute the whole spectrum of the singular values of the Jacobian (assuming certain free independence of the matrices). If we swap the limits (we shouldn't do this without justification) to $\limsup_{x_\alpha\to x_\beta}\lim_{width\to\infty} $, the problem itself is reduced to computing the derivative of the composed correlation map, which is much simpler. I think these two limits are not unchangeable in general. E.g., using the order $\limsup_{x_\alpha\to x_\beta}\lim_{width\to\infty} $, both critical gaussian and orthogonal initialization give the same answer. But using the order $\lim_{width\to\infty}\limsup_{x_\alpha\to x_\beta} $, gaussian and orthogonal initialization can give different answers, see eq (17) vs (22) in [1].   

--- 


Several Qs:

Q1:

How Theorem 1 leads to the four possible cases after it needs more discussion. In addition, what are the new insights quotient the existing ones from the order-chaotic analysis? It seems: the first case corresponds to the chaotic phase, the second case corresponds to the order phase. The third/fourth cases seem to be a finer analysis of the critical regime. 

Q2: Remark1 the critical initialization. Several works have already identified the issue of the polynomial rate convergence of the correlation to 1 for Relu and smooth functions; see Proposition 1 in [2]; sec B.3. in [3]. 

Q3: I can't find places to explain the legends "upper bound",  "largest found".

Q4: How does Thm1 imply eq (4.1)? Do you assume the operator norm is bounded by O(1)? 



[1] Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice,  https://arxiv.org/pdf/1711.04735.pdf 
[2] On the Impact of the Activation Function on Deep Neural Networks Training, https://arxiv.org/abs/1902.06853
[3] Disentangling Trainability and Generalization in Deep Neural Networks, https://arxiv.org/abs/1912.13053

Minors comments: 
1.) What is the domain of the inputs? It seems they are lying in the same sphere, not mentioned in the paper. 



   

summary_of_the_review: In sum, there seems to be a large gap in the proof of the main theorem. Even this gap has been fixed; it seems that a large portion of the paper's results overlaps with some existing work.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper stablishes a lower bound for the Jacobian of a fully connected neural network based on the maximum correlation between layers and the error in linear approximation of nonlinearities with respect to Gaussian measure. A similar lower bound is derived for residual networks. The authors then mention some remarks about implications of this bound as well as verifying them through experiments.
main_review: -- Strengths:
- The paper provides a general lower bound for norm of Jacobian of fully connected and residual networks. This bound can be used to deduce some implications for training of neural networks.

-- Main concerns:
- The authors only have included layer-wise normalizations in the definition of neural network architectures which seems to simplify the proofs, but it is not clear whether the results are applicable to commonly used networks that do not include such normalizations. The authors mention briefly that that the inclusion of such normalization does not affect the generality of the results, because the (average) normalizing factor can be factored in the nonlinearity, but I do not see how the result would change as this normalizing factor depends on the input and is not fixed for all the inputs. I would increase my score if the results are extended to a general fully connected network as it was not clear to me how it can be done without a fixed normalizing factor.

- In many of the plots, an upper bound is plotted in blue but there is no mention of what it is in the main body. I believe since it is used so often, there should be a mention of it or at least a clear pointer to the equation corresponding to this upper bound in the appendix.

- As the authors have mentioned in the literature review, there are many works that consider signal propagation through wide neural networks of different architectures in the past such as Lechao Xiao et al.
[Dynamical Isometry and a Mean Field Theory of CNNs: How to Train 10,000-Layer Vanilla Convolutional Neural Networks]
which consider convolutional networks and previous works that consider fully connected ones. At least for the case of convolutional networks it seems that very deep networks (e.g. with 10,000 layers) can be trained if initialized at the edge of chaos without any residual blocks or normalization. Is there something specific to convolutional architectures or does the same happen for fully connected networks as well? It seems like your theorems suggest that very deep fully connected networks should not be trainable but is 10,000 not deep enough?

- There are some works that derive a recursive form for computation of neural tangent kernel between two inputs for certain nonlinearities such as ReLU. See e.g. [On the Inductive Bias of Neural Tangent Kernels]. It seems to me that well-behavedness of NTK, for example if NTK is not exploding, should imply that the Jacobian is also well-behaved. It might be good to mention such works as well.

-- Minor comments/typos:
- Page 4, 2nd remark: $v$ is not defined in the definition of the cone

- Please add the assumption that $n_l\rightarrow \infty$ in Theorem 1.

- I believe that some of the effects described after the Theorem are not easy to understand the way they are written. The authors mention that at least one of these effects will necessarily occur. It would be helpful to add a description of when such effect might happen, for example in the second effect, one might say if the norm of the Jacobian is not large that necessarily means that $\rho_{max}$ is close to 1.

summary_of_the_review: The paper provides an interesting lower bound for the Jacobian of a neural network. But the main issue I have is that the neural network considered with normalization seems not to correspond to networks that are commonly used.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: A mean field analysis of deep feedforward networks is given, which demonstrates a tradeoff between gradient stability, expressivity, and the choice of activation functions. This tradeoff is experimentally explored and verified for several network architectures, and high-level insights are given for designing better deep networks.
main_review: ### Main Comments

* Beyond the mean field assumption, the bulk of the analysis in the paper is made without any additional assumptions, and applies to networks of any depth or activation function. Despite this generality, the insights derived from Theorem 1 and the subsequent experiments are quite useful, and represent fundamental tradeoffs inherent to the design of deep learning models.

* Although the main result is somewhat technical, the summarization of the proof method and subsequent takeaways of the result are done in a very reader-friendly manner. This makes the paper much more useful to a general audience than it would have been with a drier presentation.

### Minor Comments

* In Figures 2a and 3a, gradient shrinkage (as measured by 1 - \rho) relative to depth is bounded by a power-law relationship (blue curves), but in both cases the empirical behavior converges to a different power law with faster decay (black curves). The former prediction comes from Theorem 1, but could the authors elaborate on where this latter power law might come from?

* The black curves in Figure 2 are described as infinite depth in the figure caption, but infinite width in the figure legends.
summary_of_the_review: The paper gives a rigorous theoretical analysis that should be of interest to a wide swath of deep learning researchers and practitioners, and is presented in a very accessible manner. The experiments are well-chosen to complement this analysis.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: This paper studies signal propagation questions in wide fully-connected networks with layer norm. Its main result relates the correlation of preactivations as a function of layer to the input-output Jacobian of the network function, showing that the Jacobian is lower bounded by a function of the maximum correlation and the linearization error of activations. This analysis is then used to comment on the connection between the exploding gradient problem, choices of parameter initializations, and a phenomenon that the authors called "representation shrinkage."
main_review: This paper studies fully-connected networks in the infinite-width limit with layer norm applied at each layer before activations. To my knowledge this paper is first to study the signal propagation problem in networks with layer norm. Additionally, by decomposing activations in terms of an expansion of Hermite polynomials, the authors are able to consider the non-perturbative layer-to-layer evolution of the correlation of preactivations for networks consisting of nonlinear activation functions.

Unfortunately, I find the paper's notion of "representation shrinkage" to be not very intuitive, and I don't have a clear understanding of why it's not desirable. On page 4, the authors list a number of reasons for why this may be bad, but they do not provide strong evidence in support of these claims. Additionally, the authors say that it is a problematic that representation shrinkage occurs even at a polynomial rate (cf. "Remark 1"), but I would expect polynomial decay of correlations to be perfectly acceptable for the regimes of depth that are important in practice. Perhaps these claims would be stronger if they were connected -- either empirically or theoretically -- to desirable properties of trained networks, such as generalization error. Further to my point, as the authors discuss in their introduction, the infinite-width analysis will break down anyway for large enough depth. Thus, I think all these results should only be understood for intermediate regimes anyway.

I also think it would be helpful if the authors included in their high-level description of their results -- abstract, introduction, and one-sentence summary -- that they are focusing solely on fully-connected networks with layer norm.

Finally, the authors talk about gradient explosion throughout the paper, and heuristically, of course, there is a connection between preactivation correlation, input-output Jacobian, and gradients. However, I think it would be more clear to analyze the gradients directly, e.g. by studying gradient descent through the NTK formalism. With such an analysis, it becomes clear that there are a number of factors that can affect the gradient, and the preactivation correlations and input-output Jacobian are just part of the story.

One minor question: do we not have to be careful applying Theorem 1 to non-smooth activation functions such as ReLUs? For other analysis of such activation functions that I'm aware of, typically the computation of quantities detailing the correlation of different inputs requires a separate analysis.

One minor comment: the authors mention in their Discussion that their analysis supports searching for lighter models to promote green computing, but since the realistic regime that their analysis applies to already involves heavy models -- extremely wide networks such that the infinite-width approximation applies -- I don't see how this comment is really applicable.
summary_of_the_review: I recommend rejecting this paper as it is current written. In particular, I think there needs to be more clarity around the authors new idea of "representation shrinkage" for the technical result to be appreciated. If possible, it would be nice if the claims about the importance of preventing this were further supported -- either theoretically or empirically -- by an analysis of its effect on training. I also think it would be better to focus directly on gradient updates to make claims about gradient explosion rather than using proxy quantities in terms of preactivation correlations and the input-output Jacobian. Finally, I would like to see some evidence of the paper's claim that critically initialized networks with polynomial decay of correlations are problematic in practice.

### After Author Responses
While I appreciate the response and comments of the authors, I stand by my original score.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: 1. This paper studied the signal propagation in a deep fully-connected neural network using the mean-field formalism, which assumes that the network width goes to infinity and the parameters are i.i.d. Gaussian. Theorem 1 proves that in a non-linear network, increasing the depth leads to either gradient explosion or representation shrinkage. 
2. The theory is verified in the experiments, including the trade-off between non-linearity and representation shrinkage and the trade-off between gradient explosion and representation shrinkage.
3. On MNIST, it's shown that one can replace several non-linear middle layers of a deep network with a single linear transformation, and still gets a similar training trajectory. 
main_review: The tradeoff between gradient explosion and representation shrinkage is an important question, and mean-field formalism is also a nice way to analyze this problem. But I found many parts of this paper very confusing. In particular, I have the following questions:

1. In Observation 1, what is the meaning of "existence and uniqueness of an attracting fixed point of $P$"? And why cannot $\rho$ be negative?
2. On the top of page 4, what's $P'(1)?$ In general, I think the three initialization regimes here require much more explanation.
3. In Theorem 1, how does the linearization error come into the analysis? Why does the tradeoff between gradient explosion and representation shrinkage disappear when the activation is linear?
4. On the top of page 5, why is it that "either the infimum of $P_L...P_1$ is close to 1 or its derivative at 1 is large"?
5. In Remark 1, what's $\kappa$ in $1-\rho_l = O(l^{-\kappa})$? And what's the precise definition of gradient explosion and vanishing here? 
6. What does inequality 1 refer to in Section 4.1? 
7. There is no explanation for Theorem 2. It might be good to compare Theorem 2 and Theorem 1, and also explain how Theorem 2 guides the choice of initialization variance in ResNet. 
8. At the end of Section 5.2, what do you mean by "training should be started from a linear network"? 

summary_of_the_review: I do think proving the tradeoff between gradient explosion and representation shrinkage is an interesting result. But this paper is not well written and many parts are very confusing to me (as pointed out in the main review). Currently, I think this paper is marginally below the acceptance threshold. I will consider raising my score if my questions are well addressed. 

------------------------------------------------------------

Thanks for the response. After reading the response and other reviews, I decided to keep my original score.
review_title: null

Meta-review: This work performs a mean field analysis of a certain class of fully connected networks with and without layer normalization. Theory is provided which successfully predicts when some networks will exhibit either exploding gradients, or "representation shrinkage" which is similar to the extreme ordered phase discussed in prior works on signal propagation. The primary concerns raised by reviewers included, large overlap with prior works on signal propagation, a bug in the proof of the main theorem, lack of clarity, and many assumptions made in the theory which significantly limit the space of architectures for which the theory can be applied. Some of these concerns were addressed in the rebuttal period, notably major flaw in the main theorem was resolved and some concerns on clarity were addressed. However, with the remaining issues (notably overlap with prior work, and overly restrictive assumptions made) a majority of reviewers did not recommend acceptance in the end. The AC agrees with this final decision and recommends the authors look to further expand upon the contributions relative to prior work.
Decision: Reject