Neural SDE - Information propagation through thelens of diffusion processes

Stefano Peluchetti
Cogent Labs
speluchetti@cogent.co.jp

Stefano Favaro
University of Torino and Collegio Carlo Alberto
stefano.favaro@unito.it

Abstract
Very deep networks are known to suffer from a number of pathologies which limit
their expressive power. We consider parameter distributions that shrink as thenumber of layers increases in order to recover stochastic processes in the limit of
infinite total depth and show that the resulting diffusion processes exhibit desirable
properties.

1

Introduction

The focus of this work is on the function-space distribution properties of the output (the last layer)of very deep neural networks. We also focus on standard classes of neural networks, as opposed toproposing structural modifications in order to achieve the desired results. When neural networks arealso very wide, recent research [11, 12, 3] shows that i.i.d. initializations with constant variance canresult in undesirable properties, even when optimally initialized on the edge of chaos (EOC [3]), such
as: i) independence of the output from the input; ii) concentration of the output on restrictive families
including constant functions. We illustrate this problem in Fig. 1 where we plot output samples (for agiven pre-activation) over input values. In the tanh case the input has no discernible impact on the
output and the sampled functions are almost constant. This behavior holds for most smooth activation
functions [3]. In the ReLU case the input affects the variance of the output and the function samples
are piecewise linear. In both cases, the outputs corresponding to any two inputs end up perfectly
correlated.
Intuitively, the issue is the constant level of randomness introduced between subsequent layers. In thiswork we consider initializations where the parameters’ distribution shrinks as the number of layers
increases and establish the convergence of the output of residual networks (jointly over multiple
inputs) to diffusion processes as the number of layers goes to infinity. In this limit, the output satisfies
the desiderata: i) it retains dependency from the input; ii) it doesn’t suffer from the perfect correlation
constraint; iii) it doesn’t collapse to a deterministic function nor does it diverge.

2

Neural SDEs Limits

2.1

Diffusion Processes and SDEs

There are many ways to construct continuous-time stochastic processes as limiting dynamics of
discrete-time processes, and in this work we consider the simplest case where the limiting process
has continuous paths. In all the neural network architectures considered in this work each layer
depends exclusively on the previous one. These two conditions identify diffusion processes [13],which are continuous-time Markov processes with continuous paths, as natural candidates for thelimiting process. Diffusion processes arise as solutions to SDEs, stochastic versions of ODEs. A
SDE is described by:
dxt = µx (xt )dt + σx (xt )dwt
(1)
4th workshop on Bayesian Deep Learning (NeurIPS 2019), Vancouver, Canada.

tanh

ReLU

xlast, 1(x0)

2
1
0
1
2
2

1

0

x0

1

2

2

1

0

x0

1

2

Figure 1: Output samples in blue for 2 fully connected feedforward network with 500 layers of 500
units, input x0 ∈ [−2, 2], parameters on EOC; 5%, 50% and 95% quantiles in orange.
where wt is a driving Brownian motion (BM) [6] of compatible dimensionality. It is easy to give an
intuitive characterization of its dynamics. Consider the discretization of (1):√xt+1 = xt + µx (xt )∆t + σx (xt )εt ∆t,(2)where εt is a N (0, I) random vector. The discretization (2) converges (see [7]) to the solution of (1),and we notice the similarity with the Euler discretization of an ODE in the deterministic part of (2).
2.2Compatible Architectures, Parameter Distributions and Activation Functions

To obtain diffusion limits we show that it is necessary to consider a specific form of residualarchitectures. Denote with xl ∈ Rdc , l = 1, . . . , dl , the layers of a neural network of dl layers, andwith x0 its input. Let ∆t = 1/dl define an infinitesimal unit of time and let ∆xl = xl+1 − xl define
the increments of xl . Due to the continuity of the paths of the limiting diffusion process, we needPr(k∆xl k > ε|xl ) ↓ 0 as ∆t ↓ 0 for any ε > 0, i.e. we require the increments to vanish eventually.
It’s easy to see that this cannot be achieved in a feed-forward network where xl+1 = φ(Wl xl + bl ),unless φ is linear or the distribution of (Wl , bl ) depends on x. The same holds for residual networkarchitecture (ResNet) originally introduced in [4]. This leaves us with the identity ResNet [5] where:
xl+1 = xl + Fl (xl )In general, each residual block Fl can be composed of multiple stacked layers. Here we consider thecase the simplest implementation with shallow residual blocks:∆xl = Fl (xl ) = φ(Wl xl + bl )For shallow residual blocks, the vanishing increments requirement is satisfied by having the distributions of Wl and bl both concentrate around 0 provided that φ(0) = 0. We require:
Assumption 2.1 (Parameters Distribution and Scaling). For l = 0, 1, . . . let:
i.i.d.

