Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models

1 Introduction For continuous-time system identification and control , ordinary differential equations form an essential class of models , deployed in applications ranging from robotics ( Spong et al. , 2006 ) to biology ( Jones et al. , 2009 ) . Here , it is assumed that the evolution of a system is described by the evolution of continuous state variables , whose time-derivative is given by a set of parametrized equations . Often , these equations are derived from first principles , e.g. , rigid body dynamics ( Wittenburg , 2013 ) , mass action kinetics ( Ingalls , 2013 ) , or Hamiltonian dynamics ( Greydanus et al. , 2019 ) , or chosen for computational convenience ( e.g. , linear systems ( Ljung , 1998 ) ) or parametrized to facilitate system identification ( Brunton et al. , 2016 ) . Such construction methods lead to 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 as control design ( Ortega et 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 is not available , or precise parameter values are cumbersome to obtain , system identification from raw time series data becomes ∗Equal Contribution . Correspondence to trevenl @ ethz.ch , wenkph @ ethz.ch . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . necessary . In this case , one may seek more expressive nonparametric models instead ( Rackauckas et al. , 2020 ; Pillonetto et al. , 2014 ) . If the model is completely replaced by a neural network , the resulting model is called neural ODE ( Chen et al. , 2018 ) . Despite their large number of parameters , as demonstrated by Chen et al . ( 2018 ) ; Kidger et al . ( 2020 ) ; Zhuang et 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 ( LaValle , 2006 ) , probabilistic forecasting in meteorology ( Fanfarillo et al. , 2021 ) , or active learning / experimental design ( Srinivas et al. , 2010 ) . A common way of obtaining such uncertainties is via a 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-based inference schemes for Bayesian neural ODEs suffer from several serious drawbacks and thus propose sampling-based alternatives . However , as surfaced by our experiments in Section 4 , sampling-based 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 sampling-based planning ( Karaman and Frazzoli , 2011 ) or uncertainty propagation in prediction . Contributions In this work , we propose a novel approach for uncertainty quantification in nonlinear dynamical systems ( cf . Figure 1 ) . Crucially , our approach avoids explicit costly and non-robust numerical integration , by employing a probabilistic smoother of the observational data , whose representation we learn jointly across multiple trajectories . To capture dynamics , we regularize our smoother with a dynamics model . Latter captures epistemic uncertainty in the gradients of the ODE , which we match with the smoother ’ s gradients by minimizing a Wasserstein loss , hence we call our approach Distributional Gradient Matching ( DGM ) . In summary , our main contributions are : • We develop DGM , an approach2 for capturing epistemic uncertainty about nonlinear dynamical systems by jointly training a smoother and a neural dynamics model ; • We provide a computationally efficient and statistically accurate mechanism for prediction , by focusing directly on the posterior / predictive state distribution . • We experimentally demonstrate the effectiveness of our approach on learning challenging , chaotic dynamical systems , and generalizing to new unseen inital conditions . High-level overview A high-level depiction of our algorithm is shown in Figure 2 . In principle , DGM jointly learns a smoother ( S ) and a dynamics model ( D ) . The smoother model , chosen to be a Gaussian process , maps an initial condition x0 and a time t to the state distribution pS ( x ( t ) ) and state derivatives distribution pS ( ẋ ( t ) ) reached at that time . The dynamics model , chosen to be a neural network , represents an ODE that maps states x ( t ) to the derivative distribution pD ( ẋ ( t ) ) . Both models are evaluated at some training times and all its output distributions collected in the random variables XS , ẊS and ẊD . The parameters of these models are then jointly trained using a Wasserstein-distance-based objective directly on 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 3 builds on known concepts from the literature , which we 2Code is available at : https : //github.com/lenarttreven/dgm summarize in Section 2 . Finally , in Section 4 , we present the empirical study of the DGM algorithm , where we benchmark it against the state-of-the-art , uncertainty aware dynamics models . 2 Background . 2.1 Problem Statement . Consider a continuous-time dynamical system whose K-dimensional state x ∈ RK evolves according to an unknown ordinary differential equation of the form ẋ = f∗ ( x ) . ( 1 ) Here , f∗ is an arbitrary , unknown function assumed to be locally Lipschitz continuous , to guarantee existence and uniqueness of trajectories for every initial condition . In our experiment , we initialize the system at M different initial conditions xm ( 0 ) , m ∈ { 1 , . . . , 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 ) m∈ { 1 ... M } can vary from trajectory to trajectory . Thus , a trajectory m is described by its initial condition xm ( 0 ) , and the observations ym : = [ xm ( tn , m ) + ϵn , m ] n∈ { 1 ... Nm } at times tm : = [ tn , m ] n∈ { 1 ... Nm } , where the additive observation noise ϵn , m is assumed to be drawn i.i.d . from a zero mean Gaussian , whose covariance is given by Σϵ : = diag ( σ21 , . . . , σ 2 K ) . We denote by D the dataset , consisting of M initial conditions xm ( 0 ) , observation times tm , and observations ym . To model the unknown dynamical system , we choose a parametric Ansatz ẋ = f ( x , θ ) . 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 a priori ( e.g. , a neural network ) , and the key challenge is to infer a reasonable distribution over the parameters θ , conditioned on the data D. For later tasks , we are particularly interested in the predictive posterior state distribution p ( xnew ( tnew ) |D , tnew , xnew ( 0 ) ) , ( 2 ) i.e. , the posterior distribution of the states starting from a potentially unseen initial condition xnew ( 0 ) and evaluated at times tnew . This posterior would then be used by the downstream or prediction tasks described in the introduction . 2.2 Bayesian Parameter Inference . In the case of Bayesian parameter inference , an additional prior p ( θ ) is imposed on the parameters θ so that the posterior distribution of Equation ( 2 ) can be inferred . Unfortunately , this distribution is not analytically tractable for most choices of f ( x , θ ) , which is especially true when we model f with a neural network . Formally , for fixed parameters θ , initial condition x ( 0 ) and observation time t , the likelihood of an observation y is given by p ( y ( t ) |x ( 0 ) , t , θ , Σobs ) = N ( y ( t ) ∣∣∣∣x ( 0 ) + ∫ t 0 f ( x ( τ ) , θ ) dτ , Σobs ) . ( 3 ) Using the fact that all noise realizations are independent , the expression ( 3 ) can be used to calculate the likelihood of all observations in D. Most state-of-the-art parameter inference schemes use this fact to create samples θ̂s of the posterior over parameters p ( θ|D ) using various Monte Carlo methods . Given a new initial condition x ( 0 ) and observation time t , these samples θ̂s can then be turned into samples of the predictive posterior state again by numerically integrating x̂s ( t ) = x ( 0 ) + ∫ t 0 f ( x ( τ ) , θ̂s ) dτ . ( 4 ) Clearly , both training ( i.e. , obtaining the samples θ̂s ) and prediction ( i.e. , evaluating Equation ( 4 ) ) require integrating the system dynamics f many times . Especially when we model f with a neural network , this can be a huge burden , both numerically and computationally ( Kelly et al. , 2020 ) . As an alternative approach , we can approximate the posterior p ( θ|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 to evaluate the unnormalized posterior . Also , at inference time , to obtain a distribution over state x̂s ( t ) , we still need to integrate f several times . Furthermore , Dandekar et al . ( 2021 ) report poor forecasting performance by the variational approach . 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 θ̂ . This makes these approaches rather cumbersome : Both for obtaining estimates of θ̂ and for obtaining the predictive posterior over states , once θ̂ 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 θ . 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 a tuple ( t , x ( 0 ) ) consisting of a time point t and an initial condition x ( 0 ) ) 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 ( 2 ) . Albeit computationally attractive , this approach has one serious drawback . Since the smoother no 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 ( 1 ) . 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 , Ldynamics , which ensures that a trajectory predicted by the smoother is encouraged to follow some underlying system of the form of Equation ( 1 ) . The smoother is then trained with the multi-objective loss function L : = Ldata + λ · Ldynamics , ( 5 ) where , Ldata is a smoother-dependent loss function that ensures a sufficiently accurate data fit , and λ is a trade-off parameter .

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.