Efficient and Accurate Gradients for Neural SDEs

Introduction

Stochastic differential equations have seen widespread use in the mathematical modelling of random phenomena, such as particle systems , financial markets , population dynamics , and genetics . Featuring inherent randomness, then in modern machine learning parlance SDEs are generative models.

Such models have typically been constructed theoretically, and are usually relatively simple. For example the Black–Scholes equation, widely used to model asset prices in financial markets, has only two scalar parameters: a fixed drift and a fixed diffusion .

Neural stochastic differential equations

Neural stochastic differential equations offer a shift in this paradigm. By parameterising the drift and diffusion of an SDE as neural networks, then modelling capacity is greatly increased, and theoretically arbitrary SDEs may be approximated. (By the universal approximation theorem for neural networks .) Several authors have now studied or introduced Neural SDEs; amongst others.

Connections to recurrent neural networks

A numerically discretised (Neural) SDE may be interpreted as an RNN (featuring a residual connection), whose input is random noise – Brownian motion – and whose output is a generated sample. Subject to a suitable loss function between distributions, such as the KL divergence or Wasserstein distance , this may then simply be backpropagated through in the usual way.

Generative time series models

SDEs are naturally random. In modern machine learning parlance they are thus generative models. As such we treat Neural SDEs as generative time series models.

The (recurrent) neural network-like structure offers high-capacity function approximation, whilst the SDE-like structure offers strong priors on model space, memory efficiency, and deep theoretical connections to a well-understood literature. Relative to the classical SDE literature, Neural SDEs have essentially unprecedented modelling capacity.

(Generative) time series models are of classical interest, with forecasting models such as Holt–Winters , ARMA and so on. It has also attracted much recent interest with (besides Neural SDEs) the development of models such as Time Series GAN , Latent ODEs , GRU-ODE-Bayes , ODE2VAE , CTFPs , Neural ODE Processes and Neural Jump ODEs .

1 Contributions

We study backpropagation through SDE solvers, in particular to train Neural SDEs, via continuous adjoint methods. We introduce several technical innovations to improve both model performance and the speed of training: in particular to reduce numerical gradient errors to almost zero.

First, we introduce the reversible Heun method: a new SDE solver, constructed to be algebraically reversible. By matching the truncation errors of the forward and backward passes, the gradients computed via continuous adjoint method are precisely those of the numerical discretisation of the forward pass. This overcomes the typical greatest limitation of continuous adjoint methods – and to the best of our knowledge, is the first algebraically reversible SDE solver to have been developed.

After that, we introduce the Brownian Interval as a new way of sampling and reconstructing Brownian motion. It is fast, memory efficient and exact. It has an average (modal) time complexity of $\mathcal{O}\left(1\right)$, and consumes only $\mathcal{O}\left(1\right)$ GPU memory. This is contrast to previous techniques requiring either $\mathcal{O}\left(T\right)$ memory, or a choice of approximation error $\varepsilon\ll 1$ and then a time complexity of $\mathcal{O}\left(\log(1/\varepsilon)\right)$.

Finally, we demonstrate how the Lipschitz condition for the discriminator of an SDE-GAN may be imposed without gradient penalties – instead using careful clipping and the LipSwish activation function – so as to overcome their previous incompatibility with continuous adjoint methods.

Overall, multiple technical innovations provide substantial improvements over the state-of-the-art with respect to training speed, and with respect to classification, prediction, and MMD test metrics.

Background

We consider Neural SDEs as models of the form

RNNs as discretised SDEs

2 Training criteria for Neural SDEs

The Wasserstein distance may be used by constructing a discriminator and training adversarially, as in Kidger et al. . Let $F_{\phi}(Y)=m_{\phi}\cdot H_{T}$, where

for suitable neural networks $\xi_{\phi},f_{\phi},g_{\phi}$ and vector $m_{\phi}$. This is a deterministic function of the generated sample $Y$. Here $\cdot$ denotes a dot product. They then train with respect to

See Appendix B for additional details on this approach, and in particular how it generalises the classical approach to fitting (calibrating) SDEs.

Latent SDEs

Li et al. instead optimise a KL divergence. This consists of constructing an auxiliary process $\widehat{X}$ with drift $\nu_{\phi}$ parameterised by $\phi$, and optimising an expression of the form

The full construction is moderately technical; see Appendix B for further details.

3 Discretised observations

Observations of $Y_{\text{true}}$ are typically a discrete time series, rather than a true continuous-time path. This is not a serious hurdle. If training an SDE-GAN, then equation (2) may be evaluated on an interpolation $Y_{\text{true}}$ of the observed data. If training a Latent SDE, then $\nu_{\phi}$ in equation (4) may depend explicitly on the discretised $Y_{\text{true}}$.

4 Backpropagation through SDE solves

Whether the loss for our generated sample $Y$ is produced via a Latent SDE or via the discriminator of an SDE-GAN, it is still required to backpropagate from the loss to the parameters $\theta,\phi$.

Here we use the continuous adjoint method. Also known as simply ‘the adjoint method’, or ‘optimise-then-discretise’, this has recently attracted much attention in the modern literature on neural differential equations. This exploits the reversibility of a differential equation: as with invertible neural networks , intermediate computations such as $X_{t}$ for $t<T$ are reconstructed from output computations, so that they do not need to be held in memory.

Note that we assumed here that the loss $L$ acts only on $Z_{T}$, not all of $Z$. This is not an issue in practice. In both equations (3) and (4), the loss is an integral. As such it may be computed as part of $Z$ in a single SDE solve. This outputs a value at time $T$, the operation $L$ may simply extract this value from $Z_{T}$, and then backpropagation may proceeed as described here.

The main issue is that the two numerical approximations to $Z_{t}$, computed in the forward and backward passes of equation (5), are different. This means that the $Z_{t}$ used as an input in equation (6) has some discrepancy from the forward calculation, and the gradients $A_{0}$ suffer some error as a result. (Often exacerbating an already tricky training procedure, such as the adversarial training of SDE-GANs.)

See Appendix C for further discussion on how an SDE solve may be backpropagated through.

5 Alternate constructions

There are other uses for Neural SDEs, beyond our scope here. For example Song et al. combine SDEs with score-matching, and Xu et al. use SDEs to represent Bayesian uncertainty over parameters. The techniques introduced in this paper will apply to any backpropagation through an SDE solve.

Reversible Heun method

We introduce a new SDE solver, which we refer to as the reversible Heun method. Its key property is algebraic reversibility; moreover to the best of our knowledge it is the first SDE solver to exhibit this property.

To fix notation, we consider solving the Stratonovich SDE

We begin by selecting a step size $\Delta t$, and initialising $t_{0}=0$, $z_{0}=\widehat{z}_{0}=Z_{0}$, $\mu_{0}=\mu(0,Z_{0})$ and $\sigma_{0}=\sigma(0,Z_{0})$. Let $W$ denote a single sample path of Brownian motion. It is important that the same sample be used for both the forward and backward passes of the algorithm; computationally this may be accomplished by taking $W$ to be a Brownian Interval, which we will introduce in Section 4.

We then iterate Algorithm 1. Suppose $T=N\Delta t$ so that $z_{N},\widehat{z}_{N},\mu_{N},\sigma_{N}$ are the final output. Then $z_{N}\approx Z_{T}$ is returned, whilst $z_{N},\widehat{z}_{N},\mu_{N},\sigma_{N}$ are all retained for the backward pass.

Nothing else need be saved in memory for the backward pass: in particular no intermediate computations, as would otherwise be typical.

Algebraic reversibility

The key advantage of the reversible Heun method, and the motivating reason for its use alongside continuous-time adjoint methods, is that it is algebraically reversible. That is, it is possible to reconstruct $(z_{n},\widehat{z}_{n},\mu_{n},\sigma_{n})$ from $(z_{n+1},\widehat{z}_{n+1},\mu_{n+1},\sigma_{n+1})$ in closed form. (And without a fixed-point iteration.)

