Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit

Introduction

Ordinary differential equations (ODEs) and other types of continuous-time flows have always served as convenient abstractions for various deterministic iterative models and algorithms. Recently, however, several authors have started exploring the intriguing possibility of using ODEs for constructing and training very deep neural nets by considering the limiting case of composing a large number of infinitesimal nonlinear transformations (Haber and Ruthotto, 2017; Chen et al., 2018b; Li et al., 2018). In particular, Chen et al. (2018b) have introduced the framework of neural ODEs, in which the overall nonlinear transformation is represented by an ODE, and a black-box ODE solver is used as a computational primitive during end-to-end training.

These ideas naturally carry over to the domain of probabilistic modeling. Indeed, since deep probabilistic generative models can be viewed as time-inhomogeneous Markov chains, we can consider the limit of infinitely many layers as a continuous-time Markov process. Since the marginal distributions of such a process evolve through a deterministic continuous-time flow, one can use ODE techniques in this context as well (Tabak and Vanden-Eijnden, 2010; Chen et al., 2018a). However, an alternative possibility is to focus on the stochastic evolution of the sample paths of the limiting process, rather than on the deterministic evolution in the space of measures. This perspective is particularly useful when the generation of sample paths of the underlying process is more tractable than the computation of process distributions.

In this paper, we develop these ideas in the context of Deep Latent Gaussian Models (DLGMs), a flexible family of generative models introduced by Rezende et al. (2014). In these models, the latent variable is generated by a time-inhomogeneous Markov chain, where at each time step we pass the current state through a deterministic nonlinear map, such as a feedforward neural net, and add a small independent Gaussian perturbation. The observed variable is then drawn conditionally on the state of the chain after a large but finite number of steps. The iterative structure of DLGMs, together with the use of differentiable layer-to-layer transformations, is the basis of stochastic backpropagation (Kingma and Welling, 2014; Ranganath et al., 2014; Rezende et al., 2014), an efficient and scalable procedure for performing variational inference with approximate posteriors of the mean-field type. A key feature here is that all the randomness in the latent space is generated by sampling a large but finite number of independent standard Gaussian random vectors, and all other transformations are obtained by suitable differentiable reparametrizations.

If one considers the limiting regime of DLGMs, where the number of layers tends to infinity while the step size and the noise variance in layer-to-layer transformations both tend to zero, the resulting latent random object is a diffusion process of the Itô type (Bichteler, 2002; Protter, 2005), whose drift and diffusion coefficients are implemented by neural nets. We will refer to these models as neural SDEs, in analogy to the deterministic neural ODEs of Chen et al. (2018b). Generative models of this type have been considered in earlier work, first by Movellan et al. (2002) as a noisy continuous-time counterpart of recurrent neural nets, and, more recently, by Archambeau et al. (2007), Hashimoto et al. (2016), Ha et al. (2018), and Ryder et al. (2018). On the theoretical side, Tzen and Raginsky (2019) have investigated the expressive power of diffusion-based generative models and showed that they can be used to obtain approximate samples from any distribution whose Radon–Nikodym derivative w.r.t. the standard Gaussian measure can be efficiently represented by a neural net. In this paper, we leverage this expressive power and develop a framework for variational inference in neural SDEs:

We show that all the latent randomness can be generated by sampling from the standard multidimensional Wiener process, in analogy to the use of independent standard Gaussian random vectors in DLGMs. Thus, the natural latent space for neural SDEs is the Wiener space of continuous vector-valued functions on $$ equipped with the Wiener measure (the probability law of the standard Wiener process).

We derive a variational bound on the marginal log-likelihood for the observed variable using the Gibbs variational principle on the path space (Boué and Dupuis, 1998). Moreover, by Girsanov’s theorem, any variational approximation to the posterior is related to the primitive Wiener process by a mean shift. Thus, the natural neural SDE counterpart of a mean-field approximate posterior is obtained by adding an observation-dependent neural net drift to the standard Wiener process.

