Scaling ResNets in the Large-depth Regime

Pierre Marion, Adeline Fermanian, Gérard Biau, Jean-Philippe Vert

Introduction

The most common architectures have 50-150 layers, but ResNets can be trained with depths up to the order of thousand layers (He et al., 2016b). Yet, the training procedure needs to be carefully crafted to avoid vanishing or exploding gradients, particularly as the depth increases. As pointed out by, e.g., Shao et al. (2020), these instabilities are related to a shift in the magnitude of the variance of a signal as it passes through the network. In the original approach of He et al. (2016a), the issue was mitigated by adding a normalization step, called batch normalization (Ioffe and Szegedy, 2015), which rescales the output of each layer via centering and unit variance normalization. However, this normalization stage introduces practical and theoretical difficulties, among which computational overhead and strong dependence on the batch size (see Brock et al., 2021, and the references therein). A widespread alternative to stabilize training in deep models, explored for example by Yang and Schoenholz (2017), Arpit et al. (2019), Zhang et al. (2019b), and De and Smith (2020), is to incorporate a scaling factor αL\alpha_{L} in front of the residual term in (1), yielding the model

There is strong evidence that this scaling factor αL\alpha_{L} should depend on LL, without however any consensus to date on the exact form of this dependence, nor on the mathematical grounding of the approach. Thus, despite progresses on the empirical side, the mathematical forces in action behind the stability of deep ResNets are still poorly understood, although they are key to unlock training at arbitrary depth.

Our goal in the present paper is to take a step forward towards a better theoretical understanding of deep ResNets by providing a thorough probabilistic analysis of the sequence (hk)0⩽k⩽L(h_{k})_{0\leqslant k\leqslant L} at initialization when LL is large, and by leveraging a continuous-time interpretation of model (2) via the so-called neural ordinary differential equation (neural ODE, Chen et al., 2018) paradigm. In a nutshell, our results highlight the intimate connection that exists at initialization between stability of the learning process, the regularity of the weights, and the scaling factor αL\alpha_{L}. We offer in particular a proper mathematical grounding on why and how to choose the parameter αL\alpha_{L} as a function of the depth LL and the distribution of the weights.

2 Our contributions

The optimal parameters of ResNets are learnt by minimizing some empirical risk function via a gradient descent algorithm. As highlighted for example by Yang and Schoenholz (2017), Hanin and Rolnick (2018), and Arpit et al. (2019), a good parameter initialization of this learning phase plays a major role in the quality of the learnt model, in particular to avoid vanishing gradients and deadlock at initialization, or exploding gradients and quick divergence of the model parameters at the beginning of training. Moreover, a good initialization allows the use of larger learning rates, which have been shown to correlate with better generalization (Jastrzkebski et al., 2017). It is thus of great interest to study and understand the role played by scaling of deep ResNets at initialization. This is the context in which we place ourselves in the sequel.

The continuous approach.

As noticed by several authors (Chen et al., 2018; Thorpe and van Gennip, 2018; E et al., 2019), model (2) with a scaling factor αL=\nicefrac1L\alpha_{L}=\nicefrac{{1}}{{L}} (and not \nicefrac1L\nicefrac{{1}}{{\sqrt{L}}}) is formally similar to the discretization of a differential equation. Thus, when LL tends to infinity, the weights and hidden states change continuously with the layer according to the equation

Going further, our third contribution is to exhibit in Section 4 a continuous range of regimes that are controlled by the choice of αL\alpha_{L} (beyond the cases \nicefrac1L\nicefrac{{1}}{{\sqrt{L}}} and \nicefrac1L\nicefrac{{1}}{{L}}) and the distribution of (θk)1⩽k⩽L(\theta_{k})_{1\leqslant k\leqslant L} at initialization, derived from a continuous-time process Θ\Theta with a regularity different from a Brownian motion. More precisely, we show experimentally that there is a strong interplay (with the same three cases—explosion, identity mapping, non-trivial behavior) between the choice of αL\alpha_{L} and the regularity of (θk)1⩽k⩽L(\theta_{k})_{1\leqslant k\leqslant L} as a function of the layer index kk. In addition, empirical evidence suggests that this interplay impacts both the behavior and performance of the networks during training, beyond initialization.

3 Related work

The choice of scaling for ResNets has been discussed in many papers, without however reaching a clear consensus on the form this scaling factor should take. For instance, Hanin and Rolnick (2018) state that stability requires αL⩽\nicefrac1L\alpha_{L}\leqslant\nicefrac{{1}}{{L}}, while Zhang et al. (2019b) show that αL⩽\nicefrac1L\alpha_{L}\leqslant\nicefrac{{1}}{{\sqrt{L}}} is enough to ensure stability. On the other hand, Cohen et al. (2021) claim that the scaling factor observed in practice in trained ResNets is of the form \nicefrac1Lβ\nicefrac{{1}}{{L^{\beta}}} with β≈0.7\beta\approx 0.7. Other authors have proposed more complex choices for αL\alpha_{L} (e.g., Zhang et al., 2019a; Shao et al., 2020). Taking another point of view, De and Smith (2020) observe that batch normalization is empirically equivalent to taking a \nicefrac1L\nicefrac{{1}}{{\sqrt{L}}} normalization factor. Bachlechner et al. (2021) suggest to learn a scaling parameter αk\alpha_{k} that is allowed to vary from one layer to another, whereas, in (4), αL\alpha_{L} is kept constant across layers. These authors observe a great acceleration for training compared to traditional ResNets with no scaling. They also suggest a similar architecture for Transformers and then notice that αk≈\nicefrac1L\alpha_{k}\approx\nicefrac{{1}}{{L}} at the end of training.

Closest to our analysis at initialization are the papers of Arpit et al. (2019) and Zhang et al. (2019b). Arpit et al. (2019) develop a theoretical analysis based on mean field approximation that suggests that a scaling factor αL=\nicefrac1L\alpha_{L}=\nicefrac{{1}}{{\sqrt{L}}} prevents vanishing/exploding gradients at initialization, and provide experimental evidence that this approach is competitive with batch normalization. However, the authors do not provide rigorous mathematical statements for the three different cases αL≪\nicefrac1L\alpha_{L}\ll\nicefrac{{1}}{{\sqrt{L}}}, αL≈\nicefrac1L\alpha_{L}\approx\nicefrac{{1}}{{\sqrt{L}}}, and αL≫\nicefrac1L\alpha_{L}\gg\nicefrac{{1}}{{\sqrt{L}}}, nor do they highlight the connection with the continuous-time interpretation. Interestingly, the idea of exploiting the martingale structure to analyze the magnitude of the hidden states is present in Zhang et al. (2019b), who study the convergence of gradient descent for over-parameterized ResNets with different values of αL\alpha_{L}. Nevertheless, they consider a specific model with Gaussian weights, and only provide asymptotic results when both width and depth tend to infinity.

The connection between the choice of scaling and the continuous-time point of view has previously been noticed by Zhang et al. (2019c), then studied in detail by Cohen et al. (2021). The latter show that, under assumptions on the form of the weights, it is possible to derive limiting (stochastic or ordinary) differential equations for the hidden states. However, they do not discuss the transition between these two regimes, nor do they link differential equations regimes with the stability of the network.

Scaling at initialization

Our goal in this section is to study the effect of the scaling factor αL\alpha_{L} on the stability of ResNets at initialization, assuming that the weights are i.i.d. random variables. We start by making more precise the model and the learning problem introduced in (1).