This crucial property will mean that it is possible to backpropagate through the SDE solve, such that the gradients obtained via the continuous adjoint method (equation (6)) exactly match the (discretise-then-optimise) gradients obtained by autodifferentiating the numerically discretised forward pass.

In doing so, one of the greatest limitations of continuous adjoint methods is overcome.

To the best of our knowledge, the reversible Heun method is the first algebraically reversible SDE solver.

Computational efficiency

A further advantage of the reversible Heun method is computational efficiency. The method requires only a single function evaluation (of both the drift and diffusion) per step. This is in contrast to other Stratonovich solvers (such as the midpoint method or regular Heun’s method), which require two function evaluations per step.

Convergence of the solver

When applied to the Stratonovich SDE (7), the reversible Heun method exhibits strong convergence of order 0.5; the same as the usual Heun’s method.

Let $\{z_{n}\}$ denote the numerical solution of (7) obtained by Algorithm 1 with a constant step size $\Delta t$ and assume sufficient regularity of $\mu$ and $\sigma$. Then there exists a constant $C>0$ so that

for small $\Delta t$. That is, strong convergence of order 0.5. If $\sigma$ is constant, then this improves to order 1.

See Appendix D for the full proof. We additionally consider stability in the ODE setting. Whilst the method is not A-stable, we do show it has the same absolute stability region for a linear test equation as the (reversible) asynchronous leapfrog integrator proposed for Neural ODEs in Zhuang et al. .

Precise gradients

The backward pass is shown in Algorithm 2. As the same numerical solution $\{z_{n}\}$ is recovered on both the forward and backward passes – exhibiting the same truncation errors – then the computed gradients are precisely the (discretise-then-optimise) gradients of the numerical discretisation of the forward pass.

Each $\nicefrac{{\partial L(Z_{T})}}{{\partial z_{n}}}\approx A_{n\Delta t}$, where $A$ is the adjoint variable of equation (6).

This is unlike the case of solving equation (6) via standard numerical techniques, for which small or adaptive step sizes are necessary to obtain useful gradients .

1 Experiments

We validate the empirical performance of the reversible Heun method. For space, we present abbreviated details and results here. See Appendix F for details of the hyperparameter optimisation procedure, test metric definitions, and so on, and for further results on additional datasets and additional metrics.

We begin by comparing the reversible Heun method with the midpoint method, which also converges to the Stratonovich solution. We train an SDE-GAN on a dataset of weight trajectories evolving under stochastic gradient descent, and train a Latent SDE on a dataset of air quality over Beijing.

See Table 1. Due to the reduced number of vector field evaluations, we find that training speed roughly doubles ($1.98\times$) on the weights dataset, whilst its numerically precise gradients substantially improve the test metrics (comparing generated samples to a held-out test set). Similar behaviour is observed on the air quality dataset, with substantial test metric improvements and a training speed improvement of $1.25\times$.

Samples

We verify that samples from a model using reversible Heun resemble that of the original dataset: in Figure 1 we show the Latent SDE on the ozone concentration over Beijing.

Gradient error

We investigate the numerical error made in solving (6), compared to the (discretise-then-optimise) gradients of the numerically discretised forward pass. We fix a test problem (differentiating a small Neural SDE) and vary the step size and solver; see Figure 2. The error made in standard solvers is very large (but does at least decrease with step size). The reversible Heun method produces results accurate to floating point error, unattainable by any standard solver.

Brownian Interval

Numerically solving an SDE, via the reversible Heun method or via any other numerical solver, requires sampling Brownian motion: this is the input $W$ in Algorithms 1 and 2.

Brownian reconstruction

However, there are computational difficulties. On the backward pass, the same Brownian sample as the forward pass must be used, and potentially at locations other than were measured on the forward pass .

Time and memory efficiency

The simple but memory intensive approach would be to store every sample made on the forward pass, and then on the backward pass reuse these samples, or sample Brownian noise according to equation (8), as appropriate.

Li et al. instead offer a memory-efficient but time-intensive approach, by introducing the ‘Virtual Brownian Tree’. This approximates the real line by a tree of dyadic points. Samples are approximate, and demand deep (slow) traversals of the tree.

Binary tree of (interval, seed) pairs

In response to this, we introduce the ‘Brownian Interval’, which offers memory efficiency, exact samples, and fast query times, all at once. The Brownian Interval is built around a binary tree, each node of which is an interval and a random seed.

The tree starts as a stump consisting of the global interval $[0,T]$ and a randomly generated random seed. New leaf nodes are created as observations of the sample are made. For example, making a first query at $[s,t]\subseteq[0,T]$ (an operation that returns $W_{s,t}$) produces the binary tree shown in Figure 3(a). Algorithm 4 in Appendix E gives the formal definition of this procedure. Making a subsequent query at $[u,v]$ with $u<s<v<t$ produces Figure 3(b). Using a splittable PRNG , each child node has a random seed deterministically produced from the seed of its parent.

The tree is thus designed to completely encode the conditional statistics of a sample of Brownian motion: $W_{s,t},W_{t,u}$ are completely specified by $t,[s,u]$, $W_{s,u}$, equation (8), and the random seed for $[s,u]$.

In principle this now gives a way to compute $W_{s,t}$; calculating $W_{s,u}$ recursively. Naïvely this would be very slow – recursing to the root on every query – which we cover by augmenting the binary tree structure with a fixed-size Least Recently Used (LRU) cache on the computed increments $W_{s,t}$.

See Algorithm 3, where bridge denotes equation (8). The operation traverse traverses the binary tree to find or create the list of nodes whose disjoint union is the interval of interest, and is defined explicitly as Algorithm 4 in Appendix E.

Additionally see Appendix E for various technical considerations and extensions to this algorithm.

Advantages of the Brownian Interval

The LRU cache ensures that queries have an average-case (modal) time complexity of only $\mathcal{O}\left(1\right)$: in SDE solvers, subsequent queries are typically close to (and thus conditional on) previous queries. Even given cache misses all the way up the tree, the worst-case time complexity will only be $\mathcal{O}\left(\log(1/s)\right)$ in the average step size $s$ of the SDE solver. This is in contrast to the Virtual Brownian Tree, which has an (average or worst-case) time complexity of $\mathcal{O}\left(\log(1/\varepsilon)\right)$ in the approximation error $\varepsilon\ll s$.

Meanwhile the (GPU) memory cost is only $\mathcal{O}\left(1\right)$, corresponding to the fixed and constant size of the LRU cache. There is the small additional cost of storing the tree structure itself, but this is held in CPU memory, which for practical purposes is essentially infinite. This is in contrast to simply holding all the Brownian samples in memory, which has a memory cost of $\mathcal{O}\left(T\right)$.

Finally, queries are exact because the tree aligns with the query points. This is contrast to the Virtual Brownian Tree, which only produces samples up to some discretisation of the real line at resolution $\varepsilon$.

1 Experiments

We benchmark the performance of the Brownian Interval against the Virtual Brownian Tree considered in Li et al. . We include benchmarks corresponding to varying batch sizes, number of sample intervals, and access patterns. For space, just a subset of results are shown. Precise experimental details and further results are available in Appendix F.

See Table 2. We see that the Brownian Interval is uniformly faster than the Virtual Brownian Tree, ranging from $1.5\times$ faster on smaller problems to $10.6\times$ faster on larger problems. Moreover these speed gains are despite the Brownian Interval being written in Python, whilst the Virtual Brownian Tree is carefully optimised and written in C++.

Training SDE-GANs without gradient penalty

