Efficient and Accurate Gradients for Neural SDEs

1 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 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 ] . 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 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 . 2 Background . 2.1 Neural SDE construction . 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 . 2.2 Training criteria for Neural SDEs . 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 . 2.3 Discretised observations . 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 . 2.4 Backpropagation through SDE solves . Whether the loss for our generated sample Y is produced via a Latent SDE or via the discriminator of an SDE-GAN , it is still required to backpropagate from the loss to the parameters θ , φ . 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 .

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.