An important feature of model (4) is that the layer function takes the form of a matrix-vector multiplication, which will prove crucial to make use of concentration results on random matrices. We stress that this setting is standard in practice and that it encompasses many different types of ResNets. It includes for example simple ResNets where g(h,θ)=σ(h)g(h,\theta)=\sigma(h) with σ\sigma the activation function, and the original ResNets from He et al. (2016a), which have

Probabilistic setting at initialization.

It is stressed that the distribution of the parameters are assumed to be independent of the depth, so that all the dependence on LL is captured in the scaling factor αL\alpha_{L}. This model enables us to consider multiple architectures at once, via the function gg. By contrast, some authors formulate the problem of scaling as a choice of the variance at initialization (e.g., Yang and Schoenholz, 2017; Wang et al., 2022), which makes the analysis architecture-dependent. However, for a given architecture, these two approaches are essentially equivalent since Var⁡(αLVk)=αL2Var⁡(Vk)\operatorname*{Var}(\alpha_{L}V_{k})=\alpha_{L}^{2}\operatorname*{Var}(V_{k}).

The quantity ∥hL−h0∥/∥h0∥\|h_{L}-h_{0}\|/\|h_{0}\| carries key information on the behavior of the network at initialization. On the one hand, if ∥hL−h0∥≪∥h0∥\|h_{L}-h_{0}\|\ll\|h_{0}\|, the network is essentially equal to the identity function. On the other hand, if ∥hL−h0∥≫∥h0∥\|h_{L}-h_{0}\|\gg\|h_{0}\|, the output of the network explodes. An intermediate situation is when ∥hL−h0∥≈∥h0∥\|h_{L}-h_{0}\|\approx\|h_{0}\|. In addition, another source of information is provided by the gradients of the hidden states with respect to the empirical risk L\mathscr{L}. If ∥∂L∂h0−∂L∂hL∥≪∥∂L∂hL∥\|\frac{\partial\mathscr{L}}{\partial h_{0}}-\frac{\partial\mathscr{L}}{\partial h_{L}}\|\ll\|\frac{\partial\mathscr{L}}{\partial h_{L}}\|, the gradients do not change as they flow through the network, which means that the exact same information is backpropagated throughout the network. Conversely, if ∥∂L∂h0−∂L∂hL∥≫∥∂L∂hL∥\|\frac{\partial\mathscr{L}}{\partial h_{0}}-\frac{\partial\mathscr{L}}{\partial h_{L}}\|\gg\|\frac{\partial\mathscr{L}}{\partial h_{L}}\|, the gradients explode during backpropagation. By exploiting the martingale structure of (∥hk∥)0⩽k⩽L(\|h_{k}\|)_{0\leqslant k\leqslant L}, as well as state-of-the-art concentration inequalities for random matrices with sub-Gaussian entries, we provide in this section probabilistic bounds on the magnitude of these various quantities.

Assumptions.

The following assumptions will be needed throughout the section: for any 1⩽k⩽L1\leqslant k\leqslant L,

For some s⩾1s\geqslant 1, the entries of dVk\sqrt{d}V_{k} are symmetric i.i.d. s2s^{2} sub-Gaussian random variables, independent of dd and LL, with unit variance.

Assumption (A1)(A_{1}) is mild and satisfied by all initializations used in practice. For example, the classical Glorot initialization (Glorot and Bengio, 2010)—which is the default implementation in the Keras package (Chollet et al., 2015)—takes the entries of Vk+1V_{k+1} as uniform U(−3/d,3/d)\mathcal{U}(-\sqrt{3/d},\sqrt{3/d}) variables. This means that dVk+1\sqrt{d}V_{k+1} is initialized with U(−3,3)\mathcal{U}(-\sqrt{3},\sqrt{3}) random variables, which satisfy (A1)(A_{1}). Other examples include the Gaussian N(0,1/d)\mathcal{N}(0,1/d) initialization of He et al. (2015) and, for example, initialization with Rademacher variables.

The first part of Assumption (A2)(A_{2}) ensures that g(⋅,θk+1)g(\cdot,\theta_{k+1}) is not too far away from being an isometry in expectation. The second part is more technical and, roughly, allows to upper bound the deviations of the norm of g(hk,θk+1)g(h_{k},\theta_{k+1}). Our next Proposition 1 shows that most classical ResNet architectures verify Assumption (A2)(A_{2}). For the sake of readability, these models, together with their parameters, are summarized in Table 1 below.

Let res-1, res-2, and res-3 be the models defined in Table 1. Then

Assumption (A2)(A_{2}) is satisfied for res-1.

Assumption (A2)(A_{2}) is satisfied for res-2 and res-3, as soon as the entries of dWk+1,\sqrt{d}W_{k+1}, 0⩽k⩽L−1,0\leqslant k\leqslant L-1, are symmetric i.i.d. sub-Gaussian random variables, independent of dd and LL, with unit variance.

In the models res-1 and res-2, σ\sigma can be, for instance, taken as the parametric ReLU function, i.e., σ(x)=x++sx−\sigma(x)=x_{+}+sx_{-}, where x+x_{+} (resp. x−x_{-}) denotes the positive (resp. negative) part and the slope s∈[\nicefrac12,1]s\in[\nicefrac{{1}}{{\sqrt{2}}},1] is a parameter of the model. Observe also that res-2 differs from res-3 since the classical ReLU function is defined by ReLU⁡(x)=x+\operatorname*{ReLU}(x)=x_{+} and thus does not satisfy the condition ∣σ(x)∣⩾a∣x∣|\sigma(x)|\geqslant a|x|. Note that there is no bias term in these three models, as this term is commonly initialized to zero, and we are interested in the behavior at initialization.

2 Probabilistic bounds on the norm of the hidden states

The next two propositions describe how the quantity ∥hL−h0∥/∥h0∥\|h_{L}-h_{0}\|/\|h_{0}\| changes as a function of LαL2L\alpha_{L}^{2}. Proposition 2 provides a high-probability bound of interest when LαL2≪1L\alpha_{L}^{2}\ll 1. In this case, we see that, with high probability, the network acts as the identity function, directly mapping h0h_{0} to hLh_{L}. On the other hand, Proposition 3 provides information in the two cases LαL2≫1L\alpha_{L}^{2}\gg 1 and LαL2≈1L\alpha_{L}^{2}\approx 1. When LαL2≫1L\alpha_{L}^{2}\gg 1, the lower bound (i)(i) indicates an explosion with high probability of the norm of the last hidden state. On the other hand, when LαL2≈1L\alpha_{L}^{2}\approx 1, the bounds (i)(i) and (ii)(ii) show that hLh_{L} randomly varies around h0h_{0} with fluctuation sizes bounded from below and above.

Consider a ResNet (4) such that Assumptions (A1)(A_{1}) and (A2)(A_{2}) are satisfied. If LαL2⩽1L\alpha_{L}^{2}\leqslant 1, then, for any δ∈(0,1)\delta\in(0,1), with probability at least 1−δ1-\delta,

Consider a ResNet (4) such that Assumptions (A1)(A_{1}) and (A2)(A_{2}) are satisfied.

Assume that d⩾64d\geqslant 64 and αL2⩽2(Cs4+4C+16s4)d\alpha_{L}^{2}\leqslant\frac{2}{(\sqrt{C}s^{4}+4\sqrt{C}+16s^{4})d}. Then, for any δ∈(0,1)\delta\in(0,1), with probability at least 1−δ1-\delta,

