The future is log-Gaussian: ResNets and their infinite-depth-and-width limit at initialization

1 Introduction . The characterization of infinite-width dynamics of gradient descent ( GD ) in terms of the so-called Neural Tangent Kernel ( NTK ) [ 1–10 ] represented a major breakthrough in our understanding of deep learning in the large-width regime . Before the identification of infinite-width limits , the theoretical study of deep learning had long been hindered by the apparent analytical intractability of gradient descent and variants acting on the nonconvex objectives used to train neural networks . Despite this progress , evidence suggests that deep neural networks can outperform their infinite-width limits in practice [ 11 ] , particularly when the depth of the network is large . These observations motivate the study of other approximations that may close the gap . Several alternative limits have been proposed . Around the time of the discovery of the NTK limit , mean-field limits were also characterized [ 12–15 ] , and more recently have been linked with the NTK limit [ 16 ] . Yang and Hu [ 17 ] describe a family of infinite-width limits indexed by the scaling limits of initial weight variance , weight rescaling , and learning rates . This family includes both the NTK and mean field limits . One motivation for studying these alternative limits is that they yield a notion of feature learning , which provably does not occur in the NTK limit [ 17 ] . ∗Equal contribution authors . Correspondence : mufan.li @ mail.utoronto.ca ; nicam @ uoguelph.ca ; daniel.roy @ utoronto.ca . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Despite the variety of these limits , one common feature is that the depth of the network ( i.e. , the number of its layers ) is treated as a constant as the width of the network is allowed to grow . Indeed , for fixed width , n , the Gaussian approximations at initialization [ 18–23 ] worsen as the depth , d , increases . While real-world networks are fairly wide , their relative depth is not trivial . Hanin and Nica [ 24 ] were the first to compute an infinite-depth-and-width limit for fully connected networks . While the training dynamics of this limit are still not completely understood , we now know that , in the infinite-depth-and-width limit , the neural tangent kernel is random , and the derivative of the kernel is nonzero at initialization . This means training does not correspond to that of a linear model , like it does in the NTK limit [ 25 , 26 ] . In addition to these theoretical corrections to the Gaussian limit , practitioners have begun to notice that even basic properties of standard neural networks do not line up with those predicted by the infinite-width limit . Perhaps the most basic is that the gradients of the network early in training are not Gaussian , but instead are approximately log-Gaussian [ 27 ] . In fact , one should see logGaussian behaviour agrees with the theoretical predictions of Hanin and Nica [ 24 ] , although there has yet to be a careful empirical comparison made between the precise predictions coming from infinite-depth-and-width models and real world networks . In practice , however , fully connected networks are not often used without architectural modifications . In particular , residual connections , ushered into widespread use after the description of the ResNet architecture [ 28 ] , produced very deep architectures that were practically useful with optimization techniques available at the time . Initialization schemes for ResNets have been studied in the infinitewidth limit [ 19 , 29 ] or with modifications [ 30–32 ] . In this work , we consider the infinite-depth-and-width limit of fully connected architectures with residual connections ( “ Vanilla ResNets ” ) and standard initializations . Here , the analysis is complicated by the effect of skip connections , which introduce interlayer correlations that have a non-negligible effect in this limit . Surprisingly , we observe a counter-intuitive but fundamental phenomenon , whereby these skip connection cause the network to be hypoactivated , meaning that less than half of the neurons are activated on initialization . This fact undermines key assumptions that underlie other infinite-depth-and-width limit studies of architectures without residual connections or with nonstandard modifications , such as post-activation residual connections . Hypoactivation is a roadblock that all theoretical research into standard ResNet architectures must contend with—it is an unavoidable property of the standard architecture and the root of many technical difficulties . In order to sidestep this roadblock , we introduce a conjecture that bounds the size of hypoactivation and effect on interlayer correlation . The conjecture is inspired by empirical evidence from Monte Carlo simulations , as well as several simplified analyses that ignore certain technical difficulties . The conjecture introduces what we believe to be the minimal assumption necessary to allow rigorous theoretical work to proceed . It also defines an important open problem in the study of limits of residual architectures . To demonstrate the utility of the conjecture , we show that it leads to precise predictions . In particular , we prove a limit theorem characterizing the exact marginal distribution of the output at initialization in the infinite-depth-and-width limit , up to order O ( dn−2 ) . Our limit result shows that ResNets have log-Gaussian behaviour on initialization , and like fully connected networks [ 24 ] , the behaviour is determined by the depth-to-width aspect ratio . This corroborates recent empirical observations about deep ResNets [ 27 ] . Since real world networks are finite , the question of how well this approximates finite behaviour is of paramount importance . Based on Monte Carlo simulations , we find excellent agreement between our predictions and finite networks ( see Figure 1 ) . Moreover , for very deep networks ( e.g . d/n = 1 ) the infinite-depth-and-width prediction is extremely different than the infinite-width prediction . More surprisingly , however , is that even at comparatively small depth-towidth ratio ( e.g . d/n = 0.1 ) the two limits are already significantly different . Furthermore , we also observe that the effects due to hypoactivation and interlayer correlation are non-negligible ; these effects are precisely the difference between Vanilla ResNets and so-called “ Balanced Resnets ” in Figure 1 . Perhaps most importantly , we observe that real network outputs exhibit exponentially larger variance than predicted by infinite-width limits . This type of variance at initialization is known to cause exploding and vanishing gradients and other types of training failures [ 33 , 34 ] . Our result also implies that the output neurons of the network are not independent as predicted by the Gaussian infinite-width limit . See Figure 3 . In order to maintain the same skip connections from layer to layer , but render the activation patterns completely independent from layer to layer on initialization , we introduce the Balanced ResNet architecture , where the sign of each neuron ’ s activation function is randomized . We demonstrate that this exponentially decreases the variance of the network output on initialization . Moreover , this independence between neurons also makes the model more amenable to theoretical analysis and opens the door to future understanding of network behaviour . Although it is beyond the scope of this paper , a small preliminary empirical study ( Appendix C ) suggests that standard training regimes are not negatively affected by replacing standard ResNet architectures with Balanced ones . We summarize our main contributions as follows : • We identify and characterize a fundamental property of ResNets that we call hypoactivation : less than half of the ReLU neurons are activated . Based on empirical evidence , we formulate a precise minimal conjecture bounding the effect of hypoactivation that permits us to make precise , rigorous estimates for other properties of ResNets . • We prove a limit theorem which shows that the output of ResNets on initialization exhibits log-Gaussian behaviour with parameterization depending on the depth-to-width ratio d/n . • We provide empirical evidence from Monte Carlo simulations showing our theory provides more accurate predictions for simple properties of finite networks compared to the predictions made by the Gaussian infinite-width limit . • We introduce the Balanced ResNet architecture , which corrects the hypoactivation and variance due to layerwise correlation from Vanilla ResNets . We also prove that the output for this architecture is log-Gaussian on initialization with exponentially lower variance . This simple modification can be applied to any neural network that uses ReLU activations . 2 Main Results . In terms of the notation in Table 1 , a Vanilla ResNet with fully connected first/last layers and d hidden layers of width n is defined by z0 : = 1√ nin W 0x , z ` : = αz ` −1 +λ √ 2 n W ` ϕ+ ( z ` −1 ) for 1 ≤ ` ≤ d , zout : = 1√ n W outzd . ( 1 ) Notation Description Notation Description nin ∈ N Input dimension nout ∈ N Output dimension n ∈ N Hidden layer width d ∈ N Number of hidden layers ( depth ) ϕ+ ( · ) ReLU function ϕ+ ( x ) = max { x , 0 } ϕ− ( · ) “ Domain Flipped ” ReLU ϕ− ( x ) = max { −x , 0 } α ∈ R Skip connection coefficient λ ∈ R+ Feed-forward coefficient x ∈ Rnin Input W 0 ∈ Rnin×n Weight matrix at layer 0 zout ∈ Rnout Network output W out ∈ Rn×nout Weight matrix at final layer . z ` ∈ Rn Neurons ( pre-activation ) for layer 1 ≤ ` ≤ d W ` ∈ Rn×n Weight matrix at layer 1 ≤ ` ≤ d all weights initialized i.i.d . ∼ N ( 0 , 1 ) Table 1 : Notation Note that factors of √ 2n−1 in the hidden layer are equivalent to intializing according to the so-called He initialization [ 35 ] . Other intializations correspond to changing the coefficient λ . This setup is similar to that of “ Stable ResNets ” [ 29 ] , where the infinite-width limit is studied . In the infinite-depth-and-width limit , the intuition that half of the ReLU units are active ( i.e . nonzero ) because of symmetry is surprisingly not correct . We find that the following quantities play an important role . We define the average hypoactivation ( of layer ` ) and the total hypoactivation of the network by h ` : = E [ ∥∥ϕ+ ( ẑ ` ) ∥∥2 ] − 1 2 , htotal : = ∑d ` =1 h ` , ( 2 ) where ẑ ` : = z ` / ∥∥z ` ∥∥ and E means expectation over the choice of random network weights on initialization . The average hypoactivation is a measure of how many ReLU neurons are activated in layer ` ; h ` = 0 indicates roughly half of the neurons are active . Counter-intuitively , we observe that in a vanilla ResNet , h ` is negative and |h ` | = O ( 1/n ) , indicating slightly less than half the neurons are active.2 After compounding over d layers , the total hypoactivation is of order htotal = O ( d/n ) = O ( 1 ) in the infinite-depth-and-width limit . As we will see , this effect has a non-trivial contribution . At the same time , we also find the covariance between the activations of various layers does not vanish in the infinite-depth-and-width limit . This motivates the definition of the total interlayer covariance correction Itotal : = ∑ 1≤ ` 6= ` ′≤dCov ( 2 ∥∥ϕ+ ( ẑ ` ) ∥∥2 , 2‖ϕ+ ( ẑ ` ′ ) ‖2 ) . ( 3 ) As with the hypoactivation , skip connections cause this term to be non-trivial . We formulate Conjecture 5 , which contains a precise encapsulation of the behaviour of ẑ ` we observe . Conjecture 5 ( Informal ) . In expectation , the layers ẑ ` can be approximated by uniform random variables from the sphere up to a relative error of O ( 1/n ) . The conjecture is well supported by Monte-Carlo simulations ( see Figure 4 ) . We provide a more detailed discussion and a precise statement of Conjecture 5 in Section 4 . Assuming the conjecture holds , we prove a limit theorem about the distribution of zout . Informally , this says that zout is approximately a log-Gaussian scalar times an independent Gaussian vector zout ≈ ‖x‖√ nin ( α2 + λ2 ) d 2 exp ( 1 2 N ( −β 2 + 2chtotal , β + c 2Itotal ) ) ~Z , ( 4 ) where ~Z has iid N ( 0 , 1 ) entries , and β and c are defined by β : = 2 n + d n · 5λ 4 + 4α2λ2 ( α2 + λ2 ) 2 , c : = λ2 α2 + λ2 . ( 5 ) The precise statement , including asymptotic error bounds , is as follows . Theorem 1 . For any choice of hyperparameters nin , nout , n , d , α , λ , and every input x , the output zout at initialization has a marginal distribution which can be written in the form zout d = ‖x‖√ nin ( α2 + λ2 ) d 2 exp ( 1 2 G ) ~Z , ( 6 ) 2For a quantity f = f ( n , d ) whose dependence on width and depth may be implicit , we use the notation f = O ( da/nb ) to mean that , for all choice of constants α , λ , r− , r+ > 0 , there exists a constant C > 0 such that |f ( n , d ) | ≤ Cda/nb for all d , n where r− < d/n < r+ . This notation will allow us to state precise limit theorems when d , n→∞ with the ratio d/n converging to a constant . α = λ = 1/ √ 2 , d/n = 1 are fixed and width n is varied on the x-axis . The error bars indicate 95 % confidence interval ( CI ) of the Monte-Carlo simulation , with truncation at 1e−3 for plotting on log-scale . The balanced ResNet predictions fall within the Monte-Carlo CI and so are not plotted on log scale . Since d/n is fixed , Figure 2b indicates the error in the asymptotic prediction of variance ( 7 ) is O ( n−1 ) = O ( dn−2 ) as claimed . On the other hand , the error in mean ( 7 ) is O ( n−2 ) = O ( dn−3 ) ; this is one order smaller than the statement proven in Theorem 1 . A possible explanation is that the sub-leading error term in our approximation for G is mean zero . where ~Z ∈ Rnout is a Gaussian random vector with iid N ( 0 , 1 ) entries , G = G ( n , d , α , λ ) is a random variable which is independent of ~Z and whose distribution does not depend on nin , nout or x . Consider the limit where both the network depth d→∞ and hidden layer width n→∞ in such a way that the ratio d/n converges to a non-zero constant . In this limit , assuming that Conjecture 5 holds , then the random variable G ∈ R has the following asymptotic behaviour : E [ G ] = −β 2 + 2chtotal +O ( d n2 ) , Var [ G ] = β + c2Itotal +O ( d n2 ) , ( 7 ) where β and c are as in ( 5 ) , and moreover G converges in distribution to a Gaussian random variable with mean and variance given by ( 7 ) in this limit . The main ideas of the proof of Theorem 1 is given in Section 5 and the detailed proof is given in Appendix B . We also provide more explicit formulas for htotal , Itotal below . Proposition 2 . Assume Conjecture 5 is true . Then in the same infinite-depth-and-width limit as Theorem 1 , the total hypoactivation htotal and total interlayer covariance Itotal obey htotal = Cα , λ d n +O ( d n2 ) , Itotal = ∑ 1≤ ` 6= ` ′≤d J̄2 ( θ| ` ′− ` | ) − J̄2 ( π − θ| ` ′− ` | ) n +O ( d n2 ) . ( 8 ) Here Cα , λ is a constant depending on α , λ and J̄2 ( θ ) : = J2 ( θ ) /π = 3 sin ( θ ) cos ( θ ) /π + ( 1− θ/π ) ( 1 + 2 cos2 θ ) , where J2 ( θ ) first appeared in [ 36 ] , and θk is such that cos ( θk ) = αk/ ( α2 + λ2 ) k/2 . Remark 3 . The Gaussian infinite-width limit predicts that the marginals of zout should have the form of ( 6 ) withG being identically zero . As zout depends exponentially onG , the infinite-depth-and-width limit predicts the variance is exponentially larger than the infinite-width limit . See Section 3.1 for a detailed discussion and Figure 3 for verification against finite networks . Using Monte Carlo simulations , we estimate the constant Cα , λ . The result of Theorem 1 is then compared against finite networks in Figure 2 . Full proof can be found Appendix B . The idea of using a scaling parameter α = 1/ √ 2 in the skip connections has been noted in empirical papers [ 37 ] and also studied under simplified assumptions on the number of activation in each layer [ 38 ] . Our result shows that by choosing α2 + λ2 = 1 , the prefactor in ( 6 ) does not grow with depth thereby enabling deeper networks to be trained . In the case that α = 0 , λ = 1 , the architecture reduces to a fully connected network . In this case , log-Gaussian behaviour of the network was discovered in Hanin and Nica [ 24 ] . The fully connected case is simpler because the direction vectors ẑ ` are always uniformly distributed on the unit sphere and are independent from layer to layer . This means h ` = 0 and θn = 0 which greatly simplifies the result of Theorem 1 . Furthermore , the proof easily extends to the case where the coefficients α ` , λ ` vary from layer to layer ; see Appendix A.1 for the general statement . Hayou et al . [ 29 ] and Hanin and Rolnick [ 33 ] have studied ResNets where α ` = 1 in every layer , but λ ` is allowed to vary . The prefactor of our result in this case becomes ∏d ` =1 ( 1 2 + λ2 ` ) ≈ exp ( ∑d ` =1 λ 2 ` ) , which is akin to the behaviour found in those papers . Additionally , our result precisely quantifies the log-Gaussian behaviour and its dependence on the sequence λi through the parameters β and c in the infinite-depth-and-width limit .

Summary: This paper aims to deal with the problem that the gaussian process approximation of Neural network at initialization will become worse while the network become deeper due to the error accumulation. This paper proposed a central limit theorem in depth to have a log-gaussian limit of the infinite width and depth limit.  The subject of paper is pretty cool, but the assumption in the paper and conclusion doesn't convinced me the new characterization is right and essential.