Efficient and Accurate Gradients for Neural SDEs
Authors: Patrick Kidger
James Foster
Xuechen Li
Terry Lyons
Venue: NeurIPS 2021 Poster (2021)
Keywords: neural sde
neural stochastic differential equation
neural ode
neural ordinary differential equation
nde
neural differential equation
backpropagation
adjoint method
optimise-then-discretise
optimize-then-discretize
discretise-then-optimise
discretize-then-optimize

Abstract: Neural SDEs combine many of the best qualities of both RNNs and SDEs, and as such are a natural choice for modelling many types of temporal dynamics. They offer memory efficiency, high-capacity function approximation, and strong priors on model space. Neural SDEs may be trained as VAEs or as GANs; in either case it is necessary to backpropagate through the SDE solve. In particular this may be done by constructing 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 complexity, numerical errors in the SDE solve, and the cost of reconstructing Brownian motion. Here, we make several technical innovations to overcome these issues. First, we introduce the \textit{reversible Heun method}: a new SDE solver that is algebraically reversible -- which reduces numerical gradient errors to almost zero, improving several test metrics by substantial margins over state-of-the-art. Moreover it requires half as many function evaluations as comparable solvers, giving up to a $1.98\times$ speedup. Next, we introduce the \textit{Brownian interval}. This is a new and computationally efficient way of exactly sampling \textit{and reconstructing} Brownian motion; this is in contrast to previous reconstruction techniques that are both approximate and relatively slow. This gives up to a $10.6\times$ speed improvement over previous techniques. After that, 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\times$ speedup), and removes the truncation errors of the double adjoint required for gradient penalty, substantially improving several test metrics. Altogether these techniques offer substantial improvements over the state-of-the-art, with respect to both training speed and with respect to classification, prediction, and MMD test metrics. We have contributed implementations of all of our techniques to the \texttt{torchsde} library to help facilitate their adoption.

Efficient and Accurate Gradients for Neural SDEs
Patrick Kidger'
James Foster'
Xuechen Li Terry Lyons'
University of Oxford: The Alan Turing Insitute [kidger, foster, tlyons)@mathe @maths ox. ac uk 1xuechen@cs. toronto.........o edu
Stanford
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 natural choice for modelling many types of temporal dynamics. Training a Neural SDE (either as a VAE or as a GAÑ) requires backprop- agating through SDE solve. This may be done by solving backwards in-time SDE whose solution is the desired parameter gradients. However, this has previ- ously suffered from severe speed and accuracy issues, due high computational cost and numerical truncation errors. Here, we overcome these issues through several technical innovations. First, introduce the reversible Heun method. This is new SDE solver that algebraically reversible: eliminating numerical gradient errors, and the first such solver of which are aware. Moreover it requires half many function evaluations as comparable solvers, giving up to 1.98 x 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 a 10.6x speed improvement over previous techniques, which in contrast are both approximate and relatively slow. Third, when specifically training Neural SDEs as GANs (Kidger et 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 1.87x speedup) and removes the numer- ical truncation errors associated with gradient penalty. Altogether, we outperform the state-of-the-ee by substantial margins, with respect to training speed, and with respect classification, prediction, and MMD test metrics. We have contributed implementations all of our techniques to the torchsde library to help facilitate their adoption.
Introduction
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 financial markets, has only two scalar parameters: a fixed drift and fixed diffusion [5].
Neural stochastic differential equations Neural stochastic differential equations offer 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, 7J.) Several authors have now studied introduced Neural SDEs; [8, 9, 10, 11, 12, 13, 14, 15, 16] amongst others.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Connections to recurrent neural networks A numerically discretised (Neural) SDE may be interpreted as an RNN (featuring a residual connection), whose input random noise Brownian motion and whose output is generated sample. Subject to a suitable loss function between distributions, such as the KI. 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 models such as Time Series GAN [20], Latent ODEs [21], GRU-ODE-Bay [22], VAE [23], CTFPs [24], Neural ODE Processes [25] and Neural Jump ODEs [26].
1.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 Hewn 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 new way of sampling and reconstructing Brownian motion. It is fast, memory efficient and exact. It has an average (modal) time complexity of O (1), and consumes only O (1) GPU memory. This contrast to previous techniques requiring either O (T) memory, or a choice of approximation error << 1 and then a time complexity of O (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 to overcome their previous incompatibility with continuous adjoint methods.
Overall, multiple technical innovations provide substantial improvements over the state-of offte-o--hh-ar with respect to training speed, and with respect to classification, prediction, and MMD test metrics.
Background
2.1 Neural SDE construction
Certain minimal amount of structure Following Kidger al. [16], construct Neural SDEs with a certain minimal amount of structure. Let T > O be fixed and suppose we wish to model a path-valued random variable Yime: [0, 7] R The size of y is the dimensionality of the data.'
Let W: [] + R be a w-ddimnnsional Brownian motion, and let V ~ NO, Ie.) be drawn from a v-dimensional standard multivariate normal. The values w, hyperparameters describing the size of the noise. Let
Co +
#o: [0, x R R., og: [0, R
R
l RR RV,
where So- 14 and og neural networks and is affine. Collectively these parameterised by A. The dimension x is a hyperparameter describing the size of the hidden state.
'In practice we will typically observe some discretised time series sampled from Yine. For ease of presentation we will neglect this detail for now and will return it in Section 2.3.

---

