Efficient and Accurate Gradients for Neural SDEs
Year: 2022

Abstract:
Neural SDEs combine many of the best qualities of both RNNs and SDEs: memory efficient training, high-capacity function approximation, and strong priors on model space. This makes them a natural choice for modelling many types of temporal dynamics. Training a Neural SDE (either as a VAE or as a GAN) requires backpropagating through an SDE solve. This may be done by solving a backwards-in-time SDE whose solution is the desired parameter gradients. However, this has previously suffered from severe speed and accuracy issues, due to high computational cost and numerical truncation errors. Here, we overcome these issues through several technical innovations. First, we introduce the reversible Heun method. This is a new SDE solver that is algebraically reversible: eliminating numerical gradient errors, and the first such solver of which we are aware. Moreover it requires half as many function evaluations as comparable solvers, giving up to a 1.98× speedup. Second, we introduce the Brownian Interval: a new, fast, memory efficient, and exact way of sampling and reconstructing Brownian motion. With this we obtain up to a 10.6× speed improvement over previous techniques, which in contrast are both approximate and relatively slow. Third, when specifically training Neural SDEs as GANs (Kidger et al. 2021), we demonstrate how SDE-GANs may be trained through careful weight clipping and choice of activation function. This reduces computational cost (giving up to a 1.87× speedup) and removes the numerical truncation errors associated with gradient penalty. Altogether, we outperform the state-of-the-art by substantial margins, with respect to training speed, and with respect to classification, prediction, and MMD test metrics. We have contributed implementations of all of our techniques to the torchsde library to help facilitate their adoption.

heading: 1 Introduction
text: Stochastic differential equations Stochastic differential equations have seen widespread use in the mathematical modelling of random phenomena, such as particle systems [1], financial markets [2], population dynamics [3], and genetics [4]. 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 [5].
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 [6, 7].) Several authors have now studied or introduced Neural SDEs; [8, 9, 10, 11, 12, 13, 14, 15, 16] amongst others.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
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 [15] or Wasserstein distance [16], 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 [17, 18], ARMA [19] and so on. It has also attracted much recent interest with (besides Neural SDEs) the development of models such as Time Series GAN [20], Latent ODEs [21], GRU-ODE-Bayes [22], ODE2VAE [23], CTFPs [24], Neural ODE Processes [25] and Neural Jump ODEs [26].

heading: 1.1 Contributions
text: 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 ofO (1), and consumes only O (1) GPU memory. This is contrast to previous techniques requiring either O (T ) memory, or a choice of approximation error ε 1 and then a time complexity of O (log(1/ε)). 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.

heading: 2 Background

heading: 2.1 Neural SDE construction
text: Certain minimal amount of structure Following Kidger et al. [16], we construct Neural SDEs with a certain minimal amount of structure. Let T > 0 be fixed and suppose we wish to model a path-valued random variable Ytrue : [0, T ]→ Ry . The size of y is the dimensionality of the data.1
Let W : [0, T ]→ Rw be a w-dimensional Brownian motion, and let V ∼ N (0, Iv,) be drawn from a v-dimensional standard multivariate normal. The values w, v are hyperparameters describing the size of the noise. Let
ζθ : Rv → Rx, µθ : [0, T ]× Rx → Rx, σθ : [0, T ]× Rx → Rx×w, `θ : Rx → Ry,
where ζθ, µθ and σθ are neural networks and `θ is affine. Collectively these are parameterised by θ. The dimension x is a hyperparameter describing the size of the hidden state.
1In practice we will typically observe some discretised time series sampled from Ytrue. For ease of presentation we will neglect this detail for now and will return to it in Section 2.3.
We consider Neural SDEs as models of the form
X0 = ζθ(V ), dXt = µθ(t,Xt) dt+ σθ(t,Xt) ◦ dWt, Yt = `θ(Xt), (1) for t ∈ [0, T ], with X : [0, T ]→ Rx the (strong) solution to the SDE.2 The solution X is guaranteed to exist given mild conditions: that µθ, σθ are Lipschitz, and that EV [ ζθ(V ) 2 ] <∞.
We seek to train this model such that Y d≈ Ytrue. That is to say, the model Y should have approximately the same distribution as the target Ytrue, for some notion of approximate. (For example, to be similar with respect to the Wasserstein distance).
RNNs as discretised SDEs The minimal amount of structure is chosen to parallel RNNs. The solution X may be interpreted as hidden state, and the `θ maps the hidden state to the output of the model. In Appendix A we provide sample PyTorch [27] code computing a discretised SDE according to the Euler–Maruyama method. The result is an RNN consuming random noise as input.

