Neural Jump Ordinary Differential Equations: Consistent Continuous-Time Prediction and Filtering

1 INTRODUCTION . Stochastic processes are widely used in many fields to model time series that exhibit a random behaviour . In this work , we focus on processes that can be expressed as solutions of stochastic differential equations ( SDE ) of the form dXt = µ ( t , Xt ) dt+ σ ( t , Xt ) dWt , with certain assumptions on the drift µ and the diffusion σ . With respect to the L2-norm , the best prediction of a future value of the process is provided by the conditional expectation given the current value . If the drift and diffusion are known or a good estimation is available , the conditional expectation can be approximated by a Monte Carlo ( MC ) simulation . However , since µ and σ are usually unknown , this approach strongly depends on the assumptions made on their parametric form . A more flexible approach is given by neural SDEs , where the drift µ and diffusion σ are modelled by neural networks ( Tzen & Raginsky , 2019 ; Li et al. , 2020 ; Jia & Benson , 2019 ) . Nevertheless , modelling the diffusion can be avoided if one is only interested in forecasting the behaviour instead of sampling new paths . An alternative widely used approach is to use Recurrent Neural Networks ( RNN ) , where a neural network dynamically updates a latent variable with the observations of a discrete input time-series . RNNs are successfully applied to tasks for which time-series are regularly sampled , as for example speech or text recognition . However , often observations are irregularly observed in time . The standard approach of dividing the time-line into equally-sized intervals and imputing or aggregating observations might lead to a significant loss of information ( Rubanova et al. , 2019 ) . Frameworks that overcome this issue are the GRU-ODE-Bayes ( Brouwer et al. , 2019 ) and the ODE-RNN ( Rubanova et al. , 2019 ) , which combine a RNN with a neural ODE ( Chen et al. , 2018 ) . In standard RNNs , the hidden state is updated at each observation and constant in between . Conversely , in the GRU-ODEBayes and ODE-RNN framework , a neural ODE is trained to model the continuous evolution of the hidden state of the RNN between two observations . While GRU-ODE-Bayes and ODE-RNN both provide convincing empirical results , they lack thorough theoretical guarantees . Contribution . In this paper , we introduce a mathematical framework to precisely describe the problem statement of online prediction and filtering of a stochastic process with temporal irregular observations . Based on this rigorous mathematical description , we introduce the Neural Jump ODE ( NJ-ODE ) . The model architecture is very similar to the one of GRU-ODE-Bayes and ODE-RNN , however we introduce a novel training framework , which in contrast to them allows us to prove convergence guarantees for the first time . Moreover , we demonstrate empirically the capabilities of our model . Precise problem formulation . We emphasize that a precise definition of all ingredients is needed , to be able to show theoretical convergence guarantees , which is the main purpose of this work . Since the objects of interest are stochastic processes , we use tools from probability theory and stochastic calculus . To make the paper more readable and comprehensible also for readers without background in these fields , the precise formulations and demonstrations of all claims are given in the appendix , while the main part of the paper focuses on giving well understandable heuristics . 2 PROBLEM STATEMENT . The problem we consider in this work , is the online forecasting of temporal data . We assume that we make observations of a Markovian stochastic process described by the stochastic differential equation ( SDE ) dXt = µ ( t , Xt ) dt+ σ ( t , Xt ) dWt , ( 1 ) at irregularly-sampled time points . Between those observation times , we want to predict the stochastic process , based only on the observations that we made previously in time , excluding the possibility to interpolate observations . Due to the Markov property , only the last observation is needed for an optimal prediction . Hence , after each observation we extrapolate the current observation into the future until we make the next observation . The time at which the next observation will be made is random and assumed to be independent of the stochastic process itself . More precisely , we suppose to have a training set of N independent realisations of the RdX - dimensional stochastic process X defined in ( 1 ) . Each realisation j is observed at nj random observation times t ( j ) 1 , . . . , t ( j ) nj ∈ [ 0 , T ] with values x ( j ) 1 , . . . , x ( j ) nj ∈ RdX . We assume that all coordinates of the vector x ( j ) i are observed . We are interested in forecasting how a new independent realization evolves in time , such that our predictions of X minimize the expected squared distance ( L2-metric ) to the true unknown path . The optimal prediction , i.e . the L2-minimizer , is the conditional expectation . Given that the value of the new realization at time t is xt , we are therefore interested in estimating the function f ( xt , t , s ) : = E [ Xt+s|Xt = xt ] , s ≥ 0 , ( 2 ) which is the L2-optimal prediction until the next observation is made . To learn an approximation f̂ of f we make use of the N realisations of the training set . After training , f̂ is applied to the new realization . Hence , this can be interpreted as a special type of filtering problem . The following example illustrates the considered problem . Example . A complicated to measure vital parameter of patients in a hospital is measured multiple times during the first 48 hours of their stay . For each patient , this happens at different times depending on the resources , hence the observation dates are irregular and exhibit some randomness . Patient 1 has n1 = 4 measurements at hours ( t ( 1 ) 1 , t ( 1 ) 2 , t ( 1 ) 3 , t ( 1 ) 4 ) = ( 1 , 14 , 27 , 34 ) where the values ( x ( 1 ) 1 , x ( 1 ) 2 , x ( 1 ) 3 , x ( 1 ) 4 ) = ( 0.74 , 0.65 , 0.78 , 0.81 ) are measured . Patient 2 only has n2 = 2 measurements at hours ( t ( 2 ) 1 , t ( 2 ) 2 ) = ( 3 , 28 ) where the values ( x ( 2 ) 1 , x ( 2 ) 2 ) = ( 0.56 , 0.63 ) are measured . Similarly , the j-th patient has nj measurements at times ( t ( j ) 1 , . . . , t ( j ) nj ) and has the measured values ( x ( j ) 1 , . . . , x ( j ) nj ) . Based on this data , we want to forecast the vital parameter of new patients coming to the hospital . In particular , for a patient with measured values x1 at time t1 , we want to predict what his values will likely be at any time t1 + s > t1 . Importantly , we do not only focus on predicting the value at some t2 > t1 , but we want to know the entire evolution of the value . 3 BACKGROUND . Recurrent Neural Network . The input to a RNN is a discrete time series of observations { x1 , · · · , xn } . At each observation time ti+1 , a neural network , the RNNCell , updates the latent variable h using the previous latent variable hi and the input xi+1 as hi+1 : = RNNCell ( hi , xi+1 ) . Neural Ordinary Differential Equation . Neural ODEs ( Chen et al. , 2018 ) are a family of continuous-time models defining a latent variable ht : = h ( t ) to be the solution to an ODE initial-value problem ht : = h0 + ∫ t t0 f ( hs , s , θ ) ds , t ≥ t0 , ( 3 ) where f ( · , · , θ ) = fθ is a neural network with weights θ . Therefore , the latent variables can be updated continuously by solving this ODE ( 3 ) . We can emphasize the dependence of ht on a numerical ODE solver by rewriting ( 3 ) as ht : = ODESolve ( fθ , h0 , ( t0 , t ) ) . ( 4 ) ODE-RNN . ODE-RNN ( Rubanova et al. , 2019 ) is a mixture of a RNN and a neural ODE . In contrast to a standard RNN , we are not only interested in an output at the observation times ti , but also in between those times . In particular , we want to have an output stream that is generated continuously in time . This is achieved by using a neural ODE to model the latent dynamics between two observation times , i.e . for ti−1 < t < ti the latent variable is defined as in ( 3 ) and ( 4 ) , with h0 and t0 replaced by hi−1 and ti−1 . At the next observation time ti , the latent variable is updated by a RNN with the new observation xi . Fixing h0 , the entire latent process can be computed by iteratively solving an ODE followed by applying a RNN . Rubanova et al . ( 2019 ) write this as { h′i : = ODESolve ( fθ , hi−1 , ( ti−1 , ti ) ) hi : = RNNCell ( h ′ i , xi ) . ( 5 ) GRU-ODE-Bayes . The model architecture describing the latent variable in GRU-ODE-Bayes ( Brouwer et al. , 2019 ) is defined as a special case of the ODE-RNN architecture . In particular , a gated recurrent unit ( GRU ) is used for the RNN-cell and a continuous version of the GRU for the neural ODE fθ . Therefore , we focus on explaining the difference between our model architecture and the ODE-RNN architecture , in the following section . 4 PROPOSED METHOD – NEURAL JUMP ODE . Markovian paths . Our assumptions on the stochastic process X imply that it is a Markov process . In particular , the optimal prediction of a future state of X only depends on the current state rather than on the full history . Hence , the previous values do not provide any additional useful information for the prediction . JumpNN instead of RNN . Using a neural ODE between two observation has the advantage that it allows to continuously model the hidden state between two observations . But since the underlying process is Markov , there is no need to use a RNN-cell to model the updates of the hidden state at each new observation . Instead , whenever a new observation is made , the new hidden state can solely be defined from this observation . Therefore , we replace the RNN-cell used in ODE-RNN by a standard neural network mapping the observation to the hidden state . We call this the jumpNN which can be interpreted as an encoder map . Compared to ODE-RNN , this architecture is easier to train . Last observation and time increment as additional inputs for the neural ODE . The neural network fθ used in the neural ODE takes two arguments as inputs , the hidden state ht and the current time t. However , our theoretical problem analysis suggests , that instead of t the last observation time ti−1 and the time increment t− ti−1 should be used . Additionally the last observation xi−1 should also be part of the input . NJ-ODE . Combining the ODE-RNN architecture ( 5 ) with the previous considerations , we introduce the modified architecture of Neural Jump ODE ( NJ-ODE ) { h′i : = ODESolve ( fθ , ( hi−1 , xi−1 , ti−1 , t− ti−1 ) , ( ti−1 , ti ) ) hi : = jumpNN ( xi ) . ( 6 ) An implementable version of this method is presented in the Algorithm 1 . A neural ODE fθ transforms the hidden state between observations , and the hidden state jumps according to jumpNN when a new observation is available . The outputNN , a standard neural network , maps any hidden state ht to the output yt . To implement the continuous-in-time ODE evaluation , a discretization scheme is provided by the inner loop . In the training process , the weights of all three neural networks , jumpNN , the neural ODE fθ and outputNN are optimized . Algorithm 1 The NJ-ODE . A small step size ∆t is fixed and we denote tn+1 : = T . Input : Data points with timestamps { ( xi , ti ) } i=0 ... n , for i = 0 to n do hti = jumpNN ( xi ) . Update hidden state given next observation xi yti = outputNN ( hti ) . compute output s← ti while s+ ∆t ≤ ti+1 do hs+∆t = ODESolve ( fθ , hs , xi , ti , s− ti , ( s , s+ ∆t ) ) . get next hidden state ys+∆t = outputNN ( hs+∆t ) . compute output s← s+ ∆t end while end for Objective function . Our goal is to train the NJ-ODE model such that its output approximates the conditional expectation ( 2 ) , which is the optimal prediction of the target process X with respect to the L2-norm . Therefore , we define a new objective function , with which we can prove convergence . Let yi− denote the output of the NJ-ODE at ti before the jump and yi the output at ti after the jump . Note that the outputs depend on parameters θ and the previously observed xi which are inputs to the model . Then the objective function is defined as Φ̂N ( θ ) : = 1 N N∑ j=1︸ ︷︷ ︸ paths 1 nj nj∑ i=1︸ ︷︷ ︸ dates ( |x ( j ) i − y ( j ) i |︸ ︷︷ ︸ jump part at observations + |y ( j ) i − y ( j ) i− |︸ ︷︷ ︸ continuous part between two observations ) 2 . ( 7 ) We give an intuitive explanation for this definition . The “ jump part ” of the loss function forces the jumpNN to produce good updates based on new observations , while the other part forces the jump size to be small in ( the empirical ) L2-norm . Since the conditional expectation minimizes the jump size with respect to the L2-norm , this forces the neural ODE fθ to continuously transform the hidden state such that the output approximates the conditional expectation . Moreover , both parts of the loss function force the outputNN to reasonably transform the hidden state ht to the output yt .

Summary: The authors propose a method for learning the conditional expectation of stochastic process in an online fashion. The paper bears a considerable theoretical treatment, derived from the stochastic filtering literature, which is present both in the main body of the paper and the appendix. Besides the model, the paper also aims to provide a theoretical justification of the convergence of their method.