PaperlayerView exact revision on arXiv

Source-linked AI summary

SDE Matching: Scalable and Simulation-Free Training of Latent Stochastic Differential Equations

Grigory Bartosh, Dmitry Vetrov, Christian A. Naesseth

arXiv:2502.02472v3stat.MLcs.LG

TL;DR

Latent SDE training relies on costly simulation and backpropagation through approximate SDE solutions, limiting scalability. SDE Matching parameterizes the posterior process to train without simulation; experiments report comparable or better performance than adjoint sensitivity-based methods with significantly reduced computational complexity.

  • Problem

    Latent SDE training relies on costly numerical integration and backpropagation through SDE solutions, making the models difficult to scale.

    Source: §1 · p. 1

  • Method

    SDE Matching parameterizes the Latent SDE posterior process to enable simulation-free training using direct access to posterior samples and Monte Carlo estimation.

    Source: §1 · p. 2

  • Results

    Across synthetic and real-world experiments, SDE Matching achieves performance comparable to or better than adjoint sensitivity-based methods while significantly reducing computational complexity.

    Source: §8 · p. 9

  • Takeaways & Limitations

    Estimating latent trajectories without solving high-dimensional SDEs makes SDE Matching particularly suitable for large-scale time-series and sequence modeling.

    Source: §8 · p. 9

    The posterior parameterization must be smooth and invertible with respect to ε and provide access to the conditional score, restricting its flexibility and the posterior approximation.

    Source: §7 · p. 9

Abstract

from arXiv · show

The Latent Stochastic Differential Equation (SDE) is a powerful tool for time series and sequence modeling. However, training Latent SDEs typically relies on adjoint sensitivity methods, which depend on simulation and backpropagation through approximate SDE solutions, which limit scalability. In this work, we propose SDE Matching, a new simulation-free method for training Latent SDEs. Inspired by modern Score- and Flow Matching algorithms for learning generative dynamics, we extend these ideas to the domain of stochastic dynamics for time series and sequence modeling, eliminating the need for costly numerical simulations. Our results demonstrate that SDE Matching achieves performance comparable to adjoint sensitivity methods while drastically reducing computational complexity.

1. Introduction

Latent SDEs model stochastic continuous-time dynamics, but simulation-based training requires costly numerical integration and backpropagation. SDE Matching uses a parameterized posterior process to train without simulation, with comparable performance and lower computational cost.

1. Introduction

  • Latent SDEs extend continuous-time modeling to stochastic and chaotic processes, which deterministic ODE approaches limit by encoding uncertainty only in initial conditions.The paper also notes applications ranging from neuroscience to video prediction.Source: §1 · p. 1
  • Latent SDE training has resisted scaling because simulation-based methods require costly integration and backpropagation through solutions and are difficult to parallelize.The paper links this training bottleneck to the computational infeasibility of adjoint training on high-dimensional problems.Source: §1 · pp. 1–2
  • The method is motivated by a connection between diffusion models and Latent SDEs, alongside simulation-free Score and Flow Matching methods for generative dynamics.The paper establishes the connection as motivation for a more efficient training procedure.Source: §1 · pp. 1–2
  • SDE Matching parameterizes the Latent SDE posterior process to provide a simulation-free training procedure for Latent and Neural SDEs.The approach directly accesses samples from the latent posterior approximation at any time, avoiding SDE integration during training; its objective is estimated by Monte Carlo with O(1) memory and time complexity.Source: §1 · p. 2
  • SDE Matching achieves comparable performance to adjoint sensitivity methods on synthetic and real data while reducing training cost and improving convergence.It also enables Latent SDE applications to high-dimensional problems where adjoint training is computationally infeasible.Source: §1 · p. 2

2. Background

A Latent SDE generates observations by sampling a latent trajectory from a prior process and then sampling observations conditioned on that trajectory. Training uses a variational posterior process because the latent posterior is generally intractable, but estimating the objective requires costly numerical integration.

2. Background

  • The model assumes observations over [0, 1] are generated by sampling an initial latent state, evolving it with a prior SDE, and sampling each observation conditioned on its latent state.The prior SDE defines latent distributions over time through its drift, diffusion, and Wiener process.Source: §2 · p. 2
  • Because the latent posterior is generally intractable, variational training introduces an observation-conditioned posterior process with its own initial distribution and drift.The posterior shares the prior process's diffusion term despite having a different drift.Source: §2 · p. 2
  • The training objective is a variational bound on the observations' log-marginal likelihood, whose stochastic-optimization estimate requires joint numerical integration of the posterior SDE and an integral.Simulation-based integration also requires backpropagation through numerical SDE solutions and can be costly, difficult to parallelize, and numerically unstable.Source: §2 · p. 2

3. Diffusion Models as Latent SDEs

