PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Gradients for Stochastic Differential Equations

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

arXiv:2001.01328v6cs.LGmath.NAstat.ML

TL;DR

Existing SDE-gradient methods face poor scaling in time or memory. This paper introduces stochastic adjoint sensitivities with Brownian-noise reconstruction and combines them with variational inference for latent SDEs, achieving competitive dynamics-modeling performance while reducing memory use on studied tasks.

  • Problem

    Existing SDE-gradient approaches scale poorly in time or memory, limiting efficient training of stochastic dynamics models.

  • Method

    The paper derives backward stochastic adjoint dynamics, reconstructs Brownian noise memory-efficiently, and combines the method with gradient-based stochastic variational inference for latent SDEs.

  • Results

    The adjoint approach reduced memory usage by 1/2-1/3 compared to direct solver backpropagation on the considered tasks and was competitive for latent SDE dynamics modeling.

  • Takeaways & Limitations

    The framework enables time-efficient, constant-memory gradient computation for SDE solutions with high-order adaptive solvers and supports latent stochastic-dynamics modeling.

  • Takeaways & Limitations

    Convergence rates for numerical gradients with general schemes remain unknown, and it is unclear whether the results extend from strong to weak errors.

Abstract

from arXiv · show

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.

1 Introduction

SDEs model stochastic dynamics but existing gradient methods scale poorly in time or memory. The paper generalizes adjoint sensitivities to SDEs and combines them with variational inference for latent SDE models.

  • Motivation: SDEs add instantaneous noise to ODE dynamics and model phenomena driven by many small, unobserved interactions.Examples include molecular motion, gene-pool allele frequencies, and market prices.
  • Motivation: Existing SDE-gradient approaches scale poorly: pathwise differentiation grows poorly with parameters and states, while solver backpropagation grows poorly in memory.
  • Method: The stochastic adjoint sensitivity method generalizes adjoint sensitivities to SDEs, supports tens of thousands of parameters, and permits high-order adaptive solvers.
  • Method: The method derives backward Stratonovich dynamics whose solution computes the necessary gradient.
  • Method: A noise-reconstruction algorithm queries Brownian motion at arbitrary times while storing only a single random seed.
  • Applications: The approach is combined with gradient-based stochastic variational inference for latent SDEs, which generalize several existing dynamics-modeling families.

2 Background: Stochastic Flows

The background develops adjoint sensitivities, neural SDEs, Stratonovich integration, and stochastic flows. Stratonovich symmetry supports backward dynamics, while smoothness assumptions yield stochastic flows with backward SDEs.

  • Adjoint Sensitivity Method: ODE adjoint sensitivity computes gradients through a backward costate system using vector-Jacobian products without storing intermediate computations.
  • Neural SDEs: Neural SDEs use neural networks for drift and diffusion, but prior approaches explicitly evolved large Jacobians rather than cheap vector-Jacobian products.
  • Backward Stratonovich Integral: Stratonovich integration is used because its symmetry gives convenient backward dynamics, while Itô SDEs can be handled through a conversion rule.
  • Backward Stratonovich Integral: Figure 1 contrasts ODE and Stratonovich-SDE adjoint pseudocode, with differences highlighted and augmented quantities concatenated in brackets.
  • Stochastic Flow of Diffeomorphisms: Under smoothness assumptions, the SDE generates a stochastic flow of diffeomorphisms whose maps satisfy backward SDEs.
  • Stochastic Flow of Diffeomorphisms: For Itô dynamics, simply negating drift and diffusion fails to reconstruct the path backward, whereas the converted Stratonovich system does reconstruct it.

3 Sensitivity via Stochastic Adjoint

The paper derives a stochastic adjoint process whose reverse-time dynamics compute gradients for SDE solutions using vector-Jacobian products without storing intermediate quantities. It also gives convergence conditions for numerical approximations and discusses solver and diffusion limitations.

  • Stochastic Adjoint Process: The stochastic adjoint sensitivity method computes SDE gradients by solving reverse-time dynamics using cheap vector-Jacobian products without storing intermediate quantities.The formulation extends gradients from initial states to drift and diffusion parameters by augmenting the state with zero-drift, zero-diffusion parameter dynamics.
  • Stochastic Adjoint Process: The backward Stratonovich SDE addresses the mathematical challenge of reversing SDE dynamics while reconstructing the necessary gradient.The paper distinguishes the deterministic-endpoint backward SDE from the forward-flow endpoint case, which is formulated through the Itô map because adaptedness is lost.
  • Stochastic Adjoint Process: The continuous-time gradient is obtained by composing the backward SDE system with the original forward SDE.This composition evaluates the gradient associated with the terminal loss after the forward flow reaches its endpoint.
  • Numerical Approximation: Numerical solutions converge in probability when the forward and backward schemes converge in probability and the forward scheme is uniformly consistent over bounded starting positions.The paper proves the uniform-consistency condition for Euler-Maruyama, while detailed analysis of other schemes remains beyond scope.
  • The Algorithm: For general matrix-valued diffusion, strong order beyond 1/2 requires simulating costly multiple Wiener-process integrals.With diagonal noise, commutativity permits strong-order-1 schemes such as Milstein and stochastic Runge-Kutta without approximating Lévy areas.
  • The Algorithm: Analysis of high-weak-order numerical schemes is outside the paper’s current scope.The implementation includes PyTorch SDE solvers with adaptive time-stepping and a torch.autograd.Function interface.