2
Wl ∼ N (0, σw∆t)i.i.d.

bl ∼ N (0, σb2 ∆t)
Assumption 2.2 (Activation Function Regularity). The function φ : R → R satisfies: φ(0) = 0, φ iscontinuously differentiable three times on R, its second and third derivatives have at most exponential
tails growth, i.e. for some k > 0:
|φ00 (x)||φ000 (x)|
+ lim
<∞
k|x|
|x|↑∞ e|x|↑∞ ek|x|
lim

2.3

Joint Diffusion Limits

We now state the main result, which establishes the convergence in distribution of the ResNet to adiffusion process jointly over 2 inputs as dl ↑ ∞. We index with t instead of l as we need to introduce
a continuous-time interpolation (the limiting-process is a continuous-time one). We consider weaksolutions and weak uniqueness as we are interested only in the distributional properties of the limiting
SDE. The proof and more details these points are in the Appendix. We use xt and x0t to denote thejoint evolution of the resent for two inputs x0 and x00 .
2

dl = 15dl = 50dl = 150

dl = 500

dc = 15
xT, 1(x0)

5
0
5
10

dc = 500
xT, 1(x0)

5
0
5
10

2

1

0

x0

1

2

2

1

0

x0

1

2

2

1

0

x0

1

2

2

1

0

x0

1

2

Figure 2: Output samples in blue for a fully connected identity ResNet, different levels of dl anddc , input x0 ∈ [−2, 2] (i.e. x0,c has the√ same value x0 for each c ∈ dc ), parameters according toAssumption 2.1 with σb = 1, σw = 1/ dc , tanh activation, 5%, 50% and 95% quantiles in orange;
the scaling on σw has a stabilizing effect and has been used to obtain well-defined limits in [8, 11, 12];
we observe a dependency on the input and flexible function shapes (not piecewise linear) compared
to Fig. 1 for the tanh activation; moreover we observe similar distribution properties across different
orders of magnitude for dl and dc which suggests the existence of a stochastic limit as dc ↑ ∞.
Theorem 2.1. The continuous time interpolation xt , x0t of xl , x0l converges to the weakly uniqueweak solution of:
1/2
  

 2
xt
µ(x)
σ (xt , xt ) σ 2 (xt , x0t )dt + 2 0
dwt(t ∈ [0, T = 1])
(3)
d 0 =
xt
µ(x0 )σ (xt , xt ) σ 2 (x0t , x0t )12µ(x) = φ00 (0)(σb2 + σw
kxk2 )1dc
22
σ 2 (x, y) = φ0 (0)2 (σb2 + σw
hx, yi) Idc
where wt is a 2dc -dimensional BM, 1dc is a dc -dimensional unit vector, Idc is a dc × dc identity
matrix. When φ00 (0) = 0 (for instance φ = tanh) the solution is guaranteed not to explode.
An immediate consequence of Theorem 2.1 is that the bivariate distribution of the output given two
inputs converges to the transition density of the solution of (3), a stochastic distribution. If φ00 (0) = 0,
the process is non-explosive (i.e. such transition density is always well-defined). As the integration
time is finite, the dependency on the input doesn’t vanish in the limit of infinite depth and can becontrolled via σw and σb . As the diffusion matrix squared is non-singular we also don’t suffer fromthe perfect-correlation problem. This satisfies our desiderata i),ii),iii) in the introduction.
We empirically investigate the distribution properties of (3) in Fig. 2 where we plot sample outputs(for a given component c ∈ dc ) over input values. Fig. 2 can be contrasted with Fig. 1. In both figures,
each sample is generated by sampling all parameters once and computing the outputs corresponding
to each input in a range.

3

