Scalable Gradients for Stochastic Differential Equations

By Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, David Kristjanson Duvenaud. 2020

434 citations recorded by the source corpus.

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.

Subfields: Computational Mechanics, Numerical Analysis, Statistical and Nonlinear Physics, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computational Mechanics

In Numerical Analysis

In Statistical and Nonlinear Physics

In Optimization & Training

Citations in the corpus

Stochastic Normalizing Flows (2020, background): “Neural stochastic differential equations learn stochastic processes from observations along the path (Tzen & Raginsky, 2019; Jia & Benson, 2019; Liu et al., 2019; Li et al., 2020), but are not designed for marginal density estimation or asymptotically unbiased sampling.”
Maximum Likelihood Training of Score-Based Diffusion Models (2021, methodology): “…between two SDEs with shared diffusion coefﬁcients and starting points exists under assumptions (vii) (viii) (ix) (x) (xi) (see, e.g ., [52, 30]), and can be computed via the Girsanov theorem [34]: where (i) is due to Girsanov Theorem II [34, Theorem 8.6.6], and (ii) is due to the…”
A Variational Perspective on Diffusion-Based Generative Models and Score Matching (2021, methodology): “In this case, if we reparameterize Ys = μs + σs where (7)Black-box SDE solvers such as torchsde (Li et al., 2020) might not be optimized for the deterministic case, since their stochastic adjoint method scales O(L logL) in time whereas deterministic numerical solvers are usually faster.”
TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics (2020, neutral): “In future work, we would like to consider stochastic dynamics (Li et al., 2020) and learning the growth term together with the dynamics.”
Universal Differential Equations for Scientific Machine Learning (2020, methodology): “± Additionally, Supplement S ?? demonstrates a 1,600x performance advantage for the SciML ecosystem over torchsde using the geometric Brownian motion example from the torchsde documentation [65].”
Improving and generalizing flow-based generative models with minibatch optimal transport (2023, positive): “Simulation-free training is common in stochastic flow models where backpropagating through the simulation is numerically challenging and has high variance (Li et al., 2020).”
Learning Differential Equations that are Easy to Solve (2020, methodology): “While we investigated only ODEs in this paper, this approach could presumably be extended straight-forwardly to neural stochastic differential equations ﬁt by adaptive solvers (Li et al., 2020) and other ﬂavors of parametric differential equations ﬁt by gradient descent (Rackauckas et al., 2019).”
Diffusion Normalizing Flow (2021, methodology): “We note another similar approach [26] of training SDEs caches random noise dw and further takes advantage of the pseudo-random generator to save memory for the intermediate noises, resulting in constant memory consumption However, the approach can not reproduce exact trajectories due to time discret”
On Maximum Likelihood Training of Score-Based Generative Models (2021, background): “These ‘neural SDEs’ have received recent attention, but must explicitly address how to compute gradients of the solution to an SDE [Tzen and Raginsky, 2019; Li et al., 2020].”
Normalizing Flows: An Introduction and Review of Current Methods (2020, background): “Following this, Li et al. [62] derived the adjoint SDE whose solution gives the gradient of the original Neural SDE.”
Continuous-in-Depth Neural Networks (2020, background): “Other theoretical research fronts include interfacing Neural ODEs with normalizing flows [22, 79, 40, 56], graph NNs [61, 67], stochastic differential equations [52, 37, 24, 47] and RNNs [10, 65, 27, 18, 49].”
Diffusion Models for Adversarial Purification (2022, methodology): “To overcome this, we propose to use the adjoint method (Li et al., 2020) to efficiently compute full gradients of the SDE without the memory issue.”
Stochastic Normalizing Flows (2020, methodology): “This is to be compared with Itô integration which is deﬁned instead by Stratonovich diﬀerential equations were recognized in Li et al. (2020) where x ∗ is an independent copy of x , and the subscript denotes the i -th row.”
GRAND: Graph Neural Diffusion (2021, methodology): “These have been applied to problems where the PDE can be described on a graph (Li et al., 2020c).”