Gradient Explosion and Representation Shrinkage in Infinite Networks
Authors: Adam Klukowski
Venue: ICLR 2022 Submitted (2022)
Keywords: deep learning theory
mean-field approximation

Abstract: We study deep fully-connected neural networks using the mean field formalism,
and carry out a non-perturbative analysis of signal propagation. As a result, we
demonstrate that increasing the depth leads to gradient explosion or to another
undesirable phenomenon we call representation shrinkage. The appearance of at
least one of these problems is not restricted to a specific initialization scheme or
a choice of activation function, but rather is an inherent property of the fully-
connected architecture itself. Additionally, we show that many popular normal-
ization techniques fail to mitigate these problems. Our method can also be applied
to residual networks to guide the choice of initialization variances.

Under review as a conference paper at ICLR 2022
GRADIENT EXPLOSION AND REPRESENTATION SHRINKAGE IN INFINITE NETWORKS
Anonymous authors Paper under double-blind review
ABSTRACT
We study deep fully-connected neural networks with layer normalization using the mean field formalism, and carry out a non non--errtie -onttutuurettive analysis of signal propaga- tion. As result, we demonstrate that increasing the depth leads to gradient explo- sion or to another undesirable phenomenon we call representation shrinkage. The appearance of at least of these problems is not restricted to a specific initial- ization scheme or a choice of activation function, but rather is an inherent property of the fully-connected architecture itself. Additionally, we show that many popu- lar normalization techniques fail to mitigate these problems. Our method can also be applied residual networks to guide the choice of initialization variances.
INTRODUCTION
Deep learning is arguably the most successful modern machine learning method when it comes to modelling complex data. Informally, this is often attributed to hidden representations not being con- strained by the designer's ' knowledge. Empirically, deep neural networks tend to outperform shallow ones, which is explained by them learning a richer hierarchy of representations. There is a body of works making precise a number of aspects of this idea, e.g. showing that expressivity (Montufar ct al. 2014) and disentangling ability (Poole et al. 2016) grow with the network's 's depth.
Because large depth is often desired, much effort was devoted to networks with many layers. There are three main groups of techniques to aid their training: normalization, critical initialization and skip connections. Normalizations ensure the right scale of certain preactivation statisties, for ample magnitude or batch moments. Critical initialization is based on mean -ield analysis of signal propagation. It aims at picking initialization scheme that brings the singular values of the input- output Jacobian close to 1. Skip connections are used in residual blocks, which learn the difference between the target function and the identity. Such construction guarantees that the network can rep- resent the identity. A different approach with similar motivation is to use parameterized activation functions, which are linear for some value of its parameters, so that not only the identity but every linear map is guaranteed to be representable. Our results demonstrate a limitation of normalization methods and critical initialization, and naturally suggest initializing deep networks as linear maps.
A different line of research studied the physics of large networks. It is known that wide but shallow neural networks at initialization simplify to Gaussian processes (GP) (Neal 1996), and their training is shown to be equivalent to kernel methods (Jacot et al.] (018). However, realistic finite networks outperform their infinite limits (Ghorbani et al. 2020). which is explained by the lack of feature learning in large networks (Chizat et al. 2019). The situation with deep networks is qualitatively different the GP regime breaks down when depth and width are large but comparable. Therefore, a bell-behaved infinite depth and width limit would provide a promising starting point for theoretical model capable of capturing feature learning. Suitably scaled infinite residual networks (ResNets) become diffusion processes; one could hope that by appropriately combining different activation functions and normalization methods, one could obtain a similar behaviour in MLPs. We hope that by explaining why this is impossible, our result will help to illuminate most promising future directions in physics of large networks.
The main advice of our work that there are limitations to increasing the depth, even at critical initialization was also noted in Pennington et al. (2017) and Pennington al. (2018), This line of work expresses the input-output Jacobian product of random matrices, and random matrix

---

Under review as a conference paper at ICLR 2022
theory (RMT) to analyse its full spectrum. They find that the maximum and variance of its singular values grow with depth even if the mean remains ((1)
1.1 RELATED WORK
Normalizations:
Rescaling the preactivations to zero mean and unit variance across a mini-batch was introduced in loffe & Szegedy (2015); across all the neurons in single layer was done explicitly in (Lei Ba et al. 2016) and implicitly in (klambauer et al., 2017). For an overview see Huang et al. 2020). Mean-fied field analysis of batch normalization
carried out in Yang al. (019)).
Critical initialization: Mean-field formalism, describing signal propagation in wide networks, was introduced in Poole et al. (2016) based ideas from Neal (1996). It was applied in SShoonnholz al. (217)) to characterize Lyapunov exponents for signal and gradient propagation in fully-connected networks. The impact of activation function initialization this work, as well as Xiao et al. (2019), note polynomially-gue -quulkck---kkuuuuo convergence of correlation with depth at criticality, A form of correlation degeneracy was noted in Daniely et al. (2017) Initialization of convolutional networks was discussed in Xiao al. (2018).
investigated in Hayou al. (2019);
Skip connections: Residual blocks were introduced in He et al. (2016). Decreasing the variance of residual weights with depth was shown to be necessary for non-trivial kemel (Hayou et al. 2021) and good signal propagation with class of activations (Yang & Schoenholz 2017); its utility was
observed empirically in (Wang et al. 2018). Trainable activations, brought closer to linear at the initialization, were reported to help in traning very deep networks in
He al. 2015).
Limiting behaviour of large nets: Breakdown of GP regime was demonstrated in Hanin & Nica (2020). Physics of large networks was studied systematically in Roberts al. (2021). ) The corre- spondence between infinitely deep ResNets and diffusion processes was described in Peluchetti &
Favaro (2020), and generalized to doubly infinite ResNets in Peluuchetti Favaro (2021). different kind of problems with wide and deep ReLU networks was analysed in Hanin & Rolnick (2018).
1.2 CONTRIBUTION
To the best our knowledge, we present the first general non-perturbative result on signal prop- agation in infinitely wide networks. It holds for arbitrary] activation functions, and they vary between layers. Our proof does not rely on the assumption of independence between gradients and weights. On high level, our framework explains the origin of "signal distortion" in fully- connected networks, and explains why they are harder to stack than residual networks. Unlike RMT, our method can handle pairs of data points with non-infinitesimal differences.
We provide precise rrfe between having rich hidden representations and well-behaved gradi- ents in cullecconeetted and residual networks, making rigorous the intuition of competition between these goals. Unlike in previous work, result is non-asymptotic in depth and avoids formu- lation "up to a constant". We hope that improvements to our result will serve as a guidance for architectural choices.
SET-UP
2.1 NOTATION
We denote the inner product and vector norms of u, v E Rd as
d
u.e
U
|u||
Vu.u
For matrix M E Raxm it will be natural to induce the operator norm by the root-mean-squared norm
max n 2 M VER" m ||
| olop
TTchhically activations have be square-integrable with respect Gaussian weight
2