Diffusion models can be viewed as a special case of Latent SDEs, and this connection motivates extending simulation-free matching methods to latent stochastic dynamics. Their forward process has a tractable conditional marginal, enabling score learning without numerical integration.

3. Diffusion Models as Latent SDEs

  • Diffusion models use a forward SDE to perturb data with noise and a reverse SDE to generate data from a prior distribution.The forward process has a closed-form conditional marginal, and the reverse process is approximated by learning its score with denoising score matching.Source: §3 · p. 3
  • The denoising score matching objective is simulation-free because samples from the forward process's conditional marginal can be drawn directly for Monte Carlo estimation.This avoids integrating the forward SDE to estimate the objective.Source: §3 · p. 3
  • A diffusion model is a special case of a Latent SDE with a particular posterior process and a single observation at t = 1.Substituting the diffusion prior and reverse-process parameterization into the Latent SDE objective yields a reweighted denoising score matching objective.Source: §3 · p. 3
  • The connection between diffusion models and Latent SDEs, together with simulation-free diffusion training, motivates a simulation-free framework for training Latent SDEs.The paper derives the conditional forward process's reverse SDE and inverts the model's time direction to establish the connection.Source: §3 · p. 3

4. SDE Matching

SDE Matching makes simulation-free Latent SDE training possible by parameterizing posterior marginals directly and deriving a compatible posterior SDE. It retains the standard prior process while changing how the training objective is estimated.

4. SDE Matching

  • Directly parameterizing posterior marginals with a learnable function enables simulation-free Latent SDE training while leaving the prior's generative dynamics flexible.The conventional posterior marginal is intractable, preventing simulation-free training.Source: §4 · pp. 3–4

4.1. Generative Model

  • The generative model retains the standard Latent SDE prior, with observations sampled conditionally independently from their latent states.The prior drift and diffusion may use neural networks of arbitrary form.Source: §4.1 · p. 4

4.2. Posterior Process

  • The posterior process is built by defining qφ(zt|X) through a noise-to-latent function, then deriving a conditional ODE and SDE with those marginals.The SDE Matching construction makes the correspondence between its posterior SDE and marginals exact.Source: §4.2 · p. 4
  • The posterior SDE shares the prior diffusion term, which affects trajectory distributions but not posterior marginals; its drift requires a conditional score that can be costly to compute.Efficient score computation is possible for parameterizations with efficient Jacobian log-determinants, including functions linear in ε or RealNVP architectures.Source: §4.2 · p. 4
  • High-dimensional evaluation of the diffusion-related drift term can be costly, but structured diffusion choices enable efficient computation without adding restrictions beyond memory-efficient Latent SDE training.Examples include scalar-valued σθ(zt, t) times matrix-valued Γθ(t), diagonal state-dependent diffusion, or state-independent diffusion, for which the term vanishes.Source: §4.2 · p. 5

4.3. Optimization

  • SDE Matching optimizes the same variational bound as adjoint sensitivity training, estimating its losses by sampling posterior latents at selected times rather than integrating the conditional SDE.The diffusion and reconstruction losses can be estimated with arbitrary numbers of parallel samples; one sample sufficed in the reported experiments.Source: §4.3 · p. 5

4.4. Sampling

  • For forecasting from partial observations, the posterior latent state can be sampled directly at the last observed time before integrating the prior; interpolation can also use posterior dynamics.Conventional posterior parameterizations require simulating the conditional SDE up to the last observed time.Source: §4.4 · pp. 5–6
  • An ODE may preserve latent and observation marginals while distorting their joint trajectory distributions, so deterministic sampling does not generally preserve time-series structure.Sequential Monte Carlo is suggested for more general conditioning and steering, but those approaches are beyond the work's scope.Source: §4.4 · p. 6

5. Experiments

The experiments assess whether SDE Matching can train Latent SDEs efficiently, rather than outperform all time-series approaches, using synthetic and real data. They report comparable or better accuracy with lower computational cost and faster convergence, including on high-dimensional video.

5. Experiments

  • The experiments aim to demonstrate computationally efficient Latent SDE training, and the authors state that SDE Matching can be combined with existing extensions for further improvements.The experiments use the same hyperparameters for conventional Latent SDE and SDE Matching parameterizations, without annealing or additional regularization.Source: §5 · pp. 6–7
  • SDE Matching achieves comparable or better accuracy than simulation-based Latent SDE training while reducing computational cost per iteration and converging in fewer iterations.The authors suggest the findings support scaling Latent SDEs to longer and higher-dimensional time series.Source: §5 · p. 7