Kidger et al. train SDEs as GANs, as discussed in Section 2.2, using a neural CDE as a discriminator as in equation (2). They found that only gradient penalty was suitable to enforce the Lipschitz condition, given the recurrent structure of the discriminator.

However gradient penalty requires calculating second derivatives (a ‘double-backward’). This complicates the use of continuous adjoint methods: the double-continuous-adjoint introduces substantial truncation error; sufficient to obstruct training and requiring small step sizes to resolve.

Here we overcome this limitation, and moreover do so independently of the possibility of obtaining exact double-gradients via the reversible Heun method. For simplicity we now assume throughout that our discriminator vector fields $f_{\phi}$, $g_{\phi}$ are MLPs, which is also the choice we make in practice.

The key point is that the vector fields $f_{\phi}$, $g_{\phi}$ of the discriminator must not only be Lipschitz, but must have Lipschitz constant at most one.

Given vector fields with Lipschitz constant $\lambda$, then the recurrent structure of the discriminator means that the Lipschitz constant of the overall discriminator will be $\mathcal{O}\left(\lambda^{T}\right)$. Ensuring $\lambda\approx 1$ with $\lambda\leq 1$ thus enforces that the overall discriminator is Lipschitz with a reasonable Lipschitz constant.

Hard constraint

The exponential size of $\mathcal{O}\left(\lambda^{T}\right)$ means that $\lambda$ only slightly greater than one is still insufficient for stable training. We found that this ruled out enforcing $\lambda\leq 1$ via soft constraints, via either spectral normalisation or gradient penalty across just vector field evaluations.

Clipping

LipSwish activation functions

Next we must pick an activation function with Lipschitz constant at most one. It should additionally be at least twice continuously differentiable to ensure convergence of the numerical SDE solver (Appendix D). In particular this rules out the ReLU.

The overall vector fields $f_{\phi}$, $g_{\phi}$ of the discriminator consist of linear operations (which are constrained by clipping), adding biases (an operation with Lipschitz constant one), and activation functions (taken to be LipSwish). Thus the Lipschitz constant of the overall vector field is at most one, as desired.

1 Experiments

We test the SDE-GAN on a dataset of time-varying Ornstein–Uhlenbeck samples. For space only a subset of results are shown; see Appendix F for further details of the dataset, optimiser, and so on.

See Table 3 for the results. We see that the test metrics substantially improve with clipping, over gradient penalty (which struggles due to numerical errors in the double adjoint). The lack of double backward additionally implies a computational speed-up. This reduced training time from 55 hours to just 33 hours. Switching to reversible Heun additionally and substantially improves the test metrics, and further reduced training time to 29 hours; a speed improvement of $1.87\times$.

Discussion

To facilitate the use of the techniques introduced here – in particular without requiring a technical background in numerical SDEs – we have contributed implementations of both the reversible Heun method and the Brownian Interval to the open-source torchsde package. (In which the Brownian Interval has already become the default choice, due to its speed.)

2 Limitations

The reversible Heun method, Brownian Interval, and training of SDE-GANs via clipping, all appear to be strict improvements over previous techniques. Across our experiments we have observed no limitations relative to previous techniques.

3 Ethical statement

No significant negative societal impacts are anticipated as a result of this work. A positive environmental impact is anticipated, due to the reduction in compute costs implied by the techniques introduced. See Appendix G for a more in-depth discussion.

Conclusion

We have introduced several improvements over the previous state-of-the-art for Neural SDEs, with respect to both training speed and test metrics. This has been accomplished through several novel technical innovations, including a first-of-its-kind algebraically reversible SDE solver; a fast, exact, and memory efficient way of sampling and reconstructing Brownian motion; and the development of SDE-GANs via careful clipping and choice of activation function.

Acknowledgments and Disclosure of Funding

PK was supported by the EPSRC grant EP/L015811/1. PK, JF, TL were supported by the Alan Turing Institute under the EPSRC grant EP/N510129/1. PK thanks Chris Rackauckas for discussions related to the reversible Heun method.

References

Appendix A RNNs as discretised SDEs

Consider the autonomous one-dimensional Itô SDE

where $\Delta t$ is some fixed time step and $W_{n}\sim\mathcal{N}(0,\Delta t)$. This may be implemented in very few lines of PyTorch code – see Figure 4 – and subject to a suitable loss function between distributions, such as the KL divergence or Wasserstein distance , simply backpropagated through in the usual way.

In this way we see that a discretised SDE is simply an RNN consuming random noise.

In passing we note that this is a common occurrence: many popular neural network architectures may be interpreted as discretised differential equations.

Residual networks are discretisations of ODEs . RNNs are discretised controlled differential equations .

StyleGAN2 and denoising diffusion probabilistic models are both essentially discretised SDEs .

Many invertible neural networks resemble discretised differential equations; for example using an explicit Euler method on the forward pass, and recovering the intermediate computations via the implicit Euler method on the backward pass .

Appendix B Training criteria for Neural SDEs

This ensures that the model and the data behave the same with respect to the functions $F_{i}$. (Known as either ‘witness functions’ or ‘payoff functions’ depending on the field.)

Kidger et al. generalise this by replacing $F_{1},\ldots F_{N}$ with a parameterised function $F_{\phi}$ – a neural network with parameters $\phi$ – and training adversarially:

Thus making a connection to the GAN literature.

In principle $F_{\phi}$ could be parameterised as any neural network capable of operating on the path-valued $Y$. There is a natural choice: use a Neural CDE . This is a differential equation capable of acting on path-valued inputs. This means letting the discriminator be $F_{\phi}(Y)=m_{\phi}\cdot H_{T}$, where

for suitable neural networks $\xi_{\phi},f_{\phi},g_{\phi}$ and vector $m_{\phi}$. This is a deterministic function of the generated sample $Y$. Here $\cdot$ denotes a dot product.

Adding regularisation to control the derivative of $F_{\phi}$ ensures that equation (3) corresponds to the dual formulation of the Wasserstein distance, so that $Y$ is capable of perfectly matching $Y_{\text{true}}$ given enough data, training time, and model capacity .

In our experience, this approach produces models with a very high modelling capacity, but which are somewhat involved to train – GANs being notoriously hard to train stably.

Latent SDEs

Li et al. have an alternate approach. Let

Note that $\widehat{X}$ is a random variable over SDEs, as $Y_{\text{true}}$ is still a random variable.So that $\widehat{V}$, $\widehat{X}$, $\widehat{Y}$ might be better denoted $\widehat{V}|Y_{\text{true},0}$, $\widehat{X}|Y_{\text{true}}$, $\widehat{Y}|Y_{\text{true}}$ to reflect their dependence on $Y_{\text{true}}$; we elide this for simplicity of notation.

This may be derived as an evidence lower-bound (ELBO). The first two terms are simply a VAE for generating $Y_{0}$, with latent $V$. The third term and fourth term are a VAE for generating $Y$, by autoencoding $Y_{\text{true}}$ to $\widehat{Y}$, and then fitting $X$ to $\widehat{X}$.

In our experience, this produces less expressive models than the SDE-GAN approach; however the model is easier to train, due to the lack of adversarial training.For ease of presentation this section features a slight abuse of notation: writing the KL divergence between random variables rather than probability distributions. It is also a slight specialisation of Li et al. , who allow losses other than the $L^{2}$ loss between data and sample.

Appendix C Adjoints for SDEs

Recall that we wish to backpropagate from our generated sample $Y$ to the parameters $\theta,\phi$. Here we provide a more complete overview of the options for how this may be done.

One way is to simply backpropagate through the internals of every numerical solver. (Also known as ‘discretise-then-optimise’.) However this requires $\mathcal{O}\left(HT\right)$ memory, where $H$ denotes the amount of memory used to evaluate and backpropagate through each neural network once.

Optimise-then-discretise