Related Work and Conclusions[14] considers initializations for residual networks which are not encompassed yet by our analysis.
Conversely, the residual blocks in [14] cannot be shallow. Some parameters are initialized witha similar scaling, but the residual blocks are multiplied by parameters initialized at 0, hence ourdesiderata iii) is not satisfied.
The desire of obtaining flexible distributions in function space is especially relevant for Bayesianinference. For instance, a prior model that puts all the probability mass on constant functions cannot
fit non-constant functions. While our results are a "pre-requisite" to construct infinitely deep models,in order to obtain competitive performance more attention needs to be paid to architectural choices,
also at the level of input and output layers. Our results can be extended without requiring novel
theoretical developments to convolutional networks and non i.i.d. parameter distributions.
3

References[1] P. Billingsley. Convergence of Probability Measures. Wiley-Interscience, 2nd edition, 1999.
[2] S. N. Ethier and T. G. Kurtz. Markov processes: characterization and convergence. WileyInterscience, 2009.
[3] S. Hayou, A. Doucet, and J. Rousseau. On the impact of the activation function on deep neuralnetworks training. In Proceedings of the 36th International Conference on Machine Learning,
pages 2672–2680, 2019.
[4] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In
Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–
778, 2016.
[5] K. He, X. Zhang, S. Ren, and J. Sun. Identity mappings in deep residual networks. In European
conference on computer vision, pages 630–645. Springer, 2016.
[6] I. Karatzas and S. Shreve. Brownian Motion and Stochastic Calculus. Springer, 2nd edition,
1999.
[7] P. E. Kloeden and E. Platen. Numerical Solution of Stochastic Differential Equations. Springer,
corrected edition, 1992.
[8] R. M. Neal. Bayesian Learning for Neural Networks. PhD thesis, University of Toronto, 1995.
[9] D. B. Nelson. Arch models as diffusion approximations. Journal of econometrics, 45(1-2):7–38,
1990.
[10] B. Øksendal. Stochastic Differential Equations: An Introduction with Applications. Springer,
6th edition, 2003.
[11] B. Poole, S. Lahiri, M. Raghu, J. Sohl-Dickstein, and S. Ganguli. Exponential expressivity indeep neural networks through transient chaos. In Advances in Neural Information Processing
Systems 29, pages 3360–3368, 2016.
[12] S. S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein. Deep information propagation.
In International Conference on Learning Representations, 2017.
[13] D. W. Stroock and S. S. Varadhan. Multidimensional diffusion processes. Springer, 2006 edition,
2006.
[14] H. Zhang, Y. N. Dauphin, and T. Ma. Residual learning without normalization via better
initialization. In International Conference on Learning Representations, 2019.

4

A

Diffusions and Diffusion Limits

We review results which are useful in establishing convergence to diffusion limits. Let xl , l = 0, 1, . . .
be a generic d-dimensional discrete-time Markov process. Let ∆t > 0 define an infinitesimal unitof time and ∆xl = xl+1 − xl define the increments of xl . We will rely on the following condition
where it’s implicit that the distribution p(xl |xl−1 ) depends on ∆t.
Assumption A.1 (Infinitesimal Coefficients). Let xl , l = 0, 1, . . . be a d-dimensional discrete-timeMarkov process, and assume that there exist µx (x) : Rd → Rd and σx2 (x) : Rd → Rd×d such that:E[∆xl |xl ]= µx (xl )∆tV[∆xl |xl ]lim= σx2 (xl )
∆t↓0
∆t
lim

∆t↓0

2+δ
E[∆xl |xl ]
=0
∆t↓0
∆t

lim

(4)
(5)
(6)