5.1. Synthetic Datasets

  • On the synthetic Lorenz attractor, both methods learned the dynamics, but SDE Matching used one posterior-drift evaluation per iteration versus 100 simulation steps and ran about five times faster per iteration.The combined per-iteration and convergence speed-up was over 50× in this experiment.Source: §5.1 · p. 7
  • For the Lorenz experiment, adjoint-method log10 gradient norms rose from 6.26 ± 0.14 at T=1 to 8.82 ± 0.28 at T=10, while SDE Matching remained at 4.92 ± 0.24 across the tested horizons.The authors interpret this as greater stability for long time-series modeling.Source: §5.1 · p. 7

5.2. Motion Capture Dataset

  • On motion-capture data, SDE Matching achieved similar performance to the adjoint method and outperformed ARCTA, which requires about 100 latent samples and SDE evaluations per batch rather than one.The authors attribute ARCTA's weaker long-term dependency learning to its less flexible prior and posterior parameterization and reliance on observations near each time point.Source: §5.2 · pp. 7–8

5.3. Video Dataset

  • With an identical 20k training iterations, SDE Matching learned the underlying dynamics in approximately 20 minutes on a single GPU, while the adjoint sensitivity method required around 30 hours and still failed to accurately capture them.Source: §5.3 · p. 8

6. Related Work

Related work covers Neural ODEs and Latent SDEs, existing simulation-based training approaches, and newer simulation-free methods. SDE Matching is positioned as a flexible latent-space framework related to, but broader than, prior simulation-free approaches.

6. Related Work

  • Neural ODEs model continuous dynamics and have been used for time-series modeling, but their deterministic dynamics are inadequate for fundamentally stochastic processes, especially with sparse or uninformative observations.Latent SDEs are presented as better suited to inherently stochastic and chaotic processes and more robust to data shifts.Source: §6 · p. 8
  • Many Neural ODE and Latent SDE improvements still backpropagate through numerical solutions, limiting scalability, while Verma et al. (2024) instead use a site-based Gaussian approximation and fixed-point training.Earlier Latent SDE work also explored Monte Carlo training, variational inference, discretization, and stochastic optimal control.Source: §6 · p. 8
  • Course & Nair (2023) is simulation-free but restricts posterior marginals to Gaussians, excludes state-dependent prior diffusion, and relies on roughly 100 latent samples per batch.Zhang et al. (2024) learns dynamics in data space and conditions them on fixed past observations, which increases simulation cost and limits long-range dependencies.Source: §6 · p. 8
  • SDE Matching treats conventional diffusion models and several newer diffusion methods with complex or learnable noising processes as special cases of Latent SDEs.The connection follows from interpreting diffusion models as Latent SDEs, including the single-observation case.Source: §6 · p. 9

7. Limitations and Future Work

The authors identify constraints on posterior parameterization and the cost of evaluating general high-dimensional drift terms as trade-offs. They point to more expressive parameterizations and further study of latent versus data-space dynamics as future work.

7. Limitations and Future Work

  • Simulation-free training constrains Fφ to be smooth and invertible in ε and to provide the conditional score, limiting flexibility in the posterior approximation.The authors nevertheless describe it as, to their knowledge, the most flexible parameterization that enables simulation-free conditional-posterior sampling.Source: §7 · p. 9
  • Evaluating the posterior SDE drift can be computationally costly for general high-dimensional gθ(zt, t), a limitation the authors say is shared by other memory-efficient Latent SDE training methods.Source: §7 · p. 9
  • Future work includes more flexible Fφ and diffusion parameterizations and studying how flexibility affects latent-dynamics modeling versus modeling dynamics in the original data space.The authors also identify audio and video generation, simulation-free Latent ODE training, and reinforcement-learning policies as potential extensions.Source: §7 · p. 9

8. Conclusion

The conclusion presents SDE Matching as a simulation-free Latent SDE training method that parameterizes posterior marginals directly. Experiments show comparable or better performance with lower computational complexity, while posterior flexibility remains a trade-off.

8. Conclusion

  • SDE Matching eliminates costly numerical simulations by directly parameterizing marginal posterior distributions, supporting efficient training and conditional inference.Source: §8 · p. 9
  • Across synthetic and real-world experiments, SDE Matching achieved performance comparable to or better than adjoint sensitivity methods while reducing computational complexity.The authors identify estimating latent trajectories without solving high-dimensional SDEs as a reason the method is suited to large-scale time-series and sequence modeling.Source: §8 · p. 9
  • SDE Matching trades some posterior-parameterization flexibility for simulation-free latent-variable sampling, and the authors identify more expressive posterior functions as future work.Source: §8 · p. 9

A.1. Marginal Distribution

  • The marginal posterior is constructed by transforming noise ε with Fφ for each time t and observation sequence X, with the stated density derivation assuming invertibility and differentiability in ε.Source: §A.1 · p. 13

A.2. Conditional ODE

  • A smooth, time-indexed invertible Fφ defines a flow whose infinitesimal generator gives a conditional ODE matching the flow's solution trajectories.Source: §A.2 · p. 13

