SING: SDE Inference via Natural Gradients

2025-06-21

AI Summary

Existing variational inference (VI)-based latent stochastic differential equation (SDE) models suffer from slow convergence and numerical instability in posterior path inference. To address these issues, we propose a natural-gradient VI framework tailored for continuous-time dynamical systems, integrating a Gaussian process SDE prior, adaptive numerical integration, and temporal parallelization. Our method jointly infers latent state trajectories and nonlinear drift functions efficiently by explicitly leveraging the Riemannian geometry of the variational distribution—thereby enhancing optimization stability and accelerating convergence—while enabling time-parallel computation. Experiments on multiple benchmark datasets demonstrate substantial improvements in state reconstruction accuracy and drift function estimation fidelity. Notably, our approach achieves superior performance and enhanced interpretability in modeling neural dynamics from freely behaving animal data.

Technology Category

• Machine Learning: Learning with Manifolds
• Intelligent Robots: State Estimation
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Latent stochastic differential equation (SDE) models are important tools for the unsupervised discovery of dynamical systems from data, with applications ranging from engineering to neuroscience. In these complex domains, exact posterior inference of the latent state path is typically intractable, motivating the use of approximate methods such as variational inference (VI). However, existing VI methods for inference in latent SDEs often suffer from slow convergence and numerical instability. Here, we propose SDE Inference via Natural Gradients (SING), a method that leverages natural gradient VI to efficiently exploit the underlying geometry of the model and variational posterior. SING enables fast and reliable inference in latent SDE models by approximating intractable integrals and parallelizing computations in time. We provide theoretical guarantees that SING will approximately optimize the intractable, continuous-time objective of interest. Moreover, we demonstrate that better state inference enables more accurate estimation of nonlinear drift functions using, for example, Gaussian process SDE models. SING outperforms prior methods in state inference and drift estimation on a variety of datasets, including a challenging application to modeling neural dynamics in freely behaving animals. Altogether, our results illustrate the potential of SING as a tool for accurate inference in complex dynamical systems, especially those characterized by limited prior knowledge and non-conjugate structure.

Problem

• Efficient inference in latent SDE models
• Improving convergence and stability in variational inference
• Accurate estimation of nonlinear drift functions

Innovation

• Natural gradient VI for latent SDEs
• Approximates intractable integrals efficiently
• Parallelizes computations in time