heading: 2.2 Training criteria for Neural SDEs
text: Equation (1) produces a random variable Y : [0, T ]→ Ry implicitly depending on parameters θ. This model must still be fit to data. This may be done by optimising a distance between the probability distributions (laws) for Y and Ytrue.
SDE-GANs The Wasserstein distance may be used by constructing a discriminator and training adversarially, as in Kidger et al. [16]. Let Fφ(Y ) = mφ ·HT , where
H0 = ξφ(Y0), dHt = fφ(t,Ht) dt+ gφ(t,Ht) ◦ dYt, (2) for suitable neural networks ξφ, fφ, gφ and vectormφ. This is a deterministic function of the generated sample Y . Here · denotes a dot product. They then train with respect to
min θ max φ
( EY [Fφ(Y )]− EYtrue [Fφ(Ytrue)] ) . (3)
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. [15] instead optimise a KL divergence. This consists of constructing an auxiliary process X̂ with drift νφ parameterised by φ, and optimising an expression of the form
min θ,φ EW,Ytrue
[∫ T
0
(Ytrue,t − `θ(X̂t))2 + 1
2
∥∥∥(σθ(t, X̂t))−1(µθ(t, X̂t)− νφ(t, X̂t, Ytrue)) ∥∥∥ 2
2 dt
] . (4)
The full construction is moderately technical; see Appendix B for further details.

heading: 2.3 Discretised observations
text: Observations of Ytrue 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 Ytrue of the observed data. If training a Latent SDE, then νφ in equation (4) may depend explicitly on the discretised Ytrue.

heading: 2.4 Backpropagation through SDE solves
text: 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 θ, φ.
Here we use the continuous adjoint method. Also known as simply ‘the adjoint method’, or ‘optimisethen-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
2The notation ‘ ◦ dWt’ denotes Stratonovich integration. Itô is less efficient; see Appendix C.
networks [28], intermediate computations such as Xt for t < T are reconstructed from output computations, so that they do not need to be held in memory.
Given some Stratonovich SDE dZt = µ(t, Zt) dt+ σ(t, Zt) ◦ dWt for t ∈ [0, T ], (5) and a loss L : Rz → R on its terminal value ZT , then the adjoint process At = dL(ZT )/dZt ∈ Rz is a (strong) solution to
dAit = −Ajt ∂µj
∂Zi (t, Zt) dt−Ajt
∂σj,k
∂Zi (t, Zt) ◦ dW kt , (6)
which in particular uses the same Brownian motion W as on the forward pass. Equations (5) and (6) may be combined into a single SDE and solved backwards-in-time3 from t = T to t = 0, starting from ZT = ZT (computed on the forward pass of equation (5)) and AT = L(ZT )/dZT . Then A0 = dL(ZT )/Z0 is the desired backpropagated gradient.
Note that we assumed here that the loss L acts only on ZT , 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 ZT , and then backpropagation may proceeed as described here.
The main issue is that the two numerical approximations to Zt, computed in the forward and backward passes of equation (5), are different. This means that the Zt used as an input in equation (6) has some discrepancy from the forward calculation, and the gradients A0 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.

heading: 2.5 Alternate constructions
text: There are other uses for Neural SDEs, beyond our scope here. For example Song et al. [30] combine SDEs with score-matching, and Xu et al. [31] use SDEs to represent Bayesian uncertainty over parameters. The techniques introduced in this paper will apply to any backpropagation through an SDE solve.
3 Reversible Heun method
Algorithm 1: Forward pass Input: tn, zn, ẑn, µn, σn,∆t,W tn+1 = tn + ∆t
∆Wn = Wtn+1 −Wtn ẑn+1 = 2zn − ẑn + µn∆t+ σn∆Wn µn+1 = µ(tn+1, ẑn+1)
σn+1 = σ(tn+1, ẑn+1)
zn+1 = zn + 1
2 (µn + µn+1)∆t
+ 1
2 (σn + σn+1)∆Wn
Output: tn+1, zn+1, ẑn+1, µn+1, σn+1
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
dZt = µ(t, Zt) dt+ σ(t, Zt) ◦ dWt, (7) with known initial condition Z0.
Solver We begin by selecting a step size ∆t, and initialising t0 = 0, z0 = ẑ0 = Z0, µ0 = µ(0, Z0) and σ0 = σ(0, Z0). 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∆t so that zN , ẑN , µN , σN are the final output. Then zN ≈ ZT is returned, whilst zN , ẑN , µN , σ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.
3Li et al. [15] give rigorous meaning to this via two-sided filtrations; for the reader familiar with rough path theory then Kidger et al. [29, Appendix A] also give a pathwise interpretation. The reader familiar with neither of these should feel free to intuitively treat Stratonovich (but not Itô) SDEs like ODEs.
Algorithm 2: Backward pass Input: tn+1, zn+1, ẑn+1, µn+1, σn+1,∆t,W,
∂L(ZT ) ∂zn+1 , ∂L(ZT )∂ẑn+1 , ∂L(ZT ) ∂µn+1 , ∂L(ZT )∂σn+1
# Reverse step tn = tn+1 −∆t
∆Wn = Wtn+1 −Wtn ẑn = 2zn+1 − ẑn+1 − µn+1∆t− σn+1∆Wn µn = µ(tn, ẑn)
σn = σ(tn, ẑn)
zn = zn+1 − 1
2 (µn + µn+1)∆t
− 1 2 (σn + σn+1)∆Wn
# Local forward zn+1, ẑn+1, µn+1, σn+1 = Forward(tn, zn, ẑn, µn,
σn,∆t,W )
# Local backward ∂L(ZT )
∂(zn, ẑn, µn, σn) =
∂L(ZT )
∂(zn+1, ẑn+1, µn+1, σn+1)
· ∂(zn+1, ẑn+1, µn+1, σn+1) ∂(zn, ẑn, µn, σn)
Output: tn, zn, ẑn, µn, σn, ∂L(ZT ) ∂zn , ∂L(ZT )∂ẑn , ∂L(ZT ) ∂µn , ∂L(ZT )∂σn
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 (zn, ẑn, µn, σn) from (zn+1, ẑn+1, µn+1, σ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.
Theorem. Let {zn} denote the numerical solution of (7) obtained by Algorithm 1 with a constant step size ∆t and assume sufficient regularity of µ and σ. Then there exists a constant C > 0 so that
E [∥∥zN − ZT ∥∥ 2 ] ≤ C √ ∆t ,
for small ∆t. That is, strong convergence of order 0.5. If σ is constant, then this improves to order 1.
The key idea in the proof is to consider two adjacent steps of the SDE solver. Then the update ẑn 7→ ẑn+2 becomes a step of a midpoint method, whilst zn 7→ zn+1 is similar to Heun’s method. This makes it possible to show that {ẑn} and {zn} stay close together: E [ ‖zn − ẑn‖42 ] ∼ O(∆t2). With this L4 bound on z − ẑ, we can then apply standard lines of argument from the numerical SDE literature. Chaining together local mean squared error estimates, we obtain E [ ‖zN−ZT ‖22 ] ∼ O(∆t).
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. [32].
Precise gradients The backward pass is shown in Algorithm 2. As the same numerical solution {zn} 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 ∂L(ZT )/∂zn ≈ An∆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 [15].
Table 1: SDE-GAN on weights dataset; Latent SDE on air quality dataset. Mean± standard deviation averaged over three runs.
Dataset, Solver Label classificationaccuracy (%) MMD (×10 −2) Training time
Weights, Midpoint — 4.38 ± 0.67 5.12 ± 0.01 days Weights, Reversible Heun — 1.75 ± 0.3 2.59 ± 0.05 days Air quality, Midpoint 46.3 ± 5.1 0.591 ± 0.206 5.58 ± 0.54 hours Air quality, Reversible Heun 49.2 ± 0.02 0.472 ± 0.290 4.47 ± 0.31 hours
3.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.
Versus midpoint 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×) 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×.
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 (discretisethen-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.

heading: 4 Brownian Interval
text: 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 bridges Mathematically, sampling Brownian motion is straightforward. A fixed-step numerical solver may simply sample independent increments during its time stepping. An adaptive solver (which may reject steps) may use Lévy’s Brownian bridge [33] formula to generate increments with the appropriate correlations: letting Wa,b = Wb −Wa ∈ Rw, then for u < s < t,
Wu,s|Wu,t = N ( s− u t− uWu,t , (t− s)(s− u) t− u Iw ) . (8)
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 [15].
Algorithm 3: Sampling the Brownian Interval Type: Let Node denote a 5-tuple consisting of an
interval, a seed, and three optional Nodes, corresponding to the parent node, and two child nodes, respectively. (Optional as the root has no parent and leaves have no children.)
State: Binary tree with elements of type Node, with root Î = ([0, T ], ŝ, ∗, Îleft, Îright). A Node Ĵ . Input: Interval [s, t] ⊆ [0, T ]
# The returned ‘nodes’ is a list of Nodes whose # intervals partition [s, t]. Practically speaking # this will usually have only one or two elements. # Ĵ is a hint for where in the tree we are. nodes = traverse(Ĵ , [s, t] )
def sample(I : Node): if I is Î then
return N (0, T ) sampled with seed ŝ. Let I = ([a, b], s, Iparent, Ileft, Iright) Let Iparent = ([ap, bp], sp, Ipp, Ilp, Irp) Wparent = sample(Iparent) if Ii is Irp then
Wleft = bridge(ap, bp, a,Wparent, s) return Wparent −Wleft
else return bridge(ap, bp, b, Wparent, s)
sample = LRUCache(sample)
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. [15] 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] ⊆ [0, T ] (an operation that returns Ws,t) produces the binary tree shown in Figure 3a. 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 3b. Using a splittable PRNG [34, 35], 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: Ws,t,Wt,u are completely specified by t, [s, u], Ws,u, equation (8), and the random seed for [s, u].
In principle this now gives a way to compute Ws,t; calculating Ws,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 Ws,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 O (1): 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 O (log(1/s)) 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 O (log(1/ε)) in the approximation error ε s.
Meanwhile the (GPU) memory cost is only O (1), 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 O (T ). 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 ε.

heading: 4.1 Experiments
text: We benchmark the performance of the Brownian Interval against the Virtual Brownian Tree considered in Li et al. [15]. 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× faster on smaller problems to 10.6× 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++.

heading: 5 Training SDE-GANs without gradient penalty
text: Kidger et al. [16] 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 [36] 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φ, gφ are MLPs, which is also the choice we make in practice.
Lipschitz constant one The key point is that the vector fields fφ, gφ of the discriminator must not only be Lipschitz, but must have Lipschitz constant at most one.
Given vector fields with Lipschitz constant λ, then the recurrent structure of the discriminator means that the Lipschitz constant of the overall discriminator will be O ( λT ) . Ensuring λ ≈ 1 with λ ≤ 1 thus enforces that the overall discriminator is Lipschitz with a reasonable Lipschitz constant.
Hard constraint The exponential size of O ( λT )
means that λ only slightly greater than one is still insufficient for stable training. We found that this ruled out enforcing λ ≤ 1 via soft constraints, via either spectral normalisation [37] or gradient penalty across just vector field evaluations.
Clipping The first part of enforcing this Lipschitz constraint is careful clipping. Considering each linear operation from Ra → Rb as a matrix in A ∈ Ra×b, then after each gradient update we clip its entries to the region [−1/b, 1/b]. Given x ∈ Ra, then this enforces ‖Ax‖∞ ≤ ‖x‖∞.
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.
There remain several admissible choices; we use the LipSwish activation function introduced by Chen et al. [38], defined as ρ(x) = 0.909x sigmoid(x). This was carefully constructed to have Lipschitz constant one, and to be smooth. Moreover the SiLU activation function [39, 40, 41] from which it is derived has been reported as an empirically strong choice.
The overall vector fields fφ, gφ 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.

heading: 5.1 Experiments
text: 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×.

heading: 6 Discussion
text: 6.1 Available implementation in torchsde
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 [42] package. (In which the Brownian Interval has already become the default choice, due to its speed.)

heading: 6.2 Limitations
text: 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.

heading: 6.3 Ethical statement
text: 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.

heading: 7 Conclusion
text: 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.

Review summaries:
The paper aims to address issues with training Neural SDEs; in particular, it aims to overcome issues relating to backpropagation through the SDE solve step. Previously this computation was marred by speed and accuracy issues arising due to high computational complexity, numerical errors in the SDEs solve step, as well as due to the high cost of reconstructing Brownian motion. To mitigate these issues, the paper introduces a variety of constructs. It introduces the reversible Heun method: an algebraically reversible SDE solver that reduces gradient errors to (effectively) zero. It also presents a Brownian interval construction to efficiently sample from and reconstruct Brownian motion. Additionally a weight-clipping procedure is described for improved training of neural SDEs as GANs. All of these contributions considerably advance the state-of-the-art in neural SDE training, as exemplified by various test tasks in the paper.

This paper studies both neural SDE and construction of Brownian intervals. The neural SDE considered in the current paper is in a more general form comparing to other recent works where diffusion coefficients are constant. The Brownian interval seems to provide efficient sampling algorithm for Brownian motion which could be used in the discretized algorithm.

Neural SDE is an emerging tool for temporal modeling. This paper aims to improve the efficiency and accuracy of gradient calculation in the training of Neural SDEs. To this end, the authors introduce three techniques: (1) reversible Heun method that overcomes inconsistency in the forward and backward pass, (2) Brownian interval that speeds up sampling Brownian motion, and (3) weight clipping and specific activation to avoid gradient penalty in training. Through several numerical tests the authors demonstrate that these techniques altogether offer substantial improvements over the existing Neural SDE algorithms in speed and accuracy metrics.