The continuous adjoint method,Frequently abbreviated to simply ‘adjoint method’ in the modern literature, although this term is ambiguous as it is also used to refer to backpropagation-through-the-solver in other literature . also known as “optimise-then-discretise”, instead exploits the reversibility of a differential equation. This means that intermediate computations such as $X_{t}$ for $t<T$ are reconstructed from output computations, and do not need to be held in memory.

Recall equations (5) and (6): given some Stratonovich SDE

The main advantage of the continuous adjoint method is that it reduces the memory footprint to only $\mathcal{O}\left(H+T\right)$. $\mathcal{O}\left(H\right)$ to compute each vector-Jacobian product ($A^{j}\nicefrac{{\partial\mu^{i}}}{{\partial Z^{j}}}$ and $A^{j}\nicefrac{{\partial\sigma^{i,k}}}{{\partial Z^{j}}}$), and $\mathcal{O}\left(T\right)$ to hold the batch of training data.

The main disadvantage (unless using the reversible Heun method) is that the two numerical approximations to $Z_{t}$, computed in the forward and backward passes of equation (5), are different. This means that the $Z_{t}$ used as an input in equation (6) does not perfectly match what is used in the forward calculation, and the gradients $A_{0}$ suffer some error as a result. This slows and worsens training. (Often exacerbating an already tricky training procedure, such as the adversarial training of SDE-GANs.)

Itô versus Stratonovich

Note that backpropagation through an Itô SDE may be performed by first adding a $-(\sigma^{j,k}\nicefrac{{\partial\sigma^{i,k}}}{{\partial Z^{j}}})(t,Z_{t})/2$ correction term to $\mu$ in equation (5), which converts it to a Stratonovich SDE, and then applying equation (6).

This additional computational cost – including double autodifferentiation to compute derivatives of the correction term – means that we prefer to use Stratonovich SDEs throughout.

Appendix D Error analysis of the reversible Heun method

In this section, we present some of the notation, definitions and assumptions used in our error analysis.

Therefore, we can define the operator norm of $T$ recursively as

With a slight abuse of notation, we shall write $\|\cdot\|_{2}$ instead of $\|\cdot\|_{\text{op}}$ for all $k$-tensors.

We shall assume that $f,g$ are bounded and twice continuously differentiable with bounded derivatives,

for $k=0,1,2$, where $\|\cdot\|_{2}$ denotes the Euclidean operator norm on $k$-tensors given by (11).

For $N\geq 1$, we will compute numerical SDE solutions on $[0,T]$ using a constant step size $h=\frac{T}{N}$. That is, numerical solutions are obtained at times $\{t_{n}\}_{0\leq n\leq N}$ with $t_{n}:=nh$ for $n\in\{0,1,\cdots,N\}$.

We recall the reversible Heun method given by Algorithm 1.

For $N\geq 1$, we construct a numerical solution $\{Y_{n},Z_{n}\}_{0\leq n\leq N}$ for the SDE (13) by setting $Y_{0}=Z_{0}=y_{0}$ and, for each $n\in\{0,1,\cdots,N-1\}$, defining $(Y_{n+1},Z_{n+1})$ from $(Y_{n},Z_{n})$ as

Whilst it is not used in our analysis, the above numerical method is time-reversible as

To simplify notation, we shall define another numerical solution $\{\widetilde{Z}_{n}\}_{0\leq n\leq N}$ with $\widetilde{Z}_{n}:=2Y_{n}-Z_{n}$.

D.2 Approximation error between components of the reversible Heun method

The key idea underlying our analysis is to consider two steps of the numerical method, which gives,

Thus $Z$ is propagated by a midpoint method and $Y$ is propagated by a trapezoidal rule / Heun method. To prove that $\{Z_{n}\}$ and $\{Y_{n}\}$ are close together when $h$ is small, we shall use the Taylor expansions:

where the remainder terms $R_{n}^{F,1}$ and $R_{n}^{F,2}$ satisfy the following estimates for any fixed $p\geq 1$,

By Taylor’s theorem with integral remainder [52, Theorem 3.5.6], the $R_{n}^{F,i}$ terms are given by

where we also used the inequality $(a+b)^{2}\leq 2a^{2}+2b^{2}$. The result now follows from the above. ∎

With Theorem D.3, it is straightforward to derive a Taylor expansion for the difference $Y_{n+2}-Z_{n+2}\,$.

For $n\in\{0,1,\cdots,N-2\}$, the difference $Y_{n+2}-Z_{n+2}$ can be expanded as

Expanding the components (17), (18) of the reversible Heun method with Theorem D.3 gives

Since $R_{n}^{f,1},R_{n}^{f,2},R_{n}^{g,1},R_{n}^{g,2}\sim O(h)$, which was shown in Theorem D.3, the above simplifies to

Substituting the formulae (17) and (19) for $Z_{n+1}$ and $Z_{n+2}$ respectively produces

As f\big{(}Y_{n+1}\big{)} and g\big{(}Y_{n+1}\big{)} are bounded by $\|f\|_{\infty}$ and $\|g\|_{\infty}$, collecting the O\big{(}h^{\frac{3}{2}}\big{)} terms yields

We note that it is direct consequence of Theorem D.3 that g\big{(}Z_{n+1}\big{)}=g\big{(}\widetilde{Z}_{n}\big{)}+O\big{(}\sqrt{h}\,\big{)}. Thus,

Let $h_{\max}>0$ be fixed. Then there exist constants $c_{1},c_{2}>0$ such that for $n\in\{0,1,\cdots,N-2\}$,

By Theorem D.4, the second term can be expanded as

The above can be estimated using the Cauchy-Schartz inequality and Young’s inequality as

The final term in (D.2) can be estimated using Minkowski’s inequality as

Since $(a+b)^{k}\leq 2^{k-1}(a^{k}+b^{k})$ for $a,b\geq 0$ (which follows from Jensen’s inequality), this gives

By Young’s inequality, we can further estimate the terms which do not contain \big{\|}Y_{n}-Z_{n}\big{\|}_{2}^{4} as

Finally, by applying the above two estimates to the inequality (23), we arrive at

Let $h_{\max}>0$ be fixed. There exist a constant $C_{1}>0$ such that for all $n\in\{0,1,\cdots,N\}$,

By the tower property of expectations, it follows from Theorem D.5 that

which, along with the fact that $Y_{0}=Z_{0}\,$, implies that

D.3 Strong convergence of the reversible Heun method

To begin, we will Taylor expand the terms in the $Y_{n}\mapsto Y_{n+1}$ update of the reversible Heun method.

where the remainder terms $R_{n}^{F,3}$ and $R_{n}^{F,4}$ satisfy the following estimates,

By Taylor’s theorem with integral remainder [52, Theorem 3.5.6], we have that for $k\in\{0,1\}$,

By the same argument used in the proof of Theorem D.3, we can estimate the remainder term as

Using the inequality $(a+b)^{4}\leq(2a^{2}+2b^{2})^{2}\leq 8a^{4}+8b^{4}$, we have

Therefore, by the tower property of expectations, for $k\in\{0,1\}$,

We consider the $k=0$ and $k=1$ cases separately. If $k=0$, then by the tower property, we have

Slightly more care should be taken when $k=1$ since $Z_{n+1}-Y_{n}$ is not independent of $W_{n}$.

Finally, by Theorem D.6 and the Lipschitz continuity of $g$, F^{\prime}\big{(}Y_{n}\big{)}\big{(}Z_{n+1}-Y_{n}\big{)} can be expanded as

The result follows from the above estimates. ∎

We are now in a position to compute the local Taylor expansion of the numerical approximation $Y$.

By Theorem D.7, we can expand $Y_{n+1}$ as

The result now follows as the bottom line is clearly O\big{(}h^{\frac{3}{2}}\big{)} by Theorem D.7. ∎

