Robust and Scalable SDE Learning: A Functional Perspective

Stochastic differential equations provide a rich class of flexible generative models,1 capable of describing a wide range of spatio-temporal processes . A host of re-2 cent work looks to learn data-representing SDEs , using neural networks and other3 flexible function approximators . Despite these advances , learning remains compu-4 tationally expensive due to the sequential nature of SDE integrators . In this work,5 we propose an importance-sampling estimator for probabilities of observations of6 SDEs for the purposes of learning . Crucially , the approach we suggest does not7 rely on such integrators . The proposed method produces lower-variance gradient8 estimates compared to algorithms based on SDE integrators and has the added9 advantage of being embarrassingly parallelizable . This facilitates the effective use10 of large-scale parallel hardware for massive decreases in computation time.11 1 INTRODUCTION12 Stochastic differential equations ( SDEs ) are a natural extension to ordinary differential equations13 which allows modelling of noisy and uncertain driving forces . These models are particularly appealing14 due to their flexibility in expressing highly complex relationships with simple equations , while15 retaining a high degree of interpretability . Much work has been done over the last century focussing16 on understanding and modelling with SDEs , particularly in dynamical systems and quantitative17 finance ( Pavliotis , 2014 ; Malliavin & Thalmaier , 2006 ) . Examples of such work are the Langevin18 model of stochastic dynamics of particles in a fluid , stochastic turbulence ( Kumar et al. , 2020 ) , and19 the Black–Scholes model . More recently , SDE models have gathered interest in machine learning20 and current work focuses on efficient and scalable methods of inferring SDEs from data.21 1.1 THE SDE LEARNING PROBLEM22 Consider an Itô SDE ( Pavliotis , 2014 ) of the following form23 dX = f ( X , t ) dt+ g ( X , t ) dW , ( 1 ) and a set of observations x1 : N = { ( tn , xn ) } N n=1 . In the simplest case , these observations would24 be directly of a realisation of the process at discrete points in time . In many physical applications,25 the observations themselves would have some uncertainty associated with them , described by an26 observation noise model . In more complex models , the SDE would be driving some unobserved27 parameters describing further complex processes ( Li et al. , 2020 ) . For example , discrete events28 may be modeled by an inhomogenious Poisson processes , where the rate parameter is evolving29 stochastically over time . In any case , the generating process can be summarized by the graphical30 model in Figure 1 and the corresponding learning problem can be described as : “ given observations31 x1 : N , infer f and g ” .32 For the most part , the model can be trained using standard techniques , such as variational inference,33 Markov chain Monte Carlo or optimizing an appropriate loss function . Generally , this requires34 estimating an expectation over sample-paths , either to estimate the marginal probability of the obser-35 vations as a function of f and g , or to estimate the mismatch of sample-paths with the observations . To36 estimate these expectations , most current work makes use of integration algorithms . These algorithms37 generate trajectories of X for a given f and g , which can be used to estimate the required expectation.38 f g W X xn N The central idea of this paper is to instead generate trajectories from a distribution much closer to54 the path-space posterior distribution when estimating these expectations . In the simplest case where55 the observations are directly of the process X , this means sampling trajectories which pass directly56 through the observations . Alternatively , one can marginalize out the observation noise distribution —57 or efficiently estimate some expected loss function — by importance sampling58 EX [ p ( x1 : N |X ) ] = Ex̃1 : N∼q [ p ( x1 : N | x̃1 : N ) q ( x̃1 : N ) pX ( x̃1 : N ) ] , ( 4 ) where q is an importance distribution which may depend on the observations themselves . For a large59 class of observation distributions , such as additive Gaussian noise or multinomial observations , a60 natural choice for q would be the corresponding conjugate distribution , in which case the importance61 weight p/q becomes a constant . In complex graphical models , q might be parameterized and62 optimized variationally , or Gibbs sampling may be used with a collection of simpler conditional63 distributions . Taking advantage of this , we shall , without loss of generality , concentrate only on64 estimating the quantity pX ( x1 : N ) = EX [ ∏ n δ ( xn −X ( tn ) ) ] .65 In this work , we introduce a novel algorithm for efficient estimation of probabilities in SDE models66 for the purposes of learning . Our algorithm has a number of favourable charateristics:67 • Our algorithm exhibits lower gradient variance in estimated expectations than algorithms68 based on SDE solvers , since it enforces that sampled trajectories pass through the observed69 data , producing accurate results even in hundreds of dimensions.70 • Since our algorithm does not rely on an SDE solver , it completely removes all sequential71 dependence between sampled points . Our algorithm is therefore trivially parallelizable , and72 can take full advantage of modern parallel hardware such as multi-GPU training.73 • Since our estimator can be calculated in a single forward pass of the function f , proba-74 bilistic approaches such as variational Gaussian processes are viable candidates for the75 representation of f .76 1.2 RELATED WORK77 Li et al . ( 2020 ) propose a method for calculating gradients of functions of the integrated process with78 respect to parameters of f and g via a backwards SDE called the adjoint method . This approach79 requires reconstructing the Wiener process backwards in time when integrating the backward equation80 and the authors propose a memory efficient algorithm to do so . Their approach uses a constant81 amount of memory independent of the number of integration steps similar to the adjoint method for82 NeuralODEs ( Chen et al. , 2018 ) . Tzen & Raginsky ( 2019 ) discuss the computation of gradients by83 simulation of a forward SDE in the derivatives with respect to parameters of the model using the84 chain rule of stochastic calculus . This approach does not require back-propagating through the SDE85 solver . Kidger et al . ( 2021 ) propose a method of learning SDEs by treating the SDE as the generator86 in a generative adversarial network setup , using a second SDE for the discriminator . All of these87 approaches use SDE integration algorithms to estimate path-space expectations.88 Batz et al . ( 2018 ) consider a nonparametric approach to SDE learning , using Gaussian processes89 for the drift and diffusion functions . They initially consider a gradient matching approximation , in90 which the drift Gaussian process may be fit directly by conventional methods , and thereafter use an91 expectation-maximization algorithm , and estimate the transition probability for sparse observations92 using a linear approximation of the process . Yildiz et al . ( 2018 ) propose an alternative approach to93 Gaussian process-based SDE learning using the Euler–Maruyama integrator , approximating f by the94 predictive mean of a Gaussian process conditioned on a MAP estimate for a set of inducing points.95 Unfortunately , this approach completely ignores any uncertainty in the posterior over f .96 It should be noted that each of the above mentioned studies rely on SDE integrators to generate97 trajectories and therefore face the same difficulties associated with them . Furthermore , almost all98 of the experiments therein are on system of dimension between 1 and 3 , with a small number of99 experiments of dimension as high as 6 . We are unaware of any work which trains neural SDEs in100 higher dimensions , and we believe this is a testament to the inherent difficulty of learning SDEs in101 high dimensions.102 Much other work focuses on learning specific classes of SDE or ODE models , such as symplectic or103 Hamiltonian systems ( Zhong et al. , 2020 ; Greydanus et al. , 2019 ) , graph-based models ( Poli et al.,104 2021 ; Sanchez-Gonzalez et al. , 2019 ) , and controlled differential equations ( Morrill et al. , 2021 ) .105 Many of these ideas are independent of the learning algorithm and can potentially make use of our106 algorithm when applied to SDEs.107 2 PATH-SPACE IMPORTANCE SAMPLING108 It is well known that solutions of linear SDEs are Gaussian processes ( see Pavliotis , 2014 , Section 3.7 ) .109 For these processes , the joint probability density over some finite collection of points is available110 in closed form . Furthermore , the posterior distribution of the process conditioned on some set of111 observations can be sampled from exactly . Unfortunately , this is not the case in general for non-linear112 SDEs , and one has to resort to approximation methods in order to calculate probabilities . Most113 commonly , one uses an SDE integrator to estimate expectations ; however , other methods , such as114 Laplace approximations and perturbation theory are sometimes used ( Karimi & McAuley , 2016 ; 115 Dass et al. , 2017 ; Hutzenthaler & Jentzen , 2020 ) .116 Alternatively , if we are able to find a second process which follows a similar SDE to the one of117 interest , it is possible to express expectations and probabilities as importance sampling estimators.118 Linear SDEs are an attractive candidate for this due to their Gaussianity and ease of sampling.119 2.1 STATE-INDEPENDENT DIFFUSION120 For now , consider the simpler scenario in which the diffusion coefficient does not depend on the state121 variable ; i.e . g ( x , t ) = σ ( t ) for some function σ . The process defined by122 dY = σ ( t ) dW , ( 5 ) is Gaussian with mean zero and conditional variance123 EY [ ( Y ( t1 ) − Y ( t0 ) ) 2 ] = ∫ t1 t0 σ ( t ) σT ( t ) dt . ( 6 ) Sample paths of Y can be generated efficiently by simulating Brownian motion . The process Y124 conditioned on a set of observations of the process is a Brownian bridge . If we are able to express125 quantities of interest as expectations under this process instead of our original SDE , then we can126 efficiently estimate such quantities via Monte Carlo sampling . As it turns out , expectations and127 probabilities for an SDE with a general drift term can in fact be expressed as such . This relation is128 given in the following theorem.129 Theorem 1 Let X and Y be the stochastic processes generated by the following SDEs dX = f ( X , t ) dt+ σ ( t ) dW , ( 7 ) dY = σ ( t ) dW . ( 8 ) Further assume that σ ( t ) σT ( t ) is Lebesgue-almost-everywhere invertible and f is sufficiently well-130 behaved such that X ( t ) has a unique probability density for all t.131 The probability density of observations { ( tn , xn ) } N n=1 , under the process X is given by the condi-132 tional expectation133 pX ( x1 : N ) = pY ( x1 : N ) EY [ eS [ Y ] ∣∣∣ { Y ( tn ) = xn } Nn=1 ] , ( 9 ) where134 S [ Y ] = ∫ fT ( Yt , t ) ( σ ( t ) σT ( t ) ) −1 dYt − 1 2 ∫ fT ( Yt , t ) ( σ ( t ) σT ( t ) ) −1 f ( Yt , t ) dt . ( 10 ) This result follows from Girsanov ’ s theorem ( Girsanov , 1960 ; Malliavin & Thalmaier , 2006 , Sec-135 tion 1.5 ) and the definition of conditional expectation . Intuitively speaking , the first term in Equa-136 tion 10 encourages f to line up with typical paths which pass through the observations , while the137 second term regularizes the magnitude of f .138 Theorem 1 allows us to develop an importance sampling algorithm for estimating probabilities under139 the process X by simulating the process Y , where the importance weights are given by ωi = eS [ Y ( i ) ] .140 In this approach , one can generate entire sample paths of Y before calculating f ( Y , t ) , allowing the141 forward pass of f to be performed in parallel . This approach is described in Algorithm 1.142 Algorithm 1 Path Integral Importance Sampling 1 : for i = 1 , . . . , K do 2 : Sample ( Yt ) tNt=t1 , s.t . for each n , Ytn = xn . Y is a Brownian bridge 3 : ft ← f ( Yt , t ) for each t 4 : α← ∑ t f T t σ −2 ( t ) ( Yt+∆t − Yt ) 5 : β ← ∑ t f T t σ −2 ( t ) ft ∆t 6 : Si ← α− 12β 7 : end for 8 : return pY ( x1 : N ) 1K ∑ i exp ( Si ) An important property of this algorithm is that the sampled paths directly pass through observations.143 This means that the probability estimates directly relate the drift function f to plausible trajectories144 instead of allowing the particles to wander around randomly . See Figure 2 for an illustration . In the145 initial stages of training , simply integrating the SDE with drift f will typically not lead to trajectories146 which pass near the data , especially in high dimensions . This can have a significant impact on the147 variance of the gradients and can inhibit the model ’ s ability to learn . Training on paths that pass148 through the data can greatly mitigate this high variance . As with other algorithms that generate SDE149 sample paths , the discretization of the integrals introduces a small bias . However , this bias vanishes150 in the ∆t→ 0 limit.151 2.1.1 A NOTE ON PARALLELIZABILITY152 With the exception of the sampling step in Line 2 of Algorithm 1 , each of the other operations,153 including the sums , can be performed in parallel over feature and batch dimensions , as well as over154 the t and i indices . Sampling a Brownian bridge , or any conditional Gauss–Markov process , can155 be performed by first sampling the unconditional process , and then linearly interpolating with the156 observations . The linear interpolation can be performed independently , and therefore in parallel , over157 the time dimension . The only operation that is not completely parallelizable is the sampling of the158 unconditional base process — in this case standard Brownian motion . Brownian paths can be sampled159 by cumulative summation of independent normal random numbers . However , this cumulative sum is160 extremely fast , and can be performed independently between each consecutive pair of observations.161 In our tests , sampling whole Brownian trajectories was about 5 orders of magnitude faster than the162 forward pass of the neural network f , and so does not create any noticeable performance impact.163 Once the Brownian paths are sampled , the rest of Algorithm 1 may be implemented in parallel164 without difficulty . Contrast this to standard integrator-base methods ; even if the Brownian paths are165 sampled upfront , the integrator must still perform each of the forward passes of f in sequence.166 See Appendix B for more details on sampling Gaussian bridges.167 2.2 STATE-DEPENDENT DIFFUSION168 The assumption that g is independent of the state is not in fact required for the validity of Theorem 1 ; 169 the reason this assumption is required is to ensure that the process Y is Gaussian , so as to enable170 easy conditional sampling and calculation of pY ( x1 : N ) . Unfortunately , one can not simply find an171 analogue of Theorem 1 to calculate expectations of a general process as expectations with respect to172 a constant diffusion process.173 To address the question of state-dependent diffusion , it is enlightening to consider a simple example174 of such a process . Perhaps the most common example is geometric Brownian motion175 dX = µX dt+ σX dW . ( 11 ) The typical way in which one would solve this is to introduce a transformation Z = log ( X ) , and176 apply Itô ’ s lemma ( see Malliavin & Thalmaier , 2006 , Section 1.5 ; Pavliotis , 2014 , Section 3.5 ) to177 obtain the SDE178 dZ = ( µ− 12σ 2 ) dt+ σ dW . ( 12 ) Similarly , one might address a problem of a general state-dependent diffusion coefficient by consider-179 ing a transformation to a space in which the transformed equation has constant diffusion . Let T ( · , t ) 180 be a map which is invertible for all t. That is , there exists a function T−1 such that if y = T ( x , t ) ,181 then x = T−1 ( y , t ) . Transforming a process with constant diffusion σ , by the map T−1 gives a182 process with diffusion coefficient183 gi , j ( x , t ) = ∂T−1i ∂yk σk , j = ( ∂T ( x , t ) ∂x ) −1 i , k σk , j , ( 13 ) where repeated indices imply summation . With such a transformation we can represent any diffusion184 matrix that can be written as a Jacobian . While the restrictions on the diffusion matrix may be seen as185 an inherent limitation of a path-space importance sampling approach , it may be argued that constant186 or simple diffusion functions would cover the majority of cases of interest , and almost all real-world187 problems require some form of data pre-processing and , potentially , link functions.188 This transformation allows us to infer the SDE in a transformed space , with constant diffusion , and,189 if desired , transform the equation back to the original space using Itô ’ s lemma . This approach is190 described in Algorithm 2 , where we have used the symbol f̃ to represent the drift in the transformed191 space as to avoid ambiguity with the drift in the original space . One important point to note is that one192 should be careful to allow gradients to propagate through the sampling step on Line 2 in Algorithm 2,193 so as to allow learning of the function T . Algorithm 2 Transformed-State Path Integral Importance Sampling 1 : for i = 1 , . . . , K do 2 : Sample ( Yt ) tNt=t1 , s.t . for each n , Ytn = T ( xn , tn ) . Y is a Brownian bridge 3 : f̃t ← f̃ ( Yt ) for each t 4 : α← ∑ t f̃ T t σ −2 ( Yt+∆t − Yt ) 5 : β ← ∑ t f̃ T t σ −2f̃t ∆t 6 : Si ← α− 12β 7 : end for 8 : return ∏ n det ( ∂T ( xn , tn ) ∂xn ) pY ( T ( x1 : N , t1 : N ) ) 1 K ∑ i exp ( Si ) 194 One would typically apply this method to infer f̃ rather than the original f for performance reasons,195 and then only reconstruct the SDE in the original space if required.196 2.2.1 RECONSTRUCTION OF THE SDE IN THE UNTRANSFORMED SPACE197 Generation of new sample paths can be performed easily in the transformed space and then simply198 transforming back by applying T−1 . However , sometimes one needs more information for the199 purposes of analysing the learned equation . In this case the SDE can be transformed back using Itô ’ s200 lemma . The drift and diffusion coefficients are201 fi ( x , t ) = ∂T−1i ∂t + ∂T−1i ∂yk f̃k ( y , t ) + 1 2 σj , k ∂2T−1i ∂yj∂yl σl , k , ( 14 ) gi , j ( x , t ) = ∂T−1i ∂yk σk , j . ( 15 ) All expression are evaluated at y = T ( x , t ) . When T is parameterized by a neural network with202 ReLU non-linearities , the second derivative term in Equation ( 14 ) is almost-everywhere zero . The203 remaining terms are Jacobian vector products , and so can be performed efficiently even in high204 dimensions.1205 3 EXPERIMENTS206 In this section , we provide some experiments to illustrate and verify the performance and capabilities207 of our proposed algorithm . Further details on the experiments are given in Appendix C.208 3.1 THE LORENZ SYSTEM209 To compare the effect of our algorithm on learning , we follow Li et al . ( 2020 ) and train a neural210 network based SDE model on the Lorenz system . In this model the drift function is given by a211 1However , this may require forward-mode instead of backward-mode differentiation . multi-layer perceptron , and the diffusion term is a simple scalar . We generated 16 independent212 paths , consisting of 200 observations each , uniformly sampled over a time interval of [ 0 , 10 ] . We213 trained the model on this data by two different methods . In the first method , we minimize the214 mean-squared error , estimated by integrating out the SDE , and calculating gradients using the adjoint215 method ( Li et al. , 2020 ) .2 In the second method , we directly optimize the log-probability calculated216 using Algorithm 1 . In both cases we used precisely the same optimization parameters , the same217 time-step size ∆t , and the same number of sampled trajectories to calculate averages . Both models218 used the same seed for the neural network initialization , and the same number of gradient steps were219 performed in both cases . These models were trained on an Nvidia RTX 3070 , with 8Gb of memory.220 To assess the performance of these approaches , we record two metrics . The first is the mean-square221 difference between the learned drift vector field , and the Lorenz vector field , normalized by the mean-222 square magnitude of the Lorenz vector field , measured at the observations . Written mathematically:223 err = ∑ n ( fθ ( xn ) − fLor ( xn ) ) 2∑ n ( fLor ( xn ) ) 2 ( 16 ) This metric is shown in Figure 3b . The second metric we record is the mean-squared error of sampled224 trajectories compared to the observations ; this is shown in Figure 3a . However , we note that this is225 not a fair performance metric since the adjoint-based method is directly optimizing for it . On the226 x-axis we show the training time in minutes.227 In Figure 3 we have limited the range of the x-axis in order for the lines to be more clearly visible.228 The total training time for the adjoint-based MSE optimization was a factor 52 longer than the time229 taken optimizing the probability estimates . We note that , while our approach was significantly faster,230 it did use approximately four times as much memory as the adjoint-based method.3231 The poor performance of the integrator-based learning algorithm in Figure 3b — despite its low mean-232 squared-error — suggests that directly optimizing for the discrepancy of typical SDE sample-paths233 compared to observations is not a robust approach to SDE learning.234 3.2 GRADIENT VARIANCE235 We applied both our algorithm and an SDE integrator to a neural network-based SDE model on236 a real-world dataset and recorded the gradients , sampled over many independent trajectories . We237 thereafter computed the variance of the gradients individually per parameter , for various time lengths238 between observations . We used a discretization step size of 0.01 . These results are shown in Figure 4.239 2We also attempted training a model without the adjoint method and instead directly backpropagating through the SDE integrator . However , a single forward pass required more memory than available on the GPU . 3The adjoint algorithm was specifically designed to have constant memory complexity in the number of integration steps . Our algorithm does not exhibit this feature . In each case we used the same neural network , initialized with a fixed seed , and the exact same values240 for all parameters . We used the same loss functions as in Section 3.1 . The data set we used was the241 Hungarian chickenpox cases dataset from the UCI machine learning repository ; 4 it has 20 features242 per observation . We scaled the features by their standard deviations before performing the experiment243 to improve numerical stability.244 Our findings are as follows:245 • For very short times , we found that both algorithms produced comparable gradient variances:246 typically less than 10−7.247 • For longer time spans ( 3 orders of magnitude larger than the discretization step size ) , the248 gradient variance of our algorithm increased by a few orders of magnitude ; the maximum249 we observed was on the order of 10−4 . However , the SDE integrator approach resulted in250 huge increases in variances , with some as large as 1018.251 • Between the above two extremes , we found that the gradient variance of our approach did252 increase with the time between observations , but stayed below 10−5 , while the gradient253 variance from the SDE integrator exhibited very poor scaling with variances of up to 108 for254 times between observations of 1.0.255 Considering that both approaches produced comparable variances on very small timescales , we feel256 that this experiment represents a fair comparison.257 3.3 UNCERTAINTY QUANTIFICATION WITH VARIATIONAL GAUSSIAN PROCESSES258 We attempt to test the viability of using our method for an SDE model with a GP prior over the drift259 function and fitting the posterior variationally . To assess this , we used data trajectories from a van der260 Pol oscillator . For the GP prior , we used a zero mean function and a radial basis function kernel , with261 independent GP priors over each output dimension of the drift vector field . These choices are not262 necessarily optimal ; however , the purpose of this experiment is merely to illustrate the feasibility of263 using variational GP to learn the vector field of an SDE model.264 We used a grid-based inducing point approach and the kernel interpolation method of Wilson & 265 Nickisch ( 2015 ) . The combination of these techniques allows one to take advantage of the Kronecker266 product structure of the covariance matrices and the Toeplitz structure of each factor for accelerated267 sampling . The model was trained variationally , sampling the variational posterior jointly along entire268 trajectories.269 4This dataset can be found at https : //archive.ics.uci.edu/ml/datasets/Hungarian+ Chickenpox+Cases The approximate GP posterior is shown in Figure 5a compared to the ground truth vector field of the270 van der Pol oscillator in Figure 5b . Comparing the GP to the van der Pol vector field , we found that271 the GP model was fairly accurate in areas with many data points , while reverting to the prior far away272 from observations . To assess the quality of the uncertainty quantification , we further calculated the273 discrepancy between the variational posterior and the ground truth vector field . Over a dense grid,274 covering an enlarged area around the observations , we calculated the discrepancy as275 δ = fvdp ( x ) − µ ( x ) σ ( x ) . ( 17 ) A histogram is shown in Figure 5c.276 Under Gaussianity assumptions , one expects to see the quantiles line up well with some associated277 number of standard deviations . However , this will only be the case near to the observations . In278 regions where the GP reverts to the prior , which has a stationary kernel , the mismatch between the279 GP and the polynomial growth of the van der Pol oscillator results in heavy tails . We expect that a280 better choice of kernel , which accounts for this polynomial behaviour , to show residuals with lighter281 tails . We found that the distribution of these δs was fatter-tailed than a Gaussian distribution ; the282 middle 50 % ( 25 % - 75 % quantiles ) slightly passed one standard deviation on either side , as opposed283 to the 68 % of a standard Gaussian distribution . The middle 80 % ( 10 % - 90 % quantiles ) were within284 3 standard deviations . While the posterior uncertainty was somewhat overconfident , it was not wholly285 unreasonable.286 We further note that , while the complexities involved in variational GP fitting can result in longer287 training times , we still found that our method applied to a GP was significantly faster than using an288 SDE integrator with a neural network.289 4 DISCUSSION290 In this work we introduced an algorithm to efficiently estimate probabilities of observations under291 the assumption of an SDE model for the purposes of learning . The proposed approach produces292 estimates which are accurate , even in high dimensions , and can fully utilize parallel hardware for293 significant speedups.294 The proposed algorithm is capable of providing robust probability estimates which can be used in295 models both with or without an explicit observation noise component . This is in contrast to SDE296 integrator-based learning approaches , which must incorporate and estimate an explicit error term due297 to their inability to generate paths which are guaranteed to pass through the observations.298 Our algorithm produces more stable estimates than alternative integrator-based approaches , particu-299 larly when the observations are sparsely distributed over long time intervals and in high dimensions.300 This in turn results in more stable and robust learning . With the rise in applicability of large-scale301 ( neural network ) SDE modeling in science , simulation , and domains such as finance , we see our302 approach as an important contribution to reliable scaling of these methods.303 REPRODUCIBILITY304 Efforts have been made to describe all parameters and details of the algorithms used in this paper,305 either in the content or in the appendix . The code included in the supplementary material includes306 a Dockerfile and instructions for running some of the experiments . Where relevant , fixed seeds for307 random number generators can be found in the code during initialization.308 REFERENCES309 Philipp Batz , Andreas Ruttor , and Manfred Opper . Approximate Bayes learning of stochastic310 differential equations . Phys . Rev . E , 98:022109 , Aug 2018. doi : 10.1103/PhysRevE.98.022109.311 URL https : //link.aps.org/doi/10.1103/PhysRevE.98.022109.312 Ricky T. Q. Chen , Yulia Rubanova , Jesse Bettencourt , and David K Duvenaud . Neural ordinary313 differential equations . 31 , 2018 . URL https : //proceedings.neurips.cc/paper/314 2018/file/69386f6bb1dfed68692a24c8686939b9-Paper.pdf.315 Sarat C. Dass , Jaeyong Lee , Kyoungjae Lee , and Jonghun Park . Laplace based approximate posterior316 inference for differential equation models , differential equations . Statistics and Computing , 27:317 679–698 , 2017. doi : 10.1007/s11222-016-9647-0.318 I. V. Girsanov . On transforming a certain class of stochastic processes by absolutely continuous319 substitution of measures . Theory of Probability & Its Applications , 5 ( 3 ) :285–301 , 1960. doi:320 10.1137/1105027 . URL https : //doi.org/10.1137/1105027.321 Samuel Greydanus , Misko Dzamba , and Jason Yosinski . Hamiltonian Neural Networks.322 In H. Wallach , H. Larochelle , A. Beygelzimer , F. d'Alché-Buc , E. Fox , and R. Gar-323 nett ( eds . ) , Advances in Neural Information Processing Systems , volume 32 . Curran Asso-324 ciates , Inc. , 2019 . URL https : //proceedings.neurips.cc/paper/2019/file/325 26cd8ecadce0d4efd6cc8a8725cbd1f8-Paper.pdf.326 Martin Hutzenthaler and Arnulf Jentzen . On a perturbation theory and on strong convergence rates327 for stochastic ordinary and partial differential equations with nonglobally monotone coefficients.328 The Annals of Probability , 48 ( 1 ) , Jan 2020 . ISSN 0091-1798. doi : 10.1214/19-aop1345 . URL329 http : //dx.doi.org/10.1214/19-AOP1345.330 Junteng Jia and Austin R Benson . Neural jump stochastic differential equations . In331 H. Wallach , H. Larochelle , A. Beygelzimer , F. d'Alché-Buc , E. Fox , and R. Garnett332 ( eds . ) , Advances in Neural Information Processing Systems , volume 32 . Curran Asso-333 ciates , Inc. , 2019 . URL https : //proceedings.neurips.cc/paper/2019/file/334 59b1deff341edb0b76ace57820cef237-Paper.pdf.335 Hadiseh Karimi and Kimberley B. McAuley . Bayesian Estimation in Stochastic Differential336 Equation Models via Laplace Approximation . IFAC-PapersOnLine , 49 ( 7 ) :1109–1114 , 2016.337 ISSN 2405-8963. doi : https : //doi.org/10.1016/j.ifacol.2016.07.351 . URL https : //www.338 sciencedirect.com/science/article/pii/S2405896316305584 . 11th IFAC339 Symposium on Dynamics and Control of Process SystemsIncluding Biosystems DYCOPS-CAB340 2016.341 Patrick Kidger , James Foster , Xuechen Li , and Terry J Lyons . Neural SDEs as Infinite-Dimensional342 GANs . In Marina Meila and Tong Zhang ( eds . ) , Proceedings of the 38th International Conference343 on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pp . 5453–5463.344 PMLR , 18–24 Jul 2021 . URL https : //proceedings.mlr.press/v139/kidger21b.345 html.346 Prashant Kumar , Martin Schmelzer , and Richard P. Dwight . Stochastic turbulence modeling in RANS347 simulations via multilevel Monte Carlo . Computers & Fluids , 201:104420 , 2020 . ISSN 0045-7930.348 doi : https : //doi.org/10.1016/j.compfluid.2019.104420 . URL https : //www.sciencedirect.349 com/science/article/pii/S0045793019303780.350 Xuechen Li , Ting-Kam Leonard Wong , Ricky T. Q. Chen , and David Duvenaud . Scalable gradients351 for stochastic differential equations . International Conference on Artificial Intelligence and352 Statistics , 2020.353 Xuanqing Liu , Tesi Xiao , Si Si , Qin Cao , S. Kumar , and Cho-Jui Hsieh . Neural SDE : Stabilizing354 Neural ODE Networks with Stochastic Noise . ArXiv , abs/1906.02355 , 2019.355 Paul Malliavin and Anton Thalmaier . Stochastic Calculus of Variations in Mathematical Finance.356 Springer , 01 2006 . ISBN 3-540-43431-3. doi : 10.1007/3-540-30799-0.357 James Morrill , Cristopher Salvi , Patrick Kidger , and James Foster . Neural rough differential equations358 for long time series . In Marina Meila and Tong Zhang ( eds . ) , Proceedings of the 38th International359 Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pp.360 7829–7838 . PMLR , 18–24 Jul 2021 . URL https : //proceedings.mlr.press/v139/361 morrill21b.html.362 GA Pavliotis . Stochastic Processes and Applications : Diffusion Processes , the Fokker–Planck and363 Langevin Equations . Springer , 2014.364 Michael Poli , Stefano Massaroli , Junyoung Park , Atsushi Yamashita , Hajime Asama , and Jinkyoo365 Park . Graph neural ordinary differential equations . 2021.366 Alvaro Sanchez-Gonzalez , Victor Bapst , Kyle Cranmer , and Peter Battaglia . Hamiltonian Graph367 Networks with ODE Integrators , 2019.368 Belinda Tzen and M. Raginsky . Neural stochastic differential equations : Deep latent Gaussian models369 in the diffusion limit . ArXiv , abs/1905.09883 , 2019.370 Andrew Wilson and Hannes Nickisch . Kernel Interpolation for Scalable Structured Gaussian Pro-371 cesses ( KISS-GP ) . In Francis Bach and David Blei ( eds . ) , Proceedings of the 32nd International372 Conference on Machine Learning , volume 37 of Proceedings of Machine Learning Research,373 pp . 1775–1784 , Lille , France , 07–09 Jul 2015 . PMLR . URL http : //proceedings.mlr.374 press/v37/wilson15.html.375 C. Yildiz , M. Heinonen , J. Intosalmi , H. Mannerstrom , and H. Lahdesmaki . Learning stochastic376 differential equations with gaussian processes without gradient matching . In IEEE International377 Workshop on Machine Learning for Signal Processing . MLSP , 2018.378 Yaofeng Desmond Zhong , Biswadip Dey , and Amit Chakraborty . Symplectic ode-net : Learning379 hamiltonian dynamics with control . In International Conference on Learning Representations,380 2020 . URL https : //openreview.net/forum ? id=ryxmb1rKDS.381 A THEOREM 1382 We have that X and Y are processes obeying the SDEs dX = f ( X ) dt+ g ( X ) dW , ( 18 ) dY = g ( Y ) dW , ( 19 ) dropping the time argument for notational simplicity . Expectations with respect to SDEs can be expressed as path integrals . In the Itô convention , the path integral representation of an SDE has the following form EX [ F [ X ] ] = lim N→∞ ∫ F [ x ] N∏ n=1 dxn√ 2π∆tng ( xn−1 ) exp { −1 2 N∑ n=1 ( ∆xn − f ( xn−1 ) ∆tn g ( xn−1 ) √ ∆tn ) 2 } ( 20 ) = : ∫ DxF [ x ] J exp { −1 2 ∫ L ( x ( t ) , ẋ ( t ) ) dt } , ( 21 ) where J is a Jacobian factor which depends on g but not f and the Lagrangian is given by383 L ( x , ẋ ) = ( ẋ− f ( x ) g ( x ) ) 2 . ( 22 ) The derivative of x here may be thought of distributionally ; however , the expression for the measure384 Dx exp { − 12 ∫ L } need only be defined on the Cameron–Martin subspace5 ; guaranteeing its unique385 extension to the whole Wiener space . In contrast to this , the expression inside the expectation386 F [ x ] should be defined almost-everywhere , and hence care should be taken when ẋ appears in an387 expectation.388 We may expand the quadratic to separate the Lagrangian into a drift-free term plus terms dependent on the drift L = ( ẋ g ( x ) ) 2 ( 23 ) − 2fT ( x ) g−2 ( x ) ẋ ( 24 ) + fT ( x ) g−2 ( x ) f ( x ) . ( 25 ) The first term , on line 23 , is the Lagrangian for the process Y , which has no drift term . Shuffling the389 rest of the terms out of the measure and into the expectation , we have390 EX [ F [ X ] ] = ∫ DxF [ x ] eS [ x ] exp { −1 2 ∫ ( ẋ g ( x ) ) 2 dt } , ( 26 ) where S [ X ] = ∫ fT ( Xt ) g −2 ( Xt ) Ẋt dt ( 27 ) − 1 2 ∫ fT ( Xt ) g −2 ( Xt ) f ( Xt ) dt ( 28 ) = ∫ fT ( Xt ) g −2 ( Xt ) dXt ( 29 ) − 1 2 ∫ fT ( Xt ) g −2 ( Xt ) f ( Xt ) dt . ( 30 ) Note that on the third line we have used the distributional definition of Ẋ ; i.e . Ẋt dt = dXt.391 Lastly , using the definition of conditional expectation , we have that392 EA [ F [ A ] p ( B = b|A ) ] = EA [ p ( B = b|A ) ] EA [ F [ A ] ∣∣∣B = b ] . ( 31 ) By trivial substitution , we have393 pX ( x1 : N ) = pY ( x1 : N ) EY [ eS [ Y ] ∣∣∣ { Y ( tn ) = xn } Nn=1 ] . ( 32 ) B SAMPLING BRIDGES394 Let Y be a Gauss–Markov process with kernel k. Since the process is Markov , we need only sample395 the process conditioned on two end-points . Let y = ( y0 , yT ) be the known values of the process at396 times τ = ( t0 , T ) . Let Ỹ be the process defined by397 Ỹt : = Yt +K ( t , τ ) K ( τ , τ ) −1 ( y − Yτ ) , for t0 ≤ t ≤ T , ( 33 ) whereK is the matrix whose entries are k ( ti , tj ) . This distribution of Ỹ is the same as the distribution398 of Y conditioned on the values y at times τ . Note that once Yt has been sampled — which may399 require a cumulative sum — this transformation can be performed independently , and hence in400 parallel , for each t. For each consecutive pair of observations , the bridge process can be sampled401 independently.402 5The Hilbert space of functions whose derivatives are square integrable . The Wiener space is the space of continuous functions . This is not a Hilbert space , instead it is a Banach space with the supremum norm . The measure of the Cameron–Martin space is in fact zero ; the sample paths of the SDE are almost-surely nowhere differentiable . In the simplest case , where Y is a Brownian motion starting from y0 at time t0 = 0 , this reduces to Ỹt = Yt + k ( t , T ) k ( T , T ) ( yT − YT ) ( 34 ) = Yt + t T ( yT − YT ) . ( 35 ) And similarly for any other processes which can be sampled with a fixed starting point.403 In the more general case we can invert the two by two matrix to get an explicit formula for Ỹt Ỹt = Yt + k ( t , t0 ) k ( T , T ) − k ( t , T ) k ( t0 , T ) k ( t0 , t0 ) k ( T , T ) − k ( t0 , T ) 2 ( y0 − Y0 ) ( 36 ) + k ( t , T ) k ( t0 , t0 ) − k ( t , t0 ) k ( t0 , T ) k ( t0 , t0 ) k ( T , T ) − k ( t0 , T ) 2 ( yT − YT ) . ( 37 ) This may be extended straightforwardly to other cases , such as conditioning on the values of the404 integral , or derivative of the process etc.405 C EXPERIMENTS406 C.1 PROBABILITY ESTIMATES407 As a sanity check , we compare the log-probability estimates of our algorithm to the actual log-408 probabilities of an Ornstein–Uhlenbeck process . The probability density for this process is available409 in closed form for any number of observations . We estimate the probability of a set of data points410 evenly spaced over a time interval of [ 0 , 10 ] , taken from a sine function xn = sin ( π/10× tn ) . For411 these estimates we used 1000 time steps to estimate the integral S , and averaged over 100 trajectories,412 independently sampled from a Brownian bridge passing through the observations . We repeated this413 process 1000 times and report the mean and standard deviation of these estimates . Figure 6 shows414 these estimates compared to the analytic probability given by the Ornstein–Uhlenbeck process for415 various numbers of observations as a function of the dimension of the problem.416 An interesting observation is that the variance of the estimator seems to decrease with the number417 of data points . This is atypical of probability estimators , since the magnitude of the log-probability418 increases linearly with the number of observations ( not to be confused with the number of samples419 used to estimate the log-probability ) . In almost all cases , with the exception of 4 observations in420 more than 40 dimensions , our algorithm gives accurate log-probability estimates with small standard421 deviations . The error of this exception is approximately 30 % , but is significantly reduced by adding422 more observations , or using more samples . We believe this error is acceptable considering the small423 number of observations and the large scale . For 10 observations in 100 dimensions the standard424 deviation of our estimator is only 0.4 % of the magnitude of the log-probability and the error of the425 mean estimate is 2 % of the analytic log-probability . This small bias can be further reduced by using426 more samples or a smaller ∆t . A single estimate in the 100 dimensional case took less than half a427 second to compute on a laptop CPU , and is therefore feasible to use in the inner loop of a learning428 algorithm.429 C.2 LORENZ430 C.2.1 DATA SET431 We generated data by simulating a Lorenz system with the exact same parameters as Li et al.432 ( 2020 ) . We generated 16 independent trajectories , with initial values sampled from a standard normal433 distribution , and discretely sampled the values of the process over 200 equally spaced time points434 ranging from 0 to 10 . Unlike Li et al . ( 2020 ) , we did not add any observation noise to the data.435 We deliberately used a sparser data set than Li et al . ( 2020 ) , due to personal experience of the difficulty436 of fitting SDE models on sparse observations.437 C.2.2 MODEL438 The model we used to fit to the observations was439 dXt = fθ ( Xt ) dt+ σ dW , ( 38 ) where σ is a scalar and fθ is a fully-connected neural network with ReLU nonlinearities and parame-440 ters θ . We used layer sizes of 3→ 32→ 256→ 32→ 3.441 C.2.3 TRAINING PARAMETERS442 We used precisely the same parameters for each model , the only difference in the training procedures443 being the loss functions used . In the first case , the loss function per trajectory was444 Lmse ( θ ; x1 : N ) = 1 KN ∑ k , n ( X ( k ) tn − xn ) 2 , ( 39 ) where N = 200 is the number of observations per trajectory , and K is the number of sampled paths445 used to estimate the MSE , and X ( k ) is sampled using an SDE integrator starting at X ( k ) t1 = x1 . We446 then averaged this loss function over the 16 independent observed trajectories.447 In the second case , the loss function per trajectory was448 LPI ( θ ; x1 : N ) = − 1 N log p̂X ( x1 : N | θ ) , ( 40 ) where p̂X ( x1 : N | θ ) is calculated using Algorithm 1 with K trajectory samples . Again this loss was449 averaged over the 16 independent observed trajectories . In both cases we used K = 64 and a time450 step size of ∆t = 10−2.451 For each model we used the Adam optimizer with a learning rate of 10−3 and ran the optimization452 algorithm for 104 iterations . The model using our importance sampling approach completed the full453 number of iterations in 34 minutes , while the integrator-based approach took 29 hours and 46 minutes454 to complete the same number of iterations.455 Our approach consistently maintained 100 % GPU utilization , except during MSE evaluation which456 used an integrator , while the integrator based approach was unable to surpass 25 % GPU utilization in457 this experiment.458 C.3 VARIATIONAL GAUSSIAN PROCESSES459 Our approach allows more straightforward implementation of SDE models using a Gaussian process460 prior for the drift function . Some previous approaches often resort to approximations such as461 using only the predictive mean or MAP estimates ( Yildiz et al. , 2018 ) . For dense observations,462 directly discretizing the SDE and fitting the GP to the observation deltas may be viable . For sparse463 observations this is much more difficult . Our proposed method makes variational GP methods slightly464 simpler , since the values of f ( Yt ) can be sampled jointly along the path . However , it is important to465 note that various complexities still arise such as the difficulty of efficiently sampling high-dimensional466 Gaussian distributions.467 We used the kernel interpolation method of Wilson & Nickisch ( 2015 ) which stores a grid of468 inducing locations per dimension . Storing inducing locations per dimension for a Kronecker product469 kernel effectively gives a variational distribution over md inducing points with only d×m memory470 requirements , where m is the number of inducing points per dimension and d is the dimension . When471 these inducing locations are evenly spaces , the covariance factors have a Toeplitz structure which472 allows sampling in linearithmic time complexity.473 To fit the model we used gradient ascent on the ELBO which is calculated as follows:474 ELBO = Ef∼q [ log p ( x|f ) ] −KL [ q ( f ) |p ( f ) ] . ( 41 ) The KL divergence is between two Gaussian distributions calculated at the inducing locations in475 the usual way . For the likelihood term , we generated trajectories ( as in Algorithm 1 ) , followed by476 jointly sampling the variational GP posterior along these trajectories and proceed with the rest of the477 computation in Algorithm 1 using these samples . Our optimization objective is calculated478 L = − 1 K ∑ K log p̂ ( x1 : N | f ( k ) ) + KL [ q ( f ) |p ( f ) ] , ( 42 ) where each f ( k ) is an independent sample over the whole path of Y . In other words , f ( k ) ( Yt1 ) and479 f ( k ) ( Yt2 ) are correlated , but f ( k ) ( Yt1 ) and f ( j ) ( Yt2 ) are not when k 6= j.480 In our experiment we used 32 evenly spaced inducing points per dimension , and jointly sampled the481 variational posterior over 64 trajectories of 1024 time steps . The model was trained with the Adam482 optimizer with an initial learning rate of 0.1 , and an exponential decay with factor 0.999 for 105483 training steps . However , the square deviation of the predictive mean from the ground truth converged484 well before 104 iterations . Training for 104 iterations took 34 minutes , 105 iterations took about 5485 and a half hours.486

Summary: The paper presents a new approach to compute the density of a given observation sequence under arbitrary drift and diffusion functions. The method relies on Girsanov's theorem, which requires evaluating the density of the observations under Wiener process. The authors further propose an inversion trick to handle state dependent diffusion functions. The experiments show that the method indeed learns accurate drift and diffusion functions much faster than standard integration-based approaches and has much smaller gradient variance