Assume that αL2⩽1C(d+128s4)\alpha_{L}^{2}\leqslant\frac{1}{\sqrt{C}(d+128s^{4})}. Then, for any δ∈(0,1)\delta\in(0,1), with probability at least 1−δ1-\delta,

Note that the assumptions of Proposition 3 on dd and αL\alpha_{L} are mild, since in the learning tasks where deep ResNets are involved, one typically has αL=\nicefrac1Lβ\alpha_{L}=\nicefrac{{1}}{{L^{\beta}}} with β>0\beta>0, d⩾102d\geqslant 10^{2} and L⩾102L\geqslant 10^{2}. Note also that condition (5) is not severe since, when dd and LL are large, it encompasses all reasonable values of δ\delta. Propositions 2 and 3 are interesting in the sense that they provide finite-depth high-probability bounds on the behavior of the hidden states, depending on the magnitude of LαL2L\alpha_{L}^{2}. The results become clearer by letting αL=\nicefrac1Lβ\alpha_{L}=\nicefrac{{1}}{{L^{\beta}}}, with β>0\beta>0, as shown in the following corollary.

Consider a ResNet (4) such that Assumptions (A1)(A_{1}) and (A2)(A_{2}) are satisfied, and let αL=\nicefrac1Lβ\alpha_{L}=\nicefrac{{1}}{{L^{\beta}}}, with β>0\beta>0.

If β<\nicefrac12\beta<\nicefrac{{1}}{{2}} and d⩾9d\geqslant 9, then

If β=\nicefrac12\beta=\nicefrac{{1}}{{2}}, d⩾64d\geqslant 64, L⩾(12Cs4+2C+8s4)d+96Cs4L\geqslant(\frac{1}{2}\sqrt{C}s^{4}+2\sqrt{C}+8s^{4})d+96\sqrt{C}s^{4}, then, for any δ∈(0,1)\delta\in(0,1), with probability at least 1−δ1-\delta,

Corollary 4 highlights three different asymptotic behaviors for ∥hL∥\|h_{L}\|, depending on the values of β\beta. For β>\nicefrac12\beta>\nicefrac{{1}}{{2}}, statement (i)(i) tells that hLh_{L} converges towards h0h_{0} in probability, as LL tends to infinity, which means that the neural network is essentially equivalent to an identity mapping. On the other hand, for β<\nicefrac12\beta<\nicefrac{{1}}{{2}}, the norm of hLh_{L} explodes with high probability. Finally, for the critical value β=\nicefrac12\beta=\nicefrac{{1}}{{2}}, we see that hLh_{L} fluctuates around h0h_{0}, with a fluctuation size independent of LL. Observe that the lower bound in (iii)(iii) is not trivial as soon as exp⁡(\nicefrac38−\nicefrac11dδ)>1\exp(\nicefrac{{3}}{{8}}-\sqrt{\nicefrac{{11}}{{d\delta}}})>1, i.e., d>\nicefrac9964δd>\nicefrac{{99}}{{64\delta}}. The message of Corollary 4 is that the only scaling leading to a non-degenerate distribution at initialization is for β=\nicefrac12\beta=\nicefrac{{1}}{{2}}.

The three statements of Corollary 4 are illustrated in Figure 1. In this experiment, we consider model res-3, a random Gaussian observation xx in dimension nin=64n_{\textnormal{in}}=64, and parameters initialized with a uniform distribution U(−3/d,3/d)\mathcal{U}(-\sqrt{3/d},\sqrt{3/d}). We refer to Appendix E for a detailed setup of all the experiments of the paper. Figure 2(a) shows the empirical distribution of ∥hL∥/∥h0∥\|h_{L}\|/\|h_{0}\| when β=\nicefrac12\beta=\nicefrac{{1}}{{2}} for a large number of realizations. This figure illustrates in particular that our bounds are reasonably sharp, since the bounds indicate that the first quartile of the distribution is larger than 0.870.87 (whereas the first quartile of the empirical histogram is equal to 1.211.21) and the third quartile is less than 2.062.06 (whereas the third quartile of the empirical histogram is equal to 1.341.34). Determining the exact distribution of ∥hL∥/∥h0∥\|h_{L}\|/\|h_{0}\| is an interesting avenue for future research that is beyond the scope of the present article. There is however a strong indication that the ratio follows a log-normal distribution, as confirmed by a normality test on (the log of) the empirical distribution.

In a nutshell, the proofs of Propositions 2 and 3 rest upon controlling of the norm of the hidden states, which obeys the recurrence

The equalities (2.2) and (8) allow deriving without further work bounds in expectation on ∥hL∥\|h_{L}\|, as already observed by Arpit et al. (2019). However, the results we are after are stronger since they involve high-probability bounds. A finer control of the deviations of ∥Vk+1g(hk,θk+1)∥2\|V_{k+1}g(h_{k},\theta_{k+1})\|^{2} and ⟨hk,Vk+1g(hk,θk+1)⟩\langle h_{k},V_{k+1}g(h_{k},\theta_{k+1})\rangle is then needed. This involves concentration inequalities on random matrices with sub-Gaussian entries.

3 Probabilistic bounds on the gradients

Analyzing the behavior of the sequence (pk)0⩽k⩽L(p_{k})_{0\leqslant k\leqslant L} is challenging since, according to the backpropagation (or reverse-mode differentiation) formula, one has

Although the equation looks qualitatively similar to (6), it has the unpleasant feature that ∂g(hk,θk+1)∂h\frac{\partial g(h_{k},\theta_{k+1})}{\partial h} depends on hkh_{k}, hence on θ1,V1,…,θk,Vk\theta_{1},V_{1},\dots,\theta_{k},V_{k}, while pk+1p_{k+1} depends on θk+2,Vk+2,…,\theta_{k+2},V_{k+2},\dots, θL−1,VL−1\theta_{L-1},V_{L-1}. This forbids applying directly the same proof techniques as for the hidden states. Therefore, to extract useful information from this recurrence equation, one needs to characterize the dependence of the distribution of ∂g(hk,θk+1)∂h\frac{\partial g(h_{k},\theta_{k+1})}{\partial h} with respect to hkh_{k}. To do so, it is sometimes assumed that these two quantities are independent (see, e.g., Yang and Schoenholz, 2017). However, assuming independence remains a strong requirement, which is not verified for many network architectures (for example model res-1). We tackle the problem from a different point of view and propose an alternative approach based on forward-mode differentiation, valid under a much weaker assumption. The cost we pay is that we obtain results in expectation and not in high probability.

Identity (9), which is similar to (4), expresses qk+1(z)q_{k+1}(z) as a function of qk(z)q_{k}(z), and therefore respects the flow of information. Next, assuming that zz is random with a Gaussian distribution, it is possible to express one of our quantities of interest, ∥p0∥/∥pL∥\|p_{0}\|/\|p_{L}\|, as a function of the last vector qL(z)q_{L}(z). Indeed,

In summary, the recurrence (9) allows us to derive bounds on the norm of qL(z)q_{L}(z), which can then transfer to ∥p0∥/∥pL∥\|p_{0}\|/\|p_{L}\| via (10). For this, it is necessary to make the following assumption on the ratio pL/∥pL∥p_{L}/\|p_{L}\|:

It is a mild assumption, which is verified for instance if nout=1n_{\textnormal{out}}=1 with squared error (for regression) or cross-entropy (for binary classification). In these cases, pL/∥pL∥=\nicefracB⊤∥B∥Fp_{L}/\|p_{L}\|=\nicefrac{{B^{\top}}}{{\|B\|_{F}}}, where ∥⋅∥F\|\cdot\|_{F} is the Frobenius norm and BB is the weight matrix of the last layer. We finally need the following assumption, which is the equivalent of Assumption (A2)(A_{2}) for the gradients.

Assumption (A4)(A_{4}) is satisfied by all the standard architectures listed in Table 1, as shown by the next proposition.

Let res-1, res-2, and res-3 be the models defined in Table 1. Assume that (A1)(A_{1}) is satisfied and σ\sigma is almost everywhere differentiable, with a⩽σ′⩽ba\leqslant\sigma^{\prime}\leqslant b. Then

Assumption (A4)(A_{4}) is satisfied for res-1.

Assumption (A4)(A_{4}) is satisfied for res-2 and res-3, when the entries of dWk+1,0⩽k⩽L−1,\sqrt{d}W_{k+1},0\leqslant k\leqslant L-1, are symmetric i.i.d. random variables, independent of dd and LL, with unit variance.

The next two propositions are the counterparts of Proposition 2 and Proposition 3 for the gradient dynamics.

Consider a ResNet (4) such that Assumptions (A1)(A_{1})-(A4)(A_{4}) are satisfied. If LαL2⩽1L\alpha_{L}^{2}\leqslant 1, then, for any δ∈(0,1)\delta\in(0,1), with probability at least 1−δ1-\delta,

Consider a ResNet (4) such that Assumptions (A1)(A_{1})-(A4)(A_{4}) are satisfied. Then

A simple corollary of the propositions above is as follows.

Consider a ResNet (4) such that Assumptions (A1)(A_{1})-(A4)(A_{4}) are satisfied, and take αL=\nicefrac1Lβ\alpha_{L}=\nicefrac{{1}}{{L^{\beta}}}, with β>0\beta>0. Then

Corollary 8 is illustrated in Figure 3. The experimental protocol is the same as in Figure 1, but we now track p0p_{0} and pLp_{L}, the gradients of the loss L\mathscr{L} with respect to the first and the last hidden states. In accordance with our results, when β>\nicefrac12\beta>\nicefrac{{1}}{{2}}, the gradient remains the same from one layer to another (left plot). On the other hand, the middle plot clearly shows that when β<\nicefrac12\beta<\nicefrac{{1}}{{2}} the gradient explodes. Once again, the case β=\nicefrac12\beta=\nicefrac{{1}}{{2}} (right plot) is the only one for which the distribution of gradients at initialization is non-trivial. Figure 2(b) illustrates that the empirical distribution of gradients in this case also seems to be log-normal.

In summary, this and the previous subsection both point towards the same conclusion: there are three different cases, depending on the value of β\beta—explosion when β<\nicefrac12\beta<\nicefrac{{1}}{{2}}, non-degenerate limit when β=\nicefrac12\beta=\nicefrac{{1}}{{2}}, and identity when β>\nicefrac12\beta>\nicefrac{{1}}{{2}}. In the explosion case, it is well known that the network cannot be trained (Yang and Schoenholz, 2017). The theory thus points out that the value \nicefrac12\nicefrac{{1}}{{2}} plays a pivotal role. Remarkably, this value has a specific interpretation in the continuous-time point of view of ResNets, in terms of SDE. This is the topic that we address in the next section.

Scaling in the continuous-time setting

Starting with the discrete ResNet (4), it is tempting to let LL go to infinity and consider the network as the discretization of a differential equation where the layer index k∈{0,…,L}k\in\{0,\dots,L\} is replaced by the time index t∈t\in. This interpretation of deep neural networks has been popularized by Chen et al. (2018) and is often referred to as the neural ODE paradigm. Notice that this setting is different from the so-called mean-field analysis, where the width of the network is assumed to be infinite beforehand. In our setting, the width dd is assumed to be finite and fixed.

One of the main messages of Section 2 is that the standard initialization with i.i.d. parameters leads to a non-degenerate model for large values of LL only if LαL2≈1L\alpha_{L}^{2}\approx 1 (Propositions 2 and 3), or, equivalently, if β=\nicefrac12\beta=\nicefrac{{1}}{{2}} when αL=\nicefrac1Lβ\alpha_{L}=\nicefrac{{1}}{{L^{\beta}}} (Corollary 4). Remarkably, in the continuous-time limit, this regime corresponds to the discretization of an SDE. Indeed, consider for simplicity the (discrete) ResNet model res-1

where the entries of Vk+1V_{k+1} are assumed to be i.i.d. N(0,\nicefrac2d)\mathcal{N}(0,\nicefrac{{2}}{{d}}). Recall the following definition:

A one-dimensional Brownian motion (Bt)t∈(B_{t})_{t\in} is a continuous-time stochastic process with B0=0B_{0}=0, almost surely continuous, with independent increments, and such that for any 0⩽s<t⩽10\leqslant s<t\leqslant 1, Bt−Bs∼N(0,t−s)B_{t}-B_{s}\sim\mathcal{N}(0,t-s).

and the increments for different values of (i,j,k)(i,j,k) are independent. As a consequence, the recurrence (11) is equivalent in distribution to the recurrence

(Note that this is true because Vk+1V_{k+1} has the same distribution as Vk+1⊤V_{k+1}^{\top}.) We recognize the Euler-Maruyama discretization (Kloeden and Platen, 1992) on the {k/L,0⩽k⩽L}\{k/L,0\leqslant k\leqslant L\} mesh of the SDE

where the output of the network is now a function of the final value of HH, that is, H1H_{1}. The link between the discrete ResNet (11) and the SDE (12) is formalized in the next proposition.

Consider the res-1 model, where the entries of Vk+1V_{k+1} are i.i.d. Gaussian N(0,\nicefrac2d)\mathcal{N}(0,\nicefrac{{2}}{{d}}) random variables. Assume that the activation function σ\sigma is Lipschitz continuous. Then the SDE (12) has a unique solution HH and, for any 0⩽k⩽L0\leqslant k\leqslant L,

Notice that the requirement that σ\sigma is Lipschitz continuous is satisfied by most classical activation functions, including ReLU. This proposition is interesting for several reasons. First, the scaling β=\nicefrac12\beta=\nicefrac{{1}}{{2}}, which is exactly the one that yields a non-trivial dynamics at initialization, corresponds in the continuous world to a remarkably ‘simple’ model of diffusion. This shows that very deep neural networks properly initialized with i.i.d. weights are equivalent to solutions of SDE. This analogy opens interesting perspectives for training deep networks using automatic differentiation for solutions of neural SDE (Li et al., 2020).

Second, we stress that the emergence of an SDE instead of an ODE carries an important message. Several authors (including, e.g., Thorpe and van Gennip, 2018) have shown that, under appropriate assumptions, a deep ResNet converges in the large depth limit to an ODE and not an SDE. The reason why we obtain an SDE here is intrinsically connected with the choice of i.i.d. initialization for the weights, which makes a Brownian motion appear at the limit, as highlighted above. In other words, the i.i.d. initialization, the choice β=\nicefrac12\beta=\nicefrac{{1}}{{2}} (the relevant critical value exhibited in Section 2), and the emergence of an SDE are intimately linked together. On the other hand, the case β=1\beta=1 matches with an ODE if the initialization is not i.i.d., as we will see in Subsection 3.2.