We consider Neural SDEs as models of the form
Xo = Co(V, dX, = pelt, Xx) dt + a Xt) o dW,,
Y toxe,
(I)
for E [0, T), with X: [0, + R the (strong) solution to the SDE. The solution X guaranteed to exist given mild conditions: that μig 0g are Lipschitz, and that By [<o(V) 00.
We seek to train this model such that Y 2 Yuue- That is to say, the model Y should have approximately the same distribution as the target Yiue: for some notion of approximate. (For example, 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 lo maps the hidden state to the output of the model. In Appendix A we provide sample Py Torch [27] code computing discretised SDE according the Maler-Maruyama method. The result is RNN consuming random noise as input.
2.2 Training criteria for Neural SDES
Equation (I) produces a random variable Y: [0, T + RR implicitly depending on parameters O. This model must still be fit to data, This may be done by optimising distance between the probability distributions (laws) for Y and Yuue.
SDE-GANs
The Wasserstein distance may be used by constructing a discriminator and training adversarially, as in Kidger et al. [16]. Let F& (Y) my HT, where
Ho = (Yo),
dH, fo(1, H)) di + (t, H2) dY,
for suitable neural networks Eo, f Us and vector m.. This is deterministic function of the generated sample Y. Here denotes a dot product. They then train with respect to
min max (Ey [F)( Ry,,.. [Fo (Yine))))
See Appendix B for additional details this approach, and in particular how it generalises the classical approach to fitting (calibrating) SDEs.
Li al. [15] instead optimise KL divergence. This consists of constructing an auxiliary process X with drift vo parameterised by ø, and optimising an expression of the form
Latent SDEs
(Xmme:)- (YYYaee - (Z),)) +
1 | ( (mo R,) uolt, Xc, Yeree)) |) dt (4)
min Ew,yon 0.0
The full construction is moderately technical; see Appendix B for further details.
2.3 Discretised observations
Observations of Yiee are typically a discrete time series, rather than a true continuous-time path. This is not serious hurdle. If training an SDE-GAN, then equation (2) may be evaluated on an interpolation Yuue of the observed data. If training a Latent SDE, then Vo in equation (4) may depend explicitly on the discretised Yis..
2.4 Backpropagation through SDE solves
Whether the loss for our generated sample Y is produced via Latent SDE or via the discriminator of SDE-GAN, it is still required to backpropagate from the loss the parameters 0,
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 neural differ- ential equations. This exploits the reversibility of differential equation: as with invertible neural
The notation "e dW, denotes Stratonovich integration. Itô is less efficient; Appendix C.

---

networks [28]. intermediate computations such as X₁ for t < T are reconstructed from output computations, so that they do not need to be held memory.
Given some Stratonovich SDE
d2: ptt, Zt) dt + (t, Zi) dW, for t E [0,T] and loss L: RR* R on its terminal value Zr, then the adjoint process A, dl(Z))/az, RR is (strong) solution to
(5)
AA; = - Opr
Sasi
azi (t, Zt) A azi Zi) o oW
(6)
which in particular uses the same Brownian motion W as on the forward pass. Equations (5) and (6) may be combined into single SDE and solved mackwards-in-tiee from starting from ZT ZT (computed on the forward pass of equation (5)) and A7 = L(2r)/dZr. Then Ao dl(2x)/Z, Zo the desired backpropagated gradient.
Note that we assumed here that the loss I. acts only on Zr, 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 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 Z,, computed in the forward and backward passes of equation (5), different. This means that the Z₁ used as an input in equation (6) has some discrepancy from the forward calculation, and the gradients Ag suffer some error as a result. (Often exacerbating an already tricky training procedure, such the adversarial training of SDE-GANS.)
See Appendix for further discussion on how an SDE solve may be backpropagated through.
2.5 Alternate constructions
There are other uses for Neural SDEs, beyond our scope here. For example Song et al. [30] combine SDEs with ccore-matching and Xu et al. [31] use SDEs to represent Bayesian uncertainty parameters. The techniques introduced in this paper will apply to any backpropagation through an SDE solve.
3 Reversible Heun method
We introduce new SDE solver, which we refer to the reversible Heun method. Its key property is algebraic reversibility; moreover the best of our knowledge it is the first SDE solver to exhibit this property.
Algorithm 1: Forward pass
Input: tpp 2n, Ens Mus on At, W
bnt1 In + Al
AW, Wen+i - We...
To fix notation, consider solving the Stratonovich SDE
Ent1 22n 2n + pn At + AW,
d plt, 7,) dt (t,, 7,) dW,, with known initial condition Zo.
(7)
Pn+1 p((n-11 Zn+1)
aat1 =o(nnt1, 2m+11)
1
Solver We begin by selecting a step size At, and ini- tialising to = 0, 20 = 20 = Zo, Mu = p10 Zo) and
2n+1 2n + (U + Mn+1 At
2
To (0, Zo). Let W denote single sample path of Brownian motion. It is important that the same sample be used for both the forward and backward passes of the
+ (on + AWn+))AW,
1
2
Output: (441, Fnt 1, E Pn+t 1nt1
algorithm; computationally this may be accomplished by taking W to be Brownian Interval, which we will introduce in Section 4.
We then iterate Algorithm 1. Suppose T = N At so that nN, 2N, ,N,, are the final output. Then NN ZT is returned, whilst N,, 2N, ,AN, NN all retained for the backward pass.
Nothing else need be saved in memory for the backward pass: in particular no intermediate computa- tions, would otherwise be typical.
'Li 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 AJ also give a pathwise interpretation. The reader familiar with neither of these should feel free to intuitively treat Stratonovich (but not It0) SDEs like ODEs.

---

Algebraic reversibility The key advan-
Algorithm 2: Backward pass
tage of the reversible Heun method, and the motivating reason for its use alongside continuous-time adjoint methods, that
Input:
(nt1, Zn Pa+ly In+1, At, W,
aLi2rl 0L(2r)
0L(Zr) 0L(2r)
is algebraically reversible. That is. it is possible to reconstruct (zn, 2no (n: from (2n+1, In+1, Mn++1; (n+1) in closed form.
Ban+1 8En+1
Oun+i 80a+1
# Reverse step
In ln+1 Al
(And without a fixed-point iteration.)
A n = Wt,1 Wt,
This crucial property will mean that it is possible to backpropagate through the SDE
En 2zn+1 -n41 Mn+ At Onti AW,
Un = pltni 2nd
solve, such that the gradients obtained via
the continuous adjoint method (equation (6)) exactly match the piiscretise-then-ppiie gradients obtained by autodifferentiating the
Tn alhnr 2nn
2n = 2n+1
(nnn + 1x+1)
2
numerically discretised forward pass.
2 (on + +n+1))AW
In doing so, one of the greatest limitations of
continuous adjoint methods is overcome.
Local forward
To the best of our knowledge, the reversible Heun method is the first algebraically re- versible SDE solver.
2n+1, 8x+1 Mn+1, 0n+1 Forward Zn, Zno Mni
Onz Al, W)
# Local backward ÖL(Zr)
Computational efficiency A further ad- vantage of the reversible Heun method is
OL(ZT)
azn, Zno Hn, (n)
O Ent1, Mn+1: xn+1)
computational efficiency. The method re- quires only single function evaluation (of both the drift and diffusion) per step. This is in contrast other Stratonovich solvers (such as the midpoint method or regular Heun's method), which require two function evalua- tions per step.
O( M22+11 Znt1. Mn+1i mn+1)
O( 2n, Zny Mnn Tn.)
Output: Ln Zn, Zn; Pn, Ons
0L(Zr) 0L(Zr)
oL(Zr)
aL(Zr)
82
8p
Convergence of the solver When applied the Stratonovich SDE (7), the reversible Heun method exhibits strong convergence of order 0.5; the same as the usual Heun's 's method.
Theorem. Let } denote the numerical solution of (7) obtained by Algorithm I with constant step size At and assume sufficient regularity of μ and a. Then there exists a constant C' >>> 0 so that
E
|| Zr|12] s CVAL,
for small At. 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
2n + 2x becomes step of midpoint method, whilst 2n 2n+1 is similar to Heun's 's method.
This makes it possible to show that {š,} and {2x} stay close together: |E 2,, - 1 ~ (At)). With this L, bound literature. Chaining together local mean squared error estimates, we obtain || - Zz 112 ~ O(At).
2, can then apply standard lines of argument from the numerical SDE
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 has the same absolute stability region for 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 As the same numerical solution {2n} is recovered on both the forward and backward passes exhibiting the same truncation errors then the computed gradients are precisely the seseretise-then-optinee discretisation of the forward pass.
gradients of the numerical
Each 8L(Zr)/0s, R AnAt, 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].
5