---

Under review as a conference paper at ICLR 2022
Gaussian covariance will be denoted
Cl - B -r 9))
To compare activation functions we will using the Gaussian norm
| = E ( 1)) = [[o|1]
We linearize functions by truncating their Hermite expansion
x(x) do Ō1 for
where hy are normalized Hermite polynomial Let us note that they satisfy p for jointly normal zero-mean I1, x2 of variances and covariance p (for proof, see for example O' Donnell (02T)).
=
2.2
SET-UP
We will focus on an untrained fully-connected neural network of depth and Ni neurons in layer We compose from pointwise non-linearities, affine transforms and layer normalization formally the output of the i-th neuron in layer on data point Fa is given by
Nr Σ
(xa)
(Ia)
"N-ve (m.)
where
(za))
w (1
(222)-1)( (xa)) + b (2.1)
(xa))
VN-1
where weights with E N and biases 6,0) E RgN, have variances Ty Lo. respectively, and 0i IR t is the activation function of /-th layer. We will refer to 200) as preactivations. At the first layer, we (x) x and regard it as an "embedding".
An important role in our analysis will be played by correlation or cosine similarity of the preactiva- tions, defined for two input data points Fa. 2B as
N
1 (xa) (2z) = 9 (x.) D. (xa) (Z (x
PI
(2.2)
When the widths are very large N₁ oc, with Ta fixed and the parameters (W:1) (()) having ij independent entric:f distributed according to the Gaussian initialization scheme), one can use the mean field approximation (Poole et al. 2016) and treat (1) (xx, normal random variable.
Then, by the law of large numbers applied to equation 2.2] the correlation changes as
Pt 1)(p-)) del Llu [[d)p-]]]
(2.3)
21y. to <|o
Letting @ Ek=o k= duhg he be the expansion of the activation function Φ into normalized Hermite polynomials, we can work out the Gaussian covariance and obtain
P(p) =
Σ
(2.4)
Observation 1. From expression 2.4 it is apparent that P is power series with non-negative coef- ficients. Therefore, is non-decreasing and convex [0, 1). This implies existence and uniqueness of an attracting (stable) fixed point 5 of P.
First two are hz(z) = 1,)
Without centering
'Technically, we also need regularity conditions (e.g. finite fourth moment) to ensure that CLT applies.
These hold in most practical scenarios.
"A solution to P(pp)
such that (P(p) " for all in some neighbourhood of prp A sufficient
condition for the latter P' (ptp)

---

Under review as a conference paper at ICLR 2022
This is relevant when dw, Fb: O are the same across layers; then forward propagation affects the correlation by repeated application of P. and Pt approaches Ptp as 1 grows. For Po we have 1- PI (1(1)( pu), so the derivative of P at determines whether the fixed point p is have three regimes of initialization, analogous to
stable unstable. Depending the ratio the phase diagram from Schoenholz et al. (207)):
Ordered initialization for P (1) < 1. Then Pip = 1 is stable and + 1 exponentially quickly with 1. Preactivation perturbations decay, so hidden representations of all data become strongly aligned and gradients vanish.
Critical initialization for (1(1) 1. Then Pro = 1 is stable and - 1 polynomially quickly with I. Norm of preactivation perturbations changes sub-exponentially quickly so the gradients are well-behaved.
Chaotic initialization for (1(1) > 1. Then, the fixed point p = 1 of P is repelling (unstable) and there exists another fixed point Pip < 1 which is attracting. In that case P'(pt) < 1, pt + Ptp exponentially quickly with l. Hidden representations of similar data points pulled far apart, preactivation perturbations grow and gradients explode.
MAIN RESULT
Our main result, theorem quantifies the trade-off between steepness, representation ampleness, non-linearity and depth.
Let (L) be a wide fully-connected network as described in equation 2.1 having lay- ers, using layer normalization and employing activation function oi at layer I. In the limit ny oc we have the following inequality
Theorem 1.
(z(
(1
Pess
0z(0)
op
where (2!1))
B2(0) is the cuputt Jacobian, | I is the error of linearizing the activation function
and the maximal correlation Pmax is defined
Pmax
min
Nr
We delay the complete proof to appendix and in this section just sketch the main ideas. Let us start with the consequences of theorem T and give interpretations for each quantity. Theorem T tells us that at least one of the following undesirable effects will necessarily occur:
The Jacobian norm
(2(2) 1g
which measures the steepness of (L) as a function of (0), opp
is large. This means that small change in the input drastically alters the output and the network becomes overly sensitive to perturbations of data. In Bayesian interpretation, the prior hypothesis favours jagged functions. Then the gradients with respect to the first layer parameters are necessarily large, which means gradient explosion. This happens when
dunt are predominantly small, resembling the chaotic phase.
2. The maximal correlation close to one Pmax 1. This means that at some layer the hid- den representation (1) (z:) of any possible input data point x lies in the cone v.e(9(()) for some vector d. We call this scenario representation shrinkage. It may lead to range of problems. First, the representations do not utilize the whole available space but rather are confined to a small region. Second, a rich representation may be necessary for some problems. In this case getting rid of the shrinkage may be a prerequisite for learning. How- ever, it costs time. We have observed that the alignment of preactivations tends to disappear bayer-by-lay rather than at all the layers simultancously, suggesting that deeper networks require more optimization steps. Third, if the angle between preactivations becomes com- parable the machine precision, then in practice gradient descent may into numerical
Ni
Pmax
4

---

Under review as a conference paper at ICLR 2022
issues. Fourth, suppose that (x) = for any data point If the following block 2(++) ..... ((+1) of layers is reasonably smooth, then for any imput it ever encounters, the block will we well approximated by its linearization at v. Therefore, .1++))...... o((1))((1))) does little processing, but requires the full computational budget to implement. This sce- nario happens when Ste. tend to be large Hordered-phase-lee behavour).
"
3. The linearization error | -ill is small. This restricts the choice of activation functions we can use ones that are close to linear. The layers need be close linear maps. so informally they will not "process the data too much". This forces us to make the MLP into something resembling a ResNet. At zero error the network is completely linear, so forward propagation perfectly preserves norms, covariances and correlations of preactivations. With no correlation distortion the theorem becoming vacuous,
trade-off disappears, which is reflected by
4. The depth I. upper-bounded. Sum of many lower-bounded terms cannot stay upper- bounded, so we cannot stack as many layers as we want.
Sketch of proof. We examine the cosine similarity of output preactivations pr. as a function of cosine similarity of inputs Po- This functional relationship in example network is illustrated figure T For infinitely wide networks, this functional dependence is exactly described by
Pr. Py.
Pi (Pn)
where Pt are taken from equation 2.3 By observation I each Pt is convex. By adding layers we compose more and more P's, and we can show that the convexity "accumulates". Then, the "flatness near 0" and the "stcepness near I" become more pronounced and the transition between them becomes sharper" Therefore, the infimum of PL .... P1 is close to 1, or its derivative at 1 is large. In the former case every possible pair of output preactivations strongly aligned and Pmax R 1; in the latter the Jacobian norm and thus gradients are large.
15
05
44
01
18
48
02
64
25
03
(a) L
(b) L 16
(c) L 128
Figure 1: Output correlation Pr. as function of input correlation Po, for the erf activation. Empirical
(N, 1024) red, infinite-width in black. Blue is upper bound on Pt o P. for details
remark 2 appendix A.
The underlying phenomenon is the pathological behaviour of PL as a function of Po- Every 2())
has correlation with itself, so (p 1) 1. However, we can find an upper-bound (blue on figure [ that tends to Pt. .... Pi pointwise on [0, 1) as L grows.
D
Remark 1. One might ask where does gradient vanishing fit in this framework. In fact, special case of representation shrinkage. This is best illustrated on the simplified example with a.w., oit), Ot independent of Gradient vanishing happens at the ordered initialization, when representation shrinking "proceeds exponentially quickly" with the number of layers, i.e. 1- O One way to think about theorem is that criticality is not sufficient to have a well-behaved network: even though it avoids gradient explosion and vanishing, representation shrinkage still happens "at polynomial rate" (theorem I flloww to deduce the rate at least 1 -Pt O( +)).
Now let us apply our technique residual networks. We consider residual blocks con- sisting of an affine transform, normalization, pointwise activation and linear map. Af- ter adding the residual we normalize the signal. Formally, is described by the equations
'Pictorially, this resembles the behaviour of the graph of the function y = a)æ" t gets large.
When a s these functions have infimum close to 1, otherwise their slope (derivative) at is big.
By observation the infimum on [0, 1] is attained p=
5

---

Under review as a conference paper at ICLR 2022
V
U
y -10)
ymm
y())
40 (1111) ym
V
( (yme 0i (Dran y()
y(t))
20 Vo z(1) 1
21-11
Where 20) = x R the imput. The parameters are U E xm V E Ruxm;; we assume their entries independent with variances 1 and ii) respectively.
Theorem 2. For such residual network in the limit n, my - do we have
0₂ (g(L)
2
(1-Pues)
Pinss
82103
8
1+0 1of
op
where as before
20 (nw (
Pmax
max min ISI I Ba
Nr
In addition to the four undesired phenomena of the fully-connected variant theorem about ResNets permits another alternative: initialization variances a1) of weights V being small. This
situation is not problematic; in the extreme ( 0 the network is initialized the identity.
EXPERIMENTS
4.1 APPLICATIONS THEOREM
We show the trade-off between non-linearity and representation shrinkage on the example of criti- cally initialized leaky ReLU networks. For leaky ReLU network with negative slope and a₁ = 0, the inequality P predicts
1 Pmax
165
1+a
(4.1)
VT
For detailed derivation of the constants appendix We see that adding layers with constant a < (figure 2a) and decreasing a at constant depth (figure 26) decrease 1 Pmax. i.e. cause representation shrinkage at some layer. To avoid this effect, we either need to make L bounded and the network shallow or a ze and keep LReLU, close to linear.
10
Tound ugper bound orpiical
125
1.00
035
Irfebee Lanest found upper bound
10'9
10
orgrical
05
02
(a) (standard RcLU), 1 512
(b) 256,
<
Figure 2: 1 - PL as function of depth L (fig 2a) and negative slope a (fig 2b). Empirical correlations for sample data pairs red; infinite-width predictions in black; smallest correlation (largest 1 Pi) found by gradient descent data (with network parameters fixed) in green; bounds obtained by rearranging inequality 4.1 in blue. Lower means smaller
i.e. stronger shrinkage.
We illustrate the relationships between gradient explosion and representation shrinkage on the ex-
ample of er activation. With critical initialization. the operator norm of the network Jacobian (3c)
Smooth increasing function
[-1, 1] defined erf(x) 3/ S v² dy

---

Under review as a conference paper at ICLR 2022
does not grow with depth, and thus gradients (3e) well-behaved: however, the whole input space gets progressively more squeezed at deeper layers Bap. Theorem predicts that this happens at the rate at least
Pmax
VT.
(4.2)
arrttan v8 8) vs
The compression of the preactivation range does not happen in chaotic initialization. As shown on figure 33b, the correlation of preactivations on any pair of different inputs approaches Pip 1. i.c. the fixed point of the function P defined in 2.3 With the approximation Prixx ~ Pip. theorem predicts
arctan 7.
3)vs
(1 ptp)
VL
(4.3)
0z0
op
As shown in 3d norm of the Jacobian indeed grows, causing large gradients in low layers (figure 3t
Infinite width empirical
10
13
infinite width upper bound
32
empirical
10"3
32'
32
10
layer number
layer number)
(a) Shrinkage, critical
(b) (No) shrinkage, chaotic
largest smallest
singuiar values
singular values
largest
10-1
smallest
10
depth
depth
(c) Jacobian, critical
(d) Jacobian, chaotic
weight
lower bound
blas
weight
gradient norm
gradient norm
bias
-15
100 layer number
layer number
(c) Gradients, critical
(f) Gradients, chaotic
Figure 3: Representation shrinkage, singular values of the Jacobian and gradient propagation in deep
erf networks. Bounds in figures 3a Salad 3f are the predictions of inequalities 4.2 and 4.3 respectively.
4.2 REMOVING REDUNDANT LAYERS
Here we compare the training of full deep network, and a "skimmed" network obtained by replac- ing a block of middle layers with single linear transformation. We hypothesize that since all inputs
7

---

Under review as a conference paper at ICLR 2022
to the middle block lie in small region, it can be replaced with its linearization without significantly affecting the network outputs.
Figure4comparr training of full and skimmed networks MNIST. Full networks has 82 critically- initialized layers, alternating between ReLU and erf; "skimmed" one has 67 layers (a block of 16 is linearized). At initialization, preactivations entering the middle block have correlation at least a 0.97, on any possible pair of inputs; as shown on figure 4a this does not change significantly after training. Figure 44 shows that their performance improves similar rates, suggesting that the full and skimmed models follow similar trajectories. We show Pearson correlation between total displacements of the parameters the linearized block, suggesting that these layers learn similar features; this is smaller (< 0.1) but positive layers closer to the linearized block, suggesting some but not complete overlap.
each layer on figure [4b] It large (up to 0.8) further from
5 DISCUSSION
5.1 THEORETICAL AND PRACTICAL IMPLICATIONS
We have demonstrated a problem with increasing the depth of infinitely wide fully-connected net- works. It holds quite generally and appears to be universal property of the fully-connected ar- chitecture rather than a particular class of activation functions or initialization schemes. It also demonstrates that critical initialization alone does not eliminate all the problems with signal propa- gation.
Our result demonstrates that vanilla MLPs are not the right model for networks of large-but- comparable depth and width. If L. and N, grow simultaneously, then the first block of log I. layers is already problematic and the function plog r. (Po) (as in figure becomes pathological. We cannot avert it by using critical initialization or mixing different layer hyperparameters. This also hints that "weakly nonlinear networks" might make more sensible models in this setting.
Theorem suggests that for ResNet initialization, breaking the symmetry between parameters is more important than specifying good prior. From the theory side, we can deduce that well-be- havedness of deep ResNets requires scaling the weight variance as L agreeing with Hayou al. 2021 From the practical side, our result suggests choosing the variance below the threshold
which theorem becomes equality for acceptable values of [200)) and (2(0)
Pmax-
Hop
In our derivation we assumed that affine layers are followed by layer normalization. This sim- plifies the final result but does not restrict generality in network in which pointwise nonlinearity follows an affine map, we can insert layer normalization between them and rescale the activation function by the average norm of preactivations. This construction requires knowing preactivation norm network without layer normalization. which studied e.g. Poole et al. (2016).
One can note that O(L 1) rate of decay of 1 PL in equations 4. and 4.2 is worse than asymptotic O(L for ReLU and O(L-1) for erf. Moreover, the right-hand side of inequality [ is additive and not multiplicative. This raises the suspicion that our result can be improved. We leave the investigation of sharper bounds for future work.
5.2 REMEDYING STRATEGIES
The proof of theorem shows that to avoid representation shrinkage we somehow need break the convexity of P-functions. Mixing activation functions cannot completely cancel out each other's 's un- desirable behaviour, because each one only "adds convexity". It also highlights a problem common to layer normalization and self-normalizing networks those methods only normalize the diago- nal of the covariance matrix a data batch. This ensures O(1) ((1)-scaling of preactivations, but off- diagonal behaviour may still cause representation shrinkage or gradient explosion. It was showed in Yang et al. (2019) that batch normalization is not a complete solution either it necessarily leads to gradient explosion (which is slowest with linear activations).
8

---

Under review as a conference paper at ICLR 2022
Our result demonstrates limitations of normalization methods and critical initialization in improving the trainability of deep networks, and suggests using trainable nonlinearities, initialized as linear functions at the start of the training (similarly to the idea from He et al. (015
initialized
trained
12
low
middle
high
IOW
middle
high
00
number
number
(a) Shrinkage in the full model before and after training. Empirical in red, infinite- width prediction black
parameter shift nignment
low layers
high layers
weights bi Nases
weights
blases
shif
(b) Alignment of parameter displacement
training
valication
skimmed
28
25
10
3000
90.0
.001
3000
1800
361
3001
1900
training accuracy
Validation on accuracy
full
Full
skimmed
skim med
aks
150
304
(c) Training curves
Figure 4: Full and "skimmed" network on MNIST

---

Under review as a conference paper at ICLR 2022
REFERENCES
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differen- tiable programming.
Advances Neural Information Processing Systems, vol-
ume 32, 2019. URL htttss://proocedinng..eeee http://pss.pppppprr////////////
pdf
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks:
The power of initialization and dual view on expressivity, 2017. URL https:
org/abs/1602. / 05897
Alexander G. de G. Matthews, Mark Rowland, Jiri Hron, Richard E. Turner, and Zoubin Ghahra- mani. Gaussian process behaviour in wide deep neural networks, 2018. WRRltttss://aaxt
Bog/abs/Dae 1804 11271
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? CoRR, abs/2006.13409, 2020. (RRL:ttps:
org/abs/ 2006. 13409
Boris Hanin. Random neural networks in the infinite width limit as gaussian processes, 2021. URL https://arxiv.orgg
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kemel. In International Conference on Learning Representations, 2020. URL /tL.pttpps//peereeviee
net/fornn fttt/EorummiidsJjndddaa
Boris Hanin and David Rolnick. How to start training: The effect of initialization and ar- chitecture. In NeurlPS, pp. 569-579, 2018. URL http::/paperss.aeeee////
Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. On the impact of the activation function on deep neural networks training. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2672-2680, 09-15 Jun 2019. http:pttt:///socceeee
ngs. ml wttpssm.l.p....//emmmmmmm wwwress//9//haae
9a html
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable resnet. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 1324-1332, PMLR, 13-15 Apr 2021. URL https:// proceedings mlr. / hayou2aa.ala.......a.......... html
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), December 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Lei Huang, Jie Qin, Yi Zhou, Fan Zhu, Liu, and Ling Shao. Normalization techniques in training dnns: Methodology, analysis and application. CoRR, abs/2009. 2020. URL https: / wawxiv.org/aba.......o..........
12836
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. ICML, pp. 448 456, 2015. URL URL/http://proceeiin
mlr wwtppppewwsssspsssseee
html
Arthur Jacot, Clément Hongler, and Franck Gabriel. nel: Convergence and generalization in neural
Neural URL http:///apprs..nis..cc////
tangent ker- In NeurIPS,
networks.
pp. 8580-8589. 2018.
10

---

Under review as a conference paper at ICLR 2022
Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-norma neural networks. In NIPS, pp. 972-981, 2017. URL http::/ nips. cc/ wttppcccppppppp........c networks
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization. arXiv e-prints, art. arXiv:: 607,064500 July 2016.
Guido F Montufar, Razvan Pascanu, Kyunghyun Cho, and Yoshua Bengio. On the number of lin- ear regions of deep neural networks. In Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL https: http::::////oceeeigsss.o. newrips cc/ 1200 Paper. pdd.
Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pp. 29-53. Springer, 1996.
Ryan O' O'Donnell. Analysis of boolean functions. CoRR, abs/2105. 2021. URL https:
//axxiv.org/aba.. wtwpp//ppp......oo//ooo 10386
Stefano Peluchetti and Stefano Favaro. Infinitely deep neural networks as diffusion processes. In Proceedings of the Twenty Third International Conference Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 2020.
PMLR, 26-28 Aug
Stefano Peluchetti and Stefano Favaro. Doubly infinite residual neural networks: a diffusion process approach. Journal of Machine Learning Research, 22(175): 2021. URL http: jmlr. htcppaapers/v2//20222
html
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learn- ing through dynamical isometry: theory and practice. Advances in neural information processing systems, 30, 2017.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. The emergence of spectral universality in deep networks. In Amos Storkey and Fernando Perez-Cruz (eds.), Proceedings of the Twenty- First International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pp. 1924 1932. PMLR, 09-11 Apr 2018. URL https:// proceedings mlr.. htt.:eeiiin....r.peeee/////.////
8a html
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohi-Dickstein, and Surya Gan- guli. Exponential expressivity in deep neural networks through transient chaos. In Advances in Neural Information Processing Systems, volume 29. Curran Asso- ciates, Inc., 2016. URL 14851003133 9 011133344422255055555553333 b2e
pdd
Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory. arXiv e-prints, art, 1eXi:22106.10165 June 2021.
Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein, Deep information propagation. In International Conference on Learning Representations, 2017.
Xintao Wang, Ke Yu, Shixiang Wu, Jinjin Gu, Yihao Liu, Chao Dong, Chen Change Loy, Yu Qiao, and Xiaoou Tang. Esrgan: Enhanced super-resolution generative adversarial networks, 2018. URL httph:tt//// http::://axxxv.oorrr/ wwa.rg//bb//8// 00219
Lechao Xiao, Yasaman Bahri, Jascha Sohk-Dickstein, Samuel S. Schoenholz, and Jeffrey Penning- ton. Dynamical isometry and mean field theory of enns: How to train 10, 000-layer vanilla con- volutional neural networks. In /CML, pp. 5389-5398, 2018. URL
m1 wwt.peess/v//xxx/ 1tt/800/xxxll500 html
Lechao Xiao, Jeffrey Pennington, and Samuel S. Schoenholz. Disentangling trainability and gen- cralization in deep learning. CoRR, abs/1912.13055, 2019. URL /tttp://arxiii .rg/ abs/ 1912. 1305%

---

Under review as a conference paper at ICLR 2022
Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. In Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL https: https///poooeeeee.neeeeee wtwprps..c//pppr//1////11 Pap.eeeeeeeecceeeeeaaaaaaeeeeeea
Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Represenia- lions, 2019.
A PROOF DETAILS
A.1 PROOF OF THEOREM
We will analyse the propagation of preactivations for two arbitrary data points Fo Tg.. Following
(..)). Let
equation 2.2. denote the cosine similarity of preactivations at /-th layer as Pt
Ni
Pt describe the effect of (-th layer on correlation, as defined in equation 2.3 Then for layers k < I have
Pt o Pi-1 ..... PR+2 o (R+1(PE)
Here we use a "physical" argument relate Pt with the Jacobian. For mathematically rigorous discussion see section A.1.1 we let one input approach the other 2 + Ta. then
8alL)
sup 23 taa
82toy (ta)
No 2 Iaa xpl
op
On the other hand, this scenario is equivalent to Po + 1, and
N (2 (2.)) (
2-2p1. 2-2po
Pro...PP((e P-po
po (1)().))))(111
No -0.12
Therefore the Jacobian has norm at least
(1) 02(
2
(Pr.....
(A.1)
op
Maximal correlation can be expressed using as
Pmax = max min P.o... Pp())
pE|$1PE|-1.1
Let us replace this quantity with something more convenient. Define ýt = Pi o..... Pi (0) and Pmax maxisI<L Ps. By observation for e [0, we have
! Pt are power series with all coefficients non-negative, so
Pl ..... (p( P o....... Pi(0) P1
Pt .... P1 2 2P 0 ... . P1( P o. Po(p) 2 2p:
Therefore pt 2 2p: 1, and Pmax < 2(1 p.ax).
we restrict our attention to p E [0, 1]. then P1 o Pi takes values in (āt, 1]. so P141 "encounters values only from this inverval". Therefore, let us "crop" all the P's to functions P : [0. 1] 5 by defining
(1()) def P1(A-1 (--)))-))
so that P + (1- = Pt + (1)
1-9t
Then the Pt, P/ are related by affine transforms, domain and codomain of P is [o, 1] and satisfies
(())) = 0, (1) = 1. Then we have
Pt ... Pk+1 (px + (1- (1-px)t) = p: + (1- ( P Pati(t)
(A.2)
and as consequence
(1 o1- P:)(11)
(A.3)
12

---

Under review as a conference paper at ICLR 2022
that Pt is exactly "cropped" Gaussian covariance of When we substitute the formula 2.3 the activation function, i.c.
P((1) Cldi 1+(1
11--elor0 < dz p-
We will now make use the excess convexity e. is a functional that takes convex functions on [0, and returns real numbers. Its definition and all the necessary proofs are in appendix |A.3 All we need to know in this proof is that it satisfies inequalities A.4
A.5 and which follow directly
from its properties and 7
The steepness estimate (property states
P o 2 Pi)
(A.4)
The superadditivity with respect to composition (property) 5
states
o... P.) 2 1=1
e(P))
(A.5)
Finally, applying the "cropping" property yields
e( Pi)
(1-A - - 2 &
(A.6)
This is all we need from to complete the proof. Combining inequalities A.4
A.5 and
A.6
gives
(Pr..... P1)
P.).)) 2e e e oP.......PP Pi)
1- Puss
dd 5 R M
Reexpressing LHS in terms of the Jacobian norm from equation A. and remembering
Pmax
2
Puss and equation A.3 we arrive the desired inequality.
Remark 2. So presented proof omits some details needed construct the upper bound on Pf. P. presented on figure [ (blue line). This bound can be derived by combining equations A.5
and A.6 and unpacking the definition of e(Pr .... Pi) follows.
We have "cropped" the function PL o. ... Pi according to equation A.2
PL c .... P1(1) PL + PL o ..... ... Pi(1)
such that PL ...... P0) 0. Then it is enough to take care of Pt o ..... ....
P.. Equations
A.5
and A.6 provide bound the excess convexily of Pτ o
e oP..oo.. 2 Pasx
bd 1d1-5,13 &
Using the property of excess convexity we can bound Pt a...o P. from above by a piecewise lin- ear function (this gives exactly the bound on figure interpretation is presented appendix A.3
The proof of this property and an intuitive
Remark 3. Having arrived at equation A.3 we could have rewritten it using chain rule and bound each (1) straight away. There are two reasons for using the excess convexity. First, describes Pr.. - PP P, on the whole of [0, 1] and not only in the neighbourhood 1. Knowing that e(P,, .o. P1) is large allows to produce the blue upper bound on figure and thus explain why PL is nearly constant for large range of po. and rapidly jumps to 1 when the data points are very close. Second, it rules out the situation where the pathological behaviour of PL ... a P1 would be restricted to some (1 1) with c -> L grows.
13

---

Under review as a conference paper at ICLR 2022
A.1.1 JUSTIFICATION FOR EQUATION A.
To arrive at equation A.I the swapped the order of limits + and lim sup Here we
give a rigorous treatment without this exchange. This will require some understanding of the excess
convexity e.
Denote E = 1-prxx Σ1=1
| 1l and pick 0 <
2+E Equations A.5 and A.6 together imply
1
that e PL .... ... P1) 2 E.
From property
of excess convexity it follows that
PL o P1 (1 (1+ E1
-
1 E)T
plugging the relation between P1 and Pi from equation
A.2 gives
o P1( 7) S Pmax + (1 Pmnax) (1 + E)7)
=1 (1 Pmax (1 + F.)r
(A.7)
Now, take any Ta. Tig satisfying
|ral| 1|2 = | no: 2ar2z
T. Consider
(b) (so ).r (x()
|| (xa (xa)-s((i (2)) (1
1i 2. -z
.1..... nt. >> 20, this converges to
1- PL ProwooP
1- o (1-r)
Convergence probability was proved moida or ReLU activations can be deduced from Daniely et al. 2017). Remembering equation A.7 this quantity least 1- Pr....ePPPll (((1)) (1 (max (1 + E)
Hanin (2021) while almost sure convergence for sig-
Therefore as BI..... nn + 00 we have
(2) ( (1 pmax) E
(A.8)
|| 2z
By mean value theorem, there exists a point Ir, on the line segment joining Ia, satisfying
(M)
n,, (1.)
@2tay (Iz))
ng 2 =xll
Combining this with the estimate A.8 yields
@2(6) Batoy (xxT)
2
-Puus
(1 Press
d li
(A.9)
2
8
Let us summarize the precise formulation. We have actually demonstrated two very similar state- ments that differ by the exact interpretation of convergence. Using the result of Hanin (002)) we obtain: on every line segment joining &o, (that satisfy |x,, | no and Zand
no
L+E, with probability tending to 1 as 71, + 0o, we can find point x, satisfying the Employing the result of Daniely et al. 20177) the precise statement becomes; with sigmoida or ReLU activation, every line segment joining za Fg (that satisfy ||ra|? || Win and 1 satisfying the inequality A.9
inequality A.9
Ra- > I+E almost surely, for sufficiently large "1,..., UI, we can find point di
)
Yha
Satisfying || bo, | bo 1 < Clllbllx for (common) constant C
14

---

Under review as a conference paper at ICLR 2022
Remark 4. The proof is simpler in the statement theorem we replace the square Jacobian norm on the LHS by its infinite-witd width analogue
lim sup (z) ( Zg 2
Ko(Fai"u) (aag)))))))0)0))2225aaaaa
where K₁ is the NNGP kernel de G. Malthews et al.
2018
((( =
lim 1
1 (xa)2 (zz)
+00
A.2 PROOF OF RESNET VARIANT
The proof follows the strategy in the fully-connected case: write down Pj. "crop" it and recenter to [0, 1J², and apply the excess convexity functional.
We can work out the covariances
(ua)) (
(+l)ll (()))))
Pt-1
(00 () (2e)
PC- + oi [(oilmm-
Pt = P((P(-1) 2 (2.).))))
01 to cld
(A.10)
1+0 1o
Where Pj in equation A. again governs the change correlation after 1-th layer. Introduce Āt + - Similarly as in the fully-connected case, we
Pt o ...o (0) and (t( Pil have
1-Mt
(2(L)
= (1 (P ..... (1)
(A.11)
0z(0) op
Using properties and of excess convexity e from appendix A.3 we can get the bound
PP.... .))()) e .... ... P1) = (P))
o
(A.12)
Now we only need a bound on e( (PP)). Substituting the equation
A.10 to the definition of Pt yields
(t(
1-
(t)
I- 1
t0
20
1-
1-
+o
where
= ||rl - c(ob1/r--1]
g(t)
Eldilà 1+(1-f1 110 e(drl
Now we bound the excess convexity of P.. From property we can deduce a Aee(g)
(1) z I- Pt. 1
And property 7 states
(1 أo 13 [1 1
e(g) 2
24
Which gives
e(P)) 2 20 1) 11 -z lv
2(1-0: (
Finally we use (1 --:-1)1l<1133 obtain
a2 - Fil 2o, ((0)1((e)
e e(P))
Combining this with equations
A.11 A.12 and
Pmaxx
5 2(1 (max) we get
(2)
(1 Pines
agi db 2
8
T+o |o
lop
1=1
15

---

Under review as a conference paper at ICLR 2022
A.3 EXCESS CONVEXITY
Here we define the excess convexity functional and prove its properties. We will be working with non-decreasing convex functions g [0, 1] satisfying g(0) 0, g(1) Note that these condi- tions imply continuity.
Each satisfying these conditions must have an argument E (0, 1) for which g(L = t. This follows from Darboux property of g(1 =)) pictorially, the graph of g must intersect the line x + 1. We then define the excess convexity of g as
e(g) = -2
turns out that despite a seemingly arbitrary definition, excess convexity possesses a number of nice properties:
Ftot
if t 1 - 2
Property (Bounding). Ife(g) 2 C, then g(t) S
(1+0)--C C
2+C
(1+C) lies above the graph of g-
Proof. The condition e(g) 2 C means that the point
1 2+0
Therefore the line segments joining it with 0 and () lie entirely above the graph of
Property 2 (Non-negativity). e(g) 2 0, with equality only for = id(o.1)
Proof. By convexity ( S ((1)) = 1 Since (t) + t is increasing, we must have S and hence 2 2 O. Strict inequality for * id follows from property below.
Property (Coerciveness). have the inequality [|g id p.
e(g)
Ite(g)
Proof. Essentially, this proof boils down to the graph of g lying the shaded area.
If x 2 - then convexity applied to arguments 0, 1 - t, gives
= g(1 t) < x-1++ 9(0) + 1 g(z) = g(x) 2 - -
So Jr - g(x)| < - - - Similarly, for x S 1 t apply convexity to arguments x, 1 - L, 1 and obtain
2 g (x) > max (0, 21-14(1-12
again |x -((x)
T-t for [0, 1 t). This means that
x g(x)l < 1 Ite(g)
e(g)
for all x E [0,
Property (Steepness estimate). If is differentiable then g'(1) 2 1 1 e(g)
Proof. By mean value theorem, there exists E (1 - (1-t,1) satisfying
g'(s) -s 0() g(1 - 1 1 e(())
t
Since g is convex, its derivative is non-decreasing g((1) (())
Property (Superadditivity with respect to composition). For any two functions 91 92 satisfying the conditions of this subsection have e(ga o 91) 2 e(g2) + e(g1)
16

---

Under review as a conference paper at ICLR 2022
Let us first present the pictorial idea behind the proof. We know can find tę with ge - (x) = t.. We examine the points
1-4 92(t1))
P2 = 91 - (1 (1-12)
Pi= =
Pa
1/
By construction, they lie on the graph of g20g1. We intersect the segment Pi P2 with the line x + y 1 to get Q. Since 92 o 9i is convex, Q lies above the graph of 92 91 and hence gives a lower bound on e(g2 o 91).
P
Proof. Now we will make this hand-wavy proof rigorous. We start by a statement about the location of Pt. By convexity 92 we have
92 91 (1 t1) (2((1) S 1-1 92((1 (2) + 1-4162 - 92(()
4112 (A.13)
1-
1 so the weights non-negative. We can get an analogous inequality for P2: using concavity
of 91
we get
91 92 (t2) 9T 12) 2 1-4 gi (t2) + (1- 12 T-
1-11-12 1-41
t12 1-
92 o 91
t2
(A.14)
It is straightforward, if somewhat tedious, to check that the right (i.e. giving rise to something of the form 92 o g1 -t) t) f t) weights to combine inequalities A. 13 A.14 are given by
W1= = (1 (1-1 f1 to (1 t2s 22
(1- 2t, My(1 t2)
U2
[1 11-h-1)](0) (2)
f (1 2ta)
(A.15)
Adding inequalities A.13 with weight w and A. 14 with weight w2 yields
92 gi (1
with w2 f-2;) S wi 92 gi (1) w u2 92 gi (1 (2l2) S 1-
hite + w2l2
After substituting the weights from A.15 this simplifies to
92 91
[
22172
S
11 + 221122
644y
Therefore
e(g2 o og))
2
lil2
-2
- 44= + e(g2)
11712-21,42
Property 6 (Cropping power series). Take a power series g(x) Ei- caat with all coefficients non-negative C₁ 2 O. Let us "crop" it to the rectangle (a, bJ x [g(a), ((() and "recentre", formally, define
g(s) ola)((--)))-s(a) g(h) ((h)-g(a)
so that
g(a+ (b-a)s) gg(a)= (g(b) g(a)) g(s)
then e())12+e(3))
(b-a)²
2k=2 C2
[+e(1))
g1 ((a)
Proof. Lett (+c(5] so that 5(1 - t) = t. By the definition of
lg(a) + (1 ()g - y (la - (1 - ())) (o(b) s(a)) (1 1 - 5(1 1))
(o(b) g(a)) 21)
(o(b) g(a)) e(σ)
(A.16)
2+e(5)
On the other hand, expanding g yields
tg(a) + (1 -))(()) g(ta + (1 - t)b) = Σ
ta + (1 -2)b - (ta + (1 (1-t)))*
k=0
Σ
cabl
1 {()) k +(1-t) - -tt f +11-1) (A.17)
17

---

Under review as a conference paper at ICLR 2022
where the constant & = 0 and linear & = 1 terms cancel. Now, consider the function k + sz (sT + 1 - s)* for some s, E [0, 1]. We will prove that it is non-decreasing on [1. (o). Its derivative with respect to k is
-sx log (sx + 1 1) log sx+1-=
2 x.( 88 + 1 - 8) k- log ! + (sI + 1- - 8) log a0+1-s
(sx + 1 5) )k- x log x + (1 s) 1 log
(sx + +1 log(sT + 1 s)
which is non-negative by Jensen's inequality for x log x. This allows to replace all brackets in A.17 with their values at 2. giving
Σ cudt (()) (1() + (1-0) - (1. + 1-+-1 2 Σ c (())) + (1-1) [cad{(((t) 2=
+ -1-
k=2
00
= t(1 - t) 4)
chb
k=2
Substituting (= @+ and comparing to equation A.16 yields
(o(b) g(a))
c(G) 2+e(3)
1+e(f) (2+e(0))2
(b a)²
cxb
which is equivalent to the desired inequality.
Property (Cropping a Gaussian covariance). Suppose we "crop" a Gaussian covariance {(flo) [r, 1] by defining
g(t) = -(11)]-e111
Then the excess convexity of g is at least
(1 (1-1)² f 71k
(--)1ls- Fl1k
e(g) 2
2(1 iv {(/r)
f
Proof. Applying property 6 with r, b = 1 gives
20(g)
e( I+c(0)
(1-r)²
k
Where f Ex-o fa is the Hermite expansion of f. Notice that the sum is simply Σ Ilf F11,,, which completes the proof of the first inequality.
km
To get the second one we only need to bound the denominator. Convexity of Clf implies that
- -ifi1]-(iflle < (1 r) aelfial Op
From Stein's lemma we can deduce o/ c[[[l]] = c(ʃ' p). so
ac pl
11.
8p
pml
Property 8 (Combining with identity). For p.q 2 0, + q = and function we have
e(p id [o.1] + q g) 2 ge(g) I+pe(g]
Proof. Set g(1 -s) = ss) - -)((1-)) =s.((s.p11))
+ qg(1 t) = t. Write
g(l - 1) - t g - 1) s - g(1 - s) s
and since g(1 -x) x is decreasing, we must have 1 - < 1-s. This allows deduce from convexity
of g that g(1 1) < [= 8) + (0) =(1() Therefore
1-
t u-p+pp= (g 1) q
x1-4)
1-3
18

---

Under review as a conference paper at ICLR 2022
Substituting 1-s [+e(g) and rearranging yields
t S 24ol))+7o(a)
1+pe(a)
which is equivalent to
e(p- (d1,) + +--9) 9) = - 1tpe(g)
ge(g)
B HERMITE EXPANSIONS OF COMMON ACTIVATION FUNCTIONS
Leaky rectified linear unit with negative slope a is defined as I.RelU, (x) = ar + +1 -a) ReI.U(r). and has the following
e|LReLU ap + (1 0)²
+p(x-arocosp)
LRcLU(x)
1-e 1+a V2 2
Therefore
LReLUU LReLU,
2 x-2 (1 a)²
LReLUU
2 lta²
N 2
And the critical initialization requires σb 0. The error function is defined erf(2)) vx So v-
v² dy. induces
((erf)pp) 2 arctan V9-40²
erf(z) v3x
Therefore
|ef -erlll 2
(arctan J/ - 3)
|/erf
I1
( 2 -arclan 3.).
And critical initialization happens when
vB
19


Reviews:
[{"review_id": "UzANH_pexf8", "reviewer": "Reviewer_W3CL", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "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.\n\nUnfortunately, 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.\n\nI 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.\n\nFinally, 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.\n\nOne 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.\n\nOne 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_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.\"", "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.\n\n### After Author Responses\nWhile I appreciate the response and comments of the authors, I stand by my original score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "O8TebFT9umW", "reviewer": "Reviewer_ZVUC", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "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.\n\n1. Similar results have been considered in existing work, e.g. [1]. \n\n2. The proof of the main theorem is not correct (see details below). \n\nIn 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. \n\n--- \n\n[Strength]: the paper is well-written, easy to follow, and with a very clear goal.\n\n[Weakness]: \n(1) There is a large gap in the proof of Theorem 1. \n(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].  \n\n---\n\nIn what follows, I elaborate (1) and (2) since they are related. \n\nThe 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].   \n\n--- \n\n\nSeveral Qs:\n\nQ1:\n\nHow 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. \n\nQ2: 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]. \n\nQ3: I can't find places to explain the legends \"upper bound\",  \"largest found\".\n\nQ4: How does Thm1 imply eq (4.1)? Do you assume the operator norm is bounded by O(1)? \n\n\n\n[1] Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice,  https://arxiv.org/pdf/1711.04735.pdf \n[2] On the Impact of the Activation Function on Deep Neural Networks Training, https://arxiv.org/abs/1902.06853\n[3] Disentangling Trainability and Generalization in Deep Neural Networks, https://arxiv.org/abs/1912.13053\n\nMinors comments: \n1.) What is the domain of the inputs? It seems they are lying in the same sphere, not mentioned in the paper. \n\n\n\n   \n", "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. ", "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. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "MgcBBmdTgL", "reviewer": "Reviewer_Vkqm", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Main Comments\n\n* 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.\n\n* 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.\n\n### Minor Comments\n\n* 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?\n\n* 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_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.", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"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."}, "extra_text": {}}, {"review_id": "MCg4jcrXPne", "reviewer": "Reviewer_vYPS", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "-- Strengths:\n- 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.\n\n-- Main concerns:\n- 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.\n\n- 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.\n\n- 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.\n[Dynamical Isometry and a Mean Field Theory of CNNs: How to Train 10,000-Layer Vanilla Convolutional Neural Networks]\nwhich 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?\n\n- 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.\n\n-- Minor comments/typos:\n- Page 4, 2nd remark: $v$ is not defined in the definition of the cone\n\n- Please add the assumption that $n_l\\rightarrow \\infty$ in Theorem 1.\n\n- 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.\n", "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.", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LgjI0QOHd-G", "reviewer": "Reviewer_PiMb", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "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:\n\n1. In Observation 1, what is the meaning of \"existence and uniqueness of an attracting fixed point of $P$\"? And why cannot $\\rho$ be negative?\n2. On the top of page 4, what's $P'(1)?$ In general, I think the three initialization regimes here require much more explanation.\n3. 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?\n4. 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\"?\n5. 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? \n6. What does inequality 1 refer to in Section 4.1? \n7. 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. \n8. At the end of Section 5.2, what do you mean by \"training should be started from a linear network\"? \n", "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. \n2. 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.\n3. 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. ", "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. \n\n------------------------------------------------------------\n\nThanks for the response. After reading the response and other reviews, I decided to keep my original score. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

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