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

Abstract:
Mechanistic models with differential equations are a key component of scientific applications of machine learning. Inference in such models is usually computationally demanding, because it involves repeatedly solving the differential equation. The main problem here is that the numerical solver is hard to combine with standard inference techniques. Recent work in probabilistic numerics has developed a new class of solvers for ordinary differential equations (ODEs) that phrase the solution process directly in terms of Bayesian filtering. We here show that this allows such methods to be combined very directly, with conceptual and numerical ease, with latent force models in the ODE itself. It then becomes possible to perform approximate Bayesian inference on the latent force as well as the ODE solution in a single, linear complexity pass of an extended Kalman filter / smoother — that is, at the cost of computing a single ODE solution. We demonstrate the expressiveness and performance of the algorithm by training, among others, a non-parametric SIRD model on data from the COVID-19 outbreak.

text: 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.

heading: 2 Problem setting
text: 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.

heading: 3 Method
text: 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].

heading: 3.1 Gauss–Markov prior
text: 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.

heading: 3.2 Two likelihoods: for observations and for the ordinary differential equation
text: 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.

heading: 3.3 Approximate inference with an extended Kalman filter
text: 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.

heading: 3.4 Algorithm and implementation
text: The procedure is summarized in Algorithm 1. The prediction step is determined by the prior and is available in closed-form (Appendix A.2). At times at which data is observed according to the linear Gaussian measurement model in Eq. (6), the update step follows the rules of the standard Kalman filter. Before updating on pseudo-observations according to the ODE likelihood (Eq. (7)), the non-linear measurement model is linearized at the predicted mean. More details are provided
Algorithm 1 Compute the filtering distribution by conditioning on both y0:N and z0:M . Input: data y0:N , time grid T = T OBSN ∪ T ODEM , vector field f , mX, PX, mU, PU Output: Filtering distribution [Eq. (9)] Initialize X0 = N (mX, PX) and U0 = N (mU, PU) [Eq. (5)] for tj ∈ T do
Predict Xj from Xj−1 and predict Uj from Uj−1 if tj ∈ T OBSN then update Xj on yj end if [Eq. (6)] if tj ∈ T ODEM then linearize measurement model and update Xj and Uj on zj end if [Eq. (7)]
end for
in Appendix A. The filtering distribution can be turned into a smoothing posterior by running a backwards-pass with a Rauch-Tung-Striebel smoother (e.g. [27]).
The computational cost of obtaining either, the filtering or the smoothing posterior, are both linear in the number of grid points and cubic in the dimension of the state-space, i.e.O((N +M)(d3ν3 + `3)). Only a single forward-backward pass is required. If desired, the approximate Gaussian posterior can be refined iteratively by means of posterior linearization and iterated Gaussian filtering and smoothing, which yields the maximum-a-posteriori (MAP) estimate [2, 31]. The experiments presented in Section 5 show how a single forward-backward pass already approximates the MAP estimate accurately.

heading: 4 Related work
text: Latent forces and ODE solvers: The explained method closely relates to probabilistic ODE solvers and latent force models [37], especially the kind of latent force model that exploits the state-space formulation of the prior [9]. The difference is that, in the spirit of probabilistic numerical algorithms, the mechanistic knowledge in the form of an ODE is injected through the likelihood function instead of the prior. A similar approach of linking observations to mechanistic constraints has previously been used in the literature on constrained Gaussian processes [14] and gradient matching [5, 36]. Probabilistic ODE solvers have been used by Kersting et al. [17] for efficient ODE inverse problem algorithms, but their approach is different to the present algorithm, in which the need for iterated optimization or sampling is avoided altogether.
Monte Carlo methods: (Markov-chain) Monte Carlo methods are also able to infer a time-dependent ODE latent force from a set of state observations. Options that are compatible with a setup similar to the present work would include sequential Monte Carlo techniques [24], elliptical slice sampling [23], or Hamiltonian Monte Carlo [3] (for instance realized as the No-U-Turn sampler [12]). The shared disadvantage of Monte Carlo methods applied to the resulting ODE inverse problem is that the complexity of obtaining a single Monte Carlo sample is of the same order of magnitude as computing the full Gaussian approximation of the posterior distribution. In Appendix B we show results from a parametric version of the SIRD-latent force model (using the No-U-Turn sampler as provided by NumPyro [26]). This sampler requires thousands of numerical ODE solutions, compared to the single solve of our method. This fact is also reflected in the wall-clock time needed for both types of inference. While the MCMC experiment in Appendix B takes in the order of hours, each experiment with our approach takes under one minute to complete. In other words, the algorithm in the present work poses an efficient yet expressive alternative to Monte Carlo methods for approximate inference with dynamical systems.

heading: 5 Experiments
text: This section describes three blocks of experiments. The implementation is based on ProbNum [35] and all experiments use a conventional, consumer-level CPU. First, a range of artificial datasets is generated by sampling ODE parameters from a prior state-space model and simulating a solution of the corresponding ODE. Inference in such a controlled environment allows comparing to the ground truth, thereby assessing the quality of the approximate inference. We consider three ODE models to this end. Second, a COVID-19 dataset will probe the predictive performance of the probabilistic model and the resulting approximate posterior distribution. Third, some changes to the model from the COVID-19 experiments, for instance, ensuring that the number of case counts must be positive,
will improve the interpretability (for example, of the credible intervals). Controlling the range of values that the prior state-space can realize introduces additional non-linearity into the model – which can also be locally approximated by the EKF – and makes the solution more physically meaningful.

heading: 5.1 Simulated environments
text: As a first test for the capabilities of the proposed method, we consider three simulated environments. To this end, the training data is generated as follows. The starting point is always an initial value problem with dynamics defined by a vector field f and a Gauss–Markov prior over the dynamics x and the unknown parameters u of the vector field. Then, (i) we sample the time-varying parameter trajectories from the Gauss–Markov prior; (ii) we solve the ODE, as parametrized by the sampled trajectories from (i), using LSODA [11] with adaptive step sizes using SciPy [33]; (iii) we subsample the ground-truth solution on a uniform grid (which will become T OBSN ) to generate artificial state observations y0:N ; (iv) we add Gaussian i.i.d. noise to the observations.
The procedure described above generates both a ground truth to compare to and a noisy, artificially observed data set. Given such a set of observations, Algorithm 1 computes a posterior distribution over the true trajectories under appropriate model assumptions. In this posterior, we look for the proximity of the mean estimate to the underlying ground truth; the closer, the better. We measure this proximity in the root-mean-square error. Furthermore, the width of the posterior (expressed by the posterior covariance) should deliver an appropriate quantification of the mismatch. We report the χ2-statistic [1], which suggests that the posterior distribution is well-calibrated if the χ2-statistic is close to the dimension d of the ground truth. Three mechanistic models serve as examples.
Van-der-Pol: The first of three test problems is the van-der-Pol oscillator [8]. It has one parameter µ (sometimes referred to as a stiffness constant, because for large µ, the van-der-Pol system is stiff). As a prior over the dynamics we choose a twice-integrated Wiener process with diffusion intensity σ2X = 300. The stiffness parameter µ is modeled as a Matérn-3/2 process with lengthscale `U = 10 and diffusion intensity σ2U = 0.3. The posterior is computed on a grid from t0 = 0 to tmax = 25 units of time with step size ∆t = 0.025.
Lotka-Volterra: The Lotka-Volterra equations [21] describe the change in the size of two populations, predators and prey. There are four parameters, which we call a, b, c, and d, which describe the interaction and death/reproduction rates of the populations. As a prior over the dynamics we choose a twice-integrated Wiener process with diffusion intensity σ2X = 10. The four parameters are modeled as Matérn-3/2 processes with lengthscales `Ua = `Ub = `Uc = `Ud = 40. The diffusion intensities are σ2Ua = σ 2 Uc = 0.01 and σ2Ub = σ 2 Ud
= 0.001. The posterior is computed on a grid from t0 = 0 to tmax = 60 units of time with step size ∆t = 0.1.
SIRD: As detailed in Section 2, the SIRD model is governed by a contact rate β(t). Recall that we assume a time-dependent β to account for governmental measures in reaction to the spread of COVID-19. The recovery rate γ and fatality rate η are fixed at γ = 0.06 and η = 0.002, like they will be in the experiments with real data in Sections 5.2 and 5.3 below. As a prior over the dynamics we choose a twice-integrated Wiener process with diffusion intensity σ2X = 50. The contact rate β is modeled as a Matérn-3/2 process with lengthscale `U = 14 and diffusion intensity σ2U = 0.1. The posterior is computed on a grid from t0 = 0 to tmax = 100 units of time with step size ∆t = 0.1.
The model allows for straightforward restriction of parameter values by using link functions. The natural support for the SIRD-contact rate is the interval [0, 1], but U(t), as a Gauss–Markov process, takes values on the real line. A change in the basis of β(t) with a logistic sigmoid function ϑ before it enters the likelihood fixes this misspecification. Similarly, the Lotka-Volterra parameters are inferred in log-space to ensure positivity. It is an appealing aspect of the EKF that these non-linear transformations do not require significant adaptation of the algorithm. Instead, the EKF treats it as merely another level of linearization of Eq. (7). Section 5.3 extends this to the state dynamics.
The results are shown in Figure 4. On all test problems, the algorithm recovers the true states and the true latent force accurately. The recovery is not exact, which shows how the Gaussian posterior is only an approximation of the true posterior. The χ2-statistic for the van-der-Pol stiffness parameter µ is 1.11, which lies in (0.0039, 3.8415), the 90% confidence interval of the χ2 distribution with 1 degree of freedom. The root-mean-square error (RMSE) to the truth is 0.14. The χ2-statistic for the Lotka-Volterra parameters is 8.06, which lies in (0.7107, 9.4877), the 90% confidence interval of the χ2 distribution with 4 degrees of freedom. The RMSE to the truth is 0.04 in log space and 0.018 in
linear space. The χ2-statistic for the contact rate β is 0.91, which lies in (0.0039, 3.8415), the 90% confidence interval of the χ2 distribution with 1 degree of freedom. The RMSE to the truth is 0.2 in logit space and 0.033 in linear space.

heading: 5.2 COVID-19 data
text: We continue with the SIRD model introduced in Eq. (3), now using data collected in Germany over the period from January 22, 2020, to May 27, 2021. Throughout the pandemic, the German government has imposed mitigation measures of varying severity. Together with seasonal effects, summer vacations, etc., they caused a continual change in the contact rate. The next experiments aim to recover said contact rate (and the SIRD counts) from the German dataset.
The Center for Systems Science and Engineering at the Johns Hopkins University publishes daily counts of confirmed (yconfirmedn ), recovered (y recovered n ), and deceased (y deceased n ) individuals [6]. One can transform this data to suit the SIRD model
In := y confirmed n −Rn −Dn, Rn := yrecoveredn , Dn := ydeceasedn . (10)
The counts In, Rn, and Dn are available for each day, starting with January 22, 2020. Assuming a constant population over time, the numbers of susceptible individuals Sn are always evident from the other quantities, thus left out of the visualizations. We fix the population at P = 83 783 945, based on public record. We rescale the data to cases per one thousand people (CPT).
As a prior over X(t), due to its popularity in constructing probabilistic ODE solvers [32], we assume a twice-integrated Wiener process. β(t) is modelled as a Matérn-3/2 process with length scale `q = 75 and diffusion intensity σ2q = 0.05. The state-space model is straightforwardly extendable to sums and products of (more) processes [30, 28]. Inferring parameters that are constant over time, however, is not straightforward due to potentially singular transition models [27, Section 12.3.1].
As described in Section 5.1, the contact rate is inferred in logit space. We shift the logistic sigmoid function such that it fulfills ϑ(0) = 0.1 in which case the stationary mean U = 0 translates to a stationary mean ϑ(U) = β = 0.1 of the Matérn process that models the contact rate. The recovery
rate and mortality rate are considered known and fixed at γ = 0.06 and η = 0.002 to isolate the effect of the inference procedure on recovering the evolution of the contact rate U(t) = β(t). We set the mean of the Gaussian initial conditions to the first data point that is available. The diffusion intensity of the prior process X(t) is set to σ2X = 10. The latent process U and all derivatives are initialized at zero. Note that due to the logistic sigmoid transform, an initial value U0 = 0 amounts to an initial contact rate β0 = 0.1.
In the present scenario, we cannot take the SIRD model as an accurate description of the underlying data but merely as a tool that aids the inference engine in recovering physically meaningful states and forces. In order to account for this model mismatch, the Dirac likelihood from Eq. (7) is relaxed towards a Gaussian likelihood with measurement noise λ2 = 0.01. This equals the data observation noise and thus balances the respective impact of either (misspecified) source of information. Intuitively, adding ODE measurement noise reduces how strictly the vector field dynamics are enforced during inference and therefore avoids overconfident estimates of β(t).
The mesh-size of the ODE is ∆t = 1/24 days, i.e. ODE updates are computed on an hourly basis. The final 14 observations are excluded from the training set to serve as validation data for evaluating the extrapolation behavior of the proposed method. Figure 5 shows the results. The mean of the state X estimates the case counts accurately in both interpolation and extrapolation tasks. The estimated contact rate rapidly decreases around late March, remains low until fall, increases momentarily, and is dampened again soon after. This aligns with a set of political measures imposed by the government (compare Figure 5 to Table 1). The uncertainty over the estimated contact rate is high in the early beginning when the case counts are still low. It then increases again in summer and with the beginning of the extrapolation phase.
If the experiment is taken as-is, the credibility intervals of the posterior over X(t) include negative numbers (mostly where the case counts are low and the uncertainty high, and when extrapolating). Of course, in a system that models counts of people in different stages of a disease, negative numbers should be excluded altogether. The proposed method provides straightforward means to address this issue. Section 5.3 explains the details.

heading: 5.3 Non-negative state estimates
text: The following experiment evaluates how the proposed method performs in combination with a statespace model that constrains the support of the dynamics. Concretely, let X(t) model the logarithm of the SIRD dynamics and the respective derivatives. With a slight abuse of notation, we will continue writing “X” even though it lives in a different space than in the previous sections. The structure of the dynamic model is the same. The diffusion intensity of the prior process X(t) is σ2X = 0.05. The diffusion is not comparable to the value in the previous section because the state dynamics moved to log-space. Using ddt exp(x(t)) = exp(x(t))ẋ(t), the ODE likelihood becomes
Zm |XODEm , U ODEm ,∼ N (ζ1 − f(ζ2; ζ3), λ2Id), (11)
with auxiliary quantities (recall the logistic sigmoid ϑ) ζ1 := exp ( X(0)(tODEm ) ) X(1)(tODEm ), ζ2 := exp ( X(0)(tODEm ) ) , ζ3 := ϑ(U(t ODE m )). (12)
The exponential function introduces an additional non-linearity into the state-space model, which necessitates smaller step-sizes for the ODE measurements (see below).
The observed case count data y0:N is transformed into the log-space, too, in which we assume additive, i.i.d. Gaussian noise. On the one hand, transforming the measurements into log-space implies that the measurement model for the counts remains linear; on the other hand, it imposes a log-normal noise model (if viewed back in “linear space”). Log-normal noise underlines how the estimated states cannot be negative. Again, we scale the counts to cases per thousand.
As depicted in Figure 6, the reconstruction of the driving processes in this setting yields results that at first glance, look similar to the previous experiment. The states match the data points well. However, the extrapolation is more realistic in that the credible intervals encode that negative values are impossible (which is due to the log-transform). The mean of the recovered contact rate closely resembles the estimate of the previous experiment. Again, upon implementation of strict governmental measures, the uncertainty decreases, whereas in the context of relaxations, the uncertainty is high.

heading: 6 Statement on Societal Impact
text: This work performs methods research to develop an efficient numerical algorithm to infer latent forces governing ordinary differential equations. As a testbed, we use data from the COVID-19 pandemic. We do so to motivate and visualize the practical value of our methods. The results of this
algorithm, however, should not be taken as policy advice. The model used in the paper is deliberately simplistic. The presented work therefore should not be misunderstood as epidemiological research. The machine learning community has, over time, frequently used data of contemporary societal concern to motivate and test new algorithmic concepts (well-known examples from the UCI collection include the Wisconsin Breast Cancer Dataset, the mushroom classification dataset, and the German credit data set). Our work follows in this line. Of course, if this algorithm, or any competitor, is used to derive policy advice, the underlying differential equation and latent states must be carefully considered by domain experts, which we are not.

heading: 7 Conclusion
text: By coupling mechanistic and data-driven inference so directly, the algorithm builds on the core premise of probabilistic numerics – that computation itself is a data source that does not differ, formally, from observational data. Information from observations and mechanistic knowledge (in the form of an ODE) can thus be described in the same language of Bayesian filtering and smoothing. This removes the need for an outer loop over multiple forward solves and thus drastically reduces the computational cost. Our experimental evaluation corroborates that the resulting approximate posterior is close to the ground truth and drastically reduces computational cost over Monte Carlo alternatives. It faithfully captures multiple sources of uncertainty from the data, numerical (discretization) error, and epistemic uncertainty about the mechanism. We hope this framework helps empower practitioners, not just by reducing computational burden but also by providing a more flexible modelling platform.

heading: Acknowledgements
text: The authors gratefully acknowledge financial support by the European Research Council through ERC StG Action 757275 / PANAMA; the DFG Cluster of Excellence “Machine Learning - New Perspectives for Science”, EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting N. Krämer. Moreover, the authors thank Nathanael Bosch and Marvin Pförtner for valuable discussions.

Review summaries:
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.

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).