Just as in the numerical analysis of ODE solvers, we also compute Taylor expansions for the solution. In our setting, we consider the following stochastic Taylor expansion for the Stratonovich SDE (13).

Using the above theorems, we can now obtain a Taylor expansion for the difference $Y_{n+1}-y_{(n+1)h}\,$.

For $n\in\{0,1,\cdots,N-1\}$, , the difference $Y_{n+2}-Z_{n+2}$ can be expanded as

Expanding $Y_{n+1}-y_{(n+1)h}$ using Theorem D.8 and Theorem D.9 gives

Having derived Taylor expansions for the approximation and solution processes, we will establish the main results of the section (namely, local and global error estimates for the reversible Heun method).

Let $h_{\max}>0$ be fixed. Then there exist constants $c_{3},c_{4}>0$ such that for all $n\in\{0,1,\cdots,N-1\}$,

Expanding the left hand side of (24) and applying the tower property of expectations yields

A simple application of the Cauchy-Schwarz inequality then gives

where the second lines follws by the Itô–Stratonovich correction.

This gives the required $O(h)$ cancellation when we expand the final term in (25) using Theorem D.10.

Similar to the proof of Theorem D.5, we use Young’s inequality to estimate remainder terms.

It is straightforward to estimate the second term in (25) using Minkowski’s inequality as

Putting this all together, the inequality (25) becomes

Just as before, we can immediately obtain a global error estimate by chaining together local estimates.

Let $h_{\max}>0$ be fixed. Then there exists a constant $C_{2}>0$ such that for all $n\in\{0,1,\cdots,N\}$,

Since $Y_{0}=Z_{0}\,$, it follows from Theorem D.11 that

D.4 The reversible Heun method in the additive noise setting

Unsurprisingly, this simplifies the analysis and gives an $O(h)$ strong convergence rate for the method.

Expanding f\big{(}Z_{n}\big{)} and f\big{(}Z_{n+1}\big{)} at $Y_{n}$ using Taylor’s theorem [52, Theorem 3.5.6] yields

where $\widetilde{R}_{n}^{Y}$ is given by

Similar to the proofs of Theorems D.3 and D.7, we can estimate this remainder term as

Likewise, the additive noise SDE (26) admits a simpler Taylor expansion than the general SDE (13).

where $J_{n}$ denotes the time integral of Brownian motion over the interval $[nh,(n+1)h]$, that is

As $\sigma$ is constant, the terms involving second and third iterated integrals of $W$ do not appear. Therefore the result follows from more general expansions, such as [53, Proposition 5.10.1]. ∎

To simplify the error analysis, we note the following lemma.

For the $d=1$ case, the lemma was shown in [54, Definition 3.5]. When $d>1$, the result is still straightforward as each coordinate of $W$ is an independent one-dimensional Brownian motion. ∎

Using the same arguments as before, we can obtain error estimates for reversible Heun method.

Let $h_{\max}>0$ be fixed. Then there exist constants $c_{5},c_{6}>0$ such that for $n\in\{0,1,\cdots,N-1\}$,

By Theorems D.13 and D.14 along with Lemma D.15, we have

The result now follows using exactly the same arguments as in the proof of Theorem D.11. ∎

Let $h_{\max}>0$ be fixed. Then there exists a constant $C_{3}>0$ such that for all $n\in\{0,1,\cdots,N\}$,

Since $Y_{0}=Z_{0}\,$, it follows from Theorem D.16 that

It is known that Heun’s method achieves second order weak convergence for additive noise SDEs . This can make Heun’s method more appealing for SDE simulation than other two-stage methods, such as the standard midpoint method – which is first order weak convergent. Whilst understanding the weak convergence of the reversible Heun method is a topic for future work, we present numerical evidence that it has similar convergence properties as Heun’s method for SDEs with additive noise.

We apply the standard and reversible Heun methods to the following scalar anharmonic oscillator:

with $y_{0}=1$, and compute the following error estimates by standard Monte Carlo simulation:

where $\{Y_{n}\}$ denotes a numerical solution of the SDE (28) obtained with step size $h=\frac{T}{N}$ and $Y_{T}^{\text{fine}}$ is an approximation of $y_{Y}$ obtained by applying Heun’s method to (28) with a finer step size of $\frac{1}{10}h$. Both $Y_{N}$ and $Y_{T}^{\text{fine}}$ are obtained using the same Brownian sample paths and the time horizon is $T=1$.

The results of this simple numerical experiment are presented in Figures 5 and 6. From the graphs, we observe that the standard and reversible Heun methods exhibit very similar convergence rates (strong order 1.0 and weak order 2.0).

D.5 Stability properties of the reversible Heun method in the ODE setting

In this section, we present a stability result for the reversible Heun method when applied to an ODE,

Just as for the error analysis, it will be helpful to consider two steps of the reversible Heun method. In particular, the updates for the $Z$ component of the numerical solution satisfy

with the second value of $Z$ being computed using a standard Euler step as Z_{1}:=Z_{0}+f\big{(}t_{0}\,,Z_{0}\big{)}h. That is, $\{Z_{n}\}$ is precisely the numerical solution obtained by the leapfrog/midpoint method, see . The absolute stability region of this ODE solver is well-known and given below.

Suppose that we apply the leapfrog/midpoint method to obtain a numerical solution $\{Z_{n}\}_{n\geq 0}$ for the linear test equation

Using similar techniques, it is straightforward to extend this result to the reversible Heun method.

Suppose that we apply the reversible Heun method to obtain a pair of numerical solutions $\{Y_{n}\}_{n\geq 0}\,,\{Z_{n}\}_{n\geq 0}\,$ for the linear test equation

By Theorem D.18, it is enough to show that $|Y_{n}|$ is bounded for all $n\geq 0$ when $\lambda h\in[-i,i]$. It follows from the difference equation (30) and the formula for $Z_{1}$ that

For each $k\geq 0$, we have Y_{k+1}=Y_{k}+\frac{1}{2}\lambda\big{(}Z_{k}+Z_{k+1}\big{)}h\hskip 0.7113pt and so we can explicitly compute $Y_{n}$ as

Since $\lambda h\in[-i,i]$, we have $\eta_{1}=\lambda h+\sqrt{1-|\lambda h|^{2}}$ and $\eta_{2}=\lambda h-\sqrt{1-|\lambda h|^{2}}$, which implies that

The reversible Heun method is not $A$-stable for ODEs as that would require $|\eta_{i}|<1$.

Appendix E Sampling Brownian motion

We begin by providing the complete traversal and splitting algorithm needed to find or create all intervals in the Brownian Interval, as in Section 4. See Algorithm 4.

Here, List is an ordered data structure that can be appended to, and iterated over sequentially. For example a linked list would suffice. We let split_seed denote a splittable PRNG as in Salmon et al. , Claessen and Pałka . We use $*$ to denote an unfilled part of the data structure, equivalent to None in Python or a null pointer in C/C++; in particular this is used as a placeholder for the (nonexistent) children of leaf nodes. We use $=$ to denote the creation of a new local variable, and $\leftarrow$ to denote in-place modification of a variable.

E.2 Discussion

The function traverse is a depth-first tree search for locating an interval within a binary tree. The search may split into multiple (potentially parallelisable) searches if the target interval crosses the intervals of multiple existing leaf nodes. If the search’s target is not found then additional nodes are created as needed.

Sections 4 and E.1 now between them define the algorithm in technical detail.

There are some further technical considerations worth mentioning. Recall that the context we are explicitly considering is when sampling Brownian motion to solve an SDE forwards in time, then the adjoint backwards in time, and then discarding the Brownian motion. This motivates several of the choices here.