Finally, we point out that Proposition 10 states the convergence of a ResNet towards an SDE for the basic architecture res-1 and for Gaussian initialization. The extension to more general settings is an interesting direction of research, although clearly beyond the scope of the present paper (see, e.g., Peluchetti and Favaro, 2020, and Cohen et al., 2021, for results in this direction).

2 Scaling in the neural ODE setting

The basic message of our Proposition 10 is that an i.i.d. initialization, together with β=\nicefrac12\beta=\nicefrac{{1}}{{2}}, leads to an SDE rather than an ODE. A natural question is then whether a different choice of weight distributions (at initialization) and scaling can lead to a classical neural ODE.

where Vk=Vk/LV_{k}=\mathscr{V}_{k/L} and θk=Θk/L\theta_{k}=\Theta_{k/L}. Of course, it is still possible to consider (Vk)1⩽k⩽L(V_{k})_{1\leqslant k\leqslant L} (resp. (θk)1⩽k⩽L(\theta_{k})_{1\leqslant k\leqslant L}) as random variables, by letting (Vt)t∈(\mathscr{V}_{t})_{t\in} (resp. (Θt)t∈(\Theta_{t})_{t\in}) be a continuous-time stochastic process. In this model, we shall need the following assumption:

For any 0⩽k⩽L−10\leqslant k\leqslant L-1, one has Vk=Vk/LV_{k}=\mathscr{V}_{k/L} and θk=Θk/L\theta_{k}=\Theta_{k/L}, where the stochastic processes V\mathscr{V} and Θ\Theta are almost surely Lipschitz continuous and bounded.

More precisely, almost surely, there exist KV,KΘ,CV,CΘ>0K_{\mathscr{V}},K_{\Theta},C_{\mathscr{V}},C_{\Theta}>0, such that, for any s,t∈s,t\in,

We shall also need the following requirement on gg, which is satisfied by all our models as soon as σ\sigma is Lipschitz continuous:

Under Assumptions (A5)(A_{5}) and (A6)(A_{6}), the recurrence (13) almost surely converges towards the neural ODE given by

Consider model (13) such that Assumptions (A5)(A_{5}) and (A6)(A_{6}) are satisfied. Then the ODE (14) has a unique solution HH, and, almost surely, there exists some c>0c>0 such that, for any 0⩽k⩽L0\leqslant k\leqslant L,

It should be stressed that the transition from the discrete recurrence (13) to the continuous-time differential equation (14) relies on the assumptions that the weight sequences (θk)1⩽k⩽L(\theta_{k})_{1\leqslant k\leqslant L} and (Vk)1⩽k⩽L(V_{k})_{1\leqslant k\leqslant L} are the discretizations of smooth limiting processes Θ\Theta and V\mathscr{V} on the one hand, and that the scaling αL\alpha_{L} is chosen as \nicefrac1L\nicefrac{{1}}{{L}} on the other hand. From a practical perspective, Proposition 11 shows that it is possible to initialize ResNets in the ODE regime, by choosing a smooth stochastic process, discretizing it at each layer, and taking a \nicefrac1L\nicefrac{{1}}{{L}} scaling. This is in sharp contrast with the results of Sections 2 and 3.1, which show that the usual i.i.d. procedure leads to a neural SDE.

Stability and scaling.

Assuming that the weights of the network are discretizations of a smooth function (Assumption (A5)(A_{5})), it is possible to obtain stability results, depending on the value of β\beta, similarly to what has been done in Section 2. We show below that β=1\beta=1 is a critical value, by examining the hidden states, in the same way as β=\nicefrac12\beta=\nicefrac{{1}}{{2}} is a critical value in the i.i.d. setting. Similar results can be shown for the gradients. We begin by a proposition handling the cases β>1\beta>1 and β=1\beta=1.

Consider a ResNet (4) such that Assumptions (A5)(A_{5}) and (A6)(A_{6}) are satisfied. Let αL=\nicefrac1Lβ\alpha_{L}=\nicefrac{{1}}{{L^{\beta}}}, with β>0\beta>0.

If β=1\beta=1, then, almost surely, there exists some c>0c>0 such that

The explosion case (β<1\beta<1) is more delicate to deal with. We prove it for a linear model, and leave for future work the extension to more general cases.

Consider the res-1 model, taking σ\sigma as the identity function. Assume that Assumption (A5)(A_{5}) is satisfied and that V0T\mathscr{V}_{0}^{T} has a positive eigenvalue. Let αL=\nicefrac1Lβ\alpha_{L}=\nicefrac{{1}}{{L^{\beta}}}, with β∈(0,1)\beta\in(0,1). Then, almost surely,

The assumption of the existence of a positive eigenvalue for V0⊤\mathscr{V}_{0}^{\top} is mild. For instance, if the entries of V0\mathscr{V}_{0} are i.i.d. random variables with finite moments of all order, Götze and Jalowy (2021) show that such an eigenvalue exists with probability at least 1−\nicefrac1d1-\nicefrac{{1}}{{d}} for dd large enough.

In this setting, we observe experimentally a behavior of the output and of the gradients when LL grows large similar to the one explored in Section 2. This is illustrated in Figures 4 and 5, which mirror Figures 1 and 3 in Section 2. The figures clearly show that there exist three cases for the output and for the gradients: an identity case (left plots), an explosion case (middle), and a non-trivial case separating explosion and identity (right). However, the remarkable point is that the separation occurs for β=1\beta=1, and not β=\nicefrac12\beta=\nicefrac{{1}}{{2}}, as predicted by Propositions 12 and 13.

Experiments

We experimentally investigate in this section two questions. The first one is to know whether there exists a range of scaling factors β>0\beta>0 and weight initializations, beyond the i.i.d. and the smooth regimes. The second question is whether our analysis, which pertains to the initialization phase, provides insights into the training phase, beyond initialization.

In order to describe the transition between the i.i.d. and smooth cases, a possible route is to consider that the weights are increments of a γ\gamma-Hölder stochastic process. This model is interesting insofar as the Brownian motion (SDE regime) is (\nicefrac12−ε)(\nicefrac{{1}}{{2}}-\varepsilon)-Hölder (ε>0\varepsilon>0) and a Lipschitz continuous stochastic process (ODE regime) is 11-Hölder.

In line with the above, in a series of experiments, we initialize the weights as increments of a fractional Brownian motion (BtH)t∈(B^{H}_{t})_{t\in}. Recall that BHB^{H} is a continuous-time Gaussian process, starting at zero, with zero expectation for all t∈t\in, and covariance function

where H∈(0,1)H\in(0,1) is called the Hurst index. This index describes the raggedness of the process, with a higher value leading to a smoother process. When H=\nicefrac12H=\nicefrac{{1}}{{2}}, the process is a standard Brownian motion (Definition 9), whose increments are independent by construction. When H>\nicefrac12H>\nicefrac{{1}}{{2}}, the increments of the process are positively correlated, while if H<\nicefrac12H<\nicefrac{{1}}{{2}} the increments are negatively correlated. Importantly, a fractional Brownian motion with Hurst index HH is (H−ε)(H-\varepsilon)-Hölder continuous for any ε>0\varepsilon>0. In the limit when H→1H\rightarrow 1, the trajectories converge to linear functions (whose increments satisfy (A5)(A_{5})). As an illustration, Figure 6 depicts three realizations of a fractional Brownian motion with H=0.2H=0.2 (left), H=0.5H=0.5 (middle), and H=0.8H=0.8 (right).