---

Table SDE-GAN on weights dataset; Latent SDE on air quality dataset. Mean + standard deviation averaged over three runs.
Label classification accuracy (%)
Dataset, Solver
MMD (x10-3)
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 49.2 = 0.02
0.591 + 0.206
5.58 + 0.54 hours
Air quality, Reversible Heun
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 optimisa- tion 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 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 dataset of air quality over Beijing.
Time
Figure 1: Samples (red) from Latent SDE Og ozone chan- nel of air quality dataset (blue).
See Table 1. Due to the reduced number of vector field evalua- tions, we find that training speed roughly doubles (1.98x) on the weights dataset, whilst its numerically precise gradients substan- tially improve the test metrics (comparing generated samples to held-out test set). Similar behaviour is observed on the air quality dataset, with substantial test metric improvements and training speed improvement of x.25x.
10-
Relative
10"9.
Micpoont
10-1n.
Heun
102t4
2
2-
Figure 2: Relative error in gra- dient calculation.
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 the ozone concentration over Beijing.
Gradient error We investigate the numerical error made in solving (6), compared to the (discretise- then-optimise) gradients the numerically discretised forward pass. We fix a test problem (differ- entiating 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.
4 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 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 's Brownian bridge [33] formula to generate increments with the appropriate correlations: letting Wa.b = W, Wa E IR then for < <
W u,s W N u W n,t
(1 )) u) 1m)
(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].