4 Virtual Brownian Tree

The virtual Brownian tree reconstructs Wiener-process values at arbitrary query times by recursively sampling Brownian-bridge midpoints with splittable pseudorandom keys. It provides constant memory and logarithmic query-time scaling, making backward noise replay compatible with adaptive solvers.

  • Motivation: Naively storing Brownian increments consumes substantial memory and complicates adaptive backward integration because replay times can differ from forward times.The backward adjoint must query the same Wiener-process sample path used during the forward pass.
  • Virtual Brownian Tree: The virtual Brownian tree combines Brownian trees with splittable PRNGs to query Wiener-process values at arbitrary times.Each midpoint sample uses a key deterministically split from its parent node, allowing tree nodes to be reconstructed on demand.
  • Complexity: O(1) memory and logarithmic time in inverse error tolerance are the virtual Brownian tree’s stated resource costs.For a solver taking L fixed-grid steps, the tree tolerance scales as 1/L and each query has per-step time complexity log L.
  • Query Procedure: The virtual Brownian tree evaluates a query by repeatedly bisecting the interval until the midpoint is within the requested error tolerance.The figure describes unique random-generator keys determined by the path taken through the recursive tree.
  • Brownian Bridges and Brownian Trees: Brownian-bridge sampling recursively determines intermediate Wiener values between timestamps whose endpoint values are already known.Repeated midpoint evaluation constructs the Brownian tree conceptually, while the algorithm reconstructs only the requested node.

5 Latent Stochastic Differential Equations

Latent SDEs use stochastic variational inference to learn prior and posterior processes while optimizing an ELBO. The framework represents latent trajectories as SDE-generated continuous-time processes and supports encoder-based inference.

  • Maximum-likelihood fitting can overfit by driving the diffusion function to zero, motivating variational inference for latent SDEs.
  • Both the prior and approximate posterior can be parameterized as SDEs with shared diffusion, enabling finite KL estimation under mild regularity conditions.
  • The ELBO combines observation likelihoods at corresponding latent states with expectations under the approximate posterior process.
  • Gradients for prior and variational parameters can be computed through one paired forward and backward SDE solve after augmenting the forward dynamics.
  • An encoder network can map observations to variational parameters, making the model analogous to a variational autoencoder with an SDE-induced latent process.

6 Related Work

Prior approaches to differentiating and fitting SDEs face scaling, memory, discretization, or backward-simulation costs. Related work also studies Bayesian SDE learning and alternative backward-SDE formulations.

  • Pathwise methods scale poorly with parameter and state counts because each forward step requires entire Jacobian matrices.
  • Reverse-mode differentiation through SDE solver operations has high memory cost and traditionally relies on fixed Euler-Maruyama discretization.
  • Euler-discretized latent SDE variational inference differs from continuous-time SDE formulations, whose discretization limit is not established as equivalent.
  • Traditional backward SDE approaches can make simulation costly because the backward pass requires estimating conditional expectations.
  • Bayesian SDE learning includes approximate posterior processes sharing diffusion with a prior, while broader non-variational methods include extended Kalman filtering and MCMC.

7 Experiments

Experiments validate stochastic-adjoint gradients, demonstrate latent SDE learning on synthetic processes, and evaluate predictive performance on motion capture data. The studies examine accuracy, efficiency, multimodal dynamics, and real-world forecasting.

  • 7.1 Numerical Studies: Analytical-gradient comparisons show that stochastic-adjoint gradient error decreases with step size in a closed-form test problem.
  • 7.1 Numerical Studies: Lower adaptive-solver absolute tolerance generally reduces mean squared gradient error, although it tends to require more function evaluations than the ODE counterpart.
  • 7.1 Numerical Studies: For two of three test problems, fixed-step Milstein stochastic adjoints are much more time-efficient than direct backpropagation through Milstein and Euler operations.
  • 7.2 Synthetic Datasets: Latent SDEs reconstruct geometric Brownian motion and stochastic Lorenz attractor data while jointly optimizing prior, posterior, diffusion, encoder, and decoder parameters.
  • 7.2 Synthetic Datasets: The learned Lorenz prior produces bimodal samples in data and latent spaces, a behavior described as hard to achieve with a latent ODE and unimodal Gaussian initial posterior.
  • 7.3 Motion Capture Dataset: The motion-capture evaluation predicts future observations on a 50-dimensional dataset using 23 walking sequences divided into training, validation, and test sets.

8 Discussion