Finally, we show how variational inference can be carried out via automatic differentiation (AD) in Wiener space. One of the salient features of the neural ODE framework of Chen et al. (2018b) is that one can backpropagate gradients efficiently through any black-box ODE solver using the so-called method of adjoints (see, e.g., Kokotović and Heller (1967) and references therein). While there exists a counterpart of the method of adjoints for SDEs (Yong and Zhou, 1999, Chap. 3), it cannot be used to backpropagate gradients through a black-box SDE solver, as we explain in Section 5.1. Instead, one has to either derive custom backpropagation rules for each specific solver or use the less time-efficient forward-mode AD with a black-box SDE solver. For the latter, we use the theory of stochastic flows (Kunita, 1984) in order to differentiate the solutions of Itô SDEs with respect to parameters of the drift and the diffusion coefficient. These pathwise derivatives are also solutions of Itô SDEs, whose drift and diffusion coefficient can be obtained from those of the original SDE using the ordinary chain rule of multivariable calculus. Thus, the overall process can be implemented using AD and a black-box SDE solver.

Extending the neural ODE framework of Chen et al. (2018b) to the setting of SDEs is a rather natural step that has been taken by several authors. In particular, the use of SDEs to enhance the expressive power of continuous-time neural nets was proposed in a concurrent work of Peluchetti and Favaro (2019). Neural SDEs driven by stochastic processes with jumps were introduced by Jia and Benson (2019) as a generative framework for hybrid dynamical systems with both continuous and discrete behavior. Hegde et al. (2019) considered generative models built from SDEs whose drift and diffusion coefficients are samples from a Gaussian process. Liu et al. (2019) and Wang et al. (2019) have proposed using SDEs (and suitable discretizations) as a noise injection mechanism to stabilize neural nets against adversarial or stochastic input perturbations.

Background: variational inference in Deep Latent Gaussian Models

In Deep Latent Gaussian Models (DLGMs) (Rezende et al., 2014), the latent variables $X_{0},\ldots,X_{k}$ and the observed variable $Y$ are generated recursively:

where $D(\cdot\|\cdot)$ is the Kullback–Leibler divergence and the infimum is over all Borel probability measures $\nu$ on $\Omega$. If we let $\mu$ be the marginal distribution of $Z_{0},\ldots,Z_{k}$ in (2) and apply (3) to the function $F_{\theta}(z_{0},\ldots,z_{k})\mathrel{\mathop{\mathchar 58\relax}}=-\log p(y|f_{\theta}(z_{0},\ldots,z_{k}))$, we obtain the well-known variational formula

The infimum in (4) is attained by the posterior density $p_{\theta}(z_{0},\ldots,z_{k}|y)$ whose computation is also generally intractable, so one typically picks a suitable family $\nu_{\beta}(\operatorname{d\!}z_{0},\ldots,\operatorname{d\!}z_{k}|y)=q_{\beta}(z_{0},\ldots,z_{k}|y)\operatorname{d\!}z_{0}\ldots\operatorname{d\!}z_{k}$ of approximate posteriors to obtain the variational upper bound

Backpropagation with Monte Carlo:

Since the expectation in (6) is w.r.t. to a collection of i.i.d. standard Gaussian vectors, it follows that the gradients can be computed by interchanging differentiation and expectation and using reverse-mode automatic differentiation or backpropagation (Baydin et al., 2018). Unbiased estimates of $\nabla_{\bullet}\mathsf{F}_{\theta,\beta}(y)$, $\bullet\in\{\theta,\beta\}$, can then be obtained by Monte Carlo sampling.

Neural Stochastic Differential Equations as DLGMs in the diffusion limit

In this work, we consider the continuous-time limit of (1), in analogy to the neural ODE framework of Chen et al. (2018b) (which corresponds to the deterministic case $\sigma_{i}\equiv 0$). In this limit, the latent object becomes a $d$-dimensional diffusion process $X=\{X_{t}\}_{t\in}$ given by the solution of the Itô stochastic differential equation (SDE)

where $Q_{t}f(x)\mathrel{\mathop{\mathchar 58\relax}}=\mathbf{E}_{Z\sim\phi_{d}}[f(x+\sqrt{t}Z)]$. Then $X_{1}\sim q$, i.e., one can use (8) to obtain an exact sample from $q$, and this construction is information-theoretically optimal (see, e.g., Dai Pra (1991), Lehec (2013), Eldan and Lee (2018), or Tzen and Raginsky (2019)). The drift term in (8) is known as the Föllmer drift (Föllmer, 1985). Replacing $Q_{1-t}f(x)$ by a Monte Carlo estimate and $f(\cdot)$ by a suitable neural net approximation $\widehat{f}(\cdot;\theta)$, we can approximate the Föllmer drift by functions of the form

