[Paper Review] Scalable Gradients for Stochastic Differential Equations

Xuechen Li, Ting‐Kam Leonard Wong+2 more|arXiv (Cornell University)|Jan 5, 2020

Model Reduction and Neural Networks77 references108 citations

TL;DR

The paper generalizes the adjoint sensitivity method to stochastic differential equations, enabling memory-efficient, scalable gradient computation for neural SDEs and latent SDE models using high-order adaptive solvers.

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.

Motivation & Objective

• Motivate scalable gradient computation for SDEs beyond ODE adjoints.
• Extend the adjoint sensitivity framework to Stratonovich SDEs with backward dynamics.
• Develop memory-efficient techniques to replay forward noise without storing trajectories.
• Enable gradient-based inference and learning for latent SDE models with irregular observations.
• Demonstrate competitive performance on high-dimensional stochastic dynamics tasks.

Proposed method

• Generalize the adjoint sensitivity method from ODEs to SDEs using backward Stratonovich SDEs to compute gradients.
• Derive a stochastic adjoint process whose solution yields the gradient with respect to Z_T through backward dynamics.
• Propose an algorithm to query forward-path noise with a single random seed, avoiding storage of all noise realizations.
• Integrate the stochastic adjoint with gradient-based stochastic variational inference for latent SDEs.
• Allow using high-order adaptive time-stepping SDE solvers within the backward gradient computation.

Experimental results

Research questions

• RQ1 Can the adjoint sensitivity method be extended to stochastic dynamics to compute gradients efficiently with constant memory?
• RQ2 What backward dynamics (Stratonovich/SDE form) correctly recover forward trajectories for gradient computation?
• RQ3 How can one replay forward noise without storing full trajectories while maintaining exact gradient information?
• RQ4 Can the proposed stochastic adjoint framework be combined with variational inference for latent SDEs?
• RQ5 Is the approach scalable to neural SDEs with high dimensionality and irregularly sampled data?

Key findings

• The stochastic adjoint method achieves constant memory usage and O(L log L) time compared to other baselines.
• The backward Stratonovich SDE framework yields correct gradient reconstruction for stochastic dynamics.
• An efficient noise-caching mechanism allows exact gradient computation with a single random seed.
• The method supports high-order adaptive solvers for SDEs, enabling scalable training of neural SDEs.
• The approach, when combined with gradient-based variational inference, yields competitive performance on latent SDEs across datasets.
• Demonstrations include 50-dimensional motion capture dynamics.