---

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, sample Brownian noise according to equation (8), as appropriate.
Algorithm 3: Sampling the Brownian Interval
Type: Let Node denote 5-tuple consisting of an interval, a seed, and three optional Nodes, corresponding to the parent node, and two child nodes, respectively. (Optional the root has no parent and leaves have no children.)
Liet et al. [15] instead offer memory-cfficient but time-intensive approach, by introducing the "Virtual Brownian Tree'. This approxi- mates the real line by a tree of dyadic points. Samples are approximate, and demand deep (slow) traversals of the tree.
State: Binary tree with elements of type Node, with root ī = (10, T]. 3, *, Irott Fright). A
Node J.
Input: Interval (s, 1 C [0,
# 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) (])
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 inter- val and a random seed.
def' sample(l: : Node):
if is I then return N/0, T) sampled with seed 8.
The tree starts as a stump consisting of the global interval [0, T) and a randomly gen- erated random seed. New leaf nodes are created as observations of the sample are made. For example, making a first query at [s S [o, (an operation that returns W..) 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 < < < L produces Figure 3b. Us- ing a splittable PRNG [34, 35), each child node has random seed deterministically produced from the seed of its parent.
Let = (la, b). Iparent Listt Leight)
Let Lparent = ((ap. 6pl, Sp, Ip: Apps Lrp)
W parentt = -ample(4paron
if I; is Ipp then
Wleft = bridge ap, bps a, W Warrent, 8)
return W Wpareat Wieft else
return bridge(ap, hp. b, W Wparent. 8) sample = LRUCache(sample)
J - nodes[-1]
W., Eienodes sample(I) Output: Ws.t
The tree is thus designed to completely en- code the conditional statisties of sample of Brownian motion: W.,. We, are com- pletely specified by é, [s, u). W s, equation (8), and the random seed for us.
0, T
[0, T
[0,
T
[0, s)
[s. T
[0, m²]
[u. s [s,
[t,
In principle this now gives a way to compute Ws.; calculating W. recursively. Naively 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..t-
[s, {] [t, [
[s v]
[v,
(a)
(b)
Figure 3: Binary tree of intervals.
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 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.
7

---

This com- Table 2: V. Brownian Tree against Brownian Interval on speed benchmarks, over 100 subintervals.
SDE solve, speed (seconds)
Doubly sequential access, speed (seconds)
Size=1 1
Size-2560
Size-32768
Size=1
Size-256t
Size-32768
V. B. Tree
1.6x10
2.0x100
5.00 10
2.4x10-1
3.9x10-1
2.9 10
B. Interval 8.2x10-1
1.3x 10
4.7x10*
5.0x 10-2
8.0x 10-2
3.5x 10-1
Meanwhile the (GPU) memory cost 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 simply holding all the Brownian samples in memory, which has memory cost of O (T))
Finally, queries 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 e.
4.1 Experiments
We benchmark the performance of the Brownian Interval against the Virtual Brownian Tree considered Li et al. [15]. We include benchmarks corresponding to varying batch sizes, number of sample intervals, and access patterns. For space, just 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.5x faster on smaller problems 10.6x 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++.
5 Training SDE-GANs without gradient penalty
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 plicates the use of continuous adjoint methods: the double-continuous-adiioo introduces substantial truncation error; sufficient to obstruct training and requiring small step sizes resolve.
Here we overcome this limitation, and moreover do so independently of the possibility of obtaining exact couble-gradient via the reversible Heun method. For simplicity now assume throughout that our discriminator vector fields fø. Io are MLPs, which also the choice we make in practice.
Lipschitz constant one The key point is that the vector fields ʃo. Io of the discriminator must not only Lipschitz, but must have Lipschitz constant at most one.
Given vector fields with Lipschitz constant X, then the recurrent structure of the discriminator means that the Lipschitz constant of the overall discriminator will be O (A²). Ensuring À 1 with À < 1 thus enforces that the overall discriminator is Lipschitz with a reasonable Lipschitz constant.
Hard constraint
The exponential size of O (X7) 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.
The first part of enforcing this Lipschitz constraint is careful clipping. Considering each linear operation from R - R² as matrix in A Baxb then after each gradient update we clip its entries to the region [-1 1/b, 1/b). Given x E Ra then this enforces || llc I | |
Clipping
I 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.
8

---

Table 3: SDE-GAN on OU dataset. Mean + standard deviation averaged over three runs.
Test Metrics
Real/fake classification accuracy (%)
Prediction loss
MMD (x10-l)
Training time (hours)
Solver
Midpoint w/ gradient penalty
98.2 2.4
2.71 + 1.03
2.58 + 1.81
55.0 + 27.7 32.5 = 12.1
Midpoint w/ clipping
93.9 + 6.9
1.65 = 0.17
1.03 + 0.10
Reversible Heun w/ clipping
67.7 1.1
1.38 + 0.06
0.45 0.22
29.4 8.9
There remain several admissible choices; we use the LipSwish activation function introduced by Chen al. [38], defined as p(x) 0.909 (z) This was carefully constructed 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 fa. 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.
5.1 Experiments
We test the SDE-GAN on dataset of tame-varying OrnstennUUllenbeck samples. For space only subset of results are shown; see Appendix F for further details 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 hours. Switching to reversible Heun additionally and substantially improves the test metrics, and further reduced training time to 29 hours; speed improvement of 1.87x.
6 Discussion
6.1 Available implementation in torchsde
To facilitate the use of the techniques introduced here in particular without requiring 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.)
6.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.
6.3 Ethical statement
No significant negative societal impacts are anticipated as a result of this work. A positive envi- rommental impact is anticipated, due to the reduction in compute costs implied by the techniques introduced. See Appendix G for 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 metries. This has been accomplished through several novel technical innovations, including firs-oofitss-king algebraically reversible SDE solver; 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/LO15811/I PK, JF, TL were supported by the Alan Turing Institute under the EPSRC grant EP/N510119/1. PK thanks Chris Rackauckas for discussions related to the reversible Heun method.
References
[1] W. T. Coffey, Y. P. Kalmykov, and J. T. Waldron. The Langevin Equation: With Applications to Stochastic Problems Physics, Chemistry and Electrical Engineering. World Scientife, 2012.
[2] S. E. Shreve. Stochastic Calculus for Finance I: Continuous-Time Models. Springer Science & Business Media, 2004.
[3] M. Arató. A famous nonlinear stochastic equation
model with diffusion).
Mathematical and Computer Modelling. 38(7-9):709-726 2003.
[4] K.-C. Chen, T.-Y. Wang. H.-H. Tseng, C.-Y. Huang, and C.-Y. Kao. A stochastic differential equation model for quantifying transcriptional regulatory network in saccharomyces cerevisiae. 2005.
Bioinformatics, 21(12):2883-22
[5] Black and M. Scholes. The Pricing of Options and Corporate Liabilities. Journal of Political Economy, 81 (13)6637-654, 1973.
[6] A. Pinkus. Approximation theory of the MLP model in neural networks. Acta Numerica, 8: 143-195, 1999.
[7] Kidger and Lyons. Universal Approximation with Deep Narrow Networks, In Proceedings of the 33rd Conference Learning Theory, pages 2306-2327, 2020.
[8] B. Tzen and M. Raginsky. Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit. arXiv: /905:/905.000 /905009855.,8,, 2019.
[9] B. Tzen and M. Raginsky. Theoretical guarantees for sampling and inference generative models with latent diffusions. In Proceedings of the 32rd Conference on Learning Theory, pages 3084-31 11143,14, 2019.
[10] J. Jia and A. Benson. Neural Jump Stochastic Differential Equations. Advances in Neural Information Processing Systems 32, pages 9847-9858, 2019.
[11] X. Liu, T. Xiao, S. Si, Q Cao, S. Kumar, and C.-J. Hsieh. Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise. arXiv: 02355,, 2019.
[12] I. Kong, J. Sun, and C. Zhang. SDE-Net: Equipping Deep Neural Networks with Uncertainty
Estimates. In Proceedings of the 37th International Conference Machine Learning, pages 5405-5415, 2020.
[13] L. Hodgkinson, C. van der Heide, F. Roosta, and M. Mahoney. Stochastic Normalizing Flows arX1v:2002.09544, 2020.
[14] Gierjatowicz, M. Sabate-Vidales D. Šiška, Szpruch, and Ž. Žurič. Robust Pricing and Hedging via Neural SDEs, 1xXiv:2007.0415 2020.
[15] X. Li, T.-K. L. Wong, R. T. Q. Chen, and D. Duvenaud. Scalable Gradients and Variational Inference for Stochastic Differential Equations. AISTATS, 2020.
[16] Kidger, J. Foster, X. Li. H. Oberhauser, and Lyons. Neural SDEs as Infinite- Dimensional GANs. In Proceedings of the 38th International Conference on Machine Learning, 2021.
[17] C. Holt. Forecasting seasonals and trends by exponentially weighted moving averages. ONR
Research Memorandum, Carnegie Institute of Technology, 52, 1957.
[18] Winters. Forecasting sales by exponentially weighted moving averages. Management Science, 6:324-342, 1960.
[19] E. J. Hannan and J. Rissanen. Recursive Estimation of Mixed Autoregressive-Moving Average Order. Biometrika, 69:81-94, 1982.
[20] Yoon, D. Jarrett, and M. van der Schaar. Time-series Generative Adversarial Networks, In Advances in Neural Information Processing Systems 32, 2019.
10

---

[21] Y. Ruuanooa,, R. T. Q. Chen, and D. Duvenaud. Latent Ordinary Differential Equations for Aregularly-Sampple Time Series. In Advances in Neural Information Processing Systems 32, pages 5320-,330, 2019.
[22] E. De Brouwer, J. Simm, A. Arany, and Y. Moreau. GRU-ODE-Bayee Continuous Modeling of Sporadically-OOseee Observed Time Series. In Advances in Neural Information Processing Systems
32, pages 7379-7390, 2019.
[23] C. Yildiz, M. Heinonen, and H. Lahdesmaki. ODE2VAE: Deep generative second order ODEs with Bayesian neural networks. In Advances in Neural Information Processing Systems 32, 2019.
[24] R. Deng, B. Chang, M. Brubaker, G. Mori, and A. Lehrmann. Modeling Continuous Stochastic Processes with Dynamic Normalizing Flows. In Advances in Neural Information Processing Systems 33, pages 7805-7815, 2020.
[25] A. Norcliffe, C. Bodnar, B. Day, J. Moss, and P. Liò. Neural ODE Processes. In International Conference on Learning Representations, 2021.
[26] C. Herrera, F. Krach, and J. Teichmann. Neural Jump Ordinary Differential Equations: Con- sistent Continuous- Time Prediction and Filtering. In International Conference on Learning Representations, 2021.
[27] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. Vetito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang. J. Bai, and S. Chintala. PyTorch: An Imperative Style, Deep Learning Library. In Advances in Neural Information Processing Systems 32, pages 8024-8035, 2019.
High-Performant
[28] J. Behrmann, W. Grathwohl, R. T. Q. Chen, D. Duvenaud, and J.-H. Jacobsen. Invertible Residual Networks. In Proceedings of the 36th International Conference on Machine Learning,
pages 573-582, 2019.
[29] P. Kidger, J. Foster, X. Li. H. Oberhauser, and T. Lyons. Neural SDEs Made Easy: SDEs are Dimenio-Dimessional GANs. OpenReview, 2020.
[30] Y. Song, Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations, 2021.
[31] W. Xu, R. T. Q. Chen, X. Li, and D. Duvenaud. Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations. arXiv:2102.06553 2021.
[32] J. Zhuang, N. C. Dvornek, S. Tatikonda, and J. S. Duncan. MALI: A memory efficient and reverse accurate integrator for Neural ODEs. International Conference on Learning Representations, 2021.
[33] D. Revuz and M. Yor. Continuous martingales and Brownian motion, volume 293. Springer Science Business Media, 2013.
[34] J. K. Salmon, M. Moraes, O. Dror, and D. E. Shaw. Parallel random numbers: as easy as
1, 2, 3. Proc. High Performance Computing, Networking, Storage and Analysis, pages 1-12, 2011.
[35] K. Claessen and M. Palka. Splittable pseudorandom number generators using cryptographic hashing. ACM SIGPLAN Notices, 48:47-58, 2013.
[36] I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. Courville. Improved Training of Wasserstein GANs. In Advances in Neural Information Processing Systems 30, pages 5767-5777, 2017.
[37] T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida. Spectral Normalization for Generative Adversarial Networks. In International Conference on Learning Representations, 2018.
[38] R. T. Q. Chen, J. Behrmann, D. Duvenaud, and J.-H. Jacobsen. Residual Flows for Invertible Generative Modeling. In Advances in Neural Information Processing Systems 32, 2019.
[39] D. Hendrycks and K. Gimpel. Gaussian Error Linear Units (GELUs). arXiv:1606.08415 2016.
[40] S. Elfwing, E. Uchibe, and K. Doya. Weinhhted Linear Units for Neural Network Function Approximation in Reinforeement Learning, arXiv: 2017.
11

---

[41] P. Ramachandran, B. Zoph, and Q. Le.
Swish:
Gated Activation Function.
arXiv:1710.05941, 2017.
[42] X. Li. torchsde, 2019. https://githus. www.chuuu.coomeeooooommmooeeeeee google researrch/torchheeeee http://hh....ooooooeeeeeerrrrraa
[43] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. Neural Ordinary Differential
Equations. In Advances in Neural Information Processing Systems 31, 2018.
[44] M. E. Sander, P. Ablin, M. Blondel, and G. Peyré. Momentum Residual Neural Networks. arXiv:2/02.07877 2021.
[45] P. Kidger, J. Morrill, J. Foster, and T. Lyons. Neural Controlled Differential Equations for Irregular Time Series. In Advances Neural Information Processing Systems 33, pages
6696-6707, 2020.
[46] J. Morrill, C. Salvi, P. Kidger, J. Foster, and T. Lyons. Neural Rough Differential Equations for Long Time Series. In Proceedings of the 38th International Conference on Machine Learning, 2021.
[47] B. Chang, M. Chen, E. Haber, and E. H. Chi. AntisymmetricRNN: A dynamical system view on recurrent neural networks. International Conference on Learning Representations. 2019.
[48] Karras, S. Laine, M. Aittala, J. Hellsten, J. Lehtinen, and T. Aila. Analyzing and Improving the Image Quality of StyleGAN. In Proc. CVPR, 2020.
[49] P. Kidger. torchtyping, 2021. https://gitee
[50] M. Arjovsky. S. Chintala, and L. Bottou. Wasserstein Generative Adversarial Networks. In Proceedings of the 34th International Conference on Machine Learning, pages 214-223, 2017.
[51] M. Giles and Glasserman. Smoking adjoints: Fast Monte Carlo greeks. Risk, 19:88-92, 2006.
[52] R. S. Hamilton. The inverse function theorem of Nash and Moser. Bulletin of the American Mathematical Society, 7(1):65-222, 1982.
[53] E. Kloeden and E. Platen. Numerical Solution of Stochastic Differential Equations. Springer, 1992.
[54] J. Foster, H. Oberhauser, and T. Lyons. An optimal polynomial approximation of Brownian motion. SIAM Journal on Numerical Analysis, 58(3): 393-14411, 2020.
[55] B. Leimkuhler, C. Matthews, and M. V. Tretyakov. On the long-time integration stochastic gradient systems. Proceedings of the Royal Society A, 470(2170), 2014.
[56] L. F. Shampine. Stability of the leapfrog/middpint method. Applied Mathematics and Computa- tion, 208(1):293-298, 2009.
[57] A. Griewank. Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. Optimization Methods and Software, 1(1):35-54, 1992.
[58] J. Gaines and T. Lyons. Variable step size control in the numerical solution of stochastic differential equations. SIAM Journal on Applied Mathematics, 57(5):1455-14888 1997.
[59] A. Rößler. Runge-Kuut Kutta Methods for the Strong Approximation of Solutions of Stochastic Differential Equations. SIAM Journal on Numerical Analysis, 48(3):922-99) 9522995,222 2010.
[60] A. Dickinson. Optimal Approximation of the Second Iterated Integral of Brownian Motion.
Stochastic Analysis and Applications, (5(5):1109-14) 1128, 2007.
[61] Gaines and T. Lyons. Random Generation of Stochastic Area Integrals. SIAM Journal on Applied Mathematics, 133--1141111116 1994.
[62] A. Davie. KMT theory applied to approximations of SDE. In Stochastic Analysis and Applica- lions 2014, pages 185-201. Springer, 2014.
[63] G. Flint and T. Lyons. Pathwise approximation of SDEs by coupling piecewise abelian rough paths. arXiv:1505.01220 2015.
[64] J. Foster. Numerical approximations for stochastic differential equations. PhD thesis, University of Oxford, 2020.
[65] M. Wiktorsson. Joint characteristic function and simultaneous simulation of iterated Itō integrals for multiple independent Brownian motions. Annals of Applied Probability, 11(2):470-487, 2001.
12

---

[66] J. Mrongowius and A. Rößler. On the Approximation and Simulation of Iterated Stochastic Integrals and the Corresponding Lévy Areas in Terms of Multidmensional Brownian Motion. arXiv:2101 095422 2021.
[67] J. Foster and K. Habermann. Brownian bridge expansions for Lévy area approximations and particular values of the Riemann zeta function. urXiv:2102./009: 2021.
[68] A. Gretton, K. M. Borgwardr, J. Rasch, B. Schölkopf, and A. Smola. A Kernel Two-Sample Test. Journal of Machine Learning Research, 13(1):723-773, 2013.
[69] Király and H. Oberhauser. Kernels for sequentially ordered data. Journal of Machine Learning Research, 20(31):1-45, 2019.
[70] C. Toth and H. Oberhauser. Bayesian Learning from Sequential Data using Gaussian Processes with Signature Covariances. In Proceedings of the 37th International Conference on Machine Learning, pages 9548-9560, 2020.
[71] P. Bonnier, P. Kidger, Perez-Arribas C. Salvi, and T. Lyons. Deep Signature Transforms. In Advances in Neural Information Processing Systems 32, 2019.
[72] C. Salvi, T. Cass, J. Foster, T. Lyons, and W. Yang. The Signature Kernel is the solution of Goursat PDE. arXiv:2006. 14794, 2021.
[73] M. Lemercier, C. Salvi, T. Cass, E. V. Bonilla, T. Damoulas, and T. Lyons. SigGPDE: Scaling Sparse Gaussian Processes on Sequential Data. In Proceedings of the 38th International Conference on Machine Learning. 2021.
[74] P. Kidger and T. Lyons. Signatory: differentiable computations of the signature and logsignature transforms, on both CPU and GPU. International Conference on Learning Representations, watmm/pariickkttee
2021. https: https://gi //github.
haapeer//gggtato signatory.
[75] J. Morrill, A. Fermanian, P. Kidger, and T. Lyons. A generalised signature method for multivari- ate time series feature extraction. arXiv:2006. 00873, 2020.
[76] Y. Li, K. Swersky, and R. Zemel. Generative Moment Matching Networks. In Proceedings of the 32nd International Conference on Machine Learning, pages 1718-1727, 2015.
[77] C.-L. Li, W.-C. Chang, Y. Cheng, Y. Yang, and B. Poczos. MMD GAN: Towards Deeper Understanding of Moment Matching Network. In Advances in Neural Information Processing Systems 30, pages 2203-2213, 2017.
[78] Kidger. torchcde, 2020. https://github. http:://ggitbbbb.///aaarr....rie httpa///oooooaorro----deeeerrrre
[79] R. T. Q. Chen, torchdiffeq. 2018 htttp:/://ittuub.cm////gg////// httpa//www.....egcccce///oonaaa.
[80] Ax. https: // http//xx.ddvvveve
[81] D. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference
on Learning Representations, 2015.
[82] M. D. Zeiler. ADADELTA: An Adaptive Learning Rate Method. arXiv:1212.5501,, 2012.
[83] Y. YAZIcı, C.-S. Foo, S. Winkler, K.-H. Yap, G. Piliouras, and V. Chandrasekhar. The Un- usual Effectiveness of Averaging in GAN training. In International Conference on Learning Representations, 2019.
[84] P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson. Averaging Weights Leads to Wider Optima and Better Generalization. Conference on Uncertainty Artificial Intelligence, 2018.
[85] D. Dua and C. Graff. UCI Machine Learning Repository, 2017. URL http://archiie... uci edu/ml.
[86] S. Zhang, B. Guo, A. Dong. J. He, Z. Xu, and S. X. Chen. Cautionary tales air-quality improvement in Beijing. Proceedings of the Royal Society A, 473(2205), 2017.
13

---

Checklist
and believe paper conforms to 1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
[Yes] The abstract reflects the paper's contributions and scope.
(b) Did you describe the limitations of your work?
[Yes] See Section 6.2. We have found every contribution to be a strict improvement on previous work, without introducing any further limitations.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] This is discussed in Section 6.3 and in further detail in Appendix G.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] We have read the ethics review guidelines available at https://nipt: wccppblicc/tthcssuuideeeee them.
2. If you are including theoretical results...
(a) Did you state the full of assumptions all theoretical results? [Yes] Full theoretical details are provided in Appendix D.
(b) Did you include complete proofs of all theoretical results?
[Yes] Complete proofs of convergence of the SDE solver are provided in Appendix D.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)?
(Yes] All code is attached as supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Full training details (including data splits, hyperparameter optimisation, and so
on) available in Appendix
(c) Did you report error bars (e-g., with respect the random seed after running experi- ments multiple times)?
[Yes] Every experiment run multiple times, and mean and standard deviations the test metrics reported.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Compute resources and computed time are specified in Appendix F.
4. If you using existing assets (e.g., code, data, models) curating/releasing assets...
(a) If your work uses existing assets, did you cite the creators? [Ye]] Citations are included in-line.
(b) Did you mention the license of the assets?
Yes]] See Appendix G.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes] Assuming that the contributions of this paper count assets, then these are made available in the supplementary material.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating?
[Yes] See Appendix G.
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content?
[Yes] See Appendix G.
5. If you used crowdsoureing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
14

---

(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
15


Reviews:
[{"review_id": "XfnuMcB-_yJ", "reviewer": "Reviewer_e5Kh", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Positive:\n\n1. The Neural SDE considered in this paper is very general. The Stratonovich SDE and the adjoint equation are nicely used in this framework. The reversible Heun method and the algorithm seems to be well explained within the Stratonovich SDE.\n\n2. A direct discretization of Stratonovich SDE would require to solve an algebraic equation in order to repeat the iteration. The Algorithm for the forward and backward process introduce an intermediate process to avoid this step. This idea seems to be very interesting. \n\n3. The Brownian interval idea following [15] is very interesting. The memory cost of this algorithm seems to be promising. \n\n4. The experiment shows impressive and supportive results of the algorithm.\n\nMinor point:\n\n1. It would be good to provide the intuition and the reason why do you define Heun method in Definition 1. What is the relation between this definition and the discretization of Stratonovich SDE.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "Wybtfk1ZUHv", "reviewer": "Reviewer_oQSS", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n\n1. The paper's contribution of a new SDE solver that is algebraically reversible (reversible Heun method) is novel. Moreover, the solver is shown to reduce numerical gradient error to near zero, which improves several test metrics on two datasets considerably over previous solvers (as exemplified by Table 1). A theorem is provided about strong convergence of the solver; the proof techniques are standard but well executed. \n\n2. The Brownian interval construction is a well-suited data structure constructed for the task of efficiently sampling and reconstructing Brownian motion. From a data structures viewpoint, the construction is not particularly interesting (binary tree of (interval, seed) pairs), although it is carefully made to solve the problem of efficient sampling and reconstruction. Given the performance improvement showcased in Table 2 and Appendix F, I think it still makes for a good contribution to the neural SDE literature.\n\nWeaknesses\n\n1. Although the contribution regarding training of SDE-GANs via clipping as opposed to with gradient penalty did empirically help on the OU dataset as measured by a variety of test metrics (Table 3), I believe the motivation for clipping specifically is not particularly principled. I understand the need for a hard 1-Lipschitz constraint, but why not enforce the constraint via a Lie algebraic method like the one in \"Cheap Orthogonal Constraints\" [a] or a Bjorck orthonormalization approach like the one in \"Sorting out Lipschitz function approximation\" [b]? I believe both of these are more principled ways of maintaining a hard 1-Lipschitz constraint in contrast to clipping (which has had its own demonstrated set of problems since it was introduced in [50]). I would like to see a more principled discussion of the motivation for the clipping method in Section 5 and think that related work such as [a,b] should be addressed.\n\n2. The baselines in the paper demonstrated the efficacy of the introduced approach to some degree, although I would have liked to see a more thorough comparison. For example, Table 1 provides exactly one baseline: the midpoint method. Why not at least provide the midpoint method together with (non-reversible) Heun? I don't believe this result is in the appendix either. I understand the field is reasonably niche, but this isn't an excuse for at the very least running Midpoint and Heun to compare with reversible Heun.\n\n3. I think some of the discussion of results is lacking. In particular, I am thinking of Section 3.1. One the weights dataset, we see a training time improvement of roughly 2x, which is fully expected since reversible Heun requires only one function evaluation (per step) whereas midpoint and regular Heun require two. However, why does the benefit dwindle to 1.25x on the air quality dataset? I don't believe this is currently addressed.\n\nVerdict\n\nThis work introduces the novel constructs of a reversible Heun SDE solver and a Brownian interval to considerably improve speed and accuracy of neural SDE training. The reversible Heun SDE solver is shown to work well relative to a single baseline on two datasets, although a fairly natural baseline (regular Heun) is omitted. The theory provided regarding convergence is well executed though makes use of largely standard proof techniques. The Brownian interval construction is effectively a well-applied binary-tree that greatly increases efficiency of sampling and reconstructing Brownian motion. The clipping approach to improving SDE-GAN training appears reasonable (and provides empirical benefit as exemplified by Table 3) but is not well-motivated with respect to existing methods of maintaining strict 1-Lipschitzness. Overall, the paper presents several practically useful contributions to the neural SDE literature and will benefit applicability, although I cannot currently accept it given my concerns above regarding various aspects of the paper.\n\nMinor Corrections\n\nThe writing in the paper is, for the most part, good. Some minor corrections are given below:\n\nL37: \"By parameterising the drift and diffusion of an SDE as neural networks, then modelling capacity is greatly increased\" -> \"\"By parameterising the drift and diffusion of an SDE as neural networks, modelling capacity is greatly increased\"\"\n\nL62: \"computed via continuous adjoint method\" -> \"computed via the continuous adjoint method\"\n\nL67: \"This is contrast\" -> \"This is in contrast\" (perhaps your phrasing is common in British English; if so, please disregard)\n\nL250: \"memory-inefficient but time-intensive\" -> \"memory-efficient but time-intensive\" (presumably this is what you meant to say)\n\nReferences\n\n[a] Casado, M.L., & Martínez-Rubio, D. (2019). Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group. ArXiv, abs/1901.08428.\n\n[b] Anil, C., Lucas, J., & Grosse, R.B. (2019). Sorting out Lipschitz function approximation. ICML.\n\n### Post-rebuttal Update\n\nI have improved my score from a 5 to a 6.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations of their work as well as societal impact. There are separate sections in the discussion  about both, and Appendix G is entirely dedicated to this matter. It is noted that the work will likely have positive environmental impact, since it improves computational cost associated with training neural SDEs.", "review_text": "", "extra_scores": {"ethical_concerns": "This paper is reasonably theoretical in nature; I have no ethical concerns about its impact.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours"}, "extra_text": {}}, {"review_id": "NKk4zMzz99Z", "reviewer": "Reviewer_Q4br", "summary": "The submission studies neural SDEs, a more flexible version of classic SDEs with neural drifts and volatilities. Specifically, to improve the estimation/learning of such models, the submission provides two novel methods and one learning trick. ", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The contributions of the submission are solid on improving the efficiency of fitting neural SDE based on Stratonovich integrals. The motivation is straightforward aiming at directly tackling the bottleneck of the gradient updates. Thorough empirical results are provided justifying the efficiency gain of the proposed methods. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The contributions of the submission on fitting neural SDE based on Stratonovich integrals are significant. However, I am curious whether explicit Euler method for Ito integrals also have some of the properties. In general, I have three questions. \n\n1. Say we consider the SDE built on Ito integrals. Then, we use explicit Euler method as the SDE solver. With a loss function on the terminal state, can't we directly do backpropogation along the path? What is the difference between this naive method and the proposed one. Does (6) \"exactly match\" the gradient obtained by autodifferentiating the forward pass, in this naive approach? \n\n2. Is building SDEs on Stratonovich integrals instead of Ito integrals a common practice for neural SDEs? In terms of model capacity, how do the two manners differ? Further, have you tried to apply the proposed methods to SDEs based on Ito-integrals as some kind of robust analysis?\n\n3. If I understand it correctly, the method directly optimizes the drift and volatility coefficient functions. So, I am curious what the accuracy for estimating these functions might be e.g., the MSE for the drift estimation?\n\n4. It seems that the results can also be applied to estimate structured parametric SDEs like the tasks in this paper (https://www.princeton.edu/~yacine/mle.pdf). (Not optimizing neural nets but functions with specific forms.) Is that correct? Would it be possible to provide comparisons with such closed-form methods?  ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "2--W_LWxq0C", "reviewer": "Reviewer_Dmv1", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper aims to improves the efficiency and accuracy of training Neural SDEs. To this end, the authors start from the motivation and mathematical background, provide three techniques in order, and provide convincing numerical experiments to support the strengths. The whole paper is well organized and clearly written. The contribution to the open-source package will likely attract researchers and practitioners to adopt the proposed techniques and facilitate this process. Below are some concerns that should be clarified in a revision.\n\n1. The authors should prove the reversibility of the Heurn method in Algorithm 1, although it is easy to verify.\n\n2. I am confused by the result in Figure 2 where the midpoint and Heun methods \"produce large errors decreasing with step size\" (line 231). Can the authors provide some explanation? Does it conflict with the point in line 207 that \"small or adaptive step sizes are necessary to obtain useful gradients (for standard numerical techniques)\"?\n\n3. The Brownian Interval is claimed to \"offer memory efficiency, exact samples, and fast query times, all at once\". To retrieve exact samples, do we need to keep the seeds at the root node the same for each pair of the forward pass and backward pass? In terms of memory efficiency, I feel the spatial complexity of the Brownian Interval is the same as storing Brownian motion because at each time step we need to store a random seed. The memory is saved because the size of the seed is much smaller than the Brownian motion increment. Is this understanding correct? By the way, in line 250, it should be \"memory-efficient\".\n\n4. The clipping strategy seems the most ad-hoc technique in this paper. Clipping each entry to [-1/b, 1/b] is sufficient to ensure the Lipschitz constant no larger than 1, but it might be too restrictive such that approximation capacity is sacrificed. A more flexible alternative might be projecting according to the singular values. Please add more discussion on it. In addition, it is not so clear in Table 3 that clipping really helps. Could the authors provide the result of Reversible Heun with gradient penalty as well?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think the clipping technique has some limitations that might be improved further. See the fourth comment above.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: After reading the author responses and discussing with the authors, the reviewers have reached a consensus that this paper should be accepted. I concur with that judgement.  The accuracy improvements proposed by this method seem significant and useful across a range of tasks. And the mathematical insights on which the result are based will be interesting to even those in the community who aren't using neural SDEs.
Decision: Accept (Poster)