This has the following implications (see Tzen and Raginsky (2019) for a detailed analysis):

the complexity of representing the Föllmer drift by a neural net is comparable to the complexity of representing the Radon–Nikodym derivative $f=\frac{q}{\phi_{d}}$ by a neural net;

the neural net approximation to the Föllmer drift takes both the space variable $x$ and the time variable $t$ as inputs, and its weight parameters $\theta$ do not explicitly depend on time.

In some cases, this can be confirmed by direct computation. As an example, consider a stochastic deep linear neural net (Hardt and Ma, 2017) in the diffusion limit:

In this representation, the net is parametrized by the matrix-valued paths $\{A_{t}\}_{t\in}$ and $\{C_{t}\}_{t\in}$, and optimizing over the model parameters is difficult even in the deterministic case. On the other hand, the process $X$ in (9) is Gaussian (in fact, all Gaussian diffusion processes are of this form), and the probability law of $X_{1}$ can be computed in closed form (Fleming and Rishel, 1975, Chap. V, Sec. 9): $X_{1}\sim\mathcal{N}(m,\Sigma)$ with

where $\Phi_{s,t}$ (for $s<t$) is the fundamental matrix that solves the ODE

The Föllmer drift provides a more parsimonious representation that does not involve time-varying network parameters. Indeed, since $q$ is the $d$-dimensional Gaussian density with mean $m$ and covariance matrix $\Sigma$, we have

where $c$ is a normalization constant. If $\det\Sigma\neq 0$, a straightforward but tedious computation yields

where $c_{t}>0$ is a constant that does not depend on $x$, $v=(\Sigma^{-1}-I_{d})x-\Sigma^{-1}m$, and $\Sigma_{t}=((1-t)I_{d}+t\Sigma^{-1})^{-1}$. Consequently, the Föllmer drift is given by

which is an affine function of $x$ with time-invariant parameters $\Sigma^{-1}$ and $m$.

Variational inference with neural SDEs

Our objective here is to develop a variational inference framework for neural SDEs that would leverage their expressiveness and the availability of adaptive black-box solvers for SDEs (Ilie et al., 2015). We start by showing that all the building blocks of DLGMs described in Section 2 have their natural counterparts in the context of neural SDEs.

that is, for each $t$, the path $\{[f_{\theta}(W)]_{s}\}_{s\in[0,t]}$ depends only on $\{W_{s}\}_{s\in[0,t]}$. With these ingredients in place, we have the following path-space analogue of (2):

The variational representation and Girsanov reparametrization:

then $Z=\{Z_{t}\}_{t\in}\sim\boldsymbol{\nu}$ with

Conversely, any such $\boldsymbol{\nu}$ can be realized in this fashion. This leads to the Girsanov reparametrization of the variational formula (11):

where $W+\int^{\bullet}_{0}u_{s}\operatorname{d\!}s$ is shorthand for the process $\{W_{t}+\int^{t}_{0}u_{s}\operatorname{d\!}s\}_{t\in}$, and $F_{\theta}(w)\mathrel{\mathop{\mathchar 58\relax}}=-\log p(y|[f_{\theta}(w)]_{1})$.

Mean-field approximation:

and we have the mean-field variational bound

One key difference from the DLGM set-up is worth mentioning: here, the only degree of freedom we need is an additive drift that affects the mean, whereas in the DLGM case we optimize over both the mean and the covariance matrix in Eq. (6).

Automatic differentiation in Wiener space

We are now faced with the problem of computing the gradients of the variational free energy

with respect to $\theta$ and $\beta$. The gradients of the first (KL-divegence) term on the right-hand side (13) can be computed straightforwardly using automatic differentiation, so we turn to the second term. To that end, let us define, for each $\theta$ and $\beta$, the Itô process $X^{\theta,\beta}$ by

(Here, we are assuming that the function $x\mapsto\log p(y|x)$ is sufficiently well-behaved to permit interchange of differentiation and integration.)

In the remainder of this section, we first compare the problem of gradient computation in neural SDEs to its deterministic counterpart in neural ODEs and then describe two possible approaches.

Consider the following problem: We have a $d$-dimensional Itô process