First, the access patterns of SDE solvers are quite specific. Queries will be over relatively small intervals: the step that the solver is making. This means that the list of nodes populated by traverse is typically small. In our experiments we observed it usually only consisting of a single element; occasionally two. In contrast if the Brownian Interval has built up a reasonable tree of previous queries, and was then queried over $[0,s]$ for $s\gg 0$, then a long (inefficient) list would be returned. It is the fact that SDE solvers do not make such queries that means this is acceptable.

Search hints: starting from J^^𝐽\widehat{J}

Moreover, the queries are either just ahead (fixed-step solvers; accepted steps of adaptive-step solvers) or just before (rejected steps of adaptive-step solvers) previous queries. Thus in Algorithm 3, we keep track of the most recent node $\widehat{J}$, so that we begin traverse near to the correct location. This is what ensures the modal time complexity is only $\mathcal{O}\left(1\right)$, and not $\mathcal{O}\left(\log(1/s)\right)$ in the average step size $s$, which for example would be the case if searching commenced from the root on every query.

LRU cache

The fact that queries are often close to one another is also what makes the strategy of using an LRU (least recently used) cache work. Most queries will correspond to a node that have a recently-computed parent in the cache.

Backward pass

The queries are broadly made left-to-right (on the forward pass), and then right-to-left (on the backward pass). (Other than the occasional rejected adaptive step.)

Left to its own devices, the forward pass will thus build up a highly imbalanced binary tree. At any one time, the LRU cache will contain only nodes whose intervals are a subset of some contiguous subinterval $[s,t]$ of the query space $[0,T]$. Letting $n$ be the number of queries on the forward pass, then this means that the backward pass will consume $\mathcal{O}\left(n^{2}\right)$ time – each time the backward pass moves past $s$, then queries will miss the LRU cache, and a full recomputation to the root will be triggered, costing $\mathcal{O}\left(n\right)$. This will then hold only nodes whose intervals are subets of some contiguous subinterval $[u,s]$: once we move past $u$ then this $\mathcal{O}\left(n\right)$ procedure is repeated, $\mathcal{O}\left(n\right)$ times. This is clearly undesirable.

This is precisely analogous to the classical problem of optimal recomputation for performing backpropagation, whereby a dependency graph is constructed, certain values are checkpointed, and a minimal amount of recomputation is desired; see Griewank .

In principle the same solution may be applied: apply a snapshotting procedure in which specific extra nodes are held in the cache. This is a perfectly acceptable solution, but implementing it requires some additional engineering effort, carefully determining which nodes to augment the cache with.

Fortunately, we have an advantage that Griewank does not: we have some control over the dependency structure between the nodes, as we are free to prespecify any dependency structure we like. That is, we do not have to start the binary tree as just a stump. We may exploit this to produce an easier solution.

Given some estimate $\nu$ of the average step size of the SDE solver (which may be fixed and known if using a fixed step size solver), a size of the LRU cache $L$, and before a user makes any queries, then we simply make some queries of our own. These queries correspond to the intervals $[0,T/2],[T/2,T],[0,T/4],[T/4,T/2],\ldots$, so as to create a dyadic tree, such that the smallest intervals (the final ones in this sequence) are of size not more than $\nu L$. (In practice we use $\frac{4}{5}\nu L$ as an additional safety factor.)

Letting $[s,t]$ be some interval at the bottom of this dyadic tree, where $t\approx s+\frac{4}{5}\nu L$, then we are capable of holding every node within this interval in the LRU cache. Once we move past $s$ on the backward pass, then we may in turn hold the entire previous subinterval $[u,s]$ in the LRU cache, and in particular the values of the nodes whose intervals lie within $[u,s]$ may be computed in only logarithmic time, due to the dyadic tree structure.

This is now analogous to the Virtual Brownian Tree of Gaines and Lyons , Li et al. . (Up to the use of intervals rather than points.) If desired, this approach may be loosely interpreted as placing a Brownian Interval on every leaf of a shallow Virtual Brownian Tree.

Recursion errors

We find that for some problems, the recursive computations of traverse (and in principle also sample, but this is less of an issue due to the LRU cache) can occasionally grow very deep. In particular this occurs when crossing the midpoint of the pre-specified tree: for this particular query, the traversal must ascend the tree to the root, and then descend all the way down again. As such traverse should be implemented with trampolining and/or tail recursion to avoid maximum depth recursion errors.

CPU vs GPU memory

We describe this algorithm as requiring only constant memory. To be more precise, the algorithm requires only constant GPU memory, corresponding to the fixed size of the LRU cache. As the Brownian Interval receives queries then its internal tree tracking dependencies will grow, and CPU memory will increase. For deep learning models, GPU memory is usually the limiting (and so more relevant) factor.

Stochastic integrals

where $\lambda_{s,t}$ is an anti-symmetric matrix with independent entries $\lambda_{s,t}^{i,j}\sim\mathcal{N}(0,\frac{1}{12}(t-s)^{2}),i<j$.

Appendix F Experimental Details and Further Results

Several metrics were used to evaluate final model performance of the trained Latent SDEs and SDE-GANs.

A classifier was trained to distinguish real from generated data. This is trained by taking an 80%/20% split of the test data, training the classifier on the 80%, and evaluating its performance on the 20%. This produces a classification accuracy as a performance metric.

We parameterise the classifier as a Neural CDE , whose vector field is an MLP with two hidden layers each of width 32. The evolving hidden state is also of width 32. A final classification result is given by applying a learnt linear readout to the final hidden state, which produces a scalar. A sigmoid is then applied and this is trained with binary cross-entropy.

It is trained for 5000 steps using Adam with a learning rate of $10^{-4}$ and a batch size of 1024.

Smaller accuracies – indicating inability to distinguish real from generated data – are better.

Label classification (train-on-synthetic-test-on-real)

Some datasets (in particular the air quality dataset) have labelled classes associated with each sample time series. For these datasets, a classifier was trained on the generated data – possible as every model we train is trained conditional on the class label as an input – and then evaluated on the real test data. This produces a classification accuracy as a performance metric.

We parameterise the classifier as a Neural CDE, with the same architecture as before. A final classification result is given by applying a learnt linear readout to the final hidden state, which produces a vector of unnormalised class probabilities. These are normalised with a softmax and trained using cross-entropy.

It is trained for 5000 steps using Adam with a learning rate of $10^{-4}$ and a batch size of 1024.

Larger accuracies – indicating similarity of real and generated data – are better.

Prediction (train-on-synthetic-test-on-real)

A sequence-to-sequence model is trained to perform time series forecasting: given the first 80% of a time series, can the latter 20% be predicted. This is trained on the generated data, and then evaluated on the real test data. This produces a regression loss as a performance metric.

We parameterise the predictor as a sequence-to-sequence Neural CDE / Neural ODE pair. The Neural CDE is as before. The Neural ODE has a vector field which is an MLP of two hidden layers, each of width 32. Its evolving hidden state is also of width 32. An evolving prediction is given by applying a learnt linear readout to the evolving hidden state, which produces a time series of predictions. These are trained using an $L^{2}$ loss.

It is trained for 5000 steps using Adam with a learning rate of $10^{-4}$ and a batch size of 1024.

Smaller losses – indicating similarity of real and generated data – are better.

Maximum mean discrepancy

In our case, $\mathcal{X}$ corresponds to the observed time series, and we use a depth-5 signature transform as the feature map . Similarly, the untruncated signature can be used as a feature map .

(Note that MMDs may also be used as differentiable optimisation metrics provided $\psi$ is differentiable . A mistake we have seen ‘in the wild’ for training SDEs is to choose a feature map that is overly simplistic, such as taking $\psi$ to be the marginal mean and variance at all times. Such a feature map would fail to capture time-varying correlations; for example $W$ and $t\mapsto W(0)\sqrt{t}$, where $W$ is a Brownian motion, would be equivalent under this feature map.)

