# Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models

arXiv: https://arxiv.org/abs/2106.11609 · HTML: https://mlsanity.com/paper/2106.11609/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 59 references that are in the map.
- First posted on arXiv 22 Jun 2021.

## Builds on (most cited first, 23 of 23)

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Hamiltonian Neural Networks](https://mlsanity.com/paper/1906.01563/) (2019)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Safe Model-based Reinforcement Learning with Stability Guarantees](https://mlsanity.com/paper/1705.08551/) (2017)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [When Gaussian Process Meets Big Data: A Review of Scalable GPs](https://mlsanity.com/paper/1807.01065/) (2018)
- [Learning Differential Equations that are Easy to Solve](https://mlsanity.com/paper/2007.04504/) (2020)
- [Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE](https://mlsanity.com/paper/2006.02493/) (2020)
- [Neural ODE Processes](https://mlsanity.com/paper/2103.12413/) (2021)
- [Bayesian Neural Ordinary Differential Equations](https://mlsanity.com/paper/2012.07244/) (2020)
- [Linearly constrained Gaussian processes](https://mlsanity.com/paper/1703.00787/) (2017)
- [MALI: A memory efficient and reverse accurate integrator for Neural ODEs](https://mlsanity.com/paper/2102.04668/) (2021)
- [Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations](https://mlsanity.com/paper/2102.06559/) (2021)
- [Fast Gaussian Process Based Gradient Matching for Parameter Identification in Systems of Nonlinear ODEs](https://mlsanity.com/paper/1804.04378/) (2018)
- [ODIN: ODE-Informed Regression for Parameter and State Inference in Time-Continuous Dynamical Systems](https://mlsanity.com/paper/1902.06278/) (2019)
- [Differentiable Likelihoods for Fast Inversion of 'Likelihood-Free' Dynamical Systems](https://mlsanity.com/paper/2002.09301/) (2020)
- [Linear Matrix Inequalities for Physically-Consistent Inertial Parameter Identification: A Statistical Perspective on the Mass Distribution](https://mlsanity.com/paper/1701.04395/) (2017)
- [Scalable Variational Inference for Dynamical Systems](https://mlsanity.com/paper/1705.07079/) (2017)
- [Probabilistic Forecasting using Deep Generative Models](https://mlsanity.com/paper/1909.11865/) (2019)
- [AReS and MaRS - Adversarial and MMD-Minimizing Regression for SDEs](https://mlsanity.com/paper/1902.08480/) (2019)
- [SLEIPNIR: Deterministic and Provably Accurate Feature Expansion for Gaussian Process Regression with Derivatives](https://mlsanity.com/paper/2003.02658/) (2020)

## Built on by (most cited first, 1 of 1)

- [Efficient Exploration in Continuous-time Model-based Reinforcement Learning](https://mlsanity.com/paper/2310.19848/) (2023)

## Abstract (arXiv)

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.