A Probabilistic State Space Model for Joint Inference from Differential Equations and Data

1 Introduction Mechanistic models based on ordinary differential equations ( ODEs ) are popular across a wide range of scientific disciplines . To increase the descriptive power of such models , it is common to consider parametrized versions of ODEs and find a set of parameters such that the dynamics reproduce empirical observations as accurately as possible . Algorithms for this purpose typically involve repeated forward simulations in the context of , e.g. , Markov-chain Monte Carlo or optimization . The need for iterated computation of ODE solutions may demand simplifications in the model to meet limits in the computational budget . This work describes an algorithm that merges mechanistic knowledge in the form of an ODE with a non-parametric model over the parameters controlling the ODE – a latent force that represents quantities of interest . The algorithm then infers a trajectory that is informed 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . by the observations but also follows sensible dynamics , as defined by the ODE , in the absence of observations ( Figure 1 ) . The main insight enabling this approach is that if probabilistic ODE solvers use the language of ( extended ) Kalman filters , conditioning on observations and solving the ODE itself is possible in one and the same process of Bayesian filtering and smoothing . Instead of iterated computation of ODE solutions , a posterior distribution arises from a single forward simulation , which has complexity equivalent to numerically computing an ODE solution , once , with a filtering-based , probabilistic ODE solver [ 32 ] . Intuitively , one can think of this as opening up the black box ODE solver and acknowledging that each task – solving the ODE and discovering a latent force – is probabilistic inference in a state-space model . The main contribution of this work is formalizing this intuition . Several experiments empirically prove the efficiency and the expressivity of the resulting algorithm . In particular , a practical model for the dynamics of the COVID-19 pandemic is considered , in which a non-parametric latent force captures the effect of policy measures that continuously change the contact rate among the population . 2 Problem setting . Let x : [ t0 , tmax ] → Rd be a process that is observed at a discrete set of points T OBSN : = ( tOBS0 , ... , tOBSN ) through a sequence of measurements y0 : N : = ( y0 , ... , yN ) ∈ R ( N+1 ) ×k . Assume that these measurements are subject to additive i.i.d . Gaussian noise , according to the observation model yn = Hx ( tn ) + n , n ∼ N ( 0 , R ) , ( 1 ) for n = 0 , ... , N and matrices H ∈ Rk×d and R ∈ Rk×k . Further suppose that x ( t ) solves the ODE ẋ ( t ) = f ( x ( t ) ; u ( t ) ) , ( 2 ) and satisfies the initial condition x ( t0 ) = x0 ∈ Rd . The vector field f : Rd × R ` → Rd is assumed to be autonomous , which is no loss of generality ( e.g . [ 19 ] ) but simplifies the notation . The latent force u : [ t0 , tmax ] → R ` parametrizes f and shall be unknown . Susceptible Infectious Recovered Deceased β γ η is assumed to remain constant . In this context , the contact rate β ( t ) is the latent force and varies over time ( in the notation from Eq . ( 2 ) , β is u ) . A time-varying contact rate provides a model for the impact of governmental measures on the dynamics of the pandemic . The experiments in Section 5 isolate the impact of the contact rate on the course of the infection counts , by assuming that γ and η are fixed and known . The method is by no means restricted to inference over a single latent force , as will also be shown in Section 5.1 . In this SIRD setting , the goal is to infer an ( approximate ) joint posterior over β ( t ) and the dynamics of S ( t ) , I ( t ) , R ( t ) , and D ( t ) as well as to use the reconstructed dynamics to extrapolate into the future . Section 3 explains the conceptual details , Section 4 distinguishes the method from related work , and Section 5 evaluates the performance . 3 Method . This section explains how to infer the unknown process u ( t ) and the ODE solution x ( t ) in a single forward solve . Section 3.1 defines the prior model , Section 3.2 describes the probabilistic numerical ODE inference setup , and Section 3.3 describes approximate Gaussian filtering and smoothing in this context . Section 3.4 summarizes the resulting algorithm . The exposition of classic concepts here is necessarily compact . In-depth introductions can be found , e.g. , in the book by Särkkä and Solin [ 28 ] . 3.1 Gauss–Markov prior . Let ν ∈ N. Define two independent Gauss–Markov processes U : [ t0 , tmax ] → R ` and X : [ t0 , tmax ] → Rd ( ν+1 ) that solve the linear , time-invariant stochastic differential equations [ 25 ] , dU ( t ) = FUU ( t ) dt+ LU dWU ( t ) , dX ( t ) = FXX ( t ) dt+ LX dWX ( t ) , ( 4 ) with drift matrices FU ∈ R ` × ` and FX ∈ Rd ( ν+1 ) ×d ( ν+1 ) , as well as dispersion matrices LU ∈ R ` ×s and LX ∈ Rd ( ν+1 ) ×d . WU : [ t0 , tmax ] → Rs and WX : [ t0 , tmax ] → Rd are Wiener processes . U and X satisfy the Gaussian initial conditions , U ( t0 ) ∼ N ( mU , PU ) , X ( t0 ) ∼ N ( mX , PX ) , ( 5 ) defined by mU ∈ R ` , PU ∈ R ` × ` , mX ∈ Rd ( ν+1 ) , and PU ∈ Rd ( ν+1 ) ×d ( ν+1 ) . U ( t ) models the unknown function u ( t ) and can be any Gauss–Markov process that admits a representation as the solution of a linear SDE with Gaussian initial conditions . X ( t ) = ( X ( 0 ) ( t ) , ... , X ( ν ) ( t ) ) ∈ Rd ( ν+1 ) models the ODE dynamics , in light of which we require X ( i ) ( t ) = d i dti X ( 0 ) ( t ) ∈ Rd , i = 0 , ... , ν . In other words , the first element in X ( t ) is an estimate for x ( t ) , the second element is an estimate for ddtx ( t ) , et cetera . Encoding that the state X consists of a model for x ( t ) as well as its first ν derivatives imposes structure on FX and LX ( see e.g . [ 18 ] ) . Examples include the Matérn , integrated Ornstein-Uhlenbeck , and integrated Wiener processes ; the canonical choice for probabilistic ODE solvers would be integrated Wiener processes [ 29 , 32 , 4 , 19 ] . The class of Gauss–Markov priors inherits its wide generalizability from Gaussian process models ; recall that Gauss–Markov processes like U and X are Gaussian processes with the Markov property . While not every Gaussian process with one-dimensional input space is Markovian , a large number of descriptions of Gauss–Markov processes emerge by translating a covariance function into an ( approximate ) SDE representation [ 28 , Chapter 12 ] . For example , this applies to ( quasi- ) periodic , squared-exponential , or rational quadratic kernels ; in particular , sums and products of Gauss–Markov processes admit a state-space representation [ 30 , 28 ] . Recent research has considered approximate SDE representations of general Gaussian processes in one dimension [ 20 ] . With these tools , prior knowledge over U or X can be encoded straightforwardly into the model . 3.2 Two likelihoods : for observations and for the ordinary differential equation . A functional relationship between the processes U ( t ) , X ( t ) and the data y0 : N emerges by combining two likelihood functions : one for the observations y0 : N ( recall Equation ( 1 ) ) , and one for the ordinary differential equation . The present section formalizes both . Let T = T OBSN ∪ T ODEM be the union of the observation-grid T OBSN , which has been introduced in Section 2 , and an ODE-grid T ODEM : = ( tODE0 , ... , tODEM ) . The name “ ODE-grid ” expresses that this grid contains the locations on which the ODE information will enter the inference problem , as described below . T OBSN contains the locations of y0 : N , in light of which the first of two observation models is Yn | X ( tOBSn ) ∼ N ( HX ( 0 ) ( tOBSn ) , R ) , ( 6 ) for n = 0 , . . . , N . This is a reformulation of the relationship between process x and observations y0 : N in Eq . ( 1 ) in terms of X ( instead of x , which is modeled by X ( 0 ) ) . Including this first measurement model ensures that the inferred solution remains close to the data points . T ODEM contains the locations on which U ( t ) connects to X ( t ) through the ODE . Specifically , the set of random variables Z0 : M ∈ R ( M+1 ) ×d , defined as Zm | X ( tODEm ) , U ( tODEm ) ∼ δ ( X ( 1 ) ( tODEm ) − f ( X ( 0 ) ( tODEm ) ; U ( t ODE m ) ) ) , ( 7 ) where δ is the Dirac delta , describes the discrepancy between the current estimate of the derivative of the ODE solution ( i.e . X ( 1 ) ) and its desired value ( i.e . f ( X ( 0 ) ; U ) ) , as prescribed by the vector field f . If the random variables Z0 : M realize small values everywhere , X ( 0 ) solves the ODE as parametrized by U . This motivates introducing artificial data points z0 : M ∈ R ( M+1 ) ×d that are equal to zero , zm = 0 ∈ Rd , m = 0 , ... , M . Conditioning the stochastic processes X and U on attaining this ( artificial ) zero data ensures that the inferred solution follows ODE dynamics throughout the domain . Figure 3 shows the discretized state-space model . 3.3 Approximate inference with an extended Kalman filter . Both X and U enter the likelihood in Eq . ( 7 ) through a possibly non-linear vector field f . Therefore , the posterior distribution ( recall z0 : M = 0 ) p ( U ( t ) , X ( t ) | Z0 : M = z0 : M , Y0 : N = y0 : N ) ( 8 ) is intractable , but can be approximated efficiently . Even though the problem is discretized , the posterior distribution is continuous [ 28 , Chapter 10 ] . There are mainly two approaches to computing a tractable approximation of the intractable posterior distribution in Eq . ( 8 ) : approximate Gaussian filtering and smoothing [ 27 ] , which computes a cheap , Gaussian approximation of this posterior , and sequential Monte Carlo methods [ 24 ] , whose approximate posterior may be more descriptive , but also more expensive to compute . Like the literature on probabilistic ODE solvers [ 32 , 4 ] , this work uses approximate Gaussian filtering and smoothing techniques for their low computational complexity . The continuous-discrete state-space model inherits its non-linearity from the ODE vector field f . Linearizing f with a first-order Taylor series expansion creates a tractable inference problem ; more specifically , it gives rise to the extended Kalman filter ( EKF ) [ 13 , 22 ] . Loosely speaking , if the random variable Z is large in magnitude , then X and U are poor estimates for the ODE and its parameter . An EKF update , based on the first-order linearization of f , approximately corrects this misalignment . If sufficiently many ODE measurements z0 : M are available , a sequence of such updates preserves sensible ODE dynamics over time . An alternative to a Taylor-series linearization is the unscented transform , which yields the unscented Kalman filter [ 34 , 15 ] . The computational complexity of both algorithms is linear in the number of grid points and cubic in the dimension of the state-space . Detailed implementation schemes can be found , for instance , in the book by Särkkä [ 27 ] . The EKF approximates the filtering distribution p ( U ( t ) , X ( t ) | Z0 : m = z0 : m , Y0 : n = y0 : n , such that tODEm , tOBSn ≤ t ) . ( 9 ) It describes the current state of the system given all the previous measurements and allows updates in an online fashion as soon as new observations emerge . If desired , the Rauch-Tung-Striebel smoother turns the filtering distribution into an approximation of the full ( smoothing ) posterior ( in Eq . ( 8 ) ) . In doing so , all observations – that is , measurements according to both Eq . ( 6 ) and Eq . ( 7 ) – are taken into account for the posterior distribution at each location t. As special cases , this setup recovers : ( i ) a Kalman filter/Rauch-Tung-Striebel smoother [ 16 ] if the ODE likelihood ( Eq . ( 7 ) ) is omitted ; ( ii ) an ODE solver [ 32 ] , if the data likelihood ( Eq . ( 6 ) ) is omitted . In the present setting , however , both likelihoods play an important role .

Summary: The paper "A Probabilistic State Space Model for Joint Inference from Differential Equations and Data" describes a method of how to estimate latent functions in ordinary differential equations using measuremt data on observed parts of the model. This is achieved by combining a (single!) run through an ODE solver with a suitable probabilistic framework (Bayes filter).