In order to assess the effect of the scaling factor β\beta and the Hurst index HH, we initialize a neural network res-3 with d=40d=40, L=1000L=1000, various values of β∈[0.2,1.3]\beta\in[0.2,1.3], and with weights taken as increments of fractional Brownian motions with various Hurst indices H∈(0,1)H\in(0,1). Figure 7 depicts the empirical magnitude of the output and the gradients at initialization as a function of the Hurst index HH and the scaling factor β\beta. First note that we recover the two regimes (i.i.d. and smooth) discussed so far. For H=\nicefrac12H=\nicefrac{{1}}{{2}}, the i.i.d. regime kicks in, with explosion (β<\nicefrac12\beta<\nicefrac{{1}}{{2}}, orange zone), non-trivial behavior (β=\nicefrac12\beta=\nicefrac{{1}}{{2}}, black zone), and identity (β>\nicefrac12\beta>\nicefrac{{1}}{{2}}, blue zone). Likewise, we see at H=1H=1 a similar pattern in the smooth regime, with, as predicted by Proposition 12, a critical value β=1\beta=1. Beyond these two specific cases, we observe for an index HH varying in (\nicefrac12,1)(\nicefrac{{1}}{{2}},1) a whole range of intermediate situations, where the transition between identity and explosion seems to happen for a critical β=H\beta=H. Interestingly, for H<\nicefrac12H<\nicefrac{{1}}{{2}}, the transition seems to saturate at the value β=\nicefrac12\beta=\nicefrac{{1}}{{2}}.

The take-home message is that the choice of the scaling of a ResNet seems to be closely linked to the regularity of the weights as a function of the layer. More precisely, for all regimes, the critical scaling factor between explosion and identity seems to have a natural interpretation as the (Hölder) regularity of the underlying continuous-time stochastic process. We believe that the mathematical understanding of this connection, beyond the fractional Brownian motion case, is a promising research direction for the future. Finally, these experiments suggest that it is sensible to initialize a ResNet for any value of the scaling β∈(\nicefrac12,1)\beta\in(\nicefrac{{1}}{{2}},1), while avoiding the identity and explosion situations, by simulating a fractional Brownian motion of Hurst index H=βH=\beta and initializing the weights as the increments of this process.

2 Beyond initialization

At initialization, before the gradient descent, the distribution of the weights (θk)1⩽k⩽L(\theta_{k})_{1\leqslant k\leqslant L} and (Vk)1⩽k⩽L(V_{k})_{1\leqslant k\leqslant L} is chosen by the practitioner. By contrast, during and after training, control is lost on these distributions, making the picture more complex. In particular, the existence and characterization of a continuous-time stochastic process whose discretization matches the trained ResNet is an interesting but difficult problem. Attacking this question requires a fine understanding of the interaction between training dynamics and the regularity of the sequence of the weights during the gradient descent. However, there is experimental evidence that the trained weights exhibit strong structure as a function of the layer index kk (Cohen et al., 2021; Bayer et al., 2022), and that their regularity strongly depends on the choice of initialization. Figure 8 depicts this mechanism by plotting a given coordinate of θk\theta_{k} as a function of the layer index kk ranging from 11 to the depth L=1000L=1000, after training.

To investigate the link between regularity of the weights at initialization, scaling, and performance after training, we train ResNets on the datasets MNIST (Deng, 2012) and CIFAR-10 (Krizhevsky, 2009). As in Subsection 4.1, we initialize the ResNets with various scaling factors and weights that are increments of fractional Brownian motions with different regularities. Then, for each combination of weight initialization and scaling factor, the ResNet is trained using the Adam optimizer (Kingma and Ba, 2015) for 1010 epochs. The results in terms of accuracy are presented in Figure 9 (light orange = good performance, blue = bad performance). We observe a pattern similar to the one of Figure 7, however shifted downwards. This means that, for a given regularity, the network is unable to learn if it is initialized with a scaling too far below the critical value, which of course is connected with the gradient explosion issue discussed previously. On the other hand, and perhaps more surprisingly, the performance seems to be more or less stable in the identity region, with perhaps a small degradation in the case of CIFAR-10. This somewhat contrasts with the results from Yang and Schoenholz (2017), who exhibit a decrease in performance for i.i.d. initialization and a large scaling factor β\beta. Note however that, in our experiments, we adapted the learning rate of the gradient descent on a grid by cross-validation. This was done to prevent a slowdown in training when the scaling factor β\beta is too large, since, in the gradient descent, the gradients updates are also scaled by the factor \nicefrac1Lβ\nicefrac{{1}}{{L^{\beta}}}—which gets smaller as β\beta increases. The interplay between the learning rate and the scaling factor is one of the keys to better assess how the performance of the trained network is connected with the scaling.

The authors thank S. Schoenholz for fruitful discussion. P. Marion has been supported by a grant from Région Île-de-France and by a Google PhD Fellowship award.

Appendix A Proofs

Throughout the proofs, the ii-th coordinate of a vector vv is denoted by viv_{i}. Similarly, the ii-th row of a matrix MM is denoted by MiM_{i}, and its (i,j)(i,j)-th entry by MijM_{ij}.

A.2 Proof of Proposition 2

But, for LαL2⩽1L\alpha_{L}^{2}\leqslant 1, we have (1+αL2)L−1⩽exp⁡(LαL2)−1⩽2LαL2.(1+\alpha_{L}^{2})^{L}-1\leqslant\exp(L\alpha_{L}^{2})-1\leqslant 2L\alpha_{L}^{2}. Therefore,

and the result follows from Markov’s inequality.

Consider a ResNet (4) such that Assumptions (A1)(A_{1}) and (A2)(A_{2}) are satisfied. Then

Taking the squared norm of the forward update rule (4) and dividing by ∥h0∥2\|h_{0}\|^{2} yields

We deduce by Assumptions (A1)(A_{1}) and (A2)(A_{2}) that

Now, observe that hL=h0+αL∑k=0L−1Vk+1g(hk,θk+1)h_{L}=h_{0}+\alpha_{L}\sum_{k=0}^{L-1}V_{k+1}g(h_{k},\theta_{k+1}). Thus, we have

By conditioning on all random variables except Vk′+1V_{k^{\prime}+1} for k<k′k<k^{\prime} (and Vk+1V_{k+1} for k>k′k>k^{\prime}), it is easy to see that the only non-zero terms are when k=k′k=k^{\prime}. This yields

A.3 Proof of Proposition 3

Dividing (15) by ∥hk∥2\|h_{k}\|^{2} and taking the logarithm leads to

and Yk=Yk,1+Yk,2Y_{k}=Y_{k,1}+Y_{k,2}. The proof of Proposition 3 strongly relies on the following lemma, which provides technical information on the moments of Yk,1Y_{k,1} and Yk,2Y_{k,2}. For the sake of clarity, its proof is postponed to Appendix B.

Assume that Assumptions (A1)(A_{1}) and (A2)(A_{2}) are satisfied. Then