In the deterministic case, i.e., when $\sigma(\cdot)\equiv 0$, Eq. (15) is an instance of a neural ODE (Chen et al., 2018b), and the computation of $\frac{\partial J(\alpha)}{\partial\alpha}$ can be carried out efficiently using any black-box ODE solver. The key idea, based on the so-called adjoint sensitivity method (see, e.g., Kokotović and Heller (1967) and references therein), is to augment the original ODE that runs forward in time with a certain second ODE that runs backward in time. This allows one to efficiently backpropagate gradients through any black-box ODE solver.

Unfortunately, there is no straightforward way to port this construction to SDEs. In very broad strokes, this can be explained as follows: While the analogue of the method of adjoints is available for SDEs (see, e.g., Yong and Zhou (1999, Chap. 3)), the adjoint equation is an SDE that has to be solved backward in time with a terminal condition that depends on the entire Wiener path $W=\{W_{t}\}_{t\in}$, but the solution at each time $t\in$ must still be measurable only w.r.t. the “past” $\{W_{s}\}_{s\in[0,t]}$. The augmented system consisting of the original forward SDE and the adjoint backward SDE is an instance of a forward-backward SDE, or FBSDE for short (Yong and Zhou, 1999, Chap. 7). To the best of our knowledge, there are no efficient black-box schemes for solving FBSDEs with computation requirements comparable to standard SDE solvers; this stems from the fact that any procedure for solving FBSDEs must rely on a routine for solving a certain class of semilinear parabolic PDEs (Milstein and Tretyakov, 2006), which will incur considerable computational costs in high-dimensional settings. (There are, however, promising first steps in this direction by Han et al. (2018); Han and Long (2019) based on deep neural nets.)

This unfortunate complication means that we have to forgo the use of adjoint-based methods for SDEs and instead develop gradient computation procedures by other means. We describe two possible approaches in the remainder of this section. As stated earlier, we assume that the following building blocks are available:

We assume, moreover, that one can pass straight-line programs for computing $f$ and $g$ as arguments to SDE.Solve.

2 Solve-then-differentiate: the Euler backprop

The most straightforward approach is to derive custom backpropagation equations for a specific SDE solver, e.g., the Euler method. We first generate an Euler approximation of the diffusion process (14) and then estimate the gradients of $-\mathbf{E}[\log p(y|X^{\theta,\beta}_{1})]$ w.r.t. $\theta$ and $\beta$ by backpropagation through the computation graph of the Euler recursion:

the forward pass — given a time mesh $0=t_{0}<t_{1}<\ldots<t_{N}=1$, we sample $Z_{1},\ldots,Z_{N}\stackrel{{\scriptstyle{\rm i.i.d.}}}{{\sim}}\gamma_{d}$ and the desired initialization $\widehat{X}^{\theta,\beta}_{t_{0}}$, and generate the updates

for $0\leq i<N$, where $h_{i+1}\mathrel{\mathop{\mathchar 58\relax}}=t_{i+1}-t_{i}$;

the backward pass — compute the gradients of $-\log p(y|\widehat{X}^{\theta,\beta}_{t_{N}})$ w.r.t. $\theta$ and $\beta$ using reverse-mode AD.

3 Differentiate-then-solve: the pathwise differentiation method

In some settings, it may be desirable to avoid explicit discretization of the neural SDE and work with a black-box SDE solver instead. This approach amounts to differentiating through the forward pass of the solver. The computation of the pathwise derivatives of $X^{\theta,\beta}$ with respect to $\theta$ or $\beta$ can be accomplished by solving another SDE, as a consequence of the theory of stochastic flows (Kunita, 1984).

The drift $b(x,t;\theta)$ and the diffusion matrix $\sigma(x,t;\theta)$ are Lipschitz-continuous with Lipschitz-continuous Jacobians in $x$ and $\theta$, uniformly in $t\in$.

Then the pathwise derivatives of $X=X^{\theta,\beta}$ in $\theta$ and $\beta$ are given by the following Itô processes:

We will use the results of Kunita (1984) on differentiability of the solutions of Itô SDEs w.r.t. initial conditions. Consider an $m$-dimensional Itô process of the form

where $W^{1},\ldots,W^{m}$ are independent standard scalar Brownian motions, with the following assumptions:

Then (Kunita, 1984, Chap. 2, Thm. 3.1) the pathwise derivatives of $Z$ w.r.t. the initial condition $\xi$ exist and are given by the Itô processes