Smaller values – indicating similarity of real and generated data – are better.

F.2 Common details

The following details are in common to all experiments.

PyTorch was used as an autodifferentiable tensor framework .

SDEs were solved using torchsde . CDEs were solved using torchcde . ODEs were solved using torchdiffeq . Signatures were computed using Signatory .

Tensors had their shapes annotated using the torchtyping library, which helped to enforce correctness of the implementation.

Hyperparameter optimisation was performed using the Ax library .

Numerical methods

SDEs were solved using either the reversible Heun method or the midpoint method (as per the experiment), and trained using continuous adjoint methods.

The CDE used in the discriminator of an SDE-GAN was solved using either the reversible Heun method, or the midpoint method, in common with the choice made in the generator, and trained using continuous adjoint methods.

The ODEs solved for the train-on-synthetic-test-on-real prediction metric used the midpoint method, and were trained using discretise-then-optimise backpropagation. The CDEs solved for the various evaluation metrics used the midpoint method, and trained using discretise-then-optimise backpropagation. (These were essentially arbitrary choices – we merely needed to fix some choices throughout to ensure a fair comparison.)

Normalisation

Every dataset is normalised so that its initial value (at time $=0$) has mean zero and unit variance. (That is, calculate mean and variance statistics of just the initial values in the dataset, and then normalise every element of the dataset using these statistics.)

We speculate that normalising based on the initial condition produces better results than calculating mean and variance statistics over the whole trajectory, as the rest of the trajectory cannot easily be learnt unless its initial condition is well learnt first. We did not perform a thorough investigation of this topic, merely finding that this worked well enough on the problems considered here.

The times at which observations were made were normalised to have mean zero and unit range. (This is of relevance to the modelling, as the generated samples must be made over the same timespan: that is the integration variable $t$ must correspond to some parameterisation of the time at which data is actually observed.)

Dataset splits

We used 70% of the data for training, 15% for validation and hyperparameter optimisation, and 15% for testing.

Optimiser

The batch size is always taken to be 1024. The number of training steps varies by experiment, see below.

Following Kidger et al. we use Adadelta to train every SDE-GAN.

Stochastic weight averaging

When training SDE-GANs, we take a Cesàro mean over the latter 50% of the training steps, of the generator’s weights, to produce the final trained model. Known as ‘stochastic weight averaging’ this often slightly improves GAN training .

Architectures

Each of $\zeta_{\theta},\mu_{\theta},\sigma_{\theta},\xi_{\phi},f_{\phi},g_{\phi}$ were parameterised as MLPs. (From equations (1), (2), (10).)

Following Li et al. , then $\nu_{\phi}$ (of equation (10)) was parameterised as an MLP composed with a GRU.

In brief, that is $\nu_{\phi}(t,\widehat{X}_{t},Y_{\text{true}})=\nu^{1}_{\phi}(t,\widehat{X}_{t},\nu^{2}_{\phi}({\left.\kern-1.2ptY_{\text{true}}\right|_{[t,T]}}))$, where $\nu^{1}_{\phi}$ is an MLP, and $\nu^{2}_{\phi}$ is a GRU run backwards-in-time from $T$ to $t$ over whatever discretisation of ${\left.\kern-1.2ptY_{\text{true}}\right|_{[t,T]}}$ is observed.

In all cases, for simplicity, the LipSwish activation function was used throughout.

Hyperparameter optimisation

Hyperparameter optimisation used the default optimisation strategy provided by Ax (initial quasirandom Sobol sampling followed by Bayesian optimisation). The optimisation metric was the MMD evaluation metric, due to the speed at which it can be computed relative to the other optimisation metrics (which require training an auxiliary model).

The Latent SDE model was hyperoptimised on every dataset. To ensure we do not bias results in our favour, the hyperoptimised models used the midpoint method, not the reversible Heun method, throughout, and was optimised using discretise-then-optimise backpropagation.

The SDE-GAN models then used the same hyperparameters where applicable, for example on learning rate, neural network size, and so on. This fixes mosts of the hyperparameters. A few extra hyperparameters were then chosen manually.

First, the size of the initial noise $V$ and the dimensionality of the Brownian motion $W$ were arbitrarily fixed at 10.

Second, we adjusted the initialisation strategy for the parameters of the SDE. For each dataset, we picked some constants

The one exception is the Ornstein–Uhlenbeck dataset, for which the SDE-GAN had $\zeta_{\theta}$, $\mu_{\theta}$, $\sigma_{\theta},\xi_{\phi},f_{\phi},g_{\phi}$ parameterised as MLPs with one hidden layer of width 32, and had an evolving hidden state $X$ of size 32; these figures were chosen as being known to work based on early experiments.

Compute resources

Experiments were performed on an internal GPU cluster. Each experiment used only a single GPU at a time. Amount of compute time varied depending on the experiment – some took a few hours, some took a few days. Precise times given in the tables of results.

(Moreover some would likely have taken a few weeks without the algorithmic speed improvements introduced in this paper. Recall that each baseline experiment used the improvements introduced in the other sections of this paper, simply to produce tractable training times. For example the Brownian Interval was used throughout.)

GPU types varied between GeForce RTX 2080 Ti, Quadro GP100, and A100s.

F.3 Weights dataset

We consider a dataset of weights of a small convolutional network, as it is trained to classify MNIST, with training using stochastic gradient descent, as in Kidger et al. . The network was trained 10 times, and all weight trajectories, across all runs, were aggregated together to form a dataset of univariate time series.

Each time series is 50 elements long, corresponding to how the weights change over 50 epochs.

We train an SDE-GAN on this dataset. We train the generator and discriminator for 80 000 steps each.

Each MLP ($\zeta_{\theta},\mu_{\theta},\sigma_{\theta},\xi_{\phi},f_{\phi},g_{\phi}$) is parameterised as having two hidden layers each of width 67. The evolving hidden states $X,H$ are each of size 62.

The parameters of $\zeta_{\theta},\xi_{\phi}$ have a learning rate of $4.9\times 10^{-3}$. The parameters of $\mu_{\theta},\sigma_{\theta},f_{\phi},g_{\phi}$ have a learning rate of $1.3\times 10^{-3}$.

The initialisation scaling parameters $\alpha$ and $\beta$ (of equation (33)) were selected to be 4.5 and 0.25 respectively.

Results

We compare the reversible Heun method to the midpoint method on the weights dataset, by training an SDE-GAN.

We compute three test metrics: classification of real versus generated data, forecasting via train-on-synthetic-test-on-real, and a maximum mean discrepancy. We additionally report training time. See Table 4.

First and most notably, we see a dramatic reduction in training time: the training time of the reversible Heun method is roughly half that of the midpoint method. This corresponds to the reduction in vector field evaluations of the reversible Heun method.

We additionally see better performance on the test metrics, as compared to the midpoint method. This corresponds to the calculation of numerically precise gradients via the reversible Heun method.

We believe that these test metrics could be further improved (in particular the classification accuracy) given further training time.

The Brownian Interval (Section 4) is used to sample Brownian noise, and the SDE-GAN is trained using clipping as in Section 5. (Without either of which the baseline experiments would have taken infeasibly long to run.)

F.4 Air quality dataset

This is a dataset of air quality samples over Beijing, as they vary over the course of a day. Each time series is 24 elements long, corresponding to a single hour each day. We consider specifically the PM2.5 particulate matter concentration, and ozone concentration, to produce a dataset of bivariate time series. In particular the ozone channel was selected as displaying obvious non-autonomous behaviour: the latter half of the time series often includes a peak.

This dataset is available via the UCI machine learning repository .

Each time series has a label, corresponding to which of 12 different locations the measurements were made at.

We train a Latent SDE on this dataset. We train for 40 000 steps.

