Scalable Gradients for Stochastic Differential Equations

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.

AISTATS 2020; 25 pages, 6 figures in main text; clarify notation in appendix

References in corpus (16)

• Adam: A Method for Stochastic Optimization
• Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation
• Stochastic Backpropagation and Approximate Inference in Deep Generative Models
• Latent ODEs for Irregularly-Sampled Time Series
• Neural Jump Stochastic Differential Equations
• ODE${}^2$VAE: Deep generative second order ODEs with Bayesian neural networks
• Adaptive importance sampling for control and inference
• Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise
• Multi-level Residual Networks from Dynamical Systems View
• Deep Temporal Sigmoid Belief Networks for Sequence Modeling
• Augmented Neural ODEs
• Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit
• Learning unknown ODE models with Gaussian processes
• Black-box Variational Inference for Stochastic Differential Equations
• Deep learning with differential Gaussian process flows
• Adaptive Path-Integral Autoencoder: Representation Learning and Planning for Dynamical Systems

Cited by in corpus (43)

• Diffusion Models in Vision: A Survey
• Normalizing Flows: An Introduction and Review of Current Methods
• How to train your neural ODE: the world of Jacobian and kinetic regularization
• Dissecting Neural ODEs
• Feasibility Study of Neural ODE and DAE Modules for Power System Dynamic Component Modeling
• Learning Differential Equations that are Easy to Solve
• Maximum Likelihood Training of Score-Based Diffusion Models
• TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics
• Stochastic Normalizing Flows
• Quantum-tailored machine-learning characterization of a superconducting qubit
• Continuous-in-Depth Neural Networks
• Interpretable machine learning for high-dimensional trajectories of aging health
• A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depth
• Learning the temporal evolution of multivariate densities via normalizing flows
• Modeling Continuous Stochastic Processes with Dynamic Normalizing Flows
• Stochastic Normalizing Flows
• STEER: Simple Temporal Regularization For Neural ODEs
• Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations
• TorchDyn: A Neural Differential Equations Library
• PDE-Driven Spatiotemporal Disentanglement
• MALI: A memory efficient and reverse accurate integrator for Neural ODEs
• Stable Neural Flows
• Automatic structured variational inference
• Stochasticity in Neural ODEs: An Empirical Study
• Path Integral Sampler: a stochastic control approach for sampling
• Learning stochastic dynamical systems with neural networks mimicking the Euler-Maruyama scheme
• Learning the Latent dynamics of Fluid flows from High-Fidelity Numerical Simulations using Parsimonious Diffusion Maps
• Probabilistic Numeric Convolutional Neural Networks
• Latent Neural Stochastic Differential Equations for Change Point Detection
• Continuous-Time Model-Based Reinforcement Learning
• Sparse Flows: Pruning Continuous-depth Models
• Learning Continuous-Time Dynamics by Stochastic Differential Networks
• GoTube: Scalable Stochastic Verification of Continuous-Depth Models
• Gradient-based Bayesian Experimental Design for Implicit Models using Mutual Information Lower Bounds
• Continuous-Depth Neural Models for Dynamic Graph Prediction
• Meta-Solver for Neural Ordinary Differential Equations
• Latent Neural Differential Equations for Video Generation
• Deep learning and differential equations for modeling changes in individual-level latent dynamics between observation periods
• Physics-aware, probabilistic model order reduction with guaranteed stability
• State-space deep Gaussian processes with applications
• Capturing Actionable Dynamics with Structured Latent Ordinary Differential Equations
• Estimating Stochastic Poisson Intensities Using Deep Latent Models
• Enhancing Human-Machine Teaming for Medical Prognosis Through Neural Ordinary Differential Equations (NODEs)