G RADIENT E XPLOSION AND R EPRESENTATION S HRINKAGE IN I NFINITE N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
We study deep fully-connected neural networks with layer normalization using the mean ﬁeld 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 speciﬁc 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 normalization techniques fail to mitigate these problems.
Our method can also be applied to residual networks to guide the choice of initialization variances
I NTRODUCTION
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 constrained 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 et al., 2014) and disentangling ability (Poole et al., 2016) grow exponentially with the network’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.
Normalization methods ensure the right scale of certain preactivation statistics, for example their magnitude or batch moments.
Critical initialization is based on mean-ﬁeld analysis of signal propagation.
It aims at picking an 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 represent the identity function.
A different approach with similar motivation is to use a 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., 2018).
However, realistic ﬁnite networks outperform their inﬁnite 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 well-behaved inﬁnite depth and width limit would provide a promising starting point for a theoretical model capable of capturing feature learning.
Suitably scaled inﬁnite 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 noted in Pennington et al.
(2017) and Pennington et al.
This line of work expresses the Jacobian as a product of random matrices, and uses random matrix theory to analyse its full spectrum.
They ﬁnd that the maximum and variance of eigenvalues grow with depth even if the mean remains O (1) .
1. R ELATED WORK
Normalizations: Rescaling the preactivations to zero mean and unit variance across a mini-batch was introduced in Ioffe & 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.
Mean-ﬁeld analysis of batch normalization was carried out in Yang et al.
Critical initialization: Mean-ﬁeld formalism, describing signal propagation in wide networks, was introduced in Poole et al.
based on ideas from Neal (1996).
It was applied in Schoenholz et al.
(2017) to characterize Lyapunov exponents for signal and gradient propagation in fully-connected networks.
The impact of activation function on initialization was investigated in Hayou et al.
; this work, as well as Xiao et al.
(2019), note polynomially-quick convergence of correlation with depth at criticality.
Initialization of convolutional networks was discussed in Xiao et al.
Skip connections: Residual blocks were introduced in He et al.
Decreasing the variance of residual weights with depth was shown to be necessary for a non-trivial kernel (Hayou et al., 2021) and good signal propagation with a class of activations (Yang & Schoenholz, 2017).
Trainable activations brought closer to linear at the initialization (He et al., 2015), and decreasing initial variance of residual weights (Wang et al., 2018), were reported to help in training very deep networks.
Limiting behaviour of large nets: Breakdown of GP regime was demonstrated in Hanin & Nica (2020).
Physics of large networks was studied systematically in Roberts et al.
The correspondence between inﬁnitely deep ResNets and diffusion processes was described in Peluchetti &
Favaro (2020), and generalized to doubly inﬁnite ResNets in Peluchetti & Favaro (2021).
A different kind of problems with wide and deep ReLU networks was analysed in Hanin & Rolnick (2018).
1. C ONTRIBUTION
To the best of our knowledge, we present the ﬁrst general non-perturbative result on signal propagation in inﬁnitely wide networks.
It holds for arbitrary 1 activation functions, and they can vary between layers.
Our proof does not rely on the assumption of independence between gradients and weights.
On a 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 approach, our method can handle pairs of data points with non-inﬁnitesimal differences.
We provide a precise trade-off between having rich hidden representations and well-behaved gradients in fully-connected and residual networks, making rigorous the intuition of competition between these two goals.
Unlike in previous work, our result is non-asymptotic in depth and avoids formulation “up to a multiplicative constant”.
We hope that improvements to our result will serve as a guidance for architectural choices.
S ET - UP 2. N OTATION
We denote the inner product and vector norms of u, v ∈ R d as
[Equation]
For a matrix M ∈ R n × m , it will be most natural to deﬁne the operator norm as induced by the root-mean-squared norm, that is
[Equation]
Gaussian covariance will be denoted as
[Equation]
To compare activation functions we will be using the Gaussian norm
[Equation]
We linearize functions by truncating their Hermite expansion
[Equation]
where h k are normalized Hermite polynomials 2 .
Let us note that they satisfy E [ h k ( x 1 ) h k (cid:48) ( x 2 )]
= δ kk (cid:48) ρ k for jointly normal zero-mean x 1 , x 2 of variances 1 and covariance ρ (for a proof, see for example O’Donnell (2021)).
2. S ET - UP
We will focus on an untrained fully-connected neural network of depth L and N l neurons in layer l .
We compose it from pointwise non-linearities, afﬁne transforms and layer normalizations 3 ; formally the output of the i -th neuron in layer l on data point x α is given by
[Equation]
where
[Equation]
1) where weights W ( l ) ∈ R N l × N l − 1 and biases b ( l ) ∈ R N l have variances σ 2( l ) ,w , σ 2( l ) ,b respectively, and φ l : R (cid:74) ⊃ is the activation function of l -th layer.
We will refer to z ( l ) as preactivations .
At the ﬁrst layer, we set φ 1 ( x ) = x and regard it as an “embedding”.
An important role in our analysis will be played by correlation or cosine similarity of the preactivations, deﬁned for two input data points x α , x β as
[Equation]
When the widths are very large N l → ∞ , with x α ﬁxed and the parameters ( W ( l ) ij , b ( l ) i ) having independent entries 4 (e.g. distributed according to the Gaussian initialization scheme), one can use the mean ﬁeld approximation (Poole et al., 2016) and treat z ( l ) i ( x α ) as a normal random variable.
Then, by the law of large numbers applied to equation 2.2, the correlation changes as
[Equation]
Letting φ = (cid:80) ∞ k =
ˆ φ k h
k be the expansion of the activation function φ into normalized Hermite polynomials, we can work out the Gaussian covariance and obtain
[Equation]
Observation 1.
From expression 2.4 it is apparent that P is a power series with non-negative coefﬁcients.
Therefore, it is non-decreasing and convex on [0 , 1] .
This implies existence and uniqueness of an attracting (stable) ﬁxed point 5 of P .
A some neighbourhood of
A sufﬁcient latter
This is relevant when σ w , σ b , φ are the same across layers; then forward propagation affects the correlation by repeated application of P , and ρ l approaches ρ fp as l grows.
For ρ 0 ≈ 1 we have 1 − ρ l ≈ P (cid:48) (1) l (1 − ρ 0 ) , so the derivative of P determines whether the ﬁxed point ρ = 1 is stable or unstable.
Depending on the ratio σ b σ w we have three regimes of initialization, analogous to the phase transition from Schoenholz et al.
• Ordered initialization for P (cid:48) (1) < 1 .
Then ρ fp = 1 is stable and ρ l → 1 exponentially quickly with l .
Preactivation perturbations decay, so hidden representations of all data become strongly aligned and gradients vanish.
• Critical initialization for P (cid:48) (1) = 1 .
Then ρ fp = 1 is stable and ρ l → 1 polynomially quickly with l .
Norm of preactivation perturbations changes sub-exponentially quickly so the gradients are well-behaved.
• Chaotic initialization for P (cid:48) (1) > 1 .
Then, the ﬁxed point ρ = 1 of P is repelling (unstable) and there exists another attracting ﬁxed point ρ fp < 1 .
In that case P (cid:48) ( ρ fp ) < 1 , so ρ l → ρ fp exponentially quickly with l .
Hidden representations of similar data points are far apart, preactivation perturbations grow and gradients explode.
M AIN RESULT
Our main result, theorem 1, quantiﬁes the trade-off between steepness, representation ampleness, non-linearity and depth.
Theorem 1.
Let z ( L ) be a wide fully-connected network as described in equation 2.1: having L layers, using layer normalization and employing activation function φ l at layer l .
In the limit n l → ∞ we have the following inequality
[Equation]
[Table]
where ∂z ( L ) ∂z (0) is the input-output Jacobian, (cid:13)(cid:13) φ − φ (cid:13)(cid:13) 2 N is the error of linearizing the activation function φ and the maximal correlation ρ max is deﬁned as
[Equation]
l
We delay the complete proof to appendix A, and in this section just sketch the main ideas.
Let us start with the consequences of theorem 1, and give interpretations for each quantity.
Theorem 1 tells us that at least one of the following undesirable effects will necessarily occur:
The Jacobian norm (cid:13)(cid:13) (cid:13) ∂z ( L ) ∂z (0) (cid:13)(cid:13) (cid:13) op , which measures the steepness of z ( L ) as a function of z (0) , is large.
This means that a 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 ﬁrst layer parameters are necessarily large, which means gradient explosion.
This happens when σ ( l ) ,b σ ( l ) ,w are predominantly small, resembling the chaotic phase.
The maximal correlation is close to one ρ max ≈ 1 .
This means that at some layer l , the hidden representation z ( l ) ( x ) of any possible input data point x lies in the cone v .
z ( l ) ( x ) N l > ρ
max
for some vector v .
We call this scenario representation shrinkage .
It may lead to a range of problems.
First, the representations do not utilize the whole
available space but rather are conﬁned 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.
However, it costs time.
We have observed that the alignment of preactivations tends to disappear layer-by-layer rather than at all the layers simultaneously, suggesting that deeper networks require more optimization steps.
Third, if the angle between preactivations becomes comparable to the machine precision, then in practice gradient descent may run into numerical
issues.
Fourth, suppose that z ( l ) ( x ) ≈ v for any data point x .
If the following block z ( l + k ) ◦ · · · ◦ z ( l +1) of k layers is reasonably smooth, then for any input it ever encounters, the block will we well approximated by its linearization at v .
Therefore, z ( l + k ) ◦· · ·◦ z ( l +1) does little processing, but requires the full computational budget to implement.
This scenario happens when σ ( l ) ,b σ ( l ) ,w tend to be large (ordered-phase-like behavour).
The linearization error (cid:107) φ l − φ l (cid:107) N is small.
This restricts the choice of activation functions we can use to ones that are close to linear.
The layers need to be close to 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 linear, so propagation perfectly preserves norms, covariances and correlations of preactivations.
With no correlation distortion the gradient-shrinkage trade-off disappears, which is reﬂected by theorembecoming vacuous.
The depth L upper-bounded.
Sum of many lower-bounded terms cannot stay upperbounded, so we cannot stack as many layers as we want.
Sketch of proof.
We examine the cosine similarity of output preactivations ρ L as a function of cosine similarity of inputs ρ 0 .
This functional relationship in an example network is illustrated on ﬁgure 1.
For inﬁnitely wide networks, this functional dependence is exactly described by
[Equation]
where P l are taken from equation 2.3.
By observation 1, each P l is convex.
By adding layers we compose more and more P ’s, and we can show that the convexity “accumulates”.
Then, the “ﬂatness near 0” and the “steepness near 1” become more pronounced and the transition between them becomes sharper 6 .
Therefore, the inﬁmum 7 of P L ◦ · · · ◦ P 1 is close to 1 , or its derivative at 1 is large.
In the former case every possible pair of output preactivations is strongly aligned and ρ max ≈ 1 ; in the latter the Jacobian norm and thus gradients are large.
[Figure]
The underlying phenomenon is the pathological behaviour of ρ L as a function of ρ 0 .
Every z ( L ) has correlation 1 with itself, so ρ L ( ρ 0 = 1)
= 1 .
However, we can ﬁnd an upper-bound (blue on ﬁgure 1) that tends to P L ◦ · · · ◦ P 1 (0) pointwise on [0 , 1) as L grows.
[Equation]
Remark 1.
One might ask where does gradient vanishing ﬁt in this framework.
In fact, it is a special case of representation shrinkage.
This is best illustrated on the simpliﬁed example with σ ( l ) ,W , σ ( l ) ,b , φ l independent of l .
Gradient vanishing happens at the ordered initialization, when representation shrinking “proceeds exponentially quickly” with the number of layers, i.e. 1 − ρ l =
O ( e − λl ) .
One way to think about theorem 1 is that criticality is not sufﬁcient to have a well-behaved network: even though it avoids gradient explosion and vanishing, representation shrinkage still happens “at a polynomial rate” (theorem 1allows to deduce the rate at least 1 − ρ l = O ( l − 12 ) ).
Now let us apply our technique to residual networks.
We consider residual blocks consisting of an afﬁne transform, normalization, pointwise activation and a linear map.
After adding the residual we normalize the signal.
Formally, it is described by the equations
(1 − a ) x α as α gets large
When otherwise big
By observation 1, the inﬁmum on [0 , 1] is attained at ρ =
[Equation]
φ
[Equation]
Where z (0) = x ∈ R n is the input.
The parameters are U
( l ) ∈
R m l × n , V ( l ) ∈ R n × m l ; we assume their entries are independent with variances 1 and σ 2( l ) respectively.
Theorem 2.
For such residual network in the limit n, m l → ∞ we have
[Equation]
where as before
[Equation]
In addition to the four undesired phenomena of the fully-connected variant 1, theorem 2 about
ResNets permits another alternative: initialization variances σ ( l ) of weights V ( l ) being small.
This situation is not problematic; in the extreme case σ ( l ) =
the network is initialized as the identity.
E XPERIMENTS 4. A PPLICATIONS OF THEOREM
We show the trade-off between non-linearity and representation shrinkage on the example of critically initialized leaky ReLU networks.
For a leaky ReLU network with negative slope a and σ b = 0 , the inequality 1 predicts
[Equation]
For detailed derivation of the constants see
appendix
B.
We see that adding layers with constant a < 1 (ﬁgure 2a) and decreasing a at constant depth (ﬁgure 2b) decrease 1 − ρ max , 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 ≈ 1 and keep LReLU a close to linear.
[Figure]
Lower means smaller 1 − ρ , i.e. stronger shrinkage.
We illustrate the relationships between gradient explosion and representation shrinkage on the example of erf 8 activation.
With critical initialization, the operator norm of the network Jacobian (3c) 8 Smooth increasing function R (cid:16) [ − 1 , 1] deﬁned as erf ( x ) = 2 √ π (cid:82) x 0 e − y 2 dy does not grow with depth, and thus gradients (3e) are well-behaved; however, the whole input space gets progressively more squeezed at deeper layers (3a).
Theorem 1 predicts that this happens at the rate at least
[Figure]
[Equation]
[Figure]
The compression of the preactivation range does not happen in chaotic initialization.
As shown on ﬁgure 3b, the correlation of preactivations on any pair of different inputs approaches ρ fp < 1 – the ﬁxed point of the function P deﬁned in 2.3.
Making the approximation ρ max = ρ fp , theorempredicts
[Figure]
As shown in 3d, the norm of the Jacobian indeed grows, which causes large gradients in low layers (ﬁgure 3f).
[Figure]
4. R EMOVING REDUNDANT LAYERS
Here we compare the training of a full deep network, and a “skimmed” network obtained by replacing a block of middle layers with a single linear transformation.
We hypothesize that since all inputs to the middle block lie in a small region, it can be replaced with its linearization without signiﬁcantly affecting the network outputs.
Figure 4 compares training of full and skimmed networks on MNIST.
Full networks has 82 criticallyinitialized layers, alternating between ReLU and erf; “skimmed” one has 67 layers (a block ofis linearized).
At initialization, preactivations entering the middle block have correlation at least ≈ 0 .
97 , on any possible pair of inputs; as shown on ﬁgure 4a, this does not change signiﬁcantly after training.
Figure 4c shows that their performance improves at similar rates, suggesting that the full and skimmed models follow similar trajectories.
We show Pearson correlation between total displacements of the parameters in each layer on ﬁgure 4b.
It is large (up to 0 . 8 ) further from the linearized block, suggesting that these layers learn similar features; this is smaller ( ≈ 0 . 1 ) but positive in layers closer to the linearized block, suggesting some but not complete overlap.
D ISCUSSION 5. T HEORETICAL AND PRACTICAL IMPLICATIONS
We have exhibited a problem with increasing the depth of inﬁnitely wide fully-connected networks.
It holds quite generally and appears to be a universal property of the fully-connected architecture 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 propagation.
Our result demonstrates that vanilla MLPs are not the right model for networks of large-butcomparable depth and width.
If L and N l grow simultaneously, then the ﬁrst block of log L layers is already problematic and the function ρ log L ( ρ 0 ) (as in ﬁgure 1) 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 2 suggests that for ResNet initialization, breaking the symmetry between parameters is more important than specifying a good prior.
From the theory side, we can deduce that well-behavedness of deep ResNets requires scaling the weight variance as L − 1 , agreeing with Hayou et al.
From the practical side, our result suggests choosing the variance below the threshold at which theorem 2 becomes equality for acceptable (cid:13)(cid:13)(cid:13)
[Equation]
op and ρ max .
In our derivation we assumed that all afﬁne layers are followed by layer normalization.
This simpliﬁes the ﬁnal result but does not restrict generality – if we have a network in which pointwise nonlinearity follows an afﬁne map, we can insert layer normalization between them and rescale the activation function by the average magnitude of preactivations.
One can note that O
( L − 12 ) rate of decay of 1 − ρ L in equations 4.1 and 4.2 is worse than asymptotic
O ( L − 2 ) for ReLU and O
( L − 1 ) for erf.
Moreover, the right-hand side of inequality 1 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. R EMEDYING STRATEGIES
The proof of theorem 1 shows that to avoid representation shrinkage we somehow need to break the convexity of P -functions.
Mixing activation functions cannot completely cancel out each other’s undesirable 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 diagonal of the covariance matrix of a data batch.
This ensures O (1) -scaling of preactivations, but offdiagonal behaviour may still cause representation shrinkage or gradient explosion.
It was showed in Yang et al.
that batch normalization is not a complete solution either – it necessarily leads to gradient explosion (which is slowest with linear activations).
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. [Figure]
A P ROOF DETAILS A. P ROOF OF THEOREM
We will analyse the propagation of preactivations for two arbitrary data points x α , x β .
Following equation 2.2, denote the cosine similarity of preactivations at l -th layer as ρ l = z ( l ) ( x α ) .
z ( l ) ( x β ) N l .
Let
P l describe the effect of l -th layer on correlation, as deﬁned in equation 2.3.
Then for layers k < l we have
[Equation]
Here we use a “physical” argument to relate P l with the Jacobian.
For mathematically rigorous discussion see section A.1.1.
If we let one input approach the other x β → x α , then
[Equation]
On the other hand, this scenario is equivalent to ρ 0 → 1 , and
[Equation]
Therefore the Jacobian has norm at least
[Equation]
Maximal correlation can be expressed using P l as
[Equation]
Let us replace this quantity with something more convenient.
Deﬁne ˜ ρ l = P l ◦ · · · ◦ P 1 (0) and ˜ ρ max = max 1 ≤ l ≤ L ˜ ρ l .
By observation 1, P l are power series with all coefﬁcients non-negative, so for ρ ∈ [0 , 1] we have
[Equation]
Therefore ρ l ≥ 2˜
ρ l − 1 , and 1 − ρ max ≤ 2(1 − ˜ ρ max ) .
If we restrict our attention to ρ ∈ [0 , 1] , then P l ◦· · ·◦ P 1 takes values in [˜ ρ l , 1] , so P l +1 “encounters values only from this inverval”.
Therefore, let us “crop” all the P ’s to functions ˜ P : [0 , 1] (cid:74) ⊃ by deﬁning
[Equation]
so that
[Equation]
Then the P l , ˜ P l are related by afﬁne transforms, domain and codomain of ˜ P l is [0 , 1] and it satisﬁes ˜
P l (0) = 0 , ˜
P l (1) = 1 .
Then we have
[Equation]
and as a consequence
[Equation]
When we substitute the formula 2.3, we see that ˜ P l is exactly a “cropped” Gaussian covariance of the activation function, i.e.
[Equation]
We will now make use of the excess convexity e .
It is a functional that takes convex functions on [0 , 1] and returns real numbers.
Its deﬁnition and all the necessary proofs are in appendix A.3.
All we need to know in this proof is that it satisﬁes inequalities A.4, A.5 and A.6, which follow directly from its properties 4, 5 and 7.
The steepness estimate (property 4) states
[Equation]
The superadditivity with respect to composition (property 5) states
[Equation]
Finally, applying the “cropping” property 7 yields
[Equation]
This is all we need from e to complete the proof.
Combining inequalities A.4, A.5 and A.6 gives
[Equation]
Reexpressing LHS in terms of the Jacobian norm from equation A.1 and remembering 1 − ˜ ρ max ≥ 1 − ρ max and equation A.3 we arrive at the desired inequality.
Remark 2.
So presented proof omits some details needed to construct the upper bound on P L ◦ · · · ◦ P 1 presented on ﬁgure 1 (blue line).
This bound can be derived by combining equations A. and A.6, and unpacking the deﬁnition of e ( ˜ P L ◦ · · · ◦ ˜ P 1 ) as follows.
We have “cropped” the function P L ◦ · · · ◦ P 1 according to equation A.
[Equation]
such that ˜ P L ◦ · · · ◦ ˜ P 1 (0) = 0 .
Then it is enough to take care of ˜ P L ◦ · · · ◦ ˜ P 1 .
Equations A. and A.6 provide a bound on the excess convexity of ˜ P L ◦ · · · ◦ ˜ P
[Equation]
Using the property 1of excess convexity we can bound ˜ P L ◦ · · · ◦ ˜ P 1 from above by a piecewise linear function (which is exactly the bound on ﬁgure 1).
The proof of this property and an intuitive interpretation is presented in appendix A.3.
Remark 3.
Having arrived at equation A.3, we could have rewritten it using chain rule and bound each ˜ P (cid:48) l (1) straight away.
There are two reasons for using the excess convexity.
First, it describes
P L ◦· · ·◦ P 1 on the whole of [0 , 1] and not only in the neighbourhood of 1 .
Knowing that e ( ˜ P L ◦· · ·◦ ˜ P 1 ) is large allows us to produce the blue upper bound on ﬁgure 1, and thus explain why ρ L is nearly constant for large range of ρ 0 , and rapidly jumps to 1 when the data points are very close.
Second, it rules out the situation where the pathological behaviour of P L ◦ · · · ◦ P 1 would be restricted to some (1 − (cid:15), 1) with (cid:15) → 0 as L grows.
A.1. J USTIFICATION FOR EQUATION A.
To arrive at equation A.1 we the swapped the order of limits lim n l →∞ and lim sup x β → x α .
Here we give a rigorous treatment without this exchange.
This will require some understanding of the excess convexity e .
Denote E = 1 − ρ max(cid:80)
Ll =1 (cid:107) φ l − φ l (cid:107) 2 N (cid:107) φ (cid:48) l (cid:107) 2 N , and let τ = 1 2+ E .
Equations A.5 and A.6 together imply that e ( ˜ P L ◦ · · · ◦ ˜ P 1 ) ≥ E .
From property 1 of excess convexity e it follows that
[Equation]
plugging the relation between ˜ P l and P l from equation A.2 gives
[Equation]
Now, take any x α , x β satisfying
[Equation]
x
[Equation]
=
[Equation]
= 1 − τ .
Consider
[Equation]
By Hanin (2021), as n 1 , . . . , n L → ∞ , this converges in probability to
[Equation]
remembering equation A.7, this becomes
[Equation]
Therefore with probability tending to 1 as n 1 , . . . , n L → ∞ we have
[Equation]
By mean value theorem, there exists a point x γ on the line segment joining x α , x β satisfying
[Equation]
Combining this with the estimate A.8 yields
[Equation]
The precise statement takes the form: on every line segment joining x α , x β (that satisfy (cid:107) x α (cid:107)
= (cid:107) x β (cid:107) 2 = n 0 and x α . x
β n= 1 − τ ), with probability tending to 1 as n 1 , . .
. , n L → ∞ , we can ﬁnd a point x γ satisfying the inequality A.9.
Remark 4.
This statement holds not only for τ = 1 2+ E but for all 0 < τ ≤ 1 2+ E .
This follows from applying property 1 in exactly the same way.
Proof for τ = 1 2+ E involves slightly simpler expressions.
Remark 5.
The proof is simpler if in the statement of theorem 1 we replace the square of Jacobian norm on the LHS by its inﬁnite-width analogue
[Equation]
where K l is the NNGP kernel (de G. Matthews et al., 2018)
[Equation]
A. P ROOF OF R ES N ET VARIANT
The proof follows the same strategy as in the fully-connected case: we write down P l , “crop” it and recenter to [0 , 1] [Equation] and apply the excess convexity functional [Equation]
We can work out the covariances
[Equation]
Where P l in equation A.10 again governs the change of correlation after l -th layer.
Introduce ˜ ρ l =
P l ◦ · · · ◦ P 1 (0) and ˜ P l ( t ) =
P l (
[Equation]
+
[Equation]
˜
[Equation]
) − ˜
[Equation]
l
[Equation]
l .
Similarly as in the fully-connected case, we have
[Equation]
Using properties 4 and 5 of excess convexity e from appendix A.3 we can get the bound
[Equation]
Now we only need a bound on e ( ˜ P l ) .
Substituting the equation A.10 to the deﬁnition of ˜ P l yields
[Equation]
where
[Equation]
Now we bound the excess convexity of ˜ P l .
From property 8 we can deduce
[Equation]
And property 7 states
[Equation]
Which gives
[Equation]
Finally we use A ≤ (1 − ˜ ρ l − 1 ) (cid:107) φ (cid:48) l (cid:107) 2 N to obtain
[Equation]
Combining this with equations A.11, A.12 and 1 − ρ max ≤ 2(1 − ˜ ρ max ) we get
[Equation]
A. E XCESS CONVEXITY
Here we deﬁne the excess convexity functional and prove its properties.
We will be working with non-decreasing convex functions g : [0 , 1] (cid:74) ⊃ satisfying g (0) = 0 , g (1) = 1 .
Note that these conditions imply continuity.
Each g satisfying these conditions must have an argument t ∈ (0 , 1) for which g (1 − t )
=
t .
This follows from Darboux property of g (1 − x ) − x ; pictorially, the graph of g must intersect the line x + y = 1 .
We then deﬁne the excess convexity of g as
[Equation]
g
[Equation]
t
It turns out that despite a seemingly arbitrary deﬁnition, excess convexity possesses a number of nice properties:
Property 1 (Bounding) .
If e ( g ) ≥ C , then g ( t ) ≤
[Equation]
Proof.
The condition e ( g ) ≥ C means that the point
[Equation] [Equation]
lies above the graph of g .
Therefore the line segments joining it with 0 and
[Equation]
[Equation]
lie entirely above the graph of g .
Property 2 (Non-negativity) .
e ( g ) ≥ 0 , with equality only for g = id [0 , 1]
Proof.
By convexity g (cid:0) 12 (cid:1) ≤ 12 (cid:0) g
(0)+ g (1) (cid:1)
= 12 .
Since g ( t )+ t is increasing, we must have t ≤and hence 1 t − 2 ≥ 0 .
Strict inequality for g (cid:54) =
id [0 , 1] follows from property 3 below.
Property 3 (Coerciveness) .
We have the inequality (cid:107)
[Equation]
≤
[Equation]
Proof.
Essentially, this proof boils down to the graph of g lying in the shaded area.
If x ≥ 1 − t then convexity applied to arguments 0 , 1 − t, x gives
[Equation]
So | x − g ( x ) | ≤ 1 − t 1 − t = 1 − 2 t 1 − t .
Similarly, for x ≤ 1 − t we apply convexity to arguments x, 1 − t, 1 and obtain
[Equation]
so again | x − g
( x ) | ≤ 1 − 2 t 1 − t for x ∈ [0 , 1 − t ] .
This means that
[Equation]
for all x ∈ [0 , 1]
Property 4 (Steepness estimate) .
If g is differentiable then g (cid:48) (1) ≥ 1 + e ( g )
Proof.
By mean value theorem, there exists s ∈ (1 − t, 1) satisfying
[Equation]
Since g is convex, its derivative is non-decreasing so g (cid:48) (1) ≥ g (cid:48)
( s ) .
Property 5 (Superadditivity with respect to composition) .
For any two functions g 1 , g 2 satisfying the conditions of this subsection we have e ( g 2 ◦ g 1 ) ≥ e ( g 2 )
+ e ( g 1 )
Let us ﬁrst present the pictorial idea behind the proof.
We know we can ﬁnd t i with g i (1 − t i )
=
t i .
We examine the points
[Equation]
By construction, they lie on the graph of g 2 ◦ g 1 .
We intersect the segment P 1 Pwith the line x
+ y = 1 to get Q .
Since g 2 ◦ g 1 is convex, Q lies above the graph of g 2 ◦ g 1 and hence gives a lower bound on e ( g 2 ◦ g 1 ) .
[Equation]
Proof.
Now we will make this hand-wavy proof rigorous.
We start by a statement about the location of P 1 .
By convexity of g 2 we have
[Equation]
t i < 12 , so the weights are non-negative.
We can get an analogous inequality for P 2 : using concavity of g − 1 1 we get
[Equation]
It is straightforward, if somewhat tedious, to check that the right (i.e. giving rise to something of the form g 2 ◦ g 1 (1 − t ) ≤ t ) weights to combine inequalities A.13, A.14 are given by
[Equation]
Adding inequalities A.13 with weight w 1 and A.14 with weight w 2 yields
[Equation]
After substituting the weights from A.15 this simpliﬁes to
[Equation]
Therefore
[Equation]
Property 6 (Cropping a power series) .
Take a power series g ( x ) =
(cid:80) ∞ k =
c k x k with all coefﬁcients non-negative c k ≥ 0 .
Let us “crop” it to the rectangle [ a, b ] × [ g ( a ) , g ( b )] and “recentre”; formally, deﬁne
[Equation]
so that
[Equation]
then e (˜ g
[Equation] [Equation]
Proof.
Let t = 1
+ e (˜ g ) so that ˜ g (1 − t ) =
t .
By the deﬁnition of ˜ g
[Equation]
On the other hand, expanding g yields
[Equation]
where the constant k = 0 and linear k = 1 terms cancel.
Now, consider the function k
(cid:55)→ sx k − ( sx +1 − s ) k for some s, x ∈ [0 , 1] .
We will prove that it is non-decreasing on [1 , ∞ ) .
Its derivative with respect to k is
[Equation]
which is non-negative by Jensen’s inequality for x log x .
This allows to replace all brackets in A. with their values at k = 2 , giving
[Equation]
Substituting t = 1 2+ e (˜ g ) and comparing to equation A.16 yields
[Equation]
which is equivalent to the desired inequality.
Property 7 (Cropping a Gaussian covariance) .
Suppose we “crop” a Gaussian covariance C [ f |• ] to [ r, 1] by deﬁning
[Equation]
Then the excess convexity of ˜ g is at least
[Equation]
Proof.
Applying property 6 with a = r, b = 1 gives
[Equation]
Where f = (cid:80) ∞ k =
ˆ f k h k is the Hermite expansion of f .
Notice that the sum is simply (cid:80) ∞
k
=
ˆ f 2 k = (cid:107) f − f (cid:107) 2 N , which completes the proof of the ﬁrst inequality.
To get the second one we only need to bound the denominator.
Convexity of C [ f |• ] implies that
[Equation]
From Stein’s lemma we can deduce
[Equation]
ρ =1 = (cid:107) f (cid:48) (cid:107) 2 N .
Property 8 (Combining with identity) .
For p, q ≥ 0 , p + q = 1 and a function g we have
[Equation]
Proof.
Set g (1 − s )
= s, p (1 − t )
+ qg (1 − t )
=
t .
Write
[Equation]
and since g (1 − x ) − x is decreasing, we must have 1 − t ≤ 1 − s .
This allows deduce from convexity of g that g (1 − t ) ≤ 1 − t
− s g (
[Equation]
s )
[Equation]
t − s 1 − s g (0)
=
s (1 − t ) 1 − s .
Therefore
[Equation]
Substituting s 1 − s = 1 1+ e ( g ) and rearranging yields
[Equation]
which is equivalent to
[Equation]
B H ERMITE EXPANSIONS OF COMMON ACTIVATION FUNCTIONS
Leaky rectiﬁed linear unit with negative slope a is deﬁned as LReLU a ( x ) =
ax +(1 − a ) · ReLU ( x ) , and has the following
[Equation]
Therefore
[Equation]
The error function is deﬁned as erf ( x ) = 2 √ π (cid:82) x
[Equation]
e − y 2 dy .
It induces
[Equation]
Therefore
[Equation]
And critical initialization happens when (cid:16)
[Equation]
b
[Equation]
(cid:17)
[Equation]
=
[Equation]
(cid:16)
[Equation]
− arctan
[Equation]
(cid:17)
[Equation]