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: This paper presents a technique for computing approximate ordinary differential equation (ODE) solution posteriors. As shown in previous work [31], the state solutions of mechanistic ODE systems can be approximated by an extended Kalman filter. The authors further extend this methodology by conditioning the filter on observed sequences. Unlike its Monte Carlo alternatives, the resulting inference scheme computes approximate posteriors in a single forward-backward pass. The synthetic data experiments demonstrate that both the state and parameter trajectories are accurately inferred. The vanilla method and a variant that allows non-negative state estimation are tested on a COVID-19 dataset, implying interpretable findings that agree with government measures.