where $\sigma_{j}(x,t;\theta)$ is the $j$th column of $\sigma(x,t;\theta)$. These vector fields satisfy the above Lipschitz continuity condition by hypothesis. Consider the Itô process (18) with the initial condition $\xi=(0^{\top},\beta^{\top},\theta^{\top})^{\top}$. Then evidently

and Eqs. (16) and (17) follow from (19), (20), and the chain rule of multivariable calculus. ∎

using forward-mode or reverse-mode AD as needed.

for (17), we need the Jacobian of $b(x,t;\theta)$ and $\sigma(x,t;\theta)$ w.r.t. $\theta$ and $x$. The total per-iteration time complexity of generating the Jacobians using AD will be

Typically, the dimension $n$ of the latent parameter $\theta$ will be on the order of $d^{2}$ for a fully connected neural net.

The solve-then-differentiate approach of Section 5.2 will generally scale better to high-dimensional problems than the black-box pathwise approach. On the other hand, the pathwise approach is more flexible since it works with a generic SDE solver, so the overall time complexity may be reduced by using an adaptive SDE solver (Ilie et al., 2015). In addition, since the pathwise approach amounts to differentiating through the forward operation of the solver, it may incur smaller storage overhead than the Euler backprop method.

As before, the gradients of the free energy w.r.t. $\beta$ and $\theta$ can be estimated using Monte Carlo methods, by averaging multiple independent runs of the SDE solver.

Experimental results

We evaluated the performance of the forward-differentiation pathwise method of Sec. 5.3 using gradient descent on synthetic data. The code for all experiments was written in Julia using the DiffEqFlux library (Rackauckas et al., 2019) and executed on a CPU. The synthetic data were generated via a numerical SDE solution of the diffusion process

Interestingly, the improvements from discretization meshes finer than $h=1/32$ are incremental, suggesting that, at least in a simple case such as this, models with “infinitely many layers” may not offer significant practical advantage over models with “finitely many" layers, such as DLGMs, where the number of layers approaches the dimensionality of the problem. On the other hand, while the method exhibits some robustness when data are not abundant ($n\approx d$), doing not much worse than when they are ($n\gg d$), it does increasingly worse—reflected in log-likelihood attained by optimized parameters—and ultimately fails in a low-data regime, as $n$ approaches $\sqrt{d}$.

Conclusion and future directions

We have presented an analysis of neural SDEs, which can be viewed as a continuous-time limit of the DLGMs of Rezende et al. (2014) or as a stochastic version of the neural ODEs of Chen et al. (2018b). In particular, this is a best-of-both-worlds perspective that enables us to both reason about the compositional expressive power of DLGMs and the inference process in the space of measures. In addition, it allows us to draw on a variety of standard scientific computing methods developed for continuous-time stochastic processes.

We have shown that optimization for neural SDEs is far from a straightforward analogue of the ODE case, with the issue of time-adaptedness of paths complicating the use of the adjoint sensitivity method for gradient computation. The discretize-then-differentiate approach was shown to exactly recover the stochastic backpropagation method of Rezende et al. (2014) when using a simple Euler discretization, with the form of the mean-field variational approximation closely preserved; and the differentiate-then-discretize pathwise approach was demonstrated to be a single-pass method that leverages a black-box SDE solver, with little storage overhead but high computational complexity and thus limited scalability due to the use of forward-mode differentiation.

One interesting direction to examine going forward is more sophisticated discretization schemes that readily enable the use of numerical tools developed for continuous-time processes (e.g., ODE solvers), where the continuous-time process is viewed as a deterministic function of random increments, perhaps arbitrarily small as the discretization becomes increasingly fine. Another promising direction is to investigate the connection between neural SDEs and probabilistic ODE solvers that return a posterior estimate rather than a deterministic approximate solution (Conrad et al., 2017; Schober et al., 2019).

The authors would like to thank Matus Telgarsky for many enlightening discussions, and Chris Rackauckas, Markus Heinonen, and Mauricio Álvarez for their comments and constructive suggestions on the first version of this work. This work was supported in part by the NSF CAREER award CCF-1254041, in part by the Center for Science of Information (CSoI), an NSF Science and Technology Center, under grant agreement CCF-0939370, in part by the Center for Advanced Electronics through Machine Learning (CAEML) I/UCRC award no. CNS-16-24811, and in part by the Office of Naval Research under grant no. N00014-12-1-0998.

References