The discussion presents stochastic adjoints as a scalable route to SDE gradients and latent-SDE variational inference. It highlights application opportunities while identifying unresolved convergence questions and variance-reduction directions.

  • The stochastic adjoint has nearly the same time and memory complexity as directly solving the SDE, while supporting latent-SDE variational inference.
  • Future work includes reducing gradient-estimate variance and extending the method toward domain-specific structural or stationarity constraints in latent-SDE priors.
  • The framework could support differentiable SDE applications including Wright-Fisher models, derivative pricing, and infinitely-deep Bayesian neural networks.
  • Convergence rates for numerical gradients under general schemes remain unknown, and the theory is based on strong rather than weak errors.

9 Appendix

This appendix establishes convergence properties for numerical SDE solutions and their derivatives, including local uniform convergence for Euler–Maruyama under stated regularity conditions. It also identifies technical requirements for extending the argument to higher dimensions and other schemes.

  • Convergence conditions: The convergence theorem assumes sufficiently smooth SDE coefficients and numerical schemes with expectation order p1 and mean-square order p2 satisfying p1 ≥ p2 + 1/2 and p2 ≥ 1/2.The bound is uniform over discretization points and starting states under the theorem’s conditions.
  • Euler–Maruyama analysis: Euler–Maruyama provides the required derivative and discretization estimates for the one-dimensional argument, with accuracy orders (p1, p2) = (1.5, 1.0).The derivative of the Euler–Maruyama discretization is the discretization of the derivative process.
  • Local uniform convergence: The resulting Euler–Maruyama approximation converges locally uniformly to the SDE flow in probability as the mesh size h tends to zero.The proof strengthens conventional fixed-starting-point convergence using a Sobolev inequality argument.
  • Scope of the result: The argument can extend to higher dimensions and other schemes if expected W 1,p errors are bounded for p > d, but obtaining the required dependence on the starting state is nontrivial.Generic extension to other numerical schemes remains technically difficult.

9.5 Stochastic Adjoint has Commutative Noise when Original SDE has Diagonal Noise

For diagonal-noise SDEs, the stochastic adjoint has a commutative diffusion structure. This permits efficient strong-order numerical schemes without simulating costly iterated integrals or Lévy areas.

  • Augmented adjoint: Augmenting the state with model parameters yields adjoint dynamics whose gradients with respect to drift and diffusion parameters can be computed through vector-Jacobian products.The parameter components are assigned zero drift and diffusion in the augmented SDE.
  • Commutative noise: When the original SDE has diagonal noise, the backward stochastic adjoint satisfies the commutativity condition.The appendix verifies this condition for the diffusion function by considering the relevant index cases.
  • Numerical implication: Commutativity avoids simulating computationally intensive double Itô integrals and Lévy areas during numerical integration.This follows from the iterated-integral property used by commutative-noise schemes.
  • Numerical implication: Milstein and stochastic Runge–Kutta schemes of strong order 1.0 can therefore be used safely for the adjoint under diagonal noise.The implementation can require a number of vector-Jacobian products independent of the original SDE dimensionality.

9.7 Stochastic Adjoint for Latent SDE

The latent-SDE construction incorporates variational inference into the augmented stochastic-adjoint framework. Its experiments use neural latent dynamics on synthetic geometric Brownian motion and stochastic Lorenz attractor data.

  • Latent-SDE inference: The variational lower bound can be simulated by augmenting the latent SDE state with an extra variable tracking its contribution.The resulting backward system has commutative diffusion, and neither the extra variable nor its adjoint needs a separate backward solve.
  • Synthetic experiments: The geometric Brownian motion experiment trains a GRU-based latent SDE with neural drift and diffusion functions on 1024 noisy time series.The series are observed every 0.02 from time 0 to time 1.
  • Synthetic experiments: The stochastic Lorenz attractor experiment trains a latent SDE with diagonal diffusion on 1024 normalized, noisy time series.The observations are sampled every 0.025 from time 0 to time 1.

9.11 Additional Visualization

Additional visualizations show posterior reconstructions and prior samples for synthetic stochastic datasets, while the mocap model uses a structured latent-SDE architecture and the released implementation uses diagonal-noise integration.

  • Synthetic visualizations: The Lorenz visualizations compare true data and posterior reconstructions with prior samples whose initial latent states are either independently sampled or held fixed across trajectories.The three rows encode these three visualization settings.
  • Visualization findings: The synthetic-data results report good posterior reconstruction and prior behavior resembling the data, while most prior uncertainty is represented through the initial latent state.The authors leave more interpretable prior processes for future work.
  • Synthetic visualizations: For the geometric Brownian motion visualization, the first row shows true data and posterior reconstructions, with an orange contour covering 95% of 512 samples.The lower rows use independently sampled or shared fixed initial latent states.
  • Mocap architecture: The mocap latent-SDE architecture consists of encoder, decoder, prior drift, posterior drift, and diffusion components.These components are arranged across the two rows of the architecture figure.
  • Implementation: The implementation provides a diagonal-noise Euler–Maruyama integrator and a custom stochastic-adjoint autograd method that runs the backward solve over reversed time intervals.The backward pass accumulates adjoints and parameter gradients while reusing Brownian-motion queries.
Loading 2001.01328v6…