Each MLP ($\zeta_{\theta},\mu_{\theta},\sigma_{\theta},\xi_{\phi},\nu^{1}_{\phi}$) is parameterised as having a single hidden layer of width 84. The evolving hidden state $X$ is of size 63.

$\nu^{2}_{\phi}$ was parameterised as GRU with hidden size 84, whose final hidden state has a learnt affine map applied, to produce vector of size 60.

The parameters of $\zeta_{\theta}$ have learning rate $1.1\times 10^{-4}$. The parameters of $\mu_{\theta},\sigma_{\theta},\xi_{\phi},\nu^{1}_{\phi},\nu^{2}_{\phi}$ have learning rate $1.9\times 10^{-5}$.

The initialisation scaling parameters $\alpha$ and $\beta$ (of equation (33)) were selected to be 2 and 1 respectively.

Results

We compare the reversible Heun method to the midpoint method on the air quality dataset, by training a Latent SDE.

We compute four test metrics: classification of real versus generated data, classification via train-on-synthetic-test-on-real, forecasting via train-on-synthetic-test-on-real, and a maximum mean discrepancy. We additionally report training time. See Table 5.

Here, the most important metric is again the improvement in training time: whilst less dramatic than the previous SDE-GAN experiment, it is still a speed improvement of $1.2\times$.

Performance on the test metrics varies between the solvers, without a clear pattern.

It is worth noting the apparently poor real/fake classification accuracies obtained using either solver. This is typical of Latent SDEs in general, in particular as opposed to SDE-GANs. Whilst Latent SDEs are substantially quicker to train, they tend to produce less convincing samples.

F.5 Gradient error analysis

We investigate the error made in the gradient calculation using continuous adjoints. We consider using the midpoint method, Heun’s method, and the reversible Heun method, and vary the step size in decreasing powers of two.

Numerical values are shown in Table 6. Results are plotted graphically in the main text.

F.6 Brownian benchmarks

We benchmark the Brownian Interval against the Virtual Brownian Tree across several benchmarks.

We subdivide the interval $$ into a disjoint union of equal-sized intervals. Across each interval we then place a query for a small Brownian increment.

We consider subdividing $$ into different numbers of subintervals (and so of different sizes): either 10, 100 or 1000 subintervals.

Sequential access: this involves querying every interval in order, from 0 to 1. Every interval is queried precisely once. This simulates an SDE solve from 0 to 1.

Doubly sequential access: this involves querying every interval in order, from 0 to 1, and then querying them again in reverse order, from 1 to 0. Every interval is queried precisely twice. This simulates an SDE solve from 0 to 1, followed by a backpropagation via the continuous adjoint method, from 1 to 0.

Random access: this involves querying every interval precisely once, in a random order.

We consider several batch sizes: either a single Brownian simulation, a typical batch size of 2560 simultaneous Brownian simulations (corresponding to a batch of size 256, each with a vector of 10 Brownian motions), and a large batch size of 32768 simultaneous Brownian simulations.

For every such combination we report metrics for the Brownian Interval and the Virtual Brownian Tree.

For every such combination we run 32 repeats. The reported metric is the fastest (minimum time) over these repeats.Not the mean. Errors in speed benchmarks are one-sided, and so the minimum time represents the least noisy measurement.

The Brownian Interval is consistently and substantially faster than the Virtual Brownian Tree, across all access patterns, batch sizes, and number of subintervals.

On the doubly sequential access benchmark (emulating the SDE solve and backpropagation typical in practice), we see that total speed-ups vary from a factor of $2.89\times$ to a factor of $13.5\times$. That is, at minimum, speed is roughly tripled. Potentially it is improved by over an order of magnitude. Typical values are speed-ups of $6$–$8\times$.

SDE solve benchmarks

We now benchmark the Brownian Interval against the Virtual Brownian Tree on the actual task of solving and backpropagating through an SDE.

As before we consider several numbers of subintervals, several batch sizes, and run 32 repeats and take the fastest time.

Our test SDE is an Itô SDE with diagonal noise:

For $i=10,100,1000$ subintervals, we calculate a forward pass for $$using the Euler–Maruyama method, to calculate$X_{j/(i-1)}$for$j\in\{0,\ldots,i-1\}$. We then backpropagate from the vector$(X_{j/(i-1)})_{j}$to$X_{0}$, using the continuous adjoint method.

We once again see that the Brownian Interval is uniformly and substantially faster than the Virtual Brownian Tree, in all regimes. On smaller problems (batch size equal to 1 or 2560), then it is typically twice as fast; at worst it is $1.89\times$ as fast. Meanwhile on larger problems (batch size equal to 32768), then it is typically ten times as fast.

These benchmarks include the realistic overheads involved in solving an SDE (such as evaluating its vector fields), and represent typical speed-ups from using the Brownian Interval over the Virtual Brownian Tree.

F.7 Time-dependent Ornstein–Uhlenbeck dataset

This is a dataset of univariate samples of length 32 from the time-dependent Ornstein–Uhlenbeck process

with $\rho=0.02,\kappa=0.1$ and $\chi=0.4$ and $t\in$

We train an SDE-GAN on this dataset. We train the generator for 20 000 steps and the discriminator for 100 000 steps.

Each MLP ($\zeta_{\theta},\mu_{\theta},\sigma_{\theta},\xi_{\phi},f_{\phi},g_{\phi}$) is parameterised as having a single hidden layer of width 32. The evolving hidden states $X,H$ are each of size 32.

The parameters of $\zeta_{\theta},\xi_{\phi}$ have a learning rate of $1.6\times 10^{-3}$. The parameters of $\mu_{\theta},\sigma_{\theta},f_{\phi},g_{\phi}$ have a learning rate of $2.0\times 10^{-4}$.

The initialisation scaling parameters $\alpha$ and $\beta$ (of equation (33)) were selected to be 5 and 0.5 respectively.

Results

We compare training SDE-GANs by careful clipping (as in Section 5) to using gradient penalty (as in Kidger et al. ) on the OU dataset.

As our implementation of the reversible Heun method does not support a double backward, we provide two comparisons of interest: reversible Heun method with clipping against midpoint with gradient penalty, and midpoint with clipping against midpoint with gradient penalty.

We compute three test metrics: classification of real versus generated data, forecasting via train-on-synthetic-test-on-real, and a maximum mean discrepancy. We additionally report training time. See Table 11.

We see that reversible Heun with clipping dominates midpoint with clipping, which in turn dominates midpoint with gradient penalty, across all metrics.

As per Kidger et al. , the poor performance of gradient penalty is due in part to the numerical errors of a double adjoint. Switching to midpoint with clipping produces substantially better test metrics. It additionally improves training speed by $1.41\times$.

Switching from midpoint with clipping to reversible Heun with clipping then produces another substantial boost to the test metrics – most notably the real-versus-fake classification accuracy. It also improves training speed by another $1.09\times$.

Appendix G Ethical statement

SDEs are already a widely used modelling paradigm, primarily in fields such as finance and science. In this regard they are a tried-and-tested mathematical tool, with, to the best of the authors knowledge, no significant ethical concerns attached.

As this paper extends this existing methodology, then broadly speaking we expect the same to be true.

We anticipate the results of this paper as having applications to finance and the sciences. For example, to model the movement of asset prices, or to model predator-prey interactions.

As such no significant negative societal impacts are anticipated.

Environmental impacts

The primary contributions of this paper are speed improvements to existing methodologies. As such we anticipate a positive environmental impact from this paper, due to a reduction in the compute resources necessary to train model.

Dataset content

The data we are using contains no personally identifiable or offensive content.

Dataset provenance

All data used has been made publicly available, for example via the UCI machine learning repository . To the best of our knowledge this availability was voluntary, and so we believe the use of the data to be ethical and without licensing issues.