for some δ > 0, where all convergences are uniform on compacts of Rd for each vector and matrix
component. Moreover µx (x) and σx2 (x) are continuous and σx2 (x) admits a Cholesky factorization
σx (x), i.e. σx (x)σx (x)> = σx2 (x).
The infinitesimal evolution of diffusion processes is characterized by its infinitesimal mean vector(4) and infinitesimal covariance matrix (5), so the first two limits pinpoint the form of the limiting
stochastic evolution. Condition (6) is a technical one in the sense that it allows us to consider thelimits (4) and (5) instead of their truncated version [9].
Under additional assumptions, the following result establishes that in the limit xl can be embedded ina diffusion process.
Theorem A.1. Let xl , l = 0, 1, . . . be a d-dimensional discrete-time Markov process, and define thecontinuous-time process xt on t ∈ [0, T ] by continuous-on-right step-wise-constant interpolation ofxl :xt = xl 1l∆t≤t<(l+1)∆t (l = 0, . . . , dl , ∆t = T /dl )(7)
for some T > 0.
Consider the d-dimensional stochastic differential equation (SDE) with initial value x0 , drift vectorµx (x) given by (4), and diffusion matrix σx (x) obtained taking the square root of (5):dxt = µx (xt )dt + σx (xt )dwt ,(8)where wt is a d-dimensional Brownian motion (BM) with independent components. Equation (8) isshort-hand notation for:Z T
Z T
xT = x0 +µx (xt )dt +σx (xt )dwt ,00where T is the integration interval, the first integral is a standard (Riemann) integral, and the second
integral is an Ito integral.
Assume that Assumption A.1 holds and that SDE (8) admits an weakly unique and non-explosive weaksolution. Then the stochastic process defined by (7) with initial value x0 converges in law to such
solution. Moreover, this result continues to hold when x0 is an independent and square integrable
random variable, in which case the driving BM is independent of x0 . In both cases, the convergencein law is on D([0, ∞), Rd ) the space of functions from [0, ∞) that are continuous from the right withfinite left limits endowed with the Skorohod metric [1].
Proof of Theorem A.1. This is [9, Theorem 2.2]: Assumption A.1 and the postulated weakly uniqueand non-explosive weak solution satisfy all the conditions required for the application of [9, Theorem2.2]. Note that we use a stronger non-explosivity condition [10]. Alternatively, for this standard result
the reader can refer to the monograph [13] on which [9] is based; yet another reference is [2].
5The reader is referred to the monograph [10] for a gentle introduction to SDEs and Ito integration
theory. In Theorem A.1, the continuous-time interpolation xt of xl is introduced because weare seeking a continuous-time limiting process from a discrete-time process. Observe that theconvergence established in Theorem A.1 is strong in the sense that it concerns with the convergenceof the distribution of the stochastic process xt as a stochastic object on the whole time interval [0, T ]to the distribution of the diffusion limit. For instance, this convergence implies the joint convergenceof xt1 , . . . , xtn for any collection of times t1 , . . . , tn , and not only the convergence of the terminal
value xT .
In Theorem A.1 we postulate the existence and weak uniqueness of a weak solution of the limiting
SDE, and its non-explosive behavior. We consider weak solutions and weak (i.e. in law) uniqueness,
instead of strong solutions and strong (pathwise) uniqueness, as we are interested exclusively in
distributional aspects of the limiting process [9, 10]. Several assumptions exist in the literature inorder to guarantee that the additional assumptions of Theorem A.1 are satisfied. The following
assumptions suffice for our goals:
Assumption A.2 (Weak Existence and Uniqueness on Compacts). The functions µx (x) and σx (x)are twice continuously differentiable.
Assumption A.3 (Non-explosive Solution). There exist a finite C > 0 such that for each x ∈ Rd :kµx (x)k + kσx (x)k ≤ C(1 + kxk)
When Assumption A.1 and Assumption A.2 hold (as it will be the case in all the ResNet models
considered) but Assumption A.3 doesn’t, we still obtain convergence to the unique solution of (8) but
xt might diverge to infinity with positive probability as dl ↑ ∞.

B

Proof of Main Theorem

Lemma B.1. If φ satisfies Assumption 2.2, ε ∼ N (0, σ 2 ) with σ 2 ≤ σ∗2 , α > 0, then we can find
M2 (α, σ∗2 ) < ∞ and M3 (α, σ∗2 ) < ∞ such that:
00

α

2

E [|φ (ε)| ] ≤ M2 (α, σ∗ )000
α2E [|φ (ε)| ] ≤ M3 (α, σ∗ )Proof. We prove the result only for φ00 (ε), the case for φ000 (ε) being identical. Let L large enoughsuch that |φ00 (x)| ≤ K1 eK2 |x| for |x| ≥ L then:
 00

 00

00
α
α
α
E [|φ (ε)| ] = E |φ (ε)| 1|ε|≤L + E |φ (ε)| 1|ε|>L≤ sup |φ00 (x)|α + K1α E[eK2 α|ε| ]|x|≤LThe first term is finite. The fact that the second term can be bounded by a finite and increasing
function in σ 2 follows from the symmetry in law of ε and the form of its movement generating
function.
Proof of Theorem 2.1. Let d = 2dc and
 