A.3. Conditional SDE

  • The conditional SDE construction uses the equivalence between SDE and probability-flow-ODE marginals; a chosen diffusion term g determines the corresponding drift h.The cited result is stated to hold under suitable conditions on the drift and diffusion terms.Source: §A.3 · pp. 13–14

A.4. Parameterization

The appendix describes an affine parameterization of the approximate posterior and conditions for it to match the true posterior. In a linear stochastic system, the SDE Matching ELBO closely converges to the true log-marginal likelihood, with faster convergence and shorter training iterations than the adjoint method.

A.4. Parameterization

  • The posterior marginals are parameterized implicitly by an invertible transformation of noise, enabling score computation from the noise log-density and the inverse transformation’s Jacobian log-determinant.The appendix uses an affine transformation with mean μφ and matrix-valued σφ, and leaves more flexible alternatives for future research.Source: §A.4 · p. 14

B. Theoretical Considerations

  • The approximate posterior matches the true posterior when their drift terms, or their corresponding deterministic processes, match.Equivalently, Fφ must learn the solution trajectories of the ODE defined by the true posterior process for the variational bound to become tight.Source: B. Theoretical Considerations · p. 15
  • The true posterior is intractable, but a sufficiently flexible reparameterization function could in principle learn it and make the variational bound equal the negative log-marginal likelihood.The true posterior process is derived using Doob’s h-transform, with future observations denoted by Xt.Source: B. Theoretical Considerations · p. 15

C.1. Convergence to Likelihood

  • The SDE Matching ELBO closely converges to the true log-marginal likelihood and requires approximately 20× less time per training iteration than the adjoint sensitivity method.The experiment uses a linear stochastic system with coefficients F(t) = −t, L(t) = t, H(t) = 1, and R(t) = 0.01.Source: §C.1 · p. 15

C.2. Robustness of Gradients

In the linear-system robustness experiment, SDE Matching exhibits lower gradient norms and reduced gradient variance than the adjoint sensitivity method.

C.2. Robustness of Gradients

  • SDE Matching consistently exhibits lower gradient norms than the adjoint sensitivity method.The comparison uses training-time dynamics in the linear-system experiment.Source: §C.2 · p. 15, Figure 4 · p. 16
  • SDE Matching consistently exhibits reduced gradient variance compared to the adjoint sensitivity method.Figure 4 presents the linear-system training dynamics for both methods.Source: §C.2 · p. 15
  • The experiment assesses robustness by comparing gradient norms and variances over training time.It uses the setup from Appendix C.1.Source: §C.2 · p. 15

C.3. Importance of State-Dependent Diffusion Term

In the stochastic Lotka–Volterra experiment, the state-dependent volatility model captures the correct trajectory structure, whereas the state-independent model does not. The authors connect state-dependent volatility to applications across several fields and to better MOCAP performance relative to state-independent ARCTA.

C.3. Importance of State-Dependent Diffusion Term

  • The experiment motivates state-dependent diffusion as a capability absent from the earlier simulation-free approach ARCTA (Course & Nair, 2023).The authors identify the Lotka–Volterra system as an example of a biological system with state-dependent volatility.Source: §C.3 · p. 16
  • The state-dependent volatility model captures the correct Lotka–Volterra trajectory structure, whereas the state-independent model fails to do so.The experiment uses volatility coefficient σ = 0.15 and follows the setup of Course & Nair (2023).Source: §C.3 · p. 16
  • State-dependent volatility is relevant to finance, biology, neuroscience, and causal inference, and contributes to the MOCAP results.In MOCAP, state-dependent SDE Matching and Li et al. (2020) perform much better than state-independent ARCTA (Course & Nair, 2023).Source: §C.3 · p. 16

D. Implementation Details

The experiments largely follow prior setups while modifying the posterior parameterization for SDE Matching. The implementation prioritizes a fair comparison with standard Latent SDE training over optimizing model architectures.

D. Implementation Details

  • The experiments follow Li et al. (2020) in prior-process parameterizations and hyperparameters, including observation models and coordinate-wise diagonal diffusion.The posterior parameterization is modified: a time-reversal GRU aggregates observations, and an MLP predicts μφ and log σφ from its context and time.Source: D. Implementation Details · p. 17
  • For moving-pendulum videos, the posterior uses the preceding experiments’ parameterization but replaces the MLP with a convolutional encoder for image observations.Data generation and the prior-process parameterization follow Course & Nair (2023).Source: D. Implementation Details · p. 17
  • The comparison prioritizes fairness against standard Latent SDE training rather than optimizing prior and posterior architectures, which the authors believe could be parameterized more efficiently.Training uses the same hyperparameters, optimizer, and iteration count without additional regularization, reweighting, or annealing.Source: D. Implementation Details · p. 17
Loading 2502.02472v3…