R OBUST AND S CALABLE SDE L EARNING : A F UNCTIONAL P ERSPECTIVE
Anonymous authors
Paper under double-blind review
A BSTRACT
Stochastic differential equations provide a rich class of ﬂexible generative models,capable of describing a wide range of spatio-temporal processes.
A host of recent work looks to learn data-representing SDEs, using neural networks and other ﬂexible function approximators.
Despite these advances, learning remains compu- tationally expensive due to the sequential nature of SDE integrators.
In this work, we propose an importance-sampling estimator for probabilities of observations of
SDEs for the purposes of learning.
Crucially, the approach we suggest does not rely on such integrators.
The proposed method produces lower-variance gradient estimates compared to algorithms based on SDE integrators and has the added advantage of being embarrassingly parallelizable.
This facilitates the effective use of large-scale parallel hardware for massive decreases in computation time
I NTRODUCTION
Stochastic differential equations (SDEs) are a natural extension to ordinary differential equationswhich allows modelling of noisy and uncertain driving forces.
These models are particularly appealingdue to their ﬂexibility in expressing highly complex relationships with simple equations, whileretaining a high degree of interpretability.
Much work has been done over the last century focussingon understanding and modelling with SDEs, particularly in dynamical systems and quantitativeﬁnance (Pavliotis, 2014; Malliavin & Thalmaier, 2006).
Examples of such work are the Langevinmodel of stochastic dynamics of particles in a ﬂuid, stochastic turbulence (Kumar et al., 2020), andthe Black–Scholes model.
More recently, SDE models have gathered interest in machine learningand current work focuses on efﬁcient and scalable methods of inferring SDEs from data.
1. T HE SDE LEARNING PROBLEM
Consider an Itô SDE (Pavliotis, 2014) of the following form
[Equation]
and a set of observations x 1: N = (cid:8) ( t n , x n ) (cid:9) N n =1 .
In the simplest case, these observations wouldbe directly of a realisation of the process at discrete points in time.
In many physical applications,the observations themselves would have some uncertainty associated with them, described by anobservation noise model.
In more complex models, the SDE would be driving some unobservedparameters describing further complex processes (Li et al., 2020).
For example, discrete eventsmay be modeled by an inhomogenious Poisson processes, where the rate parameter is evolvingstochastically over time.
In any case, the generating process can be summarized by the graphicalmodel in Figure 1 and the corresponding learning problem can be described as: “given observationsx 1: N , infer f and g ”.
For the most part, the model can be trained using standard techniques, such as variational inference,
Markov chain Monte Carlo or optimizing an appropriate loss function.
Generally, this requiresestimating an expectation over sample-paths, either to estimate the marginal probability of the observations as a function of f and g , or to estimate the mismatch of sample-paths with the observations. Toestimate these expectations, most current work makes use of integration algorithms.
These algorithmsgenerate trajectories of X for a given f and g , which can be used to estimate the required expectation.
[Equation] For example, the Euler–Maruyama algorithm (Li et al., 2020; Jia & Benson, 2019; Liu et al., 2019)generates trajectories with the following update rule
[Equation]
The expectation of some functional F is estimated using the sample mean of its value over manyindependent trajectories X ( k )
[Equation]
Unfortunately, integrators such as Equation (2) introduce a sequential dependence between each pointand the next, which inhibits parallelization.
Furthermore, it introduces difﬁculties in probabilisticmodelling of the functions f and g ; in order to properly estimate the expectation, the sample pathsshould be generated under consistent realisations of f and g .
Models which implicitly describe f as a
Gaussian process would need to sample f ( X t , t ) conditioned on all the previously sampled values off along the trajectory, which may be computationally prohibitive.
The integrator generates trajectories sampled from the prior distribution over paths.
In the initialstages of learning, the path-space prior distribution usually has minimal overlap with the path-spaceposterior and typical trajectories do not pass anywhere near the observations, resulting in highvariance gradients.
This grows considerably worse as the dimension of the problem and the timebetween observations increases.
The central idea of this paper is to instead generate trajectories from a distribution much closer tothe path-space posterior distribution when estimating these expectations.
In the simplest case wherethe observations are directly of the process X , this means sampling trajectories which pass directlythrough the observations.
Alternatively, one can marginalize out the observation noise distribution —or efﬁciently estimate some expected loss function — by importance sampling
[Equation]
where q is an importance distribution which may depend on the observations themselves.
For a largeclass of observation distributions, such as additive Gaussian noise or multinomial observations, anatural choice for q would be the corresponding conjugate distribution, in which case the importanceweight p/q becomes a constant.
In complex graphical models, q might be parameterized andoptimized variationally, or Gibbs sampling may be used with a collection of simpler conditionaldistributions.
Taking advantage of this, we shall, without loss of generality, concentrate only onestimating the quantity p X ( x 1: N
[Equation]
= E X
[Equation] In this work, we introduce a novel algorithm for efﬁcient estimation of probabilities in SDE modelsfor the purposes of learning.
Our algorithm has a number of favourable charateristics:
• Our algorithm exhibits lower gradient variance in estimated expectations than algorithmsbased on SDE solvers, since it enforces that sampled trajectories pass through the observeddata, producing accurate results even in hundreds of dimensions.• Since our algorithm does not rely on an SDE solver, it completely removes all sequentialdependence between sampled points.
Our algorithm is therefore trivially parallelizable, andcan take full advantage of modern parallel hardware such as multi-GPU training.• Since our estimator can be calculated in a single forward pass of the function f , probabilistic approaches such as variational Gaussian processes are viable candidates for therepresentation of f .
1. R ELATED WORK
Li et al.
(2020) propose a method for calculating gradients of functions of the integrated process withrespect to parameters of f and g via a backwards SDE called the adjoint method.
This approachrequires reconstructing the Wiener process backwards in time when integrating the backward equationand the authors propose a memory efﬁcient algorithm to do so.
Their approach uses a constantamount of memory independent of the number of integration steps similar to the adjoint method for
NeuralODEs (Chen et al., 2018).
Tzen & Raginsky (2019) discuss the computation of gradients bysimulation of a forward SDE in the derivatives with respect to parameters of the model using thechain rule of stochastic calculus.
This approach does not require back-propagating through the SDEsolver.
Kidger et al.
(2021) propose a method of learning SDEs by treating the SDE as the generatorin a generative adversarial network setup, using a second SDE for the discriminator.
All of theseapproaches use SDE integration algorithms to estimate path-space expectations.
Batz et al.
(2018) consider a nonparametric approach to SDE learning, using Gaussian processesfor the drift and diffusion functions.
They initially consider a gradient matching approximation, inwhich the drift Gaussian process may be ﬁt directly by conventional methods, and thereafter use anexpectation-maximization algorithm, and estimate the transition probability for sparse observationsusing a linear approximation of the process.
Yildiz et al.
(2018) propose an alternative approach to
Gaussian process-based SDE learning using the Euler–Maruyama integrator, approximating f by thepredictive mean of a Gaussian process conditioned on a MAP estimate for a set of inducing points.
Unfortunately, this approach completely ignores any uncertainty in the posterior over f .
It should be noted that each of the above mentioned studies rely on SDE integrators to generatetrajectories and therefore face the same difﬁculties associated with them.
Furthermore, almost allof the experiments therein are on system of dimension between 1 and 3, with a small number ofexperiments of dimension as high as 6.
We are unaware of any work which trains neural SDEs inhigher dimensions, and we believe this is a testament to the inherent difﬁculty of learning SDEs inhigh dimensions.
Much other work focuses on learning speciﬁc classes of SDE or ODE models, such as symplectic or
Hamiltonian systems (Zhong et al., 2020; Greydanus et al., 2019), graph-based models (Poli et al.
,2021; Sanchez-Gonzalez et al., 2019), and controlled differential equations (Morrill et al., 2021).
Many of these ideas are independent of the learning algorithm and can potentially make use of ouralgorithm when applied to SDEs.
P ATH - SPACE IMPORTANCE SAMPLING
It is well known that solutions of linear SDEs are Gaussian processes (see Pavliotis, 2014, Section 3.7).
For these processes, the joint probability density over some ﬁnite collection of points is availablein closed form.
Furthermore, the posterior distribution of the process conditioned on some set ofobservations can be sampled from exactly.
Unfortunately, this is not the case in general for non-linear
SDEs, and one has to resort to approximation methods in order to calculate probabilities.
Mostcommonly, one uses an SDE integrator to estimate expectations; however, other methods, such as
Laplace approximations and perturbation theory are sometimes used (Karimi & McAuley, 2016;
Dass et al., 2017; Hutzenthaler & Jentzen, 2020).
Alternatively, if we are able to ﬁnd a second process which follows a similar SDE to the one ofinterest, it is possible to express expectations and probabilities as importance sampling estimators.
Linear SDEs are an attractive candidate for this due to their Gaussianity and ease of sampling.
2. S TATE - INDEPENDENT DIFFUSION
For now, consider the simpler scenario in which the diffusion coefﬁcient does not depend on the statevariable; i.e. g ( x, t )
=
σ ( t ) for some function σ .
The process deﬁned by
[Equation]
is Gaussian
[Equation]
Sample paths of Y can be generated efﬁciently by simulating Brownian motion.
The process Yconditioned on a set of observations of the process is a Brownian bridge.
If we are able to expressquantities of interest as expectations under this process instead of our original SDE, then we canefﬁciently estimate such quantities via Monte Carlo sampling.
As it turns out, expectations andprobabilities for an SDE with a general drift term can in fact be expressed as such.
This relation isgiven in the following theorem.
Theorem 1 Let X
and Y be the stochastic processes generated by the following SDEs
[Equation]
Further assume that σ ( t ) σ T ( t ) is Lebesgue-almost-everywhere invertible and f is sufﬁciently wellbehaved such that X ( t ) has a unique probability density for all t .
The probability density of observations (cid:8) ( t n , x n ) (cid:9) N n =1 , under the process X is given by the conditional expectation
[Equation]
where
[Equation] [Equation]
This result follows from Girsanov’s theorem (Girsanov, 1960; Malliavin & Thalmaier, 2006, Section 1.5) and the deﬁnition of conditional expectation.
Intuitively speaking, the ﬁrst term in Equation 10 encourages f to line up with typical paths which pass through the observations, while thesecond term regularizes the magnitude of f .
Theorem 1 allows us to develop an importance sampling algorithm for estimating probabilities underthe process X by simulating the process Y , where the importance weights are given by ω i
=
e S
[ Y ( i ) ] .
In this approach, one can generate entire sample paths of Y before calculating f ( Y, t ) , allowing theforward pass of f to be performed in parallel.
This approach is described in Algorithm 1.
1: for i = 1 , . . .
, K do 2:
Sample ( Y t ) t N t = t 1 , s.t. for each n , Y t n = x n (cid:46) Y is a Brownian bridge 3: f t ← f ( Y t , t ) for each t 4: α ← (cid:80) t f
Tt σ − 2 ( t ) ( Y t +∆ t − Y t ) 5: β ← (cid:80) t f Tt σ − 2 ( t ) f t ∆ t 6: S i ← α − 12 β 7: end for 8: return p Y ( x 1: N ) 1 K (cid:80) i exp( S i ) An important property of this algorithm is that the sampled paths directly pass through observations.
This means that the probability estimates directly relate the drift function f to plausible trajectoriesinstead of allowing the particles to wander around randomly.
See Figure 2 for an illustration.
In theinitial stages of training, simply integrating the SDE with drift f will typically not lead to trajectorieswhich pass near the data, especially in high dimensions.
This can have a signiﬁcant impact on thevariance of the gradients and can inhibit the model’s ability to learn.
Training on paths that passthrough the data can greatly mitigate this high variance.
As with other algorithms that generate SDEsample paths, the discretization of the integrals introduces a small bias.
However, this bias vanishesin the ∆ t → 0 limit.
2.1. A NOTE ON PARALLELIZABILITY
With the exception of the sampling step in Line 2 of Algorithm 1, each of the other operations,including the sums, can be performed in parallel over feature and batch dimensions, as well as over [Figure]
the t and i indices.
Sampling a Brownian bridge, or any conditional Gauss–Markov process, canbe performed by ﬁrst sampling the unconditional process, and then linearly interpolating with theobservations.
The linear interpolation can be performed independently, and therefore in parallel, overthe time dimension.
The only operation that is not completely parallelizable is the sampling of theunconditional base process — in this case standard Brownian motion.
Brownian paths can be sampledby cumulative summation of independent normal random numbers.
However, this cumulative sum isextremely fast, and can be performed independently between each consecutive pair of observations.
In our tests, sampling whole Brownian trajectories was about 5 orders of magnitude faster than theforward pass of the neural network f , and so does not create any noticeable performance impact.
Once the Brownian paths are sampled, the rest of Algorithm 1 may be implemented in parallelwithout difﬁculty.
Contrast this to standard integrator-base methods; even if the Brownian paths aresampled upfront, the integrator must still perform each of the forward passes of f in sequence.
See Appendix B for more details on sampling Gaussian bridges.
2. S TATE - DEPENDENT DIFFUSION
The assumption that g is independent of the state is not in fact required for the validity of Theorem 1;the reason this assumption is required is to ensure that the process Y is Gaussian, so as to enableeasy conditional sampling and calculation of p Y ( x 1: N ) .
Unfortunately, one cannot simply ﬁnd ananalogue of Theorem 1 to calculate expectations of a general process as expectations with respect toa constant diffusion process.
To address the question of state-dependent diffusion, it is enlightening to consider a simple exampleof such a process.
Perhaps the most common example is geometric Brownian motion
[Equation]
The typical way in which one would solve this is to introduce a transformation Z = log( X ) , andapply Itô’s lemma (see Malliavin & Thalmaier, 2006, Section 1.5; Pavliotis, 2014, Section 3.5) toobtain the SDE
[Equation]
Similarly, one might address a problem of a general state-dependent diffusion coefﬁcient by considering a transformation to a space in which the transformed equation has constant diffusion.
Let T ( · , t )be a map which is invertible for all t .
That is, there exists a function T − 1 such that if y = T ( x, t ) ,then x = T − 1 ( y, t ) .
Transforming a process with constant diffusion σ , by the map T − 1 gives aprocess with diffusion coefﬁcient
[Equation]
where repeated indices imply summation.
With such a transformation we can represent any diffusionmatrix that can be written as a Jacobian.
While the restrictions on the diffusion matrix may be seen asan inherent limitation of a path-space importance sampling approach, it may be argued that constantor simple diffusion functions would cover the majority of cases of interest, and almost all real-worldproblems require some form of data pre-processing and, potentially, link functions.
This transformation allows us to infer the SDE in a transformed space, with constant diffusion, and,if desired, transform the equation back to the original space using Itô’s lemma.
This approach isdescribed in Algorithm 2, where we have used the symbol ˜ f to represent the drift in the transformedspace as to avoid ambiguity with the drift in the original space.
One important point to note is that oneshould be careful to allow gradients to propagate through the sampling step on Line 2 in Algorithm 2,so as to allow learning of the function T .
Algorithm
1: for K do 2:
Sample bridge 3: ˜ f t ← ˜ f ( Y for each t 4: α ← (cid:80) t ˜ f T σ −
5: β ← (cid:80) t ˜ f
Tt σ f t 6:
S i ← α − 12 β 7: end for 8: return (cid:81) n det (cid:16) ∂T ( x n ,t n ) ∂x n (cid:17) p Y (cid:0) T (
[Table]
1: N , t 1: N ) (cid:1) 1 K (cid:80) i exp( S i )
One would typically apply this method to infer ˜ f rather than the original f for performance reasons,and then only reconstruct the SDE in the original space if required.
2.2. R ECONSTRUCTION OF THE SDE IN THE UNTRANSFORMED SPACE
Generation of new sample paths can be performed easily in the transformed space and then simplytransforming back by applying T − 1 .
However, sometimes one needs more information for thepurposes of analysing the learned equation.
In this case the SDE can be transformed back using Itô’slemma.
The drift and diffusion coefﬁcients are
[Equation]
[Equation]
All expression are evaluated at y = T ( x, t ) .
When T is parameterized by a neural network with
ReLU non-linearities, the second derivative term in Equation (14) is almost-everywhere zero.
Theremaining terms are Jacobian vector products, and so can be performed efﬁciently even in highdimensions.
E XPERIMENTS
In this section, we provide some experiments to illustrate and verify the performance and capabilitiesof our proposed algorithm.
Further details on the experiments are given in Appendix C.
3. T HE L ORENZ SYSTEM
To compare the effect of our algorithm on learning, we follow Li et al.
and train a neuralnetwork based SDE model on the Lorenz system.
In this model the drift function is given by a multi-layer perceptron, and the diffusion term is a simple scalar.
We generated 16 independentpaths, consisting of 200 observations each, uniformly sampled over a time interval of [0 , 10] .
Wetrained the model on this data by two different methods.
In the ﬁrst method, we minimize themean-squared error, estimated by integrating out the SDE, and calculating gradients using the adjointmethod (Li et al., 2020).
In the second method, we directly optimize the log-probability calculatedusing Algorithm 1.
In both cases we used precisely the same optimization parameters, the sametime-step size ∆ t , and the same number of sampled trajectories to calculate averages.
Both modelsused the same seed for the neural network initialization, and the same number of gradient steps wereperformed in both cases.
These models were trained on an Nvidia RTX 3070, with 8Gb of memory.
To assess the performance of these approaches, we record two metrics.
The ﬁrst is the mean-squaredifference between the learned drift vector ﬁeld, and the Lorenz vector ﬁeld, normalized by the meansquare magnitude of the Lorenz vector ﬁeld, measured at the observations.
Written mathematically:
[Equation]
This metric is shown in Figure 3b.
The second metric we record is the mean-squared error of sampledtrajectories compared to the observations; this is shown in Figure 3a.
However, we note that this isnot a fair performance metric since the adjoint-based method is directly optimizing for it.
On thex-axis we show the training time in minutes.
[Figure]
minutes) [Figure] (a)
[Figure]
(minutes) 0.
r e l a t i v e d r i f t m s e path int adjoint (b)
In Figure 3 we have limited the range of the x-axis in order for the lines to be more clearly visible.
The total training time for the adjoint-based MSE optimization was a factor 52 longer than the timetaken optimizing the probability estimates.
We note that, while our approach was signiﬁcantly faster,it did use approximately four times as much memory as the adjoint-based method.
The poor performance of the integrator-based learning algorithm in Figure 3b — despite its low meansquared-error — suggests that directly optimizing for the discrepancy of typical SDE sample-pathscompared to observations is not a robust approach to SDE learning.
3. G RADIENT VARIANCE
We applied both our algorithm and an SDE integrator to a neural network-based SDE model ona real-world dataset and recorded the gradients, sampled over many independent trajectories.
Wethereafter computed the variance of the gradients individually per parameter, for various time lengthsbetween observations.
We used a discretization step size of 0.01.
These results are shown in Figure 4.
[Figure]
In each case we used the same neural network, initialized with a ﬁxed seed, and the exact same valuesfor all parameters.
We used the same loss functions as in Section 3.1.
The data set we used was the
Hungarian chickenpox cases dataset from the UCI machine learning repository; 4 it has 20 featuresper observation.
We scaled the features by their standard deviations before performing the experimentto improve numerical stability.
Our ﬁndings are as follows:
• For very short times, we found that both algorithms produced comparable gradient variances:typically less than 10 − 7 .
• For longer time spans (3 orders of magnitude larger than the discretization step size), the
gradient variance of our algorithm increased by a few orders of magnitude; the maximumwe observed was on the order of 10 − 4 .
However, the SDE integrator approach resulted inhuge increases in variances, with some as large as 10 18 .•
Between the above two extremes, we found that the gradient variance of our approach didincrease with the time between observations, but stayed below 10 − 5 , while the gradientvariance from the SDE integrator exhibited very poor scaling with variances of up to 10 8 fortimes between observations of 1.0.
Considering that both approaches produced comparable variances on very small timescales, we feelthat this experiment represents a fair comparison.
3. U NCERTAINTY QUANTIFICATION WITH VARIATIONAL G AUSSIAN PROCESSES
We attempt to test the viability of using our method for an SDE model with a GP prior over the driftfunction and ﬁtting the posterior variationally.
To assess this, we used data trajectories from a van der Pol oscillator.
For the GP prior, we used a zero mean function and a radial basis function kernel, withindependent GP priors over each output dimension of the drift vector ﬁeld.
These choices are notnecessarily optimal; however, the purpose of this experiment is merely to illustrate the feasibility ofusing variational GP to learn the vector ﬁeld of an SDE model.
We used a grid-based inducing point approach and the kernel interpolation method of Wilson & Nickisch (2015).
The combination of these techniques allows one to take advantage of the Kroneckerproduct structure of the covariance matrices and the Toeplitz structure of each factor for acceleratedsampling.
The model was trained variationally, sampling the variational posterior jointly along entiretrajectories.
[Figure]
The approximate GP posterior is shown in Figure 5a compared to the ground truth vector ﬁeld of thevan der Pol oscillator in Figure 5b.
Comparing the GP to the van der Pol vector ﬁeld, we found thatthe GP model was fairly accurate in areas with many data points, while reverting to the prior far awayfrom observations.
To assess the quality of the uncertainty quantiﬁcation, we further calculated thediscrepancy between the variational posterior and the ground truth vector ﬁeld.
Over a dense grid,covering an enlarged area around the observations, we calculated the discrepancy as
[Equation]
A histogram is shown in Figure 5c.
Under Gaussianity assumptions, one expects to see the quantiles line up well with some associatednumber of standard deviations.
However, this will only be the case near to the observations. Inregions where the GP reverts to the prior, which has a stationary kernel, the mismatch between the GP and the polynomial growth of the van der Pol oscillator results in heavy tails. We expect that abetter choice of kernel, which accounts for this polynomial behaviour, to show residuals with lightertails.
We found that the distribution of these δ s was fatter-tailed than a Gaussian distribution; themiddle 50% (25% - 75% quantiles) slightly passed one standard deviation on either side, as opposedto the 68% of a standard Gaussian distribution.
The middle 80% (10% - 90% quantiles) were within3 standard deviations.
While the posterior uncertainty was somewhat overconﬁdent, it was not whollyunreasonable.
We further note that, while the complexities involved in variational GP ﬁtting can result in longertraining times, we still found that our method applied to a GP was signiﬁcantly faster than using
an
SDE integrator with a neural network.
D ISCUSSION
In this work we introduced an algorithm to efﬁciently estimate probabilities of observations underthe assumption of an SDE model for the purposes of learning.
The proposed approach producesestimates which are accurate, even in high dimensions, and can fully utilize parallel hardware forsigniﬁcant speedups.
The proposed algorithm is capable of providing robust probability estimates which can be used inmodels both with or without an explicit observation noise component.
This is in contrast to SDEintegrator-based learning approaches, which must incorporate and estimate an explicit error term dueto their inability to generate paths which are guaranteed to pass through the observations.
Our algorithm produces more stable estimates than alternative integrator-based approaches, particularly when the observations are sparsely distributed over long time intervals and in high dimensions.
This in turn results in more stable and robust learning.
With the rise in applicability of large-scale(neural network) SDE modeling in science, simulation, and domains such as ﬁnance, we see ourapproach as an important contribution to reliable scaling of these methods.
R EPRODUCIBILITY
Efforts have been made to describe all parameters and details of the algorithms used in this paper,either in the content or in the appendix.
The code included in the supplementary material includesa Dockerﬁle and instructions for running some of the experiments.
Where relevant, ﬁxed seeds forrandom number generators can be found in the code during initialization.
B S AMPLING B RIDGES
Let Y be a Gauss–Markov process with kernel k .
Since the process is Markov, we need only samplethe process conditioned on two end-points.
Let y =
( y 0 , y T ) be the known values of the process attimes τ =
( t 0 , T ) .
Let ˜ Y be the process deﬁned by
[Equation]
where K is the matrix whose entries are k ( t i , t j ) .
This distribution of ˜ Y is the same as the distributionof Y conditioned on the values y at times τ .
Note that once Y t has been sampled — which mayrequire a cumulative sum — this transformation can be performed independently, and hence inparallel, for each t .
For each consecutive pair of observations, the bridge process can be sampledindependently.
In the simplest case, where Y is a Brownian motion starting from y 0 at time t 0 = 0 , this reduces to
[Equation]
And similarly for any other processes which can be sampled with a ﬁxed starting point.
In the more general case we can invert the two by two matrix to get an explicit formula for ˜ Y t
[Equation]
This may be extended straightforwardly to other cases, such as conditioning on the values of theintegral, or derivative of the process etc.
C E XPERIMENTSC. P ROBABILITY ESTIMATES
As a sanity check, we compare the log-probability estimates of our algorithm to the actual logprobabilities of an Ornstein–Uhlenbeck process.
The probability density for this process is availablein closed form for any number of observations.
We estimate the probability of a set of data pointsevenly spaced over a time interval of [0 , 10] , taken from a sine function x n = sin( π / 10 × t n ) .
Forthese estimates we used 1000 time steps to estimate the integral S , and averaged over 100 trajectories,independently sampled from a Brownian bridge passing through the observations.
We repeated thisprocess 1000 times and report the mean and standard deviation of these estimates.
Figure 6 showsthese estimates compared to the analytic probability given by the Ornstein–Uhlenbeck process forvarious numbers of observations as a function of the dimension of the problem.
An interesting observation is that the variance of the estimator seems to decrease with the numberof data points.
This is atypical of probability estimators, since the magnitude of the log-probabilityincreases linearly with the number of observations (not to be confused with the number of samplesused to estimate the log-probability).
In almost all cases, with the exception of 4 observations inmore than 40 dimensions, our algorithm gives accurate log-probability estimates with small standarddeviations.
The error of this exception is approximately 30%, but is signiﬁcantly reduced by addingmore observations, or using more samples.
We believe this error is acceptable considering the smallnumber of observations and the large scale.
For 10 observations in 100 dimensions the standarddeviation of our estimator is only 0.4% of the magnitude of the log-probability and the error of the mean estimate is 2% of the analytic log-probability.
This small bias can be further reduced by usingmore samples or a smaller ∆ t .
A single estimate in the 100 dimensional case took less than half asecond to compute on a laptop CPU, and is therefore feasible to use in the inner loop of a learningalgorithm.
C. L ORENZC.2. D ATA SET
We generated data by simulating a Lorenz system with the exact same parameters as Li et al.(2020).
We generated 16 independent trajectories, with initial values sampled from a standard normaldistribution, and discretely sampled the values of the process over 200 equally spaced time pointsranging from 0 to 10.
Unlike Li et al.
(2020), we did not add any observation noise to the data.
We deliberately used a sparser data set than Li et al.
(2020), due to personal experience of the difﬁcultyof ﬁtting SDE models on sparse observations.
C.2. M ODEL
The model we used to ﬁt to the observations was
[Equation]
(38) where σ is a scalar and f θ is a fully-connected neural network with ReLU nonlinearities and parameters θ .
We used layer sizes of 3 → 32 → 256 → 32 → 3 .
C.2. T RAINING PARAMETERS
We used precisely the same parameters for each model, the only difference in the training proceduresbeing the loss functions used.
In the ﬁrst case, the loss function per trajectory was
[Equation]
(39) where N = 200 is the number of observations per trajectory, and K is the number of sampled pathsused to estimate the MSE, and X ( k ) is sampled using an SDE integrator starting at X ( k ) t 1 = x 1 .
Wethen averaged this loss function over the 16 independent observed trajectories.
In the second case, the loss function per trajectory was
[Equation]
where ˆ p X ( x 1: N | θ ) is calculated using Algorithm 1 with K trajectory samples.
Again this loss wasaveraged over the 16 independent observed trajectories.
In both cases we used K = 64 and a timestep size of ∆ t = 10 − 2 .
For each model we used the Adam optimizer with a learning rate of 10 − 3 and ran the optimizationalgorithm for 10 4 iterations.
The model using our importance sampling approach completed the fullnumber of iterations in 34 minutes, while the integrator-based approach took 29 hours and 46 minutesto complete the same number of iterations.
Our approach consistently maintained 100% GPU utilization, except during MSE evaluation whichused an integrator, while the integrator based approach was unable to surpass 25% GPU utilization
inthis experiment.
C. V ARIATIONAL G AUSSIAN PROCESSES
Our approach allows more straightforward implementation of SDE models using a Gaussian processprior for the drift function.
Some previous approaches often resort to approximations such asusing only the predictive mean or MAP estimates (Yildiz et al., 2018).
For dense observations,directly discretizing the SDE and ﬁtting the GP to the observation deltas may be viable.
For sparseobservations this is much more difﬁcult.
Our proposed method makes variational GP methods slightly simpler, since the values of f ( Y t ) can be sampled jointly along the path.
However, it is important tonote that various complexities still arise such as the difﬁculty of efﬁciently sampling high-dimensional
Gaussian distributions.
We used the kernel interpolation method of Wilson & Nickisch (2015) which stores a grid ofinducing locations per dimension.
Storing inducing locations per dimension for a Kronecker productkernel effectively gives a variational distribution over m d inducing points with only d × m memoryrequirements, where m is the number of inducing points per dimension and d is the dimension.
Whenthese inducing locations are evenly spaces, the covariance factors have a Toeplitz structure whichallows sampling in linearithmic time complexity.
To ﬁt the model we used gradient ascent on the ELBO which is calculated as follows:
[Equation]
The KL divergence is between two Gaussian distributions calculated at the inducing locations inthe usual way.
For the likelihood term, we generated trajectories (as in Algorithm 1), followed byjointly sampling the variational GP posterior along these trajectories and proceed with the rest of thecomputation in Algorithm 1 using these samples.
Our optimization objective is calculated
[Equation]
(42) where each f ( k ) is an independent sample over the whole path of Y .
In other words, f ( k ) ( Y t 1 ) andf ( k ) ( Y t 2 ) are correlated, but f ( k ) ( Y t 1 ) and f ( j ) ( Y t 2 ) are not when k (cid:54)
=
j .
In our experiment we used 32 evenly spaced inducing points per dimension, and jointly sampled thevariational posterior over 64 trajectories of 1024 time steps.
The model was trained with the Adamoptimizer with an initial learning rate of 0.1, and an exponential decay with factor 0.999 for 10 5training steps.
However, the square deviation of the predictive mean from the ground truth convergedwell before 10 4 iterations.
Training for 10 4 iterations took 34 minutes, 10 5 iterations took about 5and a half hours.