xl
∈ Rd
x0l
 
b
b = l ∈ Rdbl

Wl 0 Idc
W =∈ Rd×d0 Idc Wlx=

h = W x + b ∈ Rdµ(xl )µx (x) =µ(x0l ) 2
σ (x , x ) σ 2 (xt , x0t )
σx2 (x) = 2 t0 t
σ (xt , xt ) σ 2 (x0t , x0t )6

where we dropped the dependency on l for notational convenience and reserve subscripts to indexing.
A second order Taylor expansion of φ around 0 yields for i = 1, . . . , d:√∆xi
φ(hi ∆t)
1
1== φ0 (0)hi ∆t−1/2 + φ00 (0)h2i + φ000 (εi )h3i ∆t1/2
∆t
∆t
2
6
√
√
with εi ∈ (−hi ∆t, hi ∆t). To prove (4) we want to show that ∀R > 0:11
lim sup E φ0 (0)hi ∆t−1/2 + φ00 (0)h2i + φ000 (ε)h3i ∆t1/2 − µx (x)i = 0.
∆t↓0 kxk<R2
6
Now, hi = Wi x + bi and the distribution assumptions on W and b lead to


1 00
1
0
−1/2
2
+ φ (0)hi = φ00 (0) V[W x + b]i,i = µx (x)iE φ (0)hi ∆t2
2
It remains to show that
lim sup E φ000 (εi )h3i ∆t1/2 = 0,∆t↓0 kxk<R
for which it suffices to show that supkxk<R E φ000 (εi )h3i can be bounded by M (R) < ∞ uniformly in ∆t. By Cauchy–Schwarz:



1/2
 1/2
sup E φ000 (εi )h3i ≤ sup E φ000 (εi )2
sup E h6i
.
(9)
kxk<R

kxk<R

kxk<R

The constraint supkxk<R corresponds to a constraint on the variance of hi hence the second sup is√
finite. By Lemma B.1 the first sup is finite too and not increasing in ∆t as |εi | ≤ ∆t|hi | whichallows us to produce the desired bound M (R).
Regarding (6), following a first order Taylor expansion of φ around 0 we need to show that fori = 1, . . . , d and R > 0:"
4 #φ0 (0)hi ∆t1/2 + 12 φ00 (εi )h2i ∆t=0
lim sup E∆t↓0 kxk<R∆t√√with εi ∈ (−hi ∆t, hi ∆t). Note that The term inside the expectation is composed of a sum ofterms of the form khni φ00 (εi )m ∆tα for integers n, m ≥ 0 and reals α > 0, k ∈ R. This results fromrepeated applications of the Cauchy–Schwarz inequality and Lemma B.1 as we did previously toprove (4).
Regarding (5), we can compute E[∆x(∆x)0 ]/∆t instead of V[∆x]/∆t as in the infinitesimal limitof ∆t ↓ 0 the two quantities have to agree due to the convergence of the infinitesimal mean that wehave already established. Hence following two first order Taylor expansions of φ around 0 we needto show that for i, j = 1, . . . , d and R > 0:"#
φ0 (0)hi ∆t1/2 + 12 φ00 (εi )h2i ∆t φ0 (0)hj ∆t1/2 + 12 φ00 (εj )h2j ∆t
lim sup E∆t↓0 kxk<R∆t− σx2 (x)i,j = 0√√√√
with εi ∈ (−hi ∆t, hi ∆t), εj ∈ (−hj ∆t, hj ∆t). The only term inside the expectation notvanishing in ∆t is
0
2
0
2
2
E[φ (0) hi hj ] = φ (0) V[W x + b]i,j = σx (x)i,j
The (uniform on compacts) convergence of all terms aside from σx (x)2i,j to 0 once again follows
from repeated applications of the Cauchy–Schwarz inequality and Lemma B.1.
Finally, the continuity of µx (x) and σx (x) are a consequence of the continuity of the conditional
covariance V[W x + b], and as V[W x + b] is positive semi-definite so is σx2 (x) which satisfies theexistence of its square root matrix requirement. Again from the properties of V Assumption A.2
follows, and it can be easily verified that when φ00 (0) = 0, i.e. there is no drift, Assumption A.3 issatisfied too. This completes the proof.

7