Scalable Gradients for Stochastic Differential Equations

TL;DR

This paper extends the adjoint sensitivity method to stochastic differential equations, enabling scalable, memory-efficient gradient computation for high-dimensional stochastic dynamics, with applications in neural network-based modeling.

Contribution

It generalizes the adjoint sensitivity method to stochastic differential equations, providing a scalable, memory-efficient approach for gradient computation in stochastic systems.

Findings

Achieved competitive performance on a 50-dimensional motion capture dataset.

Derived a stochastic differential equation for gradients with convergence conditions.

Combined the method with stochastic variational inference for latent SDEs.

Abstract

The adjoint sensitivity method scalably computes gradients of solutions to ordinary differential equations. We generalize this method to stochastic differential equations, allowing time-efficient and constant-memory computation of gradients with high-order adaptive solvers. Specifically, we derive a stochastic differential equation whose solution is the gradient, a memory-efficient algorithm for caching noise, and conditions under which numerical solutions converge. In addition, we combine our method with gradient-based stochastic variational inference for latent stochastic differential equations. We use our method to fit stochastic dynamics defined by neural networks, achieving competitive performance on a 50-dimensional motion capture dataset.