Distributional Gradient Matching for Learning Uncertain Neural Dynamics Models
This work addresses the need for accurate uncertainty estimates in neural ODEs for applications such as reinforcement learning and control, representing a novel method for a known bottleneck rather than a foundational advancement.
The paper tackles the problem of learning uncertain neural ODEs for tasks like active learning and robust control by proposing distributional gradient matching (DGM), which avoids numerical integration and directly models state space uncertainties, resulting in faster training and prediction with significantly improved accuracy compared to traditional methods.
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.