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