Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models
Authors: Lenart Treven
Philippe Wenk
Florian Dörfler
Andreas Krause
Venue: NeurIPS 2021 Poster (2021)
Keywords: Gaussian Processes
Neural ODEs
Dynamical Systems
Uncertainty

Abstract: Differential equations in general and neural ODEs in particular are an essential technique in continuous-time system identification. While many deterministic learning algorithms have been designed based on numerical integration via the adjoint method, many downstream tasks such as active learning, exploration in reinforcement learning, robust control, or filtering require accurate estimates of predictive uncertainties. In this work, we propose a novel approach towards estimating epistemically uncertain neural ODEs, avoiding the numerical integration bottleneck. Instead of modeling uncertainty in the ODE parameters, we directly model  uncertainties in the state space. Our algorithm distributional gradient matching (DGM) jointly trains a smoother and a dynamics model and matches their gradients via minimizing a Wasserstein loss. Our experiments show that, compared to traditional approximate inference methods based on numerical integration, our approach is faster to train, faster at predicting previously unseen trajectories, and in the context of neural ODEs, significantly more accurate.

Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models
Lenart Treven' ETH Zürich trevenl@ethz.cm
Philippe Wenk ETH Zürich wenkph@ethz.co
Florian Dörfler ETH Zürich dorflerGethz.co
Andreas Krause ETH Zürich krausea@ethz.com
Abstract
Differential equations in general and neural ODEs in particular essential technique in continuous-time system identification. While many deterministic learning algorithms have been designed based numerical integration via the adjoint method, many downstream tasks such as active learning, exploration reinforcement learning, robust control, or filtering require accurate estimates of predictive uncertainties. In this work, we propose novel approach towards estimating epistemically uncertain neural ODEs, avoiding the numerical integration bottleneck. Instead of modeling uncertainty in the ODE parameters, we directly model uncertainties in the state space. Our algorithm distributional gradient matching (DGM) jointly trains a smoother and dynamics model and matches their gradients minimizing a Wasserstein loss. Our experiments show that, compared to traditional approximate inference methods based numerical integration, our approach is faster to train, faster at predicting previously unseen trajectories, and in the context of neural ODEs, significantly more accurate.
Introduction
Independert smorth
Indegendent smoother dynamirs
For continuous-time system identification and control, ordinary differential equations form an essential class of models, deployed in ap- plications ranging from robotics (Spong al. 2006) to biology (Jones et al., 2009). Here, it assumed that the evolution scribed by the evolution of continuous state vari- ables, whose time-derivative given by a set of parametrized equations. Often, these equations are derived from first principles, e.g., rigid body dynamics (Wittenburg. (013), mass action ki- 2013), or Hamiltonian dynamics netics 20133,,
Pestle
Joint smoother
Joint smoother Synamics
system de-
Figure Illustration of DGM: Learning joint smoother (first vs second row) across trajectories enables sharing observational data. Dynamics reg- ularization (first vs second column) substantially improves prediction accuracy joint smoother.
(Greydanus et al. 2019), chosen for computa- tional convenience (e.g.. linear systems (Ljung) 1998) parametrized to facilitate system iden- tification (Brunton et al. 2016).
Such construction methods lead intriguing properties, including guarantees on physical realizability (Wensing et al.) 2017), favorable convergence properties (Ortega et al. (2018), or a structure suitable for downstream tasks such control design (Ortega al. 2002). However, such models often capture the system dynamics only approximately. leading to a potentially significant discrepancy between the model and reality (Ljung 1999). Moreover, when expert knowledge not available, or precise parameter values are cumbersome to obtain, system identification from raw time series data becomes
'Equal Contribution. Correspondence to treven1@ethz.co ch, wenkph@ethz.cc
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

GPe,
NNo
Smoother: (t,
ps ps(a(t)..s(z(() (+((t)
Dynamics: xxt)
PD (i(())
max
maxlog (oggs((ss))))) X v.P
Figure 2: High-level depiction of DGM.
necessary, In this case, one may seek more expressive nouparametric models instead (Rackauckas)
al. 2020 Pillonetto et al., 2014). If the model is completely replaced by neural network, the resulting model is called neural ODE (Chen et al. 2018). Despite their large number of parameters,
demonstrated by Chen al. (2118 Kidger et al. (2020; Zhuang al. (2020 2021), deterministic
neural ODEs can be efficiently trained, enabling accurate deterministic trajectory predictions. For many practical applications however, accurate uncertainty estimates are essential, as they guide downstream tasks like reinforcement learning (Deisenroth and Rasmussen 2011 Schulman et al.. 2015). safety guarantees (Berkenkamp et al., 2017), robust control design (Hjalmarsson) 2005). planning under uncertainty (La 2006). probabilistic forecasting in meteorology (Fanfarillo et al. 202T). or active learning experimental design (Srinivas er al.) (010). A common way of obtaining such uncertainties is via Bayesian framework. However, as observed by Dandekar et al. (2021). Bayesian training of neural ODEs in a dynamics setting remains largely unexplored. They demonstrate that initial variational-baat based inference schemes for Bayesian neural ODEs suffer from several serious drawbacks and thus propose sampling-base -based alternatives. However, surfaced by our experiments in Section[4] sampling-bast approaches still exhibit serious challenges. These pertain both to robustness (even if highly informed priors are supplied), and reliance on frequent numerical integration of large neural networks, which poses severe computational challenges for many downstream tasks like campling-based planning (Karaman and Frazzoli 2011) uncertainty propagation in prediction.
Contributions In this work, we propose novel approach for uncertainty quantification in nonlinear dynamical systems (cf. Figure T]. Crucially, our approach avoids explicit costly and non-robust numerical integration, by employing probabilistic smoother of the observational data, whose representation we learn jointly across multiple trajectories. To capture dynamics, we regularize our smoother with dynamics model. Latter captures epistemic uncertainty in the gradients of the ODE, which we match with the smoother's gradients by minimizing Wasserstein loss, hence we call our approach Distributional Gradient Matching (DGM). In summary, our main contributions are:
We develop DGM, approact for capturing epistemic uncertainty about nonlinear dy- namical systems by jointly training a smoother and a neural dynamics model;
We provide computationally efficient and statistically accurate mechanism for prediction, by focusing directly on the posterior / predictive state distribution.
We experimentally demonstrate the effectiveness our approach on learning challenging, chaotic dynamical systems, and generalizing new unseen inital conditions.
High-level overview A high-level depiction of our algorithm is shown in Figure 2 In principle, DGM jointly learns smoother (S) and a dynamics model (D). The smoother model, chosen to be Gaussian process, maps initial condition xo and time to the state distribution (*)))))
and state derivatives distribution ( (t)) reached at that time. The dynamics model, chosen to be neural network, represents an ODE that maps states z(t) to the derivative distribution PD (i))). Both models are evaluated at some training times and all its output distributions collected in the random variables Xs. Xs and Xp. The parameters of these models are then jointly trained using a Wasserstein wiitiiddddttttbbbbaee based objective directly the level of distributions. For more details on every one of these components, we refer to Section 3 There, we introduce all components individually and then present how they interplay. Section
builds known concepts from the literature, which
"ood is available https: httpgs//ggtttb...c/////atttttttt

---

summarize in Section 2 Finally, in Section 4 we present the empirical study of the DGM algorithm,
where we benchmark against the state-of-the-ee uncertainty aware dynamics models.
Background
2.1 Problem Statement
Consider a continuous time dynamical system whose K-dimenssonaa state E RR evolves according to an unknown ordinary differential equation of the form
② = (z).
(1)
Here, f is an arbitrary, unknown function assumed to be locally Lipschitz continuous, to guar- antee existence and uniqueness of trajectories for every initial condition. In our experiment, we initialize the system at M different initial conditions Zm (0), m E (l...., M), and let it evolve to generate M trajectories. Each trajectory is then observed at discrete (but not necessarily uniformly spaced) time-points, where the number of observations (Nm)me can vary from trajectory to trajectory. Thus, a trajectory m is described by its initial condition 2m (0), and the observations Ym == (mm.)) + ne (1.....Nm at times f.m. = [nn,m)ne(1.... "1....N.)" where the additive observa- tion noise En.m is assumed to drawn i.i.d. from a zero mean Gaussian, whose covariance given by E. = diag(oî. . k We denote by D the dataset, consisting of M initial conditions 20,, (0), observation times fm- and observations ym--
To model the unknown dynamical system, we choose a parametric Ansatz & f(x, 0). Depending on the amount of expert knowledge, this parameterization can follow a white-box, gray-box, or black-box methodology (Bohlin 2006). In any case, the parametric form of f is fixed priori (e.g., a neural network), and the key challenge is to infer reasonable distribution over the parameters 0, conditioned on the data D. For later tasks, we are particularly interested in the predictive posterior state distribution
p( (fnow) ) D, tnew Fnww (0)):
i... the posterior distribution of the states starting from a potentially unseen initial condition nnew (0) and evaluated at times fnww. This posterior would then be used by the downstream or prediction tasks
described in the introduction.
2.2 Bayesian Parameter Inference
In the case Bayesian parameter inference, an additional prior p(0) is imposed on the parameters i so that the posterior distribution of Equation (2) can be inferred. Unfortunately, this distribution not analytically tractable for most choices of f(x f), which is especially true when we model f with neural network. Formally, for fixed parameters 9, initial condition (0) and observation time I, the likelihood of an observation y is given by
p(y(t) (0), 0, Eobs) =N ()())((())))) ( f(x((), 0)d7, Eots
Using the fact that all noise realizations are independent, the expression can be used to calculate the likelihood of all observations D. Most otate-of-the-art parameter inference schemes use this fact to create samples Q, of the posterior over parameters p(0) (D) using various Monte Carlo methods.
Given new initial condition (0) and observation time f, these samples 6, can then be turned into samples of the predictive posterior state again by numerically integrating
(t) x00) + 1. f(x((T)
(4)
Clearly, both training (i.c., obtaining the samples θ.) and prediction (i.c., evaluating Equation 4 require integrating the system dynamics f many times. Especially when we model f with neural network, this can be a huge burden, both numerically and computationally (Kelly et al.) 2020).
As alternative approach, we can approximate the posterior p(0)D) with variational inference (Bishop) 2006). However, we run into similar bottlenecks. While optimizing the variational objective, e.g., the ELBO, many integration steps are necessary evaluate the unnormalized posterior. Also, at inference time, to obtain a distribution over state a. (t), we still need to integrate f several times. Furthermore, Dandekar 2222 report poor forecasting performance by the variational approach.
3

---

3 Distributional Gradient Matching
In both the Monte Carlo sampling-based and variational approaches, all information about the dynamical system is stored in the estimates of the system parameters 0. This makes these approaches rather cumbersome: Both for obtaining estimates of 0 and for obtaining the predictive posterior over states, once 0 is found, we need multiple rounds of numerically integrating a potentially complicated (neural) differential equation. We thus have identified two bottlenecks limiting the performance and applicability of these algorithms: namely, numerical integration of f and inference of the system parameters 0. In our proposed algorithm, we avoid both of these bottlenecks by directly working with the posterior distribution in the state space.
To this end, we introduce a probabilistic, differentiable smoother model, that directly maps tuple (t, æ (0)) consisting of a time point t and an initial condition x0)) as input and maps it to the corresponding distribution over x(t). Thus, the smoother directly replaces the costly, numerical integration steps, needed, e.g., to evaluate Equation
Albeit computationally attractive, this approach has one serious drawback. Since the smoother longer explicitly integrates differential equations, there is no guarantee that the obtained smoother model follows any vector field. Thus, the smoother model is strictly more general than the systems described by Equation ( Unlike ODEs, it is able to capture mappings whose underlying functions violate, e.g., Lipschitz or Markovianity properties, which is clearly not desirable. To address this issue, we introduce a regularization term, Laynamies, which ensures that a trajectory predicted by the smoother is encouraged follow some underlying system of the form of Equation (T). The smoother is then trained with the multi-objective loss function
C Latta + A Ldynamics:
(5)
where, Caaa is smoother- dependent loss function that ensures sufficiently accurate data fit, and À a trade-off parameter.
3.1 Regularization by Matching Distributions over Gradients
To ultimately define Laynamies first choose a parametric dynamics model similar to f(x 0) in Equation B3). that maps states to their derivatives. Second, define a set of supporting points T with the corresponding supporting gradients X' as
{{ (tsupp.) Ssupp. (e..
x = ( (
T =
Here, the 1-th element represents the event that the dynamical system's 's derivative at time tupp. is Expp.t after being initialized at time 0 at initial condition x.)p..
Given both the smoother and the dynamics model, we have now two different ways calculate distributions over X' given some data D and supporting points T. First, we can directly leverage the differentiability and global nature of our smoother model extract a distribution ps ps(X') (D, T) from the smoother model. Second, we can first use the smoother to obtain state estimates and then plug these state estimates into the dynamies model, to obtain second distribution PD (D T). Clearly, if the solution proposed by the smoother follows the dynamics, these two distributions should match. Thus, we can regularize the smoother follow the solution of Equation (3) by defining Ldymamics encode the distance between PD (X D) and (x T) to small in some metric. By minimizing the overall loss, we thus match the distributions over the gradients of the smoother and the dynamics model.
3.2 Smoothing jointly over Trajectories with Deep Gaussian Processes
The core of DGM is formed by a smoother model. In principle, the posterior state distribution of Equation 2 could be modeled by any Bayesian regression technique. However, calculating ps T) is generally more involved. Here, the key challenge is evaluating this posterior, which already computationally challenging, e.g., for simple Bayesian neural networks. For Gaussian processes, however, this becomes straightforward, since derivatives of GPs remain GPs (Solak) et al. 2003). Thus, DGM uses a GP smoother. For scalability and simplicity, we keep K different, independent smoothers, one for each state dimension. However, if computational complexity is not a concern, our approach generalizes directly to multi-output Gaussian processes. Below, we focus on the one-dimensional case, for clarity of exposition. For notational compactness, all vectors with

---

superseript should be interpreted as vectors over time in this subsection. For example, the vector x (k) consists of all the k-th elements of the state vectors 2(tn,m), n E {l..... : Nm J. m E {l...... .......,,] M We define Gaussian process with differentiable mean function ulæm (0), bn.m) ) as well as differentiable and positive-definite kernel function KRBF mm (0), (n.m). Φi m (0), by m/ ). Here, the kernel is given by the composition of standard ARD-RBF kernel (Rasmussen 2004) and differentiable feature extractor á parametrized by a deep neural network, as introduced by Wilson cet al. 2016). Following Solak et al. (2003), given fixed Suupp.
now calculate the joint density
of (dsuppp (*)) for each state dimension k. Concatenating vectors accordingly across time and trajectories, let
(R) = th) (z)(0), t)
( (1(0)) (xrupp (0), tsupp)
01
z )(h) (zx(0), t),
(k)
= x(() (xsupp (0), tupp))
supp
(k)) = k (2(1), = )
K (A)
RBF
(k)
z ((1)),
K (b)
82
K RBF (> =supp)
supp
Ōl. 012
Then the joint density of (dsupp- y (a)) can be written as
(K)
(k) Puupp y (k)
N (mi))
xC(0)
1C(0) + ofT )
(A)
(6)
8 Here we denote by कt, the partial derivative with respect to time in the first coordinate, by e the partial derivative with respect time in the second coordinate, and with of the corresponding noise variance of Dobs
Since the conditionals of a joint Gaussian random variable are again Gaussian distributed, ps is again Gaussian, i.e., ps ( (D, T) N duuppp pas: Es) with
= ju + xk()) (K(()) + + (y ul()))
Σg = K (R) xc (K(16) - - (xx())
Here, the index A is used to highlight that this is just the distribution for one state dimension. To obtain the final (A) T). we take the product over all state dimensions k.
To fit our model to the data, we minimize the negative marginal log likelihood of our observations, neglecting purely additive terms (Rasmussen 2004 i.e.,
Laata (D)0) (y()))(()) (KC(0 + oz22
(g))) ( + 2
1
logdet (x(() + ofti) (8)
(k)
Furthermore, the predictive posterior for a new point Tist given time thest and initial condition Tisst (0) has the closed form
Ps (atedl DDx trest: Frest) = N
(k) Frestt Mpoas o opoat.k)
(9)
where
Mpusa n. (k) (Iros((0), tress) + KRBE (Zrest: 2(3)) ((C(K)) + + (K)
(k)
(k)
-p))) (10)
post, & (nross #eost) - xchh (aBB(e (=ies) 2(k)) (X((K) + +oRl) k(hl (K)
(H))
(k)) (11)
3.3 Representing Uncertainty in the Dynamics Model via the Reparametrization Trick
As described at the beginning of this section, a key bottleneck of standard Bayesian approaches is the potentially high dimensionality of the dynamics parameter vector 0. The same is true for our approach. If we were to keep track of the distributions over all parameters of our dynamics model, calculating PD (X (D, T) quickly becomes infeasible.
5

---

However, especially in the case of modeling f with a neural network, the benefits of keeping distributions directly over is unclear due overparametrization. For both the downstream tasks and our training method, we are mainly interested in the distributions in the state space. Usually, the state space significantly lower dimensional compared to the parameter space of Furthermore, since the exact posterior state distributions are generally intractable, they normally have to be approximated anyways with simpler distributions for downstream tasks (Schulman et al., 2015 Houthooft et al. 2016; Berkenkamp et al. 2017). Thus, we change the parametrization of our dynamics model as follows. Instead of working directly with ((f) = (t), 0) and keeping a distribution over 0, we model uncertainty directly on the level of the vector field as
a(t) f(x(t), +)) (z)(t
(12)
where c ~ (V, Ix) is drawn once per rollout (i.e., fixed within a trajectory) and Ep is a state- dependent and positive semi-definite matrix parametrized by a neural network. Here, & are the parameters of the new dynamics model, consisting of both the original parameters O and the weights of the neural network parametrizing Ep To keep the number of parameters reasonable, we employ a weight sharing scheme, detailed in Appendix B
In spirit, this modeling paradigm is very closely related to standard Bayesian training of NODEs. In both cases, the random distributions capture distribution over set of deterministic, ordinary differential equations. This should be seen in stark contrast to stochastic differential equations, where the randomness in the state space, i.e., diffusion, is modeled with stochastic process. In comparison to (12). the latter a time-varying disturbance added to the vector field. In that sense, our model still captures the epistemic uncertainty about our system dynamies, while SDE model captures the intrinsic process noise, i.e., aleatoric uncertainty. While this reparametrization does not allow us to directly calculate pn (D, T), we obtain a Gaussian distribution for the marginals pp (asppp (supp)
To retrieve (D T), we use the smoother model's 's predictive state posterior to obtain
pD(A'D,,T) = (dt Tauppl (D, T dx
(13)
/
po Tupp Ps (2 (Ipp T D) spp-
(14)
3.4 Comparing Gradient Distributions via the Wasserstein Distance
To compare and eventually match ( (D, and (A T), we propose to use the Wasserstein distance (Kantorovich) 1939), since it allows for an analytic, closed-fone representation, and since outperforms similar measures (like forward, backward and symmetric KI. divergence) in our exploratory experiments. The squared type-2
Wasserstein distance gives rise to the term
w? that we will later use to regularize the smoothing process. To render the calculation of this reg- ularization term computationally feasible, we introduce two approximations. First, observe that an exact calculation of the expectation in Equation (5) requires mapping multivariate Gaussian through the deterministic neural networks parametrizing f and Ep in Equation (1)). To avoid complex sampling schemes, carry out eettainty-qqiivllenn approximation of the expectation, that is, we evaluate the dynamics model on the posterior smoother mean μss supp. As a result of this approximation, observe that both ( T) and ps (A')D, T) become Gaussians. However, the covariance structure of these matrices is very different. Since we use independent GPs for different state dimensions, the smoother only models the covariance between the state values within the same dimension, across different time points. Furthermore, since e, the random variable that captures the randomness of the dynamics across all time-points, is only K-dimensiional, the covariance of pn will be degenerate. Thus, we not match the distributions directly, but instead match the marginals of each state coordinate at each time point independently at the different supporting time points. Hence,
w3 (ps() T), Eu -Par (asy (D.T) (pp (aupp )]
(15)

---

using first marginalization and then the certainty equivalence, Equation
15 reduces to
|
w3 (2sp:222
w2 (tupsupp(tsuppp)) s)sppp.)D.. T), PD ( supp (fsupp [2, 7)
(a(Ae)
K
W2
(k) ps (uupp (fsup. (D, T). (i(k) Suup (tsupp.) e) ( (usp supp)
(16)
Conveniently, the Wasserstein distance can now be calculated analytically, since for two one- dimensional Gaussians a ~ N(pa: o²) and b ~ Nus, of), we have W3la, b] = (Ma pb)2
(o. - -0,)².
3.5 Final Loss Function
As explained in the previous paragraphs, distributional gradient matching trains a smoother regularized by a dynamies model. Both the parameters of the smoother yp, consisting of the trainable parameters of the GP prior mean μe, the feature map 6, and the kernel IC, and the parameters of the dynamics model & are trained concurrently, using the same loss function. This loss consists of two terms, of which the regularization term was already described in Equation (16). While this term ensures that the smoother follows the dynamics, we need a second term ensuring that the smoother also follows the data. To this end, we follow standard GP regression literature, where it is common to learn the GP hyperparameters by maximizing the marginal log likelihood of the observations, ic. Ldata (Rasmussen 2004). Combining these terms, we obtain the final objective
K
((p,,)) = Ldata A
w2 (p (D, PD (i((t) (tsupp, l (
A=
This loss function is a multi-criteria objective, where fitting the data (via the smoother) and iden- tifying the dynamics model (by matching the marginals) regularize each other. In our preliminary experiments, we found the objective to quite robust w.r.t. different choices of X. the interest of
accounting only for the possibility of having different numbers of supporting points and observations. One special case worth mentioning is À - 0, which corresponds to conventional sequential smoothing, where the second part would be used for identification in a second step. as proposed by Pillonetto and De Nico lao (201)). However, as can be seen in Figure I the smoother fails to properly identify the system without any knowledge about the dynamics and thus fails to provide meaningful state derivative estimates. Thus, especially in the case of sparse observations, joint training is strictly superior.
simplicity, we thus set it in all our experiments in Section to a default value of À = |X
DD
In final form, unlike its pure Bayesian counterparts, DGM does not require any prior knowledge about the system dynamics. Nevertheless, if some prior knowledge is available, one could add an additional, additive term log(p())) to L(p,, 1b). should be noted however that this was not done in any of experiments, and excellent performance can be achieved without.
4 Experiments
We now compare DGM against state-of-the-art methods. In a first experiment, we demonstrate the effects of overparametrized, simple dynamics model on the performance of DGM well as traditional, MC-hased algorithms SGL.D (Stochastic Gradient Lengevin Dynamics, (Welling and Teh 201Th) and SGHMC (Stochastic Gradient Hamiltonian Monte Carlo, (Chen et al. 2014)). We select our baselines based on the results of Dandekar et al. (2021). who demonstrate that both a variational approach and NUTS (No U-Turn Sampler, Hoffman and Gelman (2014) are inferior these two. Subsequently, we will investigate and benchmark the ability of DGM to correctly identify neural dynamics models and to generalize across different initial conditions. Since SGLD and SGHMC reach their computational limits in the generalization experiments, we compare against Neural ODE Processes (NDP). Lastly, we will conclude by demonstrating the necessity of all of its components. For all comparisons, we use the julia implementations of SGLD and SGHMC provided by Dandekar al. (202T), the pytorch implementation of NDP provided by Norcliffe et al. (202T), and our own JAX (Bradbury et al. 2018) implementation of DGM.
7

---

4.1 Setup
We use known parametric systems from the literature to generate simulated, noisy trajectories. For these benchmarks, chaotic Lorenz (LO) system, a four-dimensional double pendulum (DP) and a twelve-dimensional quadrocopter (QU) model. For all systems, the exact equations and ground truth parameters are provided in the Appendix A For each system, we create two different data sets. In the first, we include just one densely observed trajectory, taking the computational limitations of the benchmarks into consideration. In the second, we include many, but sparsely observed trajectories (5 for LV and DP, 10 for LO, 15 for QU). This setting aims to study generalization over different initial conditions.
the two-dimensional Lotka Volterra (LV) system, the three-dimensional
4.2 Metric
We use the log likelihood as a metric to compare the accuracy of our probabilistic models. In the 1 I-trajectory setting, we take a grid of 100 time points equidistantly on the training trajectory. We then calculate the ground truth and evaluate its likelihood using the predictive distributions of our models. When testing for generalization, we repeat the same procedure for unseen initial conditions.
4.3 Effects Overparametrization
We first study three-dimensional, linear system of the form (t() = (t( where A a randomly chosen matrix with one stable and marginally stable eigenvalues. For the dynam- ics model, we choose a linear Ansatz f(x, 0) (t) where B is parametrized as the product of multiple matrices. The di- mension of the matrices of each factorization are captured in a string of integers of the form (3, a1..... aJ. 3). For example, (3, 3) corresponds to B being just one matrix, while (3, 6, 6, 3) corresponds B B, B2 with B, E 33x6 B2 R6x6 and B3 E TR633 All of these models can be interpreted as linear neural networks, forming simple case of the nonparametric systems we study later. Unlike general neural networks, the expressiveness the Ansatz is independent the number of pa- rameters, allowing us isolate the effects of overparametrization. In Figure 3 we show the mean and standard deviation of the log likelihood of the ground truth over 10 different noise realizations. The exact procedure for noise realization is described in the appendix. Appendix C). While SGLD runs into numerical issues after medium model complexity, the performance of SGHMC continuously disintegrates, while DGM unaffected. This foreshadows the results of the next two experiments, where we observe that the MC-based approaches are not suitable for the more complicated settings.
4.0
Likelhheed
2.0
1.5
1.-
Model
Figure 3: SGLD does not converge for strongly over- parametrized models. the perfor- mance of SGHMC deteriorates. DGM is not noticeably affected.
Table 1: Log likelihood and prediction times of 100 ground truth sample points, with mean and standard deviation taken over 10 independent noise realizations, for neural ODEs trained on a single, densely sampled trajectory.
Log Likelihood
Prediction time [ms]
DGM
SCHMC
SCLD
DGM
SCHMC
SCLD
LV 1
1.96 0.21
1.36 +: 0.0693
1.03 + 0.0581
0.68
0.04
14.98 0.23
14.59 0.15
LO
-0.57 + 0.11
-3.02 0.158
-2.67 + 0.367
0.99
0.05
98.93. 5.79
105.03 12.22
DP
+ 2.13+0.14
1.88 + 0.0506
1.85 + 0.0501
1.31 + 0.05
10.60 0.221
11.34 0.76
QU
0.64 0.07
-5.00 + 1.36
NaN
3.76 = 0.12
24.68 6.55
NaN
4.4 Single Trajectory Benchmarks
In Table we evaluate the log-likes ligelikeldhodd of the ground truth for the four benchmark systems, obtained when learning these systems using neural ODE as a dynamics model (for more details, see appendix Clearly, DGM performs the best on all systems, even though we supplied both SGLD and SGHMC with very strong priors and fine-tuned them with an extensive hyperparameter sweep (see Appendix for more details). Despite this effort, we failed to get SGLD to work on Quadrocopter 1, where it always returned NaNs. This is in stark contrast to DGM, which performs reliably without any pre-training priors.
8

---

4.5
Prediction speed
To evaluate prediction speed, we consider the task of predicting 100 points on a previously unseen trajectory. To obtain fair comparison, all algorithms' prediction routines were implemented in JAX (Bradbury et al. 2018). Furthermore, while we used 1000 MC samples when evaluating the predictive posterior for the log likelihood guarantee maximal accuracy, we only used 200 samples in Table | Here, 200 was chosen as a minimal sample size guaranteeing reasonable accuracy. following
preliminary experiment visualized in Appendix C orders of magnitudes faster, can be seen in Table This further illustrates the advantage of relying on a smoother instead of costly, numerical integration to obtain predictive posteriors in the state space.
Nevertheless, the predictions of DGM are 1-2
2.5
Table 2:
Log likelihood of 100 ground truth sample points, with mean and covariance taken over 10 inde- pendent noise realizations, for neural ODEs trained on multiple, sparsely sampled trajectory. The number fol- lowing the system name denotes the number of trajecto- ries in the training set.
Log Likellhood
DOM train
1.5
NDP test
1.0
Log Likelihood
DGM
NDP
10
12
14
16
18
20
IV 100
1.81 + 0.08
0.62 - 0.27
grid polnts nts dimensian
LO 125
-2.18 + 0.76
2288 1: 0.05
Figure 4: Log likelihood of the ground truth for Lotka Volterra for increasing number of trajectories with 5 observations each.
DP 100
1.86 1 0.05
0.88 - 0.05
QU
-0.54 0.36
-0.91 0.07
4.6 Multi- Talti-Trajectory Benchmarks
Next, we take a set of trajectories starting on an equidistant grid of the initial conditions. Each trajectory is then observed at 5 equidistant observation times for LV and DP, and 10 equidistant observation times for the chaotic Lorenz and more complicated Quadrocopter. We test generalization by randomly sampling new initial condition and evaluating the negative log likelihood of the ground truth at 100 equidistant time points. In Table 2 we compare the generalization performance of DGM against NDP, since despite serious tuning efforts, the MC methods failed to produce meaningful results in this setting. DGM clearly outperforms NDP, fact which is further exemplified in Figure 4 There, we show the test log likeliood for Lotka Volterra trained on an increasing set of trajectories. Even though the time grid is fixed and we only decrease the distance between initial condition samples, the dynamics model helps the smoother to generalize across time as well. In stark contrast, NDP fails to improve with increasing data after initial jump.
4.7 Ablation study
We next study the importance of different elements of our approach via an ablation study the Lorenz 125 dataset, shown in Figure [ Comparing the two rows, we see that joint smoothing across trajectories is essential to transfer knowledge between different training trajectories. Similarly, comparing the two columns, we see that the dynamics model enables the smoother to reduce its uncertainty in between observation points.
4.8 Computational Requirements
For the one trajectory setting, all DGM related experiments were on Nvidia RTX 2080 Ti, where the longest ones took 15 minutes. The comparison methods were given 24h, on Intel Xeon Gold 6140 CPUs. For the multi-trajectory setting, used Nvidia Titan RTX, where all experiments finished in less than hours. A more detailed run time compilation can be found in Appendix B Using careful implementation, the run time of DGM scales linearly in the number of dimensions K. However, since we use an accurate RBF kernel for all our experiments reported in this section, we have cubic run time complexity in E Nm. In principle, this can be alleviated by deploying standard feature approximation methods (Rahimi et al.) 2007 Liu et al. 2020). While this a well known fact, we nevertheless refer the interested reader more detailed discussion of the subject in Appendix D

---

5 Related work
5.1 Bayesian Parameter Inference with Gaussian Processes
The idea of matching gradients of a (spline-based) smoother and a dynamics model goes back to the work of Varah (1822). For GPs, this idea is introduced by Calderhead et al. (2009), who first fit a GP to the data and then match the parameters of the dynamics. Dondelinger et al. (2013) introduce concurrent training, while Garbach et al. (0117) introduce an efficient variational inference procedure for systems with a locally-linear parametric form. All these works claim to match the distributions of the gradients of the smoother and dynamics models, by relying on a product of experts heuristics. However, WWen et al.] (2019) demonstrate that this product of experts in fact leads to statistical independence between the observations and the dynamics parameters, and that these algorithms essentially match point estimates of the gradients instead. Thus, DGM is the first algorithm to actually match gradients on the level of distributions for ODEs. In the context of stochastic differential equations (SDEs) with constant diffusion terms, terms,[Abbati et al. (2019) deploy MMD and GANs to match their gradient distributions. However, should be noted that their algorithm treats the parameters of the dynamics model deterministically and thus, they can not provide the epistemic uncertainty estimates that we seek here. Note that our work is not related to the growing literature investigating SDE approximations of Bayesian Neural ODEs in the context of classification (Xu et al. (02T). Similarly to Chen et al. (2018), these works emphasize learning a terminal state of the ODE used for other downstream tasks.
5.2 Gaussian Processes with Operator Constraints
Gradient matching approaches mainly use the smoother as a proxy to infer dynamics parameters. This is in stark contrast to our work, where we treat the smoother as the main model used for prediction. While the regularizing properties of the dynamics on the smoother are explored by Wenk et al. (2020), Jidling al. (2017) introduce algorithm to incorporate linear operator constraints directly on the kernel level. Unlike in our work, they can provide strong guarantees that the posterior always follows these constraints. However, it remains unclear how to generalize their approach to the case of complex, nonlinear operators, potentially parametrized by neural dynamics models.
5.3 Other Related Approaches
In some sense, the smoother is mimicking a probabilistic numerical integration step, but without explicitly integrating. In spirit, this approach is similar to the solution networks used in the context of PDES, as presented by Raissi et al. (219). which however typically disregard uncertainty. In the context of classical ODE parameter inference. Kersting et al. (020)) deploy a GP to directly mimic a numerical integrator in a probabilistic, differentiable manner. Albeit promising in a classical, parametric ODE setting, it remains unclear how these methods can be scaled up, as there is still the numerical integration bottleneck. Unrelated to their work, Ghosh et al. (2021) present a variational inference scheme in the same, classical ODE setting. However, they still keep distributions over all weights of the neural network (Norcliffe et al.] 2021). A similar approach is investigated by Dandekar et al. (02T), who found it to be inferior to the MC methods we use as a benchmark. Variational inference was previously employed by Yildiz et al. (2019) in the context of latent neural ODEs parametrized by a Bayesian neural network, but their work mainly focuses on dimensionality reduction. Nevertheless, their work inspired model called Neural ODE Processes by Norcliffe et al. (202T). This work is similar to ours in the sense that avoids keeping distributions over network weights and models an ensemble of deterministic ODEs via a global context variable. Consequently, we use it as a benchmark in Section showing that it does not properly capture epistemic uncertainty in low data setting, which might be problematic for downstream tasks like reinforcement learning.
6 Conclusion
In this work, we introduced a novel, GP-based collocation method, that matches gradients of a smoother and dynamics model on the distribution level using a Wasserstein loss. Through careful parametrization of the dynamics model, we manage to train complicated, neural ODE models, where state of the art methods struggle. We then demonstrate that these models are able to accurately predict unseen trajectories, while capturing epistemic uncertainty relevant for downstream tasks. In future work, we are excited to see how our training regime can be leveraged in the context of active learning of Bayesian neural ordinary differential equation for continuous- time reinforcement learning.
10

---

Acknowledgments
This research was supported by the Max Planck ETH Center for Learning Systems. This project has received funding from the European Research Council (ERC) under the European Union's 's Horizon 2020 research and innovation programme grant agreement No 815943 as well as from the Swiss National Science Foundation under NCCR Automation, grant agreement 51NF40 180545.
References
Abbati, G., Wenk, P., Osborne, M. A., Krause, A., Schölkopf, B., and Bauer, S. (2019). Ares and mars adversarial and mmd-minimizing regression sdes. International Conference on Machine Learning, pages 1-10. PMLR.
Angelis, E., Wenk, P., Schölkopf, B., Bauer, and Krause, A. (2020). Sleipnir: Deterministic and provably accurate feature expansion for gaussian process regression with derivatives. arXiv preprint arXiv:2003.02658.
Berkenkamp. F., Turchetta, M., Schoellig. A., and Krause, A. (2017). Safe model-based reinforcement learning with stability guarantees. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
Bezanson, J., Edelman, A., Karpinski, S., and Shah, V. B. (2017). Julia: fresh approach to numerical computing. STAM Review, 59(1):65-98.
Bishop, C. M. (2006). Patern Recognition and Machine Learning (Information Science and Statistics). Springer- Berlin, Heidelberg.
Bohlin, T. P. (2006). Practical grey-box process identification: theory and applications. Springer Science & Business Media.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke,
A., VanderPlas, J., Wanderman-Milnn S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs.
Brunton, S. L., Proctor, J. L., and Kutz, J. N. (2016). Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 1 113(15):3932-393)
Calderhead, B., Girolami, M., and Lawrence, (2009). Accelerating bayesian inference over nonlinear differential equations with gaussian processes. In Advances in neural information processing systems, pages 217-224. Citeseer.
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. (2018). Neural ordinary differential equations. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.
Chen, T., Fox, E., and Guestrin, C. (2014). Stochastic gradient hamiltonian monte carlo. In International conference machine learning, pages 1683-1691. PMLR.
Dandekar, R., Chung. Dixit, V., Tarek, M., Garcia-Vann Valadez, A., Vemula, V., and Rackauckas, C. (2021). Bayesian neural ordinary differential equations. Symposium on Principles of Programming Languages, POPL.
Deisenroth, M. and Rasmussen, C. E. (2011). Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference machine learning (ICML-1I), pages 465-472. Citeseer.
Dondelinger, F., Husmeier, D., Rogers, S., and Filippone, M. M. (2013), Ode parameter inference using adaptive gradient matching with gaussian processes. In Artificial intelligence and statistics, pages 216-228. PMLR.

---

Fanfarillo, A., Roozitalab, B., Hu, W., and Cervone, G. (2021). Probabilistic forecasting using deep generative models. GeoInformatica, 25(1):127-147.
Ghosh, S., Birrell, P., and De Angelis, D. (2021). Variational inference for nonlinear ordinary differential equations. In International Conference Artificial Intelligence and Stafistics, pages 2719-2727. PMLR.
Gorbach, N. S., Bauer, S., and Buhmann, J. M. (2017). Scalable variational inference for dynamical systems. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
Greydanus, S., Dzamba, M., and Yosinski, J. (2019). Hamiltonian neural networks. In Wallach, H., Larochelle, H., Beygelzimer. A., d'Alché Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
Hjalmarsson, (2005). From experiment design to closed-loop control. Automatica, 41(3):393-443
Hoffman, M. D. and Gelman, A. (2014). The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1993-1623.
Houthooft, R., Chen, X., Chen, X., Duan, Y., Schulman, J., De Turck, F,, and Abbeel, P. (2016). Vime: Variational information maximizing exploration. In Lee, D., Sugiyama, M., Luxburg, U., Guyon, I., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc.
Ingalls, B. P. (2013). Mathematical modeling in systems biology: an introduction. MIT press.
Jidling, C., Wahlström, N., Wills, A., and Schön, T. B. (2017). Linearly constrained gaussian processes. In Guyon, I., Luxburg. U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
Jones, D. S., Plank, M., and Sleeman, B. D. (2009). Differential equations and mathematical biology. CRC press.
Kantorovich, L. V. (1939). The mathematical method of production planning and organization. Management Science, 6(4):363-42 422.
Karaman, S. and Frazzoli, E. (2011). Sampling-based algorithms for optimal motion planning. The international journal of robotics research, 30(7):846-894.
Kelly, J., Bettencourt, J., Johnson, M. J., and Duvenaud, D. K. (2020). Learning differential equations that easy solve. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H., editors, Advances in Neural Information Processing Systems, volume 33, pages 4370-4380. Curran Associates, Inc.
Kersting, H., Krämer, N., Schiegg, M., Daniel, C., Tiemann, M., and Hennig, P. (2020). Differentiable likelihoods for fast inversion of" lilliiihood--fii free dynamical systems. In International Conference on Machine Learning, pages 5198-5208, PMLR.
Kidger, P., Chen, R. T., and Lyons, T. (2020). hey, that's not an ode": Faster ode adjoints with 12 lines of code. arXiv preprint arXiv:2009.0944 09457..
Lalale S. M. (2006). Planning algorithms. Cambridge university press.
Liu, H., Ong. Y.-.., Shen, X., and Cai, J. (2020). When gaussian process meets big data: review of scalable gps. IEEE transactions on neural networks and learning systems, 31 1):4405-4423,
Ljung, L. (1998). System Identification, pages 163-173. Birkhäuser Boston, Boston, MA.
Ljung, (1999). Model validation and model error modeling. Linköping University Electronic
Press.
12

---

Norcliffe, A., Bodnar, C.. Day, B.. Moss, J., and Liò, P. (2021). Neural {ode} processes. In International Conference on Learning Representations.
Ortega, R., Praly, L., Aranovskiy, S., Yi, B., and Zhang, W. (2018). On dynamic regressor extension and mixing parameter estimators: Two luenberger observers interpretations. Automatica, 95:548- 551.
Ortega, R., Van Der Schaft, A., Maschke, B., and Escobar, G. (2002). Interconnection and damping
assignment passivity-baan bbased control of port-controlled hamiltonian systems. Automatica, 38(4):585 596.
Pillonetto, and De Nicolao, G. (2010). new kernel-based approach for linear system identification.
Automatica, 46(1))81-99)
Pillonetto, G., Dinuzzo, F., Chen, T., De Nicolao, G., and Ljung, L. (2014). Kernel methods in system identification, machine learning and function estimation: A survey. Automatica, 50(3):657-66
Rackauckas, C., Ma, Y., Martensen, J., Warner, C., Zubov, K., Supekar, R., Skinner, D., Ramadhan, A., and Edelman, A. (2020). Universal differential equations for scientific machine learning. arXiv preprint #xXiv:2001.04355
Rahimi, A., Recht, B., et al. (2007). Random features for large-scale kernel machines. NIPS, volume 3, page 5, Citeseer.
Raissi, M., Perdikaris, P., and Karniadakis, G. (2019). Physics-informet neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707.
Rasmussen, C. E. (2004). Gaussian Processes in Machine Learning, pages 63-71. Springer Berlin Heidelberg, Berlin, Heidelberg.
Schulman, J., Levine, S., Abbeel, P., Jordan, M., and Moritz, (2015). Trust region policy optimiza- tion. In International conference machine learning, pages 1889-1897. PMLR.
Solak, E., Murray-Smith, R., Leithead, W. E.,, Leith, D. J., and Rasmussen, observations in gaussian process models of dynamic systems.
(2003). Derivative
Spong, M. W., Hutchinson, S., Vidyasagar, M., al. (2006). Robot modeling and control, volume 3. wiley New York.
Srinivas, N., Krause, A., Kakade, S., and Seeger. M. (2010). Gaussian process optimization in the bandit setting: no regret and experimental design. In Proceedings of the 27th International Conference International Conference on Machine Learning, pages 1015-1022.
Varah, M. (1982). A spline least squares method for numerical parameter estimation differential equations. SIAM Journal Scientific and Statistical Computing, 3(1):28-466
Welling. M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 68 681-688. Citeseer.
Wenk, P., Abbati, G., Osborne, M. A., Schölkopf, B., Krause, A., and Bauer, S. (2020). Odin: Ode-informed regression for parameter and state inference in time-continuous dynamical systems. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 6364-6371.
Wenk, P., Gotovos, A.. Bauer, S., Gorbach, N. S., Krause, A., and Buhmann, J. M. (2019). Fast gaussian process based gradient matching for parameter identification in systems of nonlinear odes. In The 22nd International Conference Artificial Intelligence and Statistics, pages 1351-1360 PMLR.
Wensing, P. M., Kim, S., and Slotine, J.J. E. (2017). Linear matrix inequalities for physically consistent inertial parameter identification: A statistical perspective on the mass distribution. IEFE Robotics and Automation Letters. 3(1):60-67.
13

---

Wilson, A. G., Hu. Z., Salakhutdinov. R., and Xing, E. P. (2016). Deep kernel learning. In Artificial intelligence and statistics, pages 370-378. PMLR.
Wittenburg, (2013). Dynamics of systems of rigid bodies, volume 33, Springer-Vertee Verlag.
Xu, W., Chen, R. T., Li, X., and Duvenaud, D. (2021). Infinitely deep bayesian neural networks with stochastic differential equations. arXiv preprint arXiv:2102.06559
Yildiz, C., Heinonen, M., and Lähdesmäki, H. (2019). Ode2vae: Deep generative second order odes with bayesian neural networks. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Bus F, Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
Zhuang. J., Dvornek, N., Li, X., Tatikonda, S., Papademetris, X., and Duncan, J. (2020). Adaptive checkpoint adjoint method for gradient estimation neural ode. International Conference on Machine Learning, pages 11639-11649 1169.. PMLR.
Zhuang, J., Dvonek, N. C.,, Tatikonda, S., and Duncan, J. S. (2021). Mali: A memory efficient and reverse accurate integrator for neural odes. arXiv preprint arXiv:2102.0466.
14


Reviews:
[{"review_id": "qvc-IDiHqYf", "reviewer": "Reviewer_3nDp", "summary": "The paper proposes to skip the integration of continuous-time models and directly learn a model that predicts the posterior of the state conditioned on the initial state and the desired time. To learn such a model with 'good' uncertainty estimation the authors mix GP's, neural networks and the wasserstein metric (I did not understand how the authors combine these approaches together). Within the experiments the paper shows that the log-likelihood is better for the proposed model on uncontrolled toy systems, e.g. Lotka Volterra, chaotic Lorenz, double pendulum and quadcopter.\n", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First of all, I am an expert on learning dynamics models but not an expert for Bayesian models. My main problem with the paper is that I cannot understand the proposed approach. The authors described the problem and their approach. However, once the authors mix all concepts of bayesian networks, GP's, Wasserstein distance and stir it all together. I personally do not have any understanding how all of this plays together. Therefore, the paper needs a major revision that restructures the writing such that the reader gets a good understanding of the approach. A good approach is to provide the intuition of the proposed approach and only afterwards deep dive into the details. However, the authors skip the intuitive approach and directly deep dive into the theory and lose the reader in the process. Besides the writing, the performed experiments are not really impressive and only try a few toy domains. Even for the toy domains the experimental evaluation does not provide much insights except my number is bigger than yours. It would be beneficial to incorporate more qualitative evaluations and explain to the reader what to look out for. Furthermore, the authors motivate their approach using control and RL but then only model uncontrolled dynamics and never scale the proposed approach to interesting control problems. How does the model work for high-dimensional systems with contacts? Therefore, it would be necessary to evaluate the proposed approach on more complex domains such as walker or humanoid.\n\nAs the writing and structure of the paper needs a major revision and the performed experiments do not evaluate whether this model actually improves learned models for control, I cannot recommend the paper for acceptance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "dlf0NaEeor", "reviewer": "Reviewer_zfk4", "summary": "The paper tackles the problem of learning nonlinear dynamical systems from historical trajectories, while estimating predictive uncertainties. They propose a new method based on a Gaussian process smoother which is trained with maximum likelihood to fit the data of all training trajectories jointly, while constrained to follow in some sense the underlying dynamics. The latter is performed thanks to a regularizer forcing the smoother’s gradients distribution to match the  distribution of gradients of the dynamics model over a fixed set of support points. The Wasserstein-2 distance is approximated to evaluate the distribution discrepancies. Then, instead of deriving uncertainty estimates over the predicted parameters, as done in standard Bayesian methodology, uncertainty is directly estimated over the predicted states. For this, a Gaussian distribution is used to model  probability of the states derivatives conditioned on the states. Thanks to these tricks and modeling choices, the authors avoid to have to integrate the complex dynamics model, unlike previous Monte Carlo methods. The proposed approach is show-cased and compared to sampling-based methods in several experiments with simulated data, including overparametrized linear dynamics, as well as single and multi-trajectory cases of 4 well-known dynamical systems. An ablation study is also presented showing the benefits of joint smoothing and gradients distribution matching regularization on one of the four simulated datasets.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Post-rebuttal additions\n------------------------------\n\nI would like to thank the authors for their detailed answers, which address all concerns I had.\n\n-----------------------------------------------------------------------------------\n\nOriginality\n--------------\n\nThe paper seems original to me and the related work seems quite well referenced to me.\n\nSignificance\n----------------\n\nThe proposed approach is very complex, but not more than existing ones, so I believe it could be a good alternative to continuous-time nonlinear system identification requiring uncertainty estimates. It solves indeed a real bottleneck of existing approches when used with NODE models (numerical integration). Also, instead of delivering parameter uncertainties as previous methods, is focuses on direct estimation of state uncertainties, which is what matters at the end of the day. Hence I believe that the work present could be quite significant in the field, both for practitioners and for researchers who could build on these ideas.\n\nQuality\n----------\n\nOverall quality of the explanations and experiments is quite good, as the method was show cased in varied settings and compared to relevant state-of-the-art alternatives. Also, the code seems clean (+1 for reproducibility). I do have some remarks though concerning some points that do not seem explained in the paper, as well as a few questions listed below:\n\n1. While it is said in section 2 that trajectories can be sampled at non-equality spaced times, all simulated experiments use uniform time grids, which is a shame.\n2. How are the supporting points chosen and how do they influence the estimates’ quality? For example, having a set of supporting points not well spread or too small should lead to degraded performance should it not?\n3. In equation 12 and line 196, it is said that dynamics noise $epsilon$ is drawn one per trajectory, while in the problem statement, observations $y$ are supposed to have i.i.d additive noise (one per sample in each trajectory). While I realize these quantities are not the same (one is in the states and the other in their derivatives), I wonder why make these different modeling choices and not also have a different disturbance realization for each instant in the dynamics model.\n4. Concerning the ablation study, you say on line 330 that by comparing the rows of Figure 1 it can be seen that joint smoothing of multiple trajectories helps. While I do see it clearly when comparing rows of the second column (i.e. with the dynamics regularizer), I see the opposite when looking at the first column. Namely, joint smoothing seems to lead a noisier and less accurate predicted trajectory, with a larger confidence region. As this is very curious, since it contradict the text, I think it should be commented and also that the sentence should be clarified (by saying that improvement is visible on the second column for example).\n5. How many training and testing trajectories are used to compute results for figure 2?\n6. Likewise, while it is said that 5 to 10 points per trajectory are used to train the models in the multi-trajectory experiment, how many trajectories are used? Is that the meaning of the numbers next to the systems considered in the first column of Table 2? If that’s the case, please add it to the legend or paragraph text.\n7. Could you please comment the fact that SGLD and SGHMC methods outperform DGM for configurations 3.3 and 3.6.3 on figure 2?\n8. Could you please develop on the reason SGLD failed on the quadcopter single trajectory experiment, leading to NaNs? Did you try to understand what these NaNs meant? \n\nClarity\n---------\n\n1. Although the paper is well written, the proposed method is very complicated and difficult to grasp (I had to read the paper many times before I started to get it). I cannot say it is not well explained though, but it might help to put more structure in section 3 (replace paragraphs by subsections + parapgraphs for example). Also, adding a diagram summarizing how the many elements of the method interact with each other could help a lot I believe.\n2. Furthermore, I found some explanation was missing concerning the certainty-equivalence approximation of the Wasserstein distance. Maybe add that the expectation you are approximating is the one in (15) and that this approximation is carried in the second line go (16).\n3. In the paragraph Final loss (page 6), the first data fitting term of the loss is denoted quite differently than in equation (8). It took me some time to realize they were the same, so I think some referencing and more consistent notation could help with the clarity here.\n4. In line 285, you denote the model used to fit the dynamics with the same notation as in line 282, where you explain the real dynamics that generated the data. I find this confusing and believe that less ambiguous notation here would help with the clarity (by using for example the dynamics function $f$ introduced in (1)).\n5. I could not find any introduction of the abbreviations SGLD and SGHMC.\n6. On line 164, I believe $z_i:=$ should not be there, as it is defined differently in the text and in equations below line 166.\n7. Sentence at lines 119-120 has a problem as the word maps is repeated twice.\n8. Line 244: another -> each other", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I can't see any societal impact related to this publication.\n\nAlso, the authors don't seem to have addressed the limitations of the proposed approach.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "Ym8-p6-lqXA", "reviewer": "Reviewer_u7cP", "summary": "This paper presents a novel approach to model and capture the uncertainty of dynamical systems by learning from trajectory data. The method is composed of a Gaussian process smoother model, used for predictions, and a neural dynamics model, used for training. The novelty in the method comes from the use of the distributions over time-derivatives of the state vectors (via a Wasserstein-distance penalty between the distributions from the smoother and the dynamics model) during the training process. Basing the approach on integrals over the state space, instead of over the model parameters space, allows for computational efficiency and more robustness when compared to previous sampling-based approaches, as evidenced by experiments.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is mostly well written, providing a good motivation, an appropriate use of citation and interesting insights in the derivation of the method. However, there are a few issues, especially regarding methodological details, which make it hard to understand the method and to assess its soundness.\n\nMajor issues:\n- The description of the methodology is somewhat confusing in a few important aspects.\nFor instance, how does the set of \"supporting gradients\" play a role in the estimation of these two distributions? Are they random variables to be inferred as part of $\\dot{\\mathcal{X}}$? Or are they additional observations which will be used somehow in the calculation of the losses?\n- Regarding $\\mathbf{z}$ in Eq. 6, I get the idea, but it's a bit confusing at first-reading to use both $\\mathbf{z}$ and the $(\\mathbf{x}, t)$-tuples as inputs to the GP when the feature map \"phi\" is supposed to be part of the kernel. It'd be better to consider $k := k \\circ \\phi$ (composition), since the composition of a kernel with an input map is still a kernel, and then represent the inputs as simply $(\\mathbf{x},t)$, without the extra $\\mathbf{z}$ notation.\n- In line 196, it is mentioned that the (instant) state noise vector in Eq. 12 is drawn only once per rollout. Does it mean that the noise term epsilon is drawn once and repeated for all time points along a trajectory? If so, the state noise is not independent, but actually correlated across time and that would possibly not allow the model in Eq. 12 to properly capture uncertainty in the state transitions.\n- I'm not sure about the effects of this approximation in (14), since later on the difference between $p_D$ and $p_S$ will be minimised. Having $p_D$ as a function of $p_S$ might have negative side effects. Any ideas on the drawbacks?\n- The certainty-equivalence approximation (see line 220) discards the uncertainty captured by the GP model in the expectation in Eq. 15. Why not using a quadrature approximation to compute Eq. 15, like the unscented transform, or anything else which could capture the states covariance matrix on the inputs?\n- In the experiments with SGLD and SGHMC, how is the likelihood for these sampling-based methods computed? These methods don't directly produce a probability density, only empirical approximations.\n- The experimental evaluation could be more complete if they included comparisons against methods which encode (inexact) prior knowledge about the dynamics, such as simulator-based inference approaches [e.g., A, B, C, below]. The relationship to these methods is also not discussed in the related work section. In my view, the proposed approach could be combined with simulator-based models by using a physics-informed mean function for the GP smoother, which is not explored/discussed in the paper.\n\nMinor issues:\n- Parametric vs. Non-parametric: Neural networks are \"parametric\" models, though in the main paper and in the appendix they are referred to as \"non-parametric\" models. The only non-parametric model in this work is the GP.\n- The equation for $\\mathbf{z}_i$ in line 164 is missing $\\phi$.\n- I couldn't find details of this weight sharing scheme in Sec. 4, despite the mention in line 200.\n- Preliminary experiments: Are these results available? If not, it'd be better to add them to the appendix or give more details/references to back it up. For example, the Wasserstein GANs paper provides some experiments and theoretical justification on why the Wasserstein distance is more appropriate than the KL divergence for models over high-dimensional space, which usually have their data distribution concentrated on a (latent) lower-dimensional subspace.\n- Experiments: At the first mention of SGLD, SGHMC and NUTS, please spell out and/or add references to these methods for readers who are unfamiliar with the literature. Also, please, consider adding a reference to NDP in line 265.\n- Experiments: Line 306, any idea why SGLD failed with the Quadrocopter 1?\n- Please, add (foot)note on what is an \"Ansatz\". At first, I thought it was a typo, as I was unfamiliar with the term.\n\nReferences:\n\n[A] Cranmer, Kyle, Johann Brehmer, and Gilles Louppe. 2020. “The Frontier of Simulation-Based Inference.” Proceedings of the National Academy of Sciences.\n\n[B] Ramos, Fabio, Rafael Carvalhaes Possas, and Dieter Fox. 2019. “BayesSim : Adaptive Domain Randomization via Probabilistic Inference for Robotics Simulators.” In Robotics: Science and Systems (RSS). Freiburg im Breisgau, Germany.\n\n[C] Okada, Masashi, and Tadahiro Taniguchi. 2019. “Variational Inference MPC for Bayesian Model-Based Reinforcement Learning.” In 3rd Conference on Robot Learning (CoRL 2019).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Some of the scalability limitations are explored in the appendix, and no negative societal impact is foreseen by the authors.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "3p4Y79XNBB", "reviewer": "Reviewer_oFcP", "summary": "This paper develops a method called distributional gradient matching for learning unknown differential equations from data. The method is based on a combination of a neural ODE model and a Gaussian process model (with a deep covariance kernel) which is encouraged to produce solutions which follow the neural ODE model by introducing a regularisation term based on Wasserstein loss. Numerical examples show that the proposed method outperforms some existing alternatives.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I am not an expert on learning ODEs, but the literature reviews in Sections 1 and 5 do make it seem like the approach is new. It would be interesting to have more commentary on how the use of the GP model in this paper differs from prior work - I presume that most prior work does not use a deep GP model.\n\nThe numerical examples seem impressive, the quadrocopter system in particular being highly non-trivial. In these examples the proposed method clearly outperforms some existing methods, both in terms of accuracy and speed. I lack expertise to say if the two MC-based algorithms, SGLD and SGHMC, that the comparisons are made to are representative of the state-of-the-art in this setting.\n\nIn my opinion the main weakness of this paper is that the presentation of the method is often quite confusing and the notation not always properly explained:\n- Line 141: What does \\dot{\\mathcal{X}} stand for?\n- p4: It is not told what the GP model is supposed to be modelling, just that \"we define a Gaussian process\".\n- p4: The use of a mapping \\phi is not standard in GP regression (at least yet), so it would be good to make it clearer that this part of the GP model comes from Wilson et al. (2016).\n- p4: To call \\phi a \"feature map\" seems potentially confusing if one is used to defining covariance kernels via inner products of feature maps.\n- p4: Is the GP covariance kernel supposed to be potentially different for each k (and hence the notation \\mathcal{K}_k)? If so, is it intentional that there is only a single prior mean function?\n- Remarks on p6 that it has \"trainable parameters\" and in Appendix B.2 that it is \"deep\" appear to be the extent of exposition on the GP prior mean function \\mu. More should perhaps be said.\n- What are t and t_{supp} in the definitions on p4?\n- \\mathcal{L}_{data} defined in Equation (8) is never used afterwards; it is denited p_{GP}(y | \\varphi) later.\n- Is \\mu_{S,supp} on p6 equal to \\mu_S in Equation (7)?\n- It would be helpful to tell explicitly that the Wasserstein loss in Equation (16) will be used as \\mathcal{L}_{dynamics}.\n\nI am not sure if the paper contains sufficient details to replicate the experiments.\n\nOther minor comments:\n\n- The sentence on lines 119-121 is quite awkward. There is also an extra closing parenthesis on line 120.\n- Line 131: \"smother-dependent\"\n- I believe the dot over \\mu in Equation (7) is larger than that in the equation between lines 166 and 167.\n- The GP prior mean needs to be subtracted from second term in Equation (10).\n- Transpose notation is not consistent: E.g., Equations (10) and (11) use \"T\" while (6) and (7) use \"\\top\".\n- \\mathcal{D} in Equation (9) should probably have subscript k.\n- That the noise variance matrix is \\Sigma_\\epsilon given that Equation (12) has a different random variable \\epsilon.\n- Line 196: \\epsilon should probably be bold.\n- Lines 216-17: \"regularizaton\"\n- Line 223: \"indepdent\"\n- \"Lahdesmaki\" - \"Lähdesmäki\" in one of the references\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations caused by the cubic computational cost of GP regression are mentioned.", "review_text": "", "extra_scores": {"ethical_concerns": "-", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Even if there was some spread in the review scores which called for discussion among the reviewers, in the end, the reviewer consensus was in favour of accepting this paper. Please, make sure to address their concerns in the reviews in the camera-ready. The concerns related to clarity, presentation, and notation (which you also comment on in your response) are of particular importance as these help ensure impact of your work.

Decision: Accept (Poster)