---
abstract: 'We introduce stochastic normalizing flows, an extension of continuous normalizing flows for maximum likelihood estimation and variational inference (VI) using stochastic differential equations (SDEs). Using the theory of rough paths, the underlying Brownian motion is treated as a latent variable and approximated, enabling efficient training of neural SDEs as random neural ordinary differential equations. These SDEs can be used for constructing efficient Markov chains to sample from the underlying distribution of a given dataset. Furthermore, by considering families of targeted SDEs with prescribed stationary distribution, we can apply VI to the optimization of hyperparameters in stochastic MCMC.'
author:
- 'Liam Hodgkinson[^1]'
- 'Chris van der Heide[^2]'
- 'Fred Roosta[^3]'
- 'Michael W. Mahoney[^4]'
title: '<span style="font-variant:small-caps;">Stochastic Normalizing Flows</span>'
---

Introduction
============

*Normalizing flows* [@rezende2015variational] are probabilistic models constructed as a sequence of successive transformations applied to some initial distribution. A key strength of normalizing flows is their expressive power as generative models, while enjoying an explicitly computable form of the likelihood function evaluated on the transformed space. This makes them especially well-equipped for variational inference (VI). Neural networks are often used as inspiration for finding effective transformations [@dinh2014nice; @berg2018sylvester].

*Continuous normalizing flows* were later developed in [@chen2018neural] as a means to perform maximum likelihood estimation and VI for large-scale probabilistic models derived from ordinary differential equations (ODEs). The framework stems from the computation of the evolving density of an ODE with random initial value, as the solution to another ODE. The jump to continuous-time dynamics affords a few computational benefits over its discrete-time counterpart, namely the presence of a trace in place of a determinant in the evolution formulae for the density, as well as the adjoint method for memory-efficient backpropagation. Motivated by deep learning, a family of ODEs, called *neural ordinary differential equations* were constructed, whose Euler discretizations resembled layer-wise transformations of residual neural networks. Further algorithmic improvements to the framework were presented by @grathwohl2018ffjord, enabling virtually arbitrary choices of parameterized classes of ODEs. Doing all this involves some technical subtlety, and effective neural ODE architectures remain the subject of ongoing research — see for example [@dupont2019augmented; @gholami2019anode; @zhang2019anodev2].

There has also been recent interest in extending these frameworks to a stochastic scenario, that is, training probabilistic models derived from *stochastic* differential equations (SDEs). For physical models, where the evolution of a dynamical system is no longer deterministic, or microscopic fluctuations are dependent on components changing too rapidly to quantify, an SDE can be more appropriate. Stochastic extensions of neural ODEs have been considered in [@tzen2019neural; @liu2019neural; @jia2019neural; @peluchetti2019infinitely] as limits of deep latent Gaussian models, where they have been suggested to show increased robustness to noisy / adversarial data. Furthermore, unlike deterministic flows, there is a foolproof recipe for constructing a family of SDEs that are ergodic with respect to some target distribution [@ma2015complete]. This particular property guarantees the convergence of the solution of an SDE to a prescribed target distribution. Such SDEs are prime candidates for the construction of stochastic MCMC algorithms, by generating sample paths via approximate stochastic integration methods.

However, developing an analogue of the continuous normalizing flows framework for flows constructed from SDEs—in particular, one that comes with simple and rigorous mathematical theory and that does not rely on ad hoc or problem-specific assumptions—is far from trivial. A common approach for conducting VI with SDEs is to rely on Girsanov’s theorem. This allows one to estimate the Kullback-Leibler divergence between densities of solutions to two SDEs (for the prior and posterior distributions) with differing drift coefficients [@beskos2006exact; @tzen2019neural]. Following this approach, @li2020scalable developed a stochastic adjoint method which scales well to high dimensions, and enables SDEs as latent models in variational autoencoders. Theoretical justification of the method proved challenging, as stochastic calculus is ill-suited for analyzing backward (approximate) solutions to SDEs. Notable deficiencies with these previous approaches include difficulties with non-diagonal diffusion, incompatibility with higher-order adaptive SDE solvers, and a complex means of reconstructing Brownian motion paths from random number generator seeds. Furthermore, the method cannot be justifiably combined with existing approaches of density estimation for SDEs (see @hurn2007seeing).

On the other hand, recent efforts have made significant strides in applying variational and MCMC methods for idealized Bayesian computation. One of the most significant contributions in this direction is @salimans2015markov, who performed VI with respect to distributions formed from steps of a reversible Markov chain. For example, the setting of Hamiltonian Monte Carlo was examined in @wolf2016variational. More recently, @liu2016two considered optimizing step size in stochastic gradient Langevin dynamics using methods derived from kernelized Stein discrepancy. Langevin flows [@rezende2015variational] have been discussed as a potential VI framework that takes inspiration from the SDEs underlying stochastic MCMC [@ma2015complete]. Once again, implementation of Langevin flows relies on the approximation of the log-likelihood for a general class of SDEs.

Contributions {#contributions .unnumbered}
-------------

We provide a general theoretical framework (which we refer to as *stochastic normalizing flows*) for approximating generative models constructed from SDEs using continuous normalizing flows. These approximations can then be trained using existing techniques. By this process, we find that theoretical and practical developments concerning continuous normalizing flows and neural ODEs extend readily to the stochastic setting, without the need of an independent framework. The key theoretical enabler underlying our strong results and simple analysis is the *theory of rough paths* [@friz2014course], an alternative stochastic calculus that enables approximation and pathwise treatment of SDEs. Our approach

1.  enables (i) density estimation, (ii) maximum likelihood estimation, and (iii) variational approximations beyond autoencoders, for *arbitrary* SDE models; and

2.  is easily implemented using any general continuous normalizing flows implementation, such as that of @grathwohl2018ffjord.

Our framework recovers the stochastic adjoint method of @li2020scalable, but our approach is sufficiently flexible to overcome its deficiencies. Moreover, using our approach, any existing neural ODE framework (such as @zhang2019anodev2) can be extended to SDEs, simply by the addition of a few extra terms.

Following a review of background material in §\[sec:Review\], the stochastic normalizing flows framework is introduced and discussed in §\[sec:Stochastic\], with our main approximation result presented in Theorem \[thm:RoughNormFlow\]. Some numerical investigations are conducted in §\[sec:Numerics\], including an application to hyperparameter optimization in stochastic MCMC.

Background Review {#sec:Review}
=================

Continuous Normalizing Flows
----------------------------

We shall begin by reviewing the continuous normalizing flow framework for training ODE models, as our development of random and stochastic normalizing flows will build upon it. Consider a parameterized class of models $\{Z_\theta\}_{\theta \in \mathbb{R}^m}$ of the following form: for $f:\mathbb{R}^d \times [0,T] \times \mathbb{R}^m \to \mathbb{R}^d$, let $Z = Z_\theta \in \mathbb{R}^d$ satisfy the ODE with random initial condition (often called a *random ordinary differential equation*) $$\label{eq:ODEModel}
\frac{{\mathrm{d}}}{{\mathrm{d}}t} Z(t) = f(Z(t), t, \theta),\quad 
Z(0) \sim p_0(\theta).$$ In a general machine learning context, one might choose $f$ such that the Euler discretization of (\[eq:ODEModel\]) resembles layer-wise updates of a residual neural network [@lu2017beyond; @chen2018neural], or one may parameterize $f$ as a neural network itself [@grathwohl2018ffjord]. The resulting ODEs constitute the class of so-called *neural ordinary differential equations*. The following theorem is a consequence of the Liouville equation (equivalently, Fokker-Planck equation) applied to the solution $Z(t)$ of the random ODE (\[eq:ODEModel\]), and it yields an ODE for the log density of $Z(t)$ evaluated at $Z(t)$.

\[thm:ContNormFlow\] Suppose that $Z(t)$ satisfies (\[eq:ODEModel\]). The distribution of $Z(t)$ is absolutely continuous with respect to Lebesgue measure, with probability density $p_t$ satisfying $$\label{eq:ContNormFlow}
\frac{{\mathrm{d}}}{{\mathrm{d}}t} \log p_t(Z(t)) = -\nabla_z \cdot f(Z(t), t, \theta) $$

Naively computing the divergence in (\[eq:ContNormFlow\]) with automatic differentiation is of quadratic complexity in the dimension $d$. As pointed out by @grathwohl2018ffjord, this can be improved to linear complexity using a trace estimator [@roosta2015improved]: $$\label{eq:TraceEst}
\nabla_z \cdot f(z) = {\mathrm{tr}}\left(\frac{\partial f}{\partial z}\right) \approx \frac1n \sum_{k=1}^{n} \epsilon_k^\top \frac{\partial f}{\partial z} \epsilon_k,$$ where each $\epsilon_k$ is an independent and identically distributed copy of a random vector $\epsilon \in \mathbb{R}^d$ with zero mean and $\mathbb{E}[\epsilon \epsilon^\top] = I$. Common choices for $\epsilon_k$ include standard normal and Rademacher random vectors.

The Adjoint Method
------------------

Training continuous normalizing flows often involves minimizing a scalar loss function involving $Z$ and/or the log-density computed via Theorem \[thm:ContNormFlow\] with respect to the parameters $\theta$. For this, we require gradients of $Z(t)$ with respect to $\theta$ for $t \in [0,T]$. The most obvious approach is to directly backpropagate through a numerical integration scheme such as in @ryder2018black, but this does not scale well in $T$. The superior alternative is the *adjoint method*, which computes derivatives of a scalar loss function by solving another differential equation in reverse time. Letting $L$ denote a scalar loss depending on $Z(T)$, the *adjoint* given by $a(t) = \frac{\partial L}{\partial Z(t)}$, as well as the gradient of $L$ in $\theta$, satisfy [@pontryagin2018mathematical §12]

\[eq:AdjointMethod\] $$\begin{aligned}
\frac{{\mathrm{d}}}{{\mathrm{d}}t} a(t) &= -\nabla_z f(Z(t),t,\theta) a(t), \\
\nabla_\theta L &= \int_0^T \nabla_\theta f(Z(t), t, \theta)a(t) {\mathrm{d}}t.\end{aligned}$$

Together with (\[eq:ODEModel\]), the equations (\[eq:AdjointMethod\]) are solved in reverse time, starting from the terminal values $Z(T)$ and $\nabla L(Z(T))$. By augmenting $Z(t)$ together with (\[eq:ContNormFlow\]), this method also allows for loss functions depending on $p_T(Z(T))$.

Solving (\[eq:ODEModel\]), (\[eq:ContNormFlow\]), and (\[eq:AdjointMethod\]) can be achieved using off-the-shelf numerical integrators. Adaptive solvers prove particularly effective, although, as pointed out in @gholami2019anode, the backward solve (\[eq:AdjointMethod\]) can often run into stability issues, suggesting a Rosenbrock or other implicit approach [@hairer1996solving]. We point out that the same is also true in the stochastic setting; see @hodgkinson2019implicit, for example. For further implementation details concerning continuous normalizing flows, we refer to @grathwohl2018ffjord.

Rough Path Theory
-----------------

The theory of rough paths was first introduced in [@lyons1998differential] to provide a supporting pathwise theory for SDEs. It has since flourished into a coherent pathwise alternative to stochastic calculus, facilitating direct stochastic generalizations of results from the theory of ODEs — we refer to @friz2014course for a gentle introduction, and @friz2010multidimensional for a thorough treatment of the topic. Suppose that we would like to prescribe meaning to the infinitesimal limit of the sequence of iterates $$\label{eq:RDELimit}
Z_{t+h} = Z_t + f(Z_t)(X_{t+h} - X_t),\quad \mbox{as }h \to 0^+.$$ In the case of SDEs, $X_t$ is a sample path of Brownian motion, so that each $X_{t+h} - X_t$ is a realization of a normal random vector with zero mean and covariance $h I$. Unfortunately, a strong limit of (\[eq:RDELimit\]) fails to exist if $X_t$ is too “rough”. In particular, suppose that $X_t$ is $\alpha$-Hölder continuous for $\alpha \in (0,1)$, that is, there exists some $C > 0$ such that $\|X_s - X_t\| \leq C|s - t|^{\alpha}$ for any $s,t \geq 0$. Since the limit (\[eq:RDELimit\]) is only well-defined if $\alpha \geq 1/2$ [@young1936inequality], a function on $[0,T]$ is *rough* if it is Hölder-continuous only for $\alpha < 1/2$. Sample paths of Brownian motion constitute rough paths under this definition. The problem is that the discretization (\[eq:RDELimit\]) invokes the *zeroth-order* approximation $f(Z_{t+s}) \approx f(Z_t)$ for $0 \leq s \leq h$, which proves too poor. By instead taking a *first-order* approximation $$\begin{aligned}
f(Z_{t+s}) &\approx f(Z_t) + \nabla_z f(Z_t)(Z_{t+s}-Z_t) \\&\approx f(Z_t) + \nabla_z f(Z_t)f(Z_t)(X_{t+s}-X_t),\end{aligned}$$ we arrive at the *Davie scheme* [@davie2008differential] $$\label{eq:DavieApprox}
Z_{t+h} = Z_t + f(Z_t) (X_{t+h} - X_t) + \nabla_z f(Z_t)f(Z_t)\mathbb{X}_{t,t+h},$$ where $\mathbb{X}_{s,t}$ represents the “integral” $\int_s^t X_r {\mathrm{d}}X_r^\top$. Once again, we cannot uniquely define $\mathbb{X}$ from the path $X$ itself, so instead we prescribe it. In fact, each choice of $\mathbb{X}$ satisfying Chen’s relations $$\mathbb{X}_{s,t} - \mathbb{X}_{s,u} - \mathbb{X}_{u,t} = (X_s - X_u)(X_t - X_u)^\top,$$ for any $s,u,t \geq 0$, will reveal a *different* limit for (\[eq:DavieApprox\]) as $h \to 0^+$, provided $\alpha \geq 1/3$ (for smaller $\alpha$, higher-order approximations are necessary). The pair $\boldsymbol{X} = (X,\mathbb{X})$ is referred to as a *rough path*, and the limit of (\[eq:DavieApprox\]) as $h \to 0^+$ is the solution to the *rough differential equation* (RDE) $$\label{eq:RoughDiffEq}
{\mathrm{d}}\boldsymbol{Z}_t = f(Z_t) {\mathrm{d}}\boldsymbol{X}_t.$$ Hölder continuity is critical to rough path theory — in the sequel, we equip the space of $\alpha$-Hölder functions with the $\alpha$-Hölder norm, defined by $$\|X\|_\alpha \coloneqq \sup_{t \in [0,T]}\|X_t\| + \sup_{\substack{s,t\in [0,T]\\s \neq t}}\frac{\|X_t - X_s\|}{|t - s|^{\alpha}}.$$ This definition extends to the iterated integral $\mathbb{X}$ by replacing $X_t$ and $X_t - X_s$ with $\mathbb{X}_{0,t}$ and $\mathbb{X}_{s,t}$, respectively.

It is useful to identify a calculus which satisfies the usual chain and product rules. This occurs when the rough path $\boldsymbol{X}$ is *geometric*, that is, $$\label{eq:GeoRoughPath}
\mathbb{X}_{s,t} - \mathbb{X}_{t,s} = \tfrac12 (X_t - X_s) (X_t - X_s)^\top,\quad \forall s,t \geq 0.$$ Every continuous and piecewise differentiable function $X$ is canonically lifted to a geometric rough path by taking $\mathbb{X}_{s,t} = \int_s^t X_r \frac{{\mathrm{d}}}{{\mathrm{d}}r} X_r^\top {\mathrm{d}}r$, where the derivative is interpreted in the weak sense. In these cases, (\[eq:RoughDiffEq\]) equates to the ODE $\frac{{\mathrm{d}}}{{\mathrm{d}}t}Z_t = f(Z_t) \frac{{\mathrm{d}}}{{\mathrm{d}}t} X_t$.

Geometric rough paths have two key properties of interest:

1.  The canonical lifts of any sequence of smooth approximations $X^{(n)}$ which converge to $X$ as $n \to \infty$ in the $\alpha$-Hölder norm, also converge in the $\alpha$-Hölder rough path metric $$\varrho_{\alpha}((X,\mathbb{X}),(Y,\mathbb{Y})) = \|X - Y\|_{\alpha} + \|\mathbb{X} - \mathbb{Y}\|_{2\alpha},$$ to a geometric rough path $(X,\mathbb{X})$. Conversely, any geometric rough path can be approximated by some sequence of smooth paths [@friz2014course Proposition 2.5].

2.  The reverse-time process $\tilde{Z}_t = Z_{T-t}$ of a solution $Z_t$ to any rough differential equation (\[eq:RoughDiffEq\]) with Lipschitz $f$, itself satisfies the reversed rough differential equation ${\mathrm{d}}\tilde{\bm{Z}}_t = -f(T-t,\tilde{Z}_t){\mathrm{d}}\bm{X}_{T-t}$ *if and only if* $\bm{X}$ is geometric.

By property I, any solution to RDEs driven by a geometric rough path can be approximated by solutions to ODEs. Property II, which follows readily from the definition (\[eq:GeoRoughPath\]) in the limit (\[eq:DavieApprox\]), enables the adjoint method for rough differential equations driven by a geometric rough path.

Stochastic Normalizing Flows {#sec:Stochastic}
============================

Let $Z_t$ satisfy the Itô SDE $$\label{eq:MainSDE}
{\mathrm{d}}Z_t = \mu_t(Z_t, \theta) {\mathrm{d}}t + \sigma_t(Z_t, \theta) {\mathrm{d}}B_t, \quad Z_0 \sim p_0(\theta),$$ where $B_t$ is an $m$-dimensional Brownian motion, and $\mu_t:\mathbb{R}^d \to \mathbb{R}^d$, $\sigma_t:\mathbb{R}^d \to \mathbb{R}^{d \times m}$ are the drift, and diffusion coefficients, respectively. Analogous to neural ODEs, neural SDEs choose $\mu_t$ to resemble a single layer of a neural network [@tzen2019neural]. The dropout-inspired construction of @liu2019neural suggests taking $\sigma_t \propto \mathrm{diag}(\mu_t)$. Alternatively, one can parameterize both $\mu_t$ and $\sigma_t$ by multi-layer neural networks.

The reliance of stochastic calculus on non-anticipating processes as well as the lack of continuity for solution maps of Itô SDEs necessitates complicated and delicate arguments for extending each piece of the continuous normalizing flow framework from §\[sec:Review\] to SDEs. We bypass the intricacies of existing theoretical treatments of neural SDEs by an approximation argument: for a smooth approximation $\tilde{B}_t$ of Brownian motion $B_t$, we estimate solutions of an SDE by a random ODE involving $\tilde{B}_t$. One must take great care with such approximations. For example, geometric Brownian motion, that is, the solution to ${\mathrm{d}}Z_t = \sigma Z_t {\mathrm{d}}B_t$, has the explicit expression $Z_t = Z_0 \exp(-\frac{\sigma^2}{2} t + \sigma B_t)$, which is not well-approximated by the solution $\tilde{Z}_t = Z_0 \exp(\sigma \tilde{B}_t)$ to $\frac{{\mathrm{d}}}{{\mathrm{d}}t} \tilde{Z}_t = \sigma \tilde{Z}_t \frac{{\mathrm{d}}\tilde{B}_t}{{\mathrm{d}}t}$. Theoretical verification of this approach is challenging using traditional stochastic calculus due to the irregularity of solution maps. Instead, we rely on rough path theory — particularly properties I and II of geometric rough paths.

In the rough path framework, one can reconstruct the Itô stochastic calculus via the rough path $\bm{B}^{{\text{It\^{o}}}} = (B, \mathbb{B}^{{\text{It\^{o}}}})$, where $\mathbb{B}^{{\text{It\^{o}}}}_{s,t} = B_t (B_t - B_s)^\top - \frac{t-s}2 I$. Indeed, by @friz2014course [Theorem 9.1], letting $\bm{B}^{{\text{It\^{o}}}}(\omega)$ denote a realization of the Itô Brownian motion rough path, the solution to the rough differential equation $$\label{eq:ItoRDE}
{\mathrm{d}}\bm{Z}_t = \mu_t(Z_t, \theta) {\mathrm{d}}t + \sigma_t(Z_t, \theta) {\mathrm{d}}\bm{B}^{{\text{It\^{o}}}}_t(\omega)$$ is a realization of the strong solution to (\[eq:MainSDE\]). Likewise, the Davie scheme (\[eq:DavieApprox\]) corresponds to the Milstein integrator for SDEs [@kloeden2013numerical §10.3].

Unfortunately, $\bm{B}^{{\text{It\^{o}}}}(\omega)$ is not a geometric rough path, and so Theorem \[thm:RoughNormFlow\] cannot be directly applied. Instead, we shall proceed according to the following steps:

1.  Convert the Itô SDE to a Stratonovich SDE (§\[sec:Strat\]).

2.  Interpret the Stratonovich SDE pathwise as an RDE driven by a geometric rough path $\bm{B}^{{\text{Strat}}}$ (\[eq:StratRDE\]).

3.  Approximate the pathwise Stratonovich RDE by a random ODE (§\[sec:WongZakai\]).

4.  Train the random ODE as a continuous normalizing flow with added latent variables (§\[sec:RandCNF\]).

Consequently, the RDE (\[eq:ItoRDE\]) is estimated by the ODE $\frac{{\mathrm{d}}Z_t(\omega)}{{\mathrm{d}}t} = F_\omega(Z_t(\omega),t,\theta)$ where $$F_\omega(z, t, \theta) = \underset{\text{Stratonovich drift}}{\underbrace{\tilde{\mu}_t(z, \theta)}} + \;\;\sigma_t(z,\theta)\underset{\text{approximation}}{\underbrace{\frac{{\mathrm{d}}B_t(\omega)}{{\mathrm{d}}t}}}.$$

Stratonovich calculus {#sec:Strat}
---------------------

The unique geometric rough path formed from Brownian motion yields the Stratonovich calculus: $\mathbb{B}^{{\text{Strat}}}_{s,t} = B_t(B_t - B_s)^\top$. A Stratonovich differential equation is commonly written in the form ${\mathrm{d}}Z_t = \mu_t(Z_t) {\mathrm{d}}t + \sigma_t(Z_t) \circ {\mathrm{d}}B_t$, where $\circ$ denotes Stratonovich integration: for a process $Y_t$ adapted to the filtration generated by $B_t$, $$\int_s^t Y_t \circ {\mathrm{d}}B_t = \lim_{|\mathcal{P}|\to 0}\sum_{k=1}^N \frac12(Y_{t_k} + Y_{t_{k-1}}) (B_{t_k} - B_{t_{k-1}}),$$ where $\mathcal{P} = \{0 = t_0 < \cdots < t_N = T\}$ is a partition with mesh size $|\mathcal{P}| = \max_k |t_k - t_{k-1}|$, and the limit is in $L^2$. This is to be compared with Itô integration which is defined instead by $$\int_s^t Y_t \,{\mathrm{d}}B_t = \lim_{|\mathcal{P}|\to 0}\sum_{k=1}^N Y_{t_{k-1}} (B_{t_k} - B_{t_{k-1}}).$$ Stratonovich differential equations were recognized in @li2020scalable to be the correct setting for extending the adjoint method to SDEs. However, the adherence to classical stochastic calculus, which relies on adaptedness, somewhat complicates the argument. In our setting, the advantages of Stratonovich differential equations are clear. Because Stratonovich differential equations can be arbitrarily well-approximated by random ODEs, *all* methods of training continuous normalizing flows extend to them, including the adjoint method. Any Itô SDE can be converted into a Stratonovich SDE by adjusting the drift [@evans2012introduction p. 123], a fact readily seen by comparing limits of (\[eq:DavieApprox\]) with $\mathbb{B}^{{\text{It\^{o}}}}$ and $\mathbb{B}^{{\text{Strat}}}$. The following formula is particularly amenable to implementation with automatic differentiation: the Itô SDE ${\mathrm{d}}Z_t = \mu_t(Z_t) {\mathrm{d}}t + \sigma_t(Z_t) {\mathrm{d}}B_t$ is equivalent to the Stratonovich SDE ${\mathrm{d}}Z_t = \tilde{\mu}_t(Z_t) {\mathrm{d}}t + \sigma_t(Z_t) \circ {\mathrm{d}}B_t$ provided that for each $i=1,\dots,d$, $$\label{eq:ItoCorrection}
\tilde{\mu}_t^i(x) = \mu_t^i(x) - \tfrac12\nabla_x \cdot (\sigma_t(x)\sigma_t^\top(x^{\ast}))_i,$$ where $x^\ast$ is an independent copy of $x$, and the subscript denotes the $i$-th row. Once again, we can make use of the trace estimator (\[eq:TraceEst\]) to increase performance in higher dimensions. In the rough path theory, Stratonovich SDEs are interpreted pathwise according to the RDE $$\label{eq:StratRDE}
{\mathrm{d}}\bm{Z}_t = \tilde{\mu}_t(Z_t, \theta) {\mathrm{d}}t + \sigma_t(Z_t,\theta) {\mathrm{d}}\bm{B}_t^{{\text{Strat}}}(\omega),$$ which is equivalent to (\[eq:ItoRDE\]).

Wong–Zakai approximations {#sec:WongZakai}
-------------------------

A random ODE $\frac{{\mathrm{d}}}{{\mathrm{d}}t} Z_t^{(n)} = \mu_t(Z_t^{(n)}) + \sigma_t(Z_t^{(n)}) \frac{{\mathrm{d}}B_t^{(n)}}{{\mathrm{d}}t}$ estimating a Stratonovich SDE ${\mathrm{d}}Z_t = \mu_t(Z_t) {\mathrm{d}}t + \sigma_t(Z_t) \circ {\mathrm{d}}B_t$ is commonly referred to as a *Wong–Zakai approximation* [@twardowska1996wong], after the authors of the seminal paper [@wong1965convergence], who first illustrated this concept for one-dimensional Brownian motion. We shall consider two types of Wong–Zakai approximation: a Karhunen-Loève expansion, and a piecewise linear function. These approximations are compared in Figure \[fig:WienerApprox\]. In practice, we have found that the Karhunen-Loève expansion with $4 \leq n \leq 10$ terms works well for training, while the piecewise linear approximation is preferable for testing.

![Karhunen-Loeve (left) and piecewise linear (right) approximations of a Brownian motion sample path with $n = 6$ and $\Delta t = \frac{1}{6}$ respectively.[]{data-label="fig:WienerApprox"}](wiener_approximations.pdf){width="\columnwidth"}

### Piecewise linear

Easily the most common approximation of Brownian motion involves exact simulation on a discrete set of times $\{0,t_1,t_2,\dots,t_n\}$, followed by linear interpolation. More precisely, letting $\Delta t_k = t_{k+1}-t_k$, for each $k=0,\dots,n-1$, we let $$B_{t_{k+1}}^{(n)} = B_{t_k}^{(n)} + \sqrt{\Delta t_k}\omega_k,\quad \omega_k \sim \mathcal{N}(0,1),$$ and consider the approximation $$B_t^{(n)} = B_{t_k}^{(n)} + \frac{t - t_k}{t_{k+1}-t_k} (B_{t_{k+1}}^{(n)} - B_{t_k}^{(n)}),\quad t \in [t_k,t_{k+1}].$$ Integrating the resulting Wong–Zakai approximation using Euler’s method on the same set of time points is equivalent to performing the Euler–Maruyama method for solving the Stratonovich SDE. By @friz2010multidimensional [Theorem 15.45], as the mesh size $\max_k \Delta t_k \to 0$, the piecewise linear approximation converges almost surely to Brownian motion in the $\alpha$-Hölder norm for any $\alpha < 1/2$.

### Karhunen-Loève expansion

For any zero-mean Gaussian process $X_t$ on $\mathbb{R}^d$ with $t\in [0,T]$, the covariance function $K(s,t) = \mathbb{E}[X_s X_t^\top]$ is a positive-definite kernel. If $K$ is also continuous, Mercer’s theorem guarantees the existence of an orthonormal basis on $L^2([0,T],\mathbb{R}^d)$ of eigenfunctions $\{e_k\}_{k=1}^{\infty}$ with corresponding positive eigenvalues $\{\lambda_k\}_{k=1}^{\infty}$ such that $K(s,t)=\sum_{j=1}^{\infty} \lambda_j e_j(s) e_j(t)$. The process $X_t$ can be expanded in terms of these eigenfunctions as $$X_t = \sum_{k=1}^{\infty} \sqrt{\lambda_k} \omega_k e_k(t), \quad \omega_k \sim \mathcal{N}(0,1),$$ where each $\omega_k$ is independent. This is called the *Karhunen-Loève expansion* of $X$. Truncating the series after $n$ terms yields the $n$-th order Karhunen-Loève approximation, and has the smallest mean squared error over all expansions with $n$ orthogonal basis elements. Recalling that we are primarily interested in the endpoints of the solution, instead of expanding Brownian motion itself, we consider an approximation $B_t^{(n)}$ derived from the Karhunen-Loève expansion of the Brownian bridge $B_t - B_T \frac{t}{T}$: $$B_t^{(n)} = \omega_0 \frac{t}{\sqrt{T}} + \sum_{k=1}^{n - 1} \omega_k \frac{\sqrt{2T} \sin(k\pi t  / T)}{k\pi},\quad n = 1,2,\dots.$$ Using this approximation ensures that the terminal density for SDEs with constant drift and diffusion coefficients is computed exactly. By @friz2010multidimensional [Theorem 15.51], $(B_t^{(n)})_{t \in [0,T]}$ converges almost surely as $n \to \infty$ to Brownian motion in the $\alpha$-Hölder norm for any $\alpha < 1/2$. Furthermore, since $B_t^{(n)}$ is smooth, Wong–Zakai approximations involving $\frac{{\mathrm{d}}B_t^{(n)}}{{\mathrm{d}}t}$ may be readily solved using adaptive ODE solvers.

Main result
-----------

Using Wong–Zakai approximations, a Stratonovich SDE can be uniformly approximated in Hölder norm by random ODEs. In Theorem \[thm:RoughNormFlow\], we show that the log-densities and loss function gradients for these random ODEs also converge appropriately. More generally, geometric rough paths (including the Stratonovich paths (\[eq:StratRDE\])) with random initial conditions can be approximately trained as random ODEs.

\[thm:RoughNormFlow\] Let $\bm{X} = (X, \mathbb{X})$ be an $\alpha$-Hölder geometric rough path, and $\{X^{(n)}\}_{n=1}^{\infty}$ a sequence of piecewise differentiable functions on $[0,T]$ that approximate $X$ under the $\beta$-Hölder norm for $\beta \in (\frac13, \frac12)$, that is, $\|X^{(n)} - X\|_\beta \to 0$ as $n \to \infty$. Let $\bm{Z},Z^1,Z^2,\dots$ be solutions to the differential equations

\[eq:RoughFlow\] $$\begin{aligned}
{\mathrm{d}}\bm Z_t &= f(Z_t, t, \theta) {\mathrm{d}}\bm{X}_t, &&& Z_0 &\sim p_0,\\ \frac{{\mathrm{d}}Z_t^{(n)}}{{\mathrm{d}}t} &= f(Z_t^{(n)}, t, \theta) \frac{{\mathrm{d}}X_t^{(n)}}{{\mathrm{d}}t} &&& Z_0^{(n)} &= Z_0,\end{aligned}$$

where $f \in \mathcal{C}_b^4(\mathbb{R}^d \times [0,T] \times \mathbb{R}^m)$ and $p_0$ is a density on $\mathbb{R}^d$ such that $\log p_0$ is continuous. Let $p_t^{(n)}$ denote the probability density of $Z_t^{(n)}$ at time $t$, given by (\[eq:ContNormFlow\]). The distribution of $Z_t$ is absolutely continuous with respect to Lebesgue measure with corresponding continuous density $p_t$ satisfying:

1.  For any $x \in \mathbb{R}^d$, $\displaystyle{\sup_{t \in [0,T]}}|\log p_t^{(n)}(x) - \log p_t(x)| \to 0$ as $n \to \infty$.

2.  The path $t \mapsto \log p_t(Z_t)$ is the unique solution to the rough differential equation $$\label{eq:RoughNF}
    {\mathrm{d}}\log p_t(Z_t) = -\nabla_z \cdot (f(Z_t, t, \theta) {\mathrm{d}}\bm X_t). $$

3.  For any smooth loss function $L:\mathbb{R}^{d+1}\to\mathbb{R}$ and $t \geq 0$, as $n \to \infty$, $$\label{eq:ApproxAdjoint}
    \nabla_\theta L(Z_t^{(n)},\log p_t^{(n)}(Z_t^{(n)})) \to \nabla_\theta L(Z_t,\log p_t(Z_t)).$$

Recall that each $X^{(n)}$ can be lifted canonically to a rough path $\bm{X}^{(n)}$ such that $\rho_{\beta}(\bm{X}^{(n)},\bm{X}) \to 0$ as $n \to \infty$. For an arbitrary rough path $\bm{Y}$, we let $\Phi_t(\boldsymbol{Y}, \xi)$ and $\Psi_t(\bm{Y},\ell)$ denote the solution maps for the rough differential equations ${\mathrm{d}}\bm{Z}_t = f(Z_t,t,\theta){\mathrm{d}}\bm{Y}_t$, $Z_0 = \xi$ and ${\mathrm{d}}\bm{L}_t = -\nabla_z\cdot f(Z_t,t,\theta) {\mathrm{d}}\bm{Y}_t$, $L_0 = \ell$, respectively. By @friz2014course [Theorem 8.10], $\Phi_t(\boldsymbol{Y},\cdot)$ is a $\mathcal{C}^1$-diffeomorphism, and hence, for $Z_0 \sim p_0(\theta)$ and any $t \in [0,T]$, $Z_t = \Phi_t(\boldsymbol{X}, Z_0)$ is an absolutely continuous random variable, whose corresponding density we denote by $p_t$. In fact, denoting by $\Phi_{-t}(\boldsymbol{Y},\cdot)$ the inverse of $\Phi_t(\bm{Y},\cdot)$, $$\begin{aligned}
\label{eq:DensityNCoM}
p_t^{(n)}(x) &= p_0(\Phi_{-t}(\bm{X}^{(n)},x))\left|\det \frac{\partial \Phi_{-t}(\bm{X}^{(n)},x)}{\partial x}\right|,\\
\label{eq:DensityCoM}
p_t(x) &= p_0(\Phi_{-t}(\bm{X},x))\left|\det \frac{\partial \Phi_{-t}(\bm{X},x)}{\partial x}\right|,\end{aligned}$$ and so both $p_t^{(n)}$ and $p_t$ are continuous. Furthermore, by @friz2014course [Theorem 8.5], for any $\frac13 < \gamma < \beta$, there exist constants $C_\gamma^{\Phi}$ and $C_\gamma^{\Psi}$ such that for any $\beta$-Hölder continuous rough paths $\boldsymbol{X}$, $\boldsymbol{Y}$ and $\xi, \tilde{\xi} \in \mathbb{R}^d$, $\ell, \tilde{\ell} \in \mathbb{R}_+$, $$\begin{aligned}
\label{eq:PhiRegularity}
\|\Phi(\boldsymbol{X}, \xi) - \Phi(\boldsymbol{Y}, \tilde{\xi})\|_{\gamma} &\leq C_\gamma^\Phi (\|\xi - \tilde{\xi}\| + \varrho_{\beta}(\boldsymbol{X}, \boldsymbol{Y}))\\
\label{eq:PsiRegularity}
\|\Psi(\bm{X}, \ell) - \Psi(\bm{Y}, \tilde{\ell})\|_{\gamma} &\leq C_\gamma^\Psi (|\ell - \tilde{\ell}| + \varrho_{\beta}(\bm{X},\bm{Y})).\end{aligned}$$ We deduce the following for any $t \in [0,T]$ and $x \in \mathbb{R}^d$: **(i)** $\|Z^{(n)} - Z\|_\gamma \to 0$ by (\[eq:PhiRegularity\]); **(ii)** using (i) and continuity of $p_t$, $\log p_t(Z_t^{(n)}) \to \log p_t(Z_t)$; **(iii)** as a consequence of (\[eq:DensityNCoM\]), (\[eq:DensityCoM\]), (\[eq:PhiRegularity\]), and @friz2014course [Theorem 8.10], $p_t^{(n)}(x)\to p_t(x)$; **(iv)** combining (ii) and (iii), $\log p_t^{(n)}(Z_t^{(n)}) \to \log p_t(Z_t)$. Since $\Psi_t(\bm{X}^{(n)},\log p_0(Z_0)) = \log p_t^{(n)}(Z_t^{(n)})$ by Theorem \[thm:ContNormFlow\], (iv) and (\[eq:PsiRegularity\]) imply $\log p_t(Z_t) = \Psi(\bm{X}, \log p_0(Z_0))$ and hence (\[eq:RoughNF\]). Let $x \in \mathbb{R}^d$ be arbitrary. To show that $\log p_t^{(n)}(x)$ converges uniformly in $t$, observe that $$\log p_t(x) = \Psi(\bm{X}, \log p_0(\Phi_{-t}(\bm{X},x))),$$ and similarly for $\log p_t^{(n)}(x)$. Together with property II of geometric rough paths, inequality (\[eq:PhiRegularity\]) with $\bm{Y} \equiv \bm{0}$ reveals that $\Phi_{-t}(\bm{X}^{(n)},x)$ and $\Phi_{-t}(\bm{X},x)$ are uniformly bounded in $t \in [0,T]$. Since $\log p_0$ is continuous, $\log p_0(\Phi_{-t}(\bm{X}^{(n)},x))$ converges to $\log p_0(\Phi_{-t}(\bm{X},x))$ uniformly in $t \in [0,T]$. Applying (\[eq:PsiRegularity\]), $$\begin{gathered}
\sup_{t\in[0,T]}|\log p_t^{(n)}(x) - \log p_t(x)| \leq C_{\gamma}^\Psi(\varrho_\beta(\bm{X}^{(n)}, \bm{X}) + |\log p_0(\Phi_{-t}(\bm{X}^{(n)}, x)) - \log p_0(\Phi_{-t}(\bm{X}, x))| \to 0.\end{gathered}$$ Finally, to prove (\[eq:ApproxAdjoint\]), by @friz2014course [Proposition 5.6], we can write $(\theta,L)$ as the solution to the rough differential equation

\[eq:LossSDE\] $$\begin{aligned}
&{\mathrm{d}}\theta = 0 \\
&{\mathrm{d}}L(Z_t,\log p_t(Z_t)) = \nabla_z L(Z_t, \log p_t(Z_t)) \cdot {\mathrm{d}}Z_t + \nabla_\ell L(Z_t, \log p_t(Z_t)) {\mathrm{d}}\log p_t(Z_t)\end{aligned}$$

and similarly for $Z_t^{(n)}$ and $\log p_t^{(n)}(Z_t^{(n)})$, where $\nabla_z$ and $\nabla_\ell$ denote the gradients with respect to $Z_t$ and $\log p_t(Z_t)$, respectively. The derivative of $L$ with respect to $\theta$ is a derivative of (\[eq:LossSDE\]) with respect to its initial condition, and hence (\[eq:ApproxAdjoint\]) follows from @friz2014course [Theorem 8.10].

Random continuous normalizing flows {#sec:RandCNF}
-----------------------------------

By a conditioning argument, any random ODE, such as a Wong-Zakai approximation, may be treated as a continuous normalizing flow. Let $Z_t$ be the solution to a random ODE of the form $$\label{eq:RandomODE}
\frac{{\mathrm{d}}}{{\mathrm{d}}t}Z_t = f(Z_t, \omega, t, \theta),$$ where $\omega = (\omega_1,\dots,\omega_n) \sim q(\omega)$ is a random vector independent of $Z_t$, $t$, and $\theta$. The reduction of a random ODE to this form is in keeping with the reparameterization trick [@xu2019variance]. In particular, for the piecewise linear and Karhunen-Loève approximations, each $\omega_i \sim \mathcal{N}(0,1)$. After conditioning on $\omega$, Theorem \[thm:ContNormFlow\] applied to (\[eq:RandomODE\]) provides a means of computing $\log p_t(Z_t\vert\omega)$, after sampling $Z_0 \sim p_0$. The density $p_t(Z_t)$ can be computed using a naive Monte Carlo estimator $$\label{eq:DensityEstNaive}
p_t^\theta(Z_t) = \int p_t^\theta(Z_t \vert \omega)q(\omega) {\mathrm{d}}\omega \approx \frac1N \sum_{i=1}^{n} p_t^\theta(Z_t \vert \omega_i),$$ where the dependence on $\theta$ has been made explicit, and can be optimized over using the adjoint method. Analogously to @chen2018neural [@grathwohl2018ffjord], the density of data ${\boldsymbol{x}}$ may be estimated along a single sample path $B_t(\omega)$ (denoted $p({\boldsymbol{x}}\vert \omega)$) in the following way: letting $\Delta \log p_t^\omega = \log p_t(Z_t\vert \omega) - \log p({\boldsymbol{x}} \vert \omega)$, we see that $\Delta \log p_t^\omega$ also satisfies (\[eq:ContNormFlow\]). By solving (\[eq:RandomODE\]) and the corresponding (\[eq:ContNormFlow\]) in reverse time from the initial conditions $Z_T = {\boldsymbol{x}}$ and $\Delta \log p_T^\omega = 0$, we obtain $Z_0$ and $\Delta \log p_0^\omega$, and compute $\log p({\boldsymbol{x}} \vert \omega) = \log p_0(Z_0) - \Delta \log p_0^\omega$. This is shown in Algorithm \[alg:DensityEst\], which depends on an ODE solver <span style="font-variant:small-caps;">odesolve</span>, and yields a density estimation procedure for stochastic normalizing flows when paired with (\[eq:DensityEstNaive\]). Note that by comparison to @grathwohl2018ffjord [Algorithm 1] which encompasses steps 6–12 of our Algorithm \[alg:DensityEst\], we see much of the density estimation procedure can be accomplished using an existing continuous normalizing flow implementation. In variational settings where $\log p_T(Z_T)$ is required, the same procedure applies, where $\mathbf{x}$ becomes $Z_T$ and is generated by the SDE as well.

drift function $\mu$, diffusion function $\sigma$, an initial distribution $p_0$, final time $T$, minibatch of samples ${\boldsymbol{x}}$, sample path $\tilde{B}_t(\omega)$ of Brownian motion approximation. an estimate of $\log p({\boldsymbol{x}} \vert \omega)$

Generate $\epsilon = (\epsilon_1,\dots,\epsilon_d)$ for (\[eq:TraceEst\]). <span style="font-variant:small-caps;">odefunc</span>$(z,t)$ Compute $\tilde{\mu}(z,t)$ via (\[eq:ItoCorrection\]).$\triangleright$ Itô correction $\tilde{\mu}(z,t) + \sigma(z,t) \frac{{\mathrm{d}}\tilde{B}_t(\omega)}{{\mathrm{d}}t}$. <span style="font-variant:small-caps;">aug</span>$((z,\log p_t),t)$ $f_t \gets $ <span style="font-variant:small-caps;">odefunc</span>($z$, $t$, $\omega$) $J_t \gets -\nabla_z(\epsilon \cdot f_t) \cdot\epsilon$$\triangleright$ Trace estimator (\[eq:TraceEst\]); $n = 1$ $(f_t, J_t)$ $(z,\Delta\log p_t^\omega) \gets $ <span style="font-variant:small-caps;">odesolve</span>(<span style="font-variant:small-caps;">aug</span>,$(x,0)$,$0$,$T$) $\log p_0(z) - \Delta \log p_t^\omega$

\[alg:DensityEst\]

A number of techniques exist for debiasing the logarithm of (\[eq:DensityEstNaive\]) — see @rhee2015unbiased and @rischard2018unbiased, for example. Alternatively, we lie in the setting of *semi-implicit variational inference* seen in @yin2018semi and @titsias2018unbiased, and those techniques directly extend to our case as well. Naturally, it would be easiest to instead optimize the upper bound $$-\log p_t^\theta(Z_t) \leq -\mathbb{E}_{\omega} \log p_t^\theta(Z_t\vert \omega),$$ and in many cases we have found this to be effective. Observing that $$\label{eq:ELBO}
\log p_t^\theta(Z_t) - {D_{\mathrm{KL}}}(q \Vert p_{\omega \vert Z_t}^{\theta}) = \mathbb{E}_\omega \log p_t^\theta(Z_t \vert \omega),$$ minimizing $-\mathbb{E}_{\omega}\log p_t^{\theta}(Z_t\vert \omega)$ maximizes the true log-likelihood regularized by the KL-divergence between the prior and posterior distributions for $\omega$, which reduces the effect of noise on the model. At the same time, parameterizations of the diffusion coefficient that allow $\|\sigma\|$ to shrink to zero will often do so, and should be avoided to remain distinct from a continuous normalizing flow.

Numerical experiments {#sec:Numerics}
=====================

Samplers and density estimation from data
-----------------------------------------

For our first experiments, we train a stochastic normalizing flow (\[eq:MainSDE\]) — using Algorithm \[alg:DensityEst\] with the upper bound (\[eq:ELBO\]) — to data generated from a specified target distribution. For our drift function, we adopt the same architecture used in the toy examples of @grathwohl2018ffjord; a four-layer fully-connected neural network with 64 hidden units in each layer. Dependence on time is removed to ensure a time-homogeneous, and hence, potentially ergodic SDE after training. All networks were trained using Adagrad [@duchi2011adaptive], with $p_0 \sim \mathcal{N}(0,I)$ and a batch size of 1000 samples.

### A two-dimensional toy example {#sec:Toy}

In our first example, our data is generated from the banana-shaped distribution $$p(x,y) \propto \exp(-\tfrac12 (x^2 + \tfrac12 (x^2 + y)^2)).$$ Two choices of diffusion coefficient are considered: the first, where $\sigma = I$, yields a neural SDE that can be trained using the techniques of @li2020scalable. For the second, we choose $$\label{eq:SigmaChoice}
\sigma(x) = \lambda \left(\begin{matrix} 1 & \sigma_1(x) \\ \sigma_2(x) & 1 \end{matrix}\right),$$ with $\lambda = 1$, and parameterize $(\sigma_1,\sigma_2)$ by a two-layered neural network with 64 hidden units. This SDE can only be trained using our method. After training, to emulate the application of these SDEs as approximate samplers, a single sample path with 10,000 steps was simulated for each model using the Euler–Maruyama method. The resulting paths are compared in Figure \[fig:FixedVsVariable\]. From data alone, both models constructed recurrent processes. The addition of a trainable diffusion coefficient led to improved adaptation of the sampler to the underlying curvature.

![Sample paths from SDEs trained as stochastic normalizing flows to a banana-shaped distribution.[]{data-label="fig:FixedVsVariable"}](fixed_vs_variable_sigma.png){width="0.6\columnwidth"}

### Visualizing regularization

As discussed in @liu2019neural, the stochastic noise injection in SDEs is a natural form of regularization, that can potentially improve robustness to noisy or adversarial data. We visualize this effect by considering the same stochastic normalizing flows treated in §\[sec:Toy\] with diffusion coefficient (\[eq:SigmaChoice\]), and adjusting the parameter $\lambda > 0$. Our data is generated in polar coordinates from a ten-pointed star-shaped distribution by $$\theta \sim \mathrm{Unif}(-\pi,\pi),\quad r\vert \theta \sim \mathcal{N}(\tfrac{2}{\sqrt{1 + \frac12 \sin(10 \theta)}}, \tfrac{9}{400}).$$ In Figure \[fig:RegDensity\], we plot the densities for $\lambda \in \{0, \tfrac{1}{10}, \tfrac12, 1\}$ computed using Algorithm \[alg:DensityEst\], noting that the $\lambda = 0$ case corresponds to a continuous normalizing flow. Increasing $\lambda$ reveals generative models with expectedly higher variance, but with improved capacity to smooth out minor (potentially, unwanted) details.

![Density plots of stochastic normalizing flows trained to a star-shaped distribution with varying diffusion coefficients.[]{data-label="fig:RegDensity"}](reg_density.png){width="0.5\columnwidth"}

Optimizing stochastic MCMC
--------------------------

An interesting class of SDE models for approximating a target distribution $p$ are *targeted diffusions*, solutions to SDEs that are $p$-ergodic. A convenient representation of such diffusions are known [@ma2015complete]. Because these diffusions are frequently used in MCMC algorithms, in a sense, conducting VI with respect to targeted diffusions is analogous to optimizing the convergence rate of stochastic MCMC algorithms.

To illustrate the potential applications of stochastic normalizing flows for finding and examining optimal stochastic MCMC algorithms for a particular target distribution, we consider a basic setup, where $p$ is the one-dimensional Cauchy distribution $p(x) \propto (1 + x^2)^{-1}$. All $p$-ergodic SDEs are of the form $$\label{eq:CauchySampler}
{\mathrm{d}}Z_t = (-2 \sigma(Z_t)^2 Z_t / (1 + Z_t^2) + \tfrac12 \sigma'(Z_t)) {\mathrm{d}}t + \sigma(Z_t) {\mathrm{d}}B_t,$$ and we may choose $\sigma$ arbitrarily. *A priori*, an optimal choice of $\sigma$ (up to constants) to ensure rapid mixing of (\[eq:CauchySampler\]) does not appear obvious. The present rule of thumb from second-order methods takes $\sigma \approx (\log p)''$ [@girolami2011riemann]. We train a stochastic normalizing flow for (\[eq:CauchySampler\]) with $\sigma$ parameterized by a four-layer neural network with 32 hidden units in each layer. The corresponding loss function is taken to be the Kullback-Leibler divergence $\log p_T(Z_T) - \log p(Z_T)$, estimated using Algorithm \[alg:DensityEst\], with an $L^1$ penalty term $10^{-4}\|{\boldsymbol{w}}\|_1$ over the weights ${\boldsymbol{w}}$ of the neural network, to prevent taking $|\sigma| \to +\infty$. The results are presented in Figure \[fig:CauchyPlot\].

![Stochastic normalizing flow targeting Cauchy distribution. *Left*: trained diffusion coefficient (blue) compared to $\sqrt{1 + x^2}$ (orange dashed). *Right*: Histogram of 20,000 generated samples (blue) and Cauchy density (orange). []{data-label="fig:CauchyPlot"}](cauchy.pdf){width="0.7\columnwidth"}

Curiously, after training, we found the apparent “optimal” choice is approximately $\sigma(x) \propto \sqrt{1 + x^2} \propto \sqrt{\frac{(\log \varphi)'}{(\log p)'}}$, where $\varphi$ is the density of the standard normal distribution.

Conclusion {#sec:Conclusion}
==========

We have extended the continuous normalizing flows framework to generative models involving SDEs. Justified by rough path theory, our framework enables practitioners of neural ODEs to apply their existing implementation for training neural SDEs. This is advantageous, as neural SDEs have been suggested to be more robust than neural ODEs in high-dimensional real-world examples [@liu2019neural; @li2020scalable]. Stochastic normalizing flows can be implemented as a device for investigating “optimal” hyperparameters in stochastic MCMC, which could prove useful for informing future research, but they may require implementational improvements for high-dimensional cases, e.g., variance reduction techniques and improved loss estimators.

#### Acknowledgements.

We would like to acknowledge DARPA, NSF, and ONR for providing partial support of this work.

[41]{} \[1\][\#1]{} \[1\][`#1`]{} urlstyle \[1\][doi: \#1]{}

Beskos, A., Papaspiliopoulos, O., Roberts, G. O., and Fearnhead, P. Exact and computationally efficient likelihood-based estimation for discretely observed diffusion processes. *Journal of the Royal Statistical Society: Series B (Statistical Methodology)*, 680 (3):0 333–382, 2006.

Chen, T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. In *Advances in Neural Information Processing Systems*, pp. 6571–6583, 2018.

Davie, A. M. Differential equations driven by rough paths: an approach via discrete approximation. *Applied Mathematics Research eXpress*, 2008, 2008.

Dinh, L., Krueger, D., and Bengio, Y. . In *Proceedings of the 3rd International Conference on Learning Representations (ICLR)*, 2015.

Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. *Journal of Machine Learning Research*, 12:0 2121–2159, 2011.

Dupont, E., Doucet, A., and Teh, Y. W. Augmented neural [ODEs]{}. In *Advances in Neural Information Processing Systems*, 2019.

Evans, L. C. *An introduction to stochastic differential equations*, volume 82. American Mathematical Society, 2012.

Friz, P. and Hairer, M. *A Course on Rough Paths*. Springer International Publishing, 2014.

Friz, P. K. and Victoir, N. B. *Multidimensional stochastic processes as rough paths: theory and applications*, volume 120. Cambridge University Press, 2010.

Gholami, A., Keutzer, K., and Biros, G. . In *Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence (IJCAI-19)*, 2019.

Girolami, M. and Calderhead, B. . *Journal of the Royal Statistical Society: Series B (Statistical Methodology)*, 730 (2):0 123–214, 2011.

Grathwohl, W., Chen, R. T., Betterncourt, J., Sutskever, I., and Duvenaud, D. . *arXiv preprint arXiv:1810.01367*, 2018.

Hairer, E. and Wanner, G. *[Solving Ordinary Differential Equations II]{}*, volume 14 of *Springer Series in Computational Mathematics*. Springer-Verlag Berlin Heidelberg, 1996.

Hodgkinson, L., Salomone, R., and Roosta, F. . *arXiv preprint arXiv:1903.12322*, 2019.

Hurn, A. S., Jeisman, J., and Lindsay, K. A. . *Journal of Financial Econometrics*, 50 (3):0 390–455, 2007.

Jia, J. and Benson, A. R. . In *Proceedings of the 33rd Conference on Neural Information Processing Systems (NeurIPS 2019)*, 2019.

Kloeden, P. E. and Platen, E. *Numerical solution of stochastic differential equations*, volume 23. Springer Science & Business Media, 2013.

Li, X., Wong, T.-K. L., Chen, R. T., and Duvenaud, D. Scalable gradients for stochastic differential equations. *arXiv preprint arXiv:2001.01328*, 2020.

Liu, Q. and Feng, Y. Two methods for wild variational inference. *arXiv preprint arXiv:1612.00081*, 2016.

Liu, X., Si, S., Cao, Q., Kumar, S., and Hsieh, C.-J. . *arXiv preprint arXiv:1906.02355*, 2019.

Lu, Y., Zhong, A., Li, Q., and Dong, B. . *arXiv preprint arXiv:1710.10121*, 2017.

Lyons, T. J. Differential equations driven by rough signals. *Revista Matem[á]{}tica Iberoamericana*, 140 (2):0 215–310, 1998.

Ma, Y.-A., Chen, T., and Fox, E. . In *Advances in Neural Information Processing Systems*, pp. 2917–2925, 2015.

Peluchetti, S. and Favaro, S. Infinitely deep neural networks as diffusion processes. *arXiv preprint arXiv:1905.11065*, 2019.

Pontryagin, L. S. *Mathematical theory of optimal processes*. Routledge, 2018.

Rezende, D. J. and Mohamed, S. Variational inference with normalizing flows. In *Proceedings of the 32nd International Conference on Machine Learning*, 2015.

Rhee, C.-h. and Glynn, P. W. . *Operations Research*, 630 (5):0 1026–1043, 2015.

Rischard, M., Jacob, P. E., and Pillai, N. . *arXiv preprint arXiv:1810.01382*, 2018.

Roosta, F. and Ascher, U. Improved bounds on sample size for implicit matrix trace estimators. *Foundations of Computational Mathematics*, 150 (5):0 1187–1212, 2015.

Ryder, T., Golightly, A., McGough, A. S., and Prangle, D. Black-box variational inference for stochastic differential equations. *arXiv preprint arXiv:1802.03335*, 2018.

Salimans, T., Kingma, D., and Welling, M. . In *International Conference on Machine Learning*, pp. 1218–1226, 2015.

Titsias, M. K. and Ruiz, F. J. Unbiased implicit variational inference. In *Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS)*, 2019.

Twardowska, K. . *Acta Applicandae Mathematica*, 430 (3):0 317–359, 1996.

Tzen, B. and Raginsky, M. . *arXiv preprint arXiv:1905.09883*, 2019.

van den Berg, R., Hasenclever, L., Tomczak, J., and Welling, M. Sylvester normalizing flows for variational inference. In *Proceedings of the Conference on Uncertainty in Artificial Intelligence (UAI)*, 2018.

Wolf, C., Karl, M., and van der Smagt, P. . *arXiv preprint arXiv:1609.08203*, 2016.

Wong, E. and Zakai, M. On the convergence of ordinary integrals to stochastic integrals. *The Annals of Mathematical Statistics*, 360 (5):0 1560–1564, 1965.

Xu, M., Quiroz, M., Kohn, R., and Sisson, S. A. Variance reduction properties of the reparameterization trick. In *The 22nd International Conference on Artificial Intelligence and Statistics*, pp.  2711–2720, 2019.

Yin, M. and Zhou, M. Semi-implicit variational inference. In *Proceedings of the 35th International Conference on Machine Learning*, 2018.

Young, L. C. . *Acta Mathematica*, 67:0 251–282, 1936.

Zhang, T., Yao, Z., Gholami, A., Keutzer, K., Gonzalez, J., Biros, G., and Mahoney, M. W. . *Advances in Neural Information Processing Systems*, 2019.

[^1]: Department of Statistics, University of California at Berkeley, USA, and International Computer Science Institute, Berkeley, CA, USA. Email: `liam.hodgkinson@berkeley.edu`

[^2]: School of Mathematics and Physics, University of Queensland, Australia. Email:

[^3]: School of Mathematics and Physics, University of Queensland, Australia, and International Computer Science Institute, Berkeley, CA, USA. Email:

[^4]: Department of Statistics, University of California at Berkeley, USA, and International Computer Science Institute, Berkeley, CA, USA. Email: `mmahoney@stat.berkeley.edu`