The last inequality is true for αL2⩽1C(d+128s4)\alpha_{L}^{2}\leqslant\frac{1}{\sqrt{C}(d+128s^{4})}. Therefore, by inequality (17), we obtain, for c>exp⁡(LαL2)c>\exp(L\alpha_{L}^{2}),

We conclude that, for any δ∈(0,1)\delta\in(0,1), with probability at least 1−δ1-\delta,

This shows statement (ii)(ii) of the proposition.

Next, to prove statement (i)(i), observe that c>0c>0,

Using the inequality ln⁡(1+x)⩾x−x2\ln(1+x)\geqslant x-x^{2} for x⩾−\nicefrac12x\geqslant-\nicefrac{{1}}{{2}}, we obtain

where the last inequality is valid for d⩾64d\geqslant 64 and αL2⩽116C(2s4+1)\alpha_{L}^{2}\leqslant\frac{1}{16\sqrt{C}(2s^{4}+1)}. Hence, for 0<c<exp⁡(\nicefrac3LαL28)0<c<\exp(\nicefrac{{3L\alpha_{L}^{2}}}{{8}}),

Using the crc_{r}-inequality (a+b)n⩽2n−1(an+bn)(a+b)^{n}\leqslant 2^{n-1}(a^{n}+b^{n}) respectively for n=2n=2 and n=4n=4, we see that

By (E3)(E_{3})-(E7)(E_{7}), it is easy to verify that, for d⩾64d\geqslant 64 and αL2⩽1(Cs4/16+2C+8s4)d\alpha_{L}^{2}\leqslant\frac{1}{(\sqrt{C}s^{4}/16+2\sqrt{C}+8s^{4})d},

This shows that, for c<exp⁡(\nicefrac3LαL28)c<\exp(\nicefrac{{3L\alpha_{L}^{2}}}{{8}}),

To conclude the proof, it remains to upper bound the second term of inequality (18). According to inequality (22) in the proof of Lemma 15 (with t=\nicefrac12t=\nicefrac{{1}}{{2}}), one has

Putting everything together, we are led to

Take δ∈(0,1)\delta\in(0,1). Then, if 2L\exp\big{(}-\frac{d}{64\alpha_{L}^{2}s^{2}}\big{)}\leqslant\frac{\delta}{11}, with probability at least 1−δ1-\delta,

Notice that this inequality is valid under the assumption αL2⩽2(Cs4+4C+16s4)d\alpha_{L}^{2}\leqslant\frac{2}{(\sqrt{C}s^{4}+4\sqrt{C}+16s^{4})d}.

A.4 Proof of Corollary 4

Statement (i)(i) is a consequence of Proposition 2, whereas (ii)(ii) is a consequence of Proposition 3 (i)(i). The latter is valid under the conditions d⩾64d\geqslant 64 and αL⩽2(Cs4+4C+16s4)d\alpha_{L}\leqslant\frac{2}{(\sqrt{C}s^{4}+4\sqrt{C}+16s^{4})d}, which is automatically satisfied for all LL large enough. Furthermore, an inspection of the proof of Proposition 3 reveals that the divergence in high probability of ∥hL∥\|h_{L}\| can be proved under the relaxed assumption d⩾9d\geqslant 9. Indeed, the main constraint on dd comes from the lower bound (19), where one needs to make sure that LαL22−4LαL2d>0\frac{L\alpha_{L}^{2}}{2}-4\frac{L\alpha_{L}^{2}}{d}>0, which is the case for d=9d=9.

To prove (iii)(iii), we use a union bound on both statements of Proposition 3.

A.5 Proof of Proposition 5

The first claim follows from the observation that

from (A1)(A_{1}), and from the assumption on σ′\sigma^{\prime}.

Denote by DD the matrix in the middle of the right-hand side. Then

Since the (Wij)1⩽i,j⩽d(W_{ij})_{1\leqslant i,j\leqslant d} are symmetric random variables, we conclude that

A.6 Proof of Proposition 6

Letting b=pL/∥pL∥b=p_{L}/\|p_{L}\|, as in Assumption (A3)(A_{3}), and taking expectation in (10), we obtain

The rest of the proof is similar to the proof of Proposition 2. From (9), we have

By independence of Vk+1V_{k+1} from qk(z)q_{k}(z) and ∂g(hk,θk+1)∂h\frac{\partial g(h_{k},\theta_{k+1})}{\partial h},

Using arguments similar to (20), we may write

Now, upon noting that qL(z)−z=qL(z)−q0(z)=αL∑k=0L−1Vk+1∂g(hk,θk+1)∂hqk(z)q_{L}(z)-z=q_{L}(z)-q_{0}(z)=\alpha_{L}\sum_{k=0}^{L-1}V_{k+1}\frac{\partial g(h_{k},\theta_{k+1})}{\partial h}q_{k}(z),

for LαL2⩽1L\alpha_{L}^{2}\leqslant 1. Note that the second equality is obtained by conditioning on every random variable except Vk′+1V_{k^{\prime}+1} for k<k′k<k^{\prime} (and Vk+1V_{k+1} for k>k′k>k^{\prime}). Finally, by using Markov’s inequality, we conclude that, for any ε>0\varepsilon>0,

A.7 Proof of Proposition 7

A.8 Proof of Corollary 8

The first statement is an immediate consequence of Proposition 6. The second one is a consequence of Proposition 7 and the fact that, for β<1\beta<1,

Finally, (iii)(iii) follows from Proposition 7.

A.9 Proof of Proposition 10

The proposition is a consequence of Kloeden and Platen (1992, Theorems 4.5.3 and 10.2.2) for the SDE

Letting a(h,t)=0a(h,t)=0 and b(h,t)=d2σ(h)b(h,t)=\sqrt{\frac{d}{2}}\sigma(h), we need to check the following assumptions:

Assumptions (H1)(H_{1}), (H4)(H_{4}), and (H5)(H_{5}) readily follow from the definitions. Assumption (H2)(H_{2}) is true since σ\sigma is Lipschitz continuous, and (H3)(H_{3}) follows from

A.10 Proof of Proposition 11

In addition, it is continuous in its second one. Thus, according to the Picard-Lindelöf theorem (Theorem 22 in Appendix D), this is enough to show that the neural ODE (14) has a unique solution on $.Notethatthesolution. Note that the solutionHiscontinuousonis continuous onandisthereforeboundedbyaconstantand is therefore bounded by a constantM>0$.

In order to prove the approximation bound of Proposition 11, we start by proving that both ψ\psi and HH are Lipschitz continuous in tt. Under (A5)(A_{5}) and (A6)(A_{6}), this is clear for ψ\psi since HH is bounded. Moreover, for any [s,t]⊂[s,t]\subset, we have

Now, let K1K_{1} and K2K_{2} denote the Lipschitz constants of ψ\psi (in both arguments) and HH respectively, and, for any 0⩽k⩽L0\leqslant k\leqslant L, let tk=\nicefrackLt_{k}=\nicefrac{{k}}{{L}}. Then we have, for k⩾1k\geqslant 1,

A.11 Proof of Proposition 12

Starting from (4) and using Assumption (A6)(A_{6}), one easily obtains the existence of C1C_{1} and C2C_{2} (whose values depend on the realization of V\mathscr{V} and Θ\Theta) such that

Hence, using αL⩽\nicefrac1L\alpha_{L}\leqslant\nicefrac{{1}}{{L}},

since we showed that each term in the sum is bounded by some constant C3>0C_{3}>0, independent of LL and kk. Hence we have that

yielding the results depending on the value of β\beta.

A.12 Proof of Proposition 13

Take yy a unit-norm eigenvector of V0⊤\mathscr{V}_{0}^{\top} with associated eigenvalue λ>0\lambda>0. Then

Since V\mathscr{V} is Lipschitz and Vk+1=Vk+1/LV_{k+1}=\mathscr{V}_{k+1/L}, there exists cc such that ∥Vk+1−V0∥⩽ck+1L\|V_{k+1}-\mathscr{V}_{0}\|\leqslant c\frac{k+1}{L}. Hence

Let M=∣⟨h0,y⟩∣2cαLM=\frac{|\langle h_{0},y\rangle|}{2c\alpha_{L}}, and suppose that ∥hk∥⩽M\|h_{k}\|\leqslant M for all 0⩽k⩽L0\leqslant k\leqslant L. Then

Then, for λαL⩽1\lambda\alpha_{L}\leqslant 1,

Thus, since LαL=L1−βL\alpha_{L}=L^{1-\beta}, we have that ∥hL∥→∞\|h_{L}\|\rightarrow\infty, which contradicts our assumption that ∥hk∥⩽M\|h_{k}\|\leqslant M for all 0⩽k⩽L0\leqslant k\leqslant L. We deduce that, for all LL large enough,

Appendix B Technical results

Proof The first part is a consequence of the assumption on σ\sigma. To prove the equality, let Xi=∑j=1dWijxjX_{i}=\sum_{j=1}^{d}W_{ij}x_{j}. Then

The result follows by summing over all i∈{1,…,d}i\in\{1,\dots,d\}.

B.2 Proof of Lemma 15

(E1)(E_{1}) and (E2)(E_{2}) are simple consequences of Assumptions (A1)(A_{1}) and (A2)(A_{2}).

To prove (E3)(E_{3}), let f(hk,θk+1)=Vk+1g(hk,θk+1)f(h_{k},\theta_{k+1})=V_{k+1}g(h_{k},\theta_{k+1}). Then

It is easy to verify that, under Assumption (A1)(A_{1}), each term of the sum above has zero expectation. This shows (E3)(E_{3}).

To establish (E4)(E_{4}), we start by noting that

To prove (E5)(E_{5}), let φ=⟨Vk+1g(hk,θk+1),hk⟩∥g(hk,θk+1)∥∥hk∥\varphi=\frac{\langle V_{k+1}g(h_{k},\theta_{k+1}),h_{k}\rangle}{\|g(h_{k},\theta_{k+1})\|\|h_{k}\|}. Then, for any t>0t>0,

by Jensen’s inequality. Finally, using Assumption (A2)(A_{2}), we deduce that

In particular, for all q⩾1q\geqslant 1 (see, e.g., Pauwels, 2020),

Finally, (E6)(E_{6}) and (E7)(E_{7}) are consequences of Lemma 21 in Appendix C.

Appendix C Concentration of sub-Gaussian random matrices

In this appendix, we are interested in the concentration of linear and quadratic forms of sub-Gaussian matrices (Lemma 20 and Lemma 21). These two propositions are byproducts of the main result of Kontorovich (2014), which generalizes McDiarmid’s inequality to sub-Gaussian variables. We start by a technical result regarding the sub-Gaussian diameter introduced by Kontorovich (2014), whose definition is recalled below.

Let XX be a real-valued random variable, X′X^{\prime} an independent copy of XX, and ε\varepsilon a Rademacher random variable, independent of XX and X′X^{\prime}. Then the sub-Gaussian diameter of XX is defined as the smallest tt such that ε∣X−X′∣\varepsilon|X-X^{\prime}| is t2t^{2} sub-Gaussian.

Let XX be a s2s^{2} sub-Gaussian symmetric random variable. Then the sub-Gaussian diameter of XX is less than 2s\sqrt{2}s.

where the last equality is a consequence of the symmetry of XX.

We are now ready to prove the two main results of this appendix.

By the triangle inequality, φ\varphi is a Lipschitz continuous function, with Lipschitz constant equal to 11. Observe also that XijX_{ij} is a \nicefracxi2s2yj2d∥x∥2∥y∥2\nicefrac{{x_{i}^{2}s^{2}y_{j}^{2}}}{{d\|x\|^{2}\|y\|^{2}}} sub-Gaussian. Thus, according to Lemma 19, the sub-Gaussian diameter of XijX_{ij} is less than \nicefrac2xisyjd∥x∥∥y∥\nicefrac{{\sqrt{2}x_{i}sy_{j}}}{{\sqrt{d}\|x\|\|y\|}}. By Kontorovich (2014, Theorem 1), for any t>0t>0, one has

Each function φi\varphi_{i} is a Lipschitz continuous function, with Lipschitz constant equal to 11. Observe now that the random variable XijX_{ij} is \nicefracxj2s2d∥x∥2\nicefrac{{x_{j}^{2}s^{2}}}{{d\|x\|^{2}}} sub-Gaussian. Thus, according to Lemma 19, the sub-Gaussian diameter of XijX_{ij} is less than \nicefrac2xjsd∥x∥\nicefrac{{\sqrt{2}x_{j}s}}{{\sqrt{d}\|x\|}}. Therefore, according to Kontorovich (2014, Theorem 1), for any t>0t>0,

From identity (21) in the proof of technical Lemma 17, given in Appendix B, we obtain that, for q=1q=1,

which is an improvement by a factor 8s28s^{2} over the previous upper bound. To conclude, it remains to conclude ∥Vx∥4\|Vx\|^{4} and ∥Vx∥8\|Vx\|^{8} with the ⟨Vi,x⟩\langle V_{i},x\rangle. To do so, observe that

Hence, by independence of the (Vi)1⩽i⩽d(V_{i})_{1\leqslant i\leqslant d},

Appendix D A version of the Picard-Lindelöf theorem

Appendix E Detailed experimental setting

Our code is available at https://github.com/PierreMarion23/scaling-resnets.

To obtain Figures 1 to 3, we initialize ResNets from res-3 with the hyper-parameters of Table 2.

Each experiment is repeated 5050 times, with independent data and weight sampling.

For Figures 4 and 5, we take the same hyper-parameters except for β\beta, which now takes values in {0.5,1,2}\{0.5,1,2\}, and for the weight distribution. The weights are now initialized as discretizations of a Gaussian process. More precisely, each entry of V\mathscr{V} and Θ\Theta is an independent Gaussian process with zero mean and an RBF kernel of variance 10−210^{-2}.

To obtain Figure 7, we take the hyper-parameters of Table 3.

More precisely, for each 1⩽i,j⩽d1\leqslant i,j\leqslant d, we let (Vk+1,i,j)0⩽k⩽L−1(V_{k+1,i,j})_{0\leqslant k\leqslant L-1} be the increments of a fractional Brownian motion (fBm), where the various fBm involved are independent. The procedure is the same for θ\theta.

In Figure 9, we use res-1, with the hyper-parameters of Table 4.

We train on MNISThttp://yann.lecun.com/exdb/mnist and CIFAR-10https://www.cs.toronto.edu/~kriz/cifar.html using the Adam optimizer (Kingma and Ba, 2015) for 1010 epochs. The learning rate is divided by 1010 after 55 epochs. The best performance on the learning rate grid is reported in the figure.

Figure 8 is obtained by plotting a random coordinate of θk\theta_{k}, after training on MNIST.

References