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

We are delighted to see that the reviewer finds our approach Innovative, efficient and note the reduction in computational cost of training. Below, we address the questions raised in the review.

Questions:

In all our experiments, including the 3D stochastic Lorenz attractor, the motion capture dataset, and additional experiments (see rebuttal to TN6r, Question 3), with the same parameterization and training hyperparameters (such as the number of training steps), SDE Matching requires approximately $5$ times less computation time than the adjoint sensitivity method per iteration. Asymptotically, a single training iteration of SDE Matching takes $\mathcal{O}(1)$ time, while the adjoint sensitivity method takes $\mathcal{O}(L)$, where $L$ is the number of steps in the simulation of the posterior SDE (Eq. 2). However, in practice, the exact training time depends not only on the length of simulation, but also on the evaluation cost of other components, such as the calculation of the prior and reconstruction losses (Eqs. 22 and 24), which affects the actual ratio.

Additionally, as demonstrated in Figure 2, SDE Matching may exhibits significantly faster convergence compared to the adjoint sensitivity method, further accelerating the overall training procedure. When combining the per-iteration improvement and the convergence speed improvement SDE Matching is ~500 times faster than the baseline in the experimental setup from Section 4.1.

At the same time, in SDE Matching, the generative model defined by the prior process (Section 4.1) is exactly the same as in training the Latent SDE using the adjoint sensitivity method. Therefore, the inference time for unconditional sampling is identical for both approaches. However, as discussed in Section 4.4, in the case of forecasting, if we have access to partial observations $x_{t_1}, \dots, x_{t_N}$, SDE Matching enables sampling of the latent state $z_{t_N}$ at the time of the last observation $t_N$ in $\mathcal{O}(1)$ time. In contrast, the conventional parameterization of the posterior process requires simulating the conditional SDE up to $t_N$. This property of SDE Matching can significantly reduce inference time when observations span a long period. However, in our experiments, we only have access to $3$ samples, all close to $t = 0$. Consequently, most of the inference time is spent simulating the prior process from $t_N$, making the forecasting inference time nearly equal for SDE Matching and the adjoint sensitivity method.

We would like to note that, in principle, SDE Matching allows simulation-free sampling of latent states $z_t$ at arbitrary time steps, potentially enabling fully simulation-free interpolation and forecasting in $\mathcal{O}(1)$ time. However, such inference of the posterior process would require a corresponding training procedure in which the model observes only a few early observations. Otherwise, the posterior process may produce biased samples.

We are pleased that the reviewer finds our claims to be well supported theoretically, the experiments well designed, and the discussions insightful. We are especially grateful that the reviewer took the time to reproduce our experiments. We sincerely thank the reviewer for such attention to our paper and for the valuable feedback. In response to the comments and questions raised, we provide the following clarifications and commitments.

Other Comments

1. Indeed, in SDE Matching we do not explicitly solve any SDE during training. Instead, we parameterize the posterior process in such a way that allows us to estimate the objective function in $\mathcal{O}(1)$ space and time. The purpose of Table 1 is primarily to reflect the computational budget in terms of number of drift/diffusion term evaluations per training iteration.

   Nevertheless, we would like to note that in SDE Matching, when sampling latent variables $z_t$ to estimate the diffusion loss (Eq. 23), we sample them from the posterior marginals $q_\phi(z_t|X)$, which, by design, are solutions of the posterior SDE (Eq. 19). Thus, during optimization, we are effectively computing gradients with respect to the solutions of the posterior SDE.

2. In the Latent SDE model, training essentially requires two things: the ability to sample from the posterior marginals $q_\phi(z_t|X)$ and access to the posterior SDE. With the conventional parameterization, we directly define the SDE via its drift, which necessitates numerical simulation in order to sample from the unknown marginals.

   The high-level idea of SDE Matching is to define the posterior process in a different order: we first parameterize a sampler of the marginal distribution and, based on it, derive the posterior SDE. Therefore, we may say that sampling from $q_\phi(z_t|X)$ implicitly solves the posterior SDE.

   It is also important to emphasize that the correspondence between the posterior SDE and its marginals in SDE Matching is exact, not approximate.

Questions:

1. There is a connection with PINNs in the tools used to define the model. Like PINNs, we use automatic differentiation to compute time derivatives (Eq. 18) and score functions (Eq. 26). We also employ a partial differential equation (the Fokker–Planck equation) to establish connections between elements of the dynamical system.

   However, unlike typical PINNs that solve PDEs, we use similar mechanics to compute coefficients of an SDE based on its solutions. So, while we believe there are conceptual connections, these are significantly different methods solving different problems.

2. The difficulty of solving the posterior SDE is reflected in the flexibility of the reparameterization function $F_\phi(\epsilon,t,X)$ (Eq. 16). As discussed in Section 7, the SDE Matching parameterization limits the flexibility of the posterior process.

   If the true posterior process is simple and has tractable marginals, solving it should not be computationally expensive. However, if the posterior is complex, an insufficiently flexible parameterization may not accurately capture the underlying dynamics.

   How complex the posterior latent process needs to be is an open question. Especially in the presence of a flexible observation model.

3. In this work, we chose a simple parameterization for the posterior process in order to stay close to the setup of Li et al. However, the reviewer’s observations motivate us to explore more expressive parameterizations, and we plan to include those results in a future revision.

   In general, we believe it may be more challenging to marginalize posterior processes of more complex forms, and that the results could be more sensitive to architectural choices in the parameterization network.

   We believe that with denser observations, the posterior process may be better approximated by Gaussian marginals. In contrast, when observations are sparse, intermediate distributions tend to resemble the unconditional marginals, which may have arbitrary shapes.

4. Small values of $g_\theta$ may introduce numerical instability, however this issue is not specific to SDE Matching, as the term involving $g_\theta^{-1}$ appears in the ELBO of Latent SDE models in general.

5. Yes, Eq. 27 implies that the posterior marginals $q_\phi(z_t|X)$ are conditionally Gaussian.

6. Yes, providing time steps to the reparameterization function $F_\phi$ makes perfect sense when working with observations at arbitrary time points.

7. In practice, the terms in the objective (Eq. 21) may indeed have different scales (e.g., if the observation model has very low variance). However, the total variational bound remains valid, and all terms should still be minimized to match the true dynamics.

We trust that the clarifications and additional discussions will strengthen your support for acceptance! If accepted, we will update the camera-ready version to reflect this discussion.

We appreciate the reviewer’s positive feedback regarding the speed, simplicity, and sensibility of our method, as well as the discussion of its connections to diffusion models. Below, we address the questions and comments raised in the review.

References:

• We will cite the Nature paper by Course & Nair in the camera-ready version if accepted.

Weaknesses:

1. To demonstrate the scalability of SDE Matching, we provide additional experiments (please refer to rebuttal to TN6r, Question 3).

2. We would like to point out that, compared to the work of Course & Nair, in addition to allowing more flexible marginals of the posterior process, SDE Matching also supports a state-dependent diffusion term $g_\theta(z_t,t)$. This is important in cases where the underlying dynamics exhibit state-dependent volatility.

   To demonstrate this difference, we designed an experiment similar to the setup in Section 4.1, using a stochastic Lotka–Volterra system with a dynamic of the following form:

   $dx=(\alpha x-\beta xy)dt+\sigma xdw$, $dy=(\delta xy-\gamma y)dt+\sigma ydw$

   We then applied SDE Matching to train models with both state-independent and state-dependent volatility functions. Visualisations of the learned trajectories are available at the anonymised link: https://imgur.com/a/jh8sM0Q. It is evident that the model with state-independent volatility fails to capture the correct form of the trajectories.

3. Please refer to the rebuttal to JkpK for the discussion on training and clarifications on inference (including interpolation).

Other Comments:

We would like to highlight that, in contrast to methods such as Flow Matching, both SDE Matching and the adjoint sensitivity method do not match marginals. Instead, they match the distributions of trajectories of the posterior and prior processes. SDE Matching essentially proposes an alternative parameterization of the posterior process that, by design, enables training of the Latent SDE model by minimising the same objective, but in a simulation-free manner.

Nevertheless, the question of the limitations on the tightness of the ELBO in SDE Matching is interesting. Due to the limited space, we may not include rigorous derivations, but we outline the theoretical aspects of this question below.

In the case of the Latent SDE, the variational bound becomes tight when the posterior process matches the true posterior of the prior process.

The prior process (Eq. 15) is given by:

$dz_t=h_\theta(z_t,t) dt+g_\theta(z_t,t)dw$

Given a series of observations $X$, the true posterior process can be derived via Doob’s $h$-transform, yielding: dz_t=\Big[\overbrace{h_\theta(z_t,t)+g_\theta(z_t,t)g^\top_\theta(z_t,t)\nabla_{z_t}\log p_\theta\left(X_t|z_t\right) }^{h_\\theta(z_t,t,X)} \Big] dt+g_\theta(z_t,t)dw,

where $X_t=\\{x_s:x_s\in X,t\leq s\\}$ denotes the set of future observations from time $t$ onward.

We can also consider the corresponding deterministic process that follows the true posterior marginals $p_\theta(z_t|X)$:

$\frac{dz_t}{dt}=\bar{h}^\theta(z_t,t,X)=h_\theta(z_t, t)+g_\theta(z_t,t)g^\top_\theta(z_t,t)\left[ \nabla_{z_t}\log p_\theta\left(X_t|z_t\right)-\frac{1}{2}\nabla_{z_t}\log p_\theta\left(z_t|X\right)\right]-\frac{1}{2}\nabla_{z_t}\cdot\left[g_\theta(z_t,t)g^\top_\theta(z_t,t)\right]$

Now, consider the approximate posterior process (Eq. 19):

$d z_t=f_{\theta,\phi}(z_t,t,X)dt+g_\theta(z_t,t)dw$

The approximate process matches the true posterior process if the drift terms are equal: $f_{\theta,\phi}(z_t,t,X)\equiv h_\theta(z_t,t,X)$, or equivalently, if the corresponding deterministic processes (Eq. 17) match: $\bar{f}^{\theta,\phi}(z_t,t,X)\equiv \bar{h}^\theta(z_t,t,X)$. In terms of the reparameterization function $F_\phi(\epsilon,t,X)$ (Eq. 16) that defines the approximate posterior process, we may say that the variational bound becomes tight, and the approximate posterior process matches the true posterior, if $F_\phi$ learns the solution trajectories of the ODE defined by $\bar{h}_\theta$.

Therefore, with a sufficiently flexible reparameterization function $F_\phi$, the SDE Matching approach should, in principle, be capable of making the variational bound tight.

Questions:

SDE Matching shows similar performance to Li et al. as there is an overlap of the confidence intervals. We attribute the slightly better performance of Li et al. to a suboptimal parameterization of the posterior process. In this work, we intentionally kept the parameterization as close as possible to the conventional setup from Li et al. to ensure a fair comparison. In principle, with a sufficiently flexible parameterization, SDE Matching should demonstrate comparable performance.

We will include additional discussions and experimental results with detailed setup explanations in the camera-ready version. We trust these clarifications, highlighting the novelty and scalability of our work, will enhance your support for acceptance!

We thank the reviewers for their valuable feedback, which helps us to improve the paper. Below, we address the questions and comments raised in the review.

Other Comments:

We apologise for any confusion regarding terminology. Allow us to provide a clearer definition of the term “simulation-free”.

In the diffusion literature, simulation-free implies that explicit simulation of dynamics (numerical integration) is not required during training. For example, the adjoint sensitivity method [1] requires numerical simulation of the posterior SDE (Eq. 2) to sample latent states $z_t \sim q_\phi(z_t|X)$ to evaluate the loss function (Eq. 5). In contrast, thanks to our alternative parameterization of the posterior process, SDE Matching allows, by design, direct sampling of $z_t$ via inference of the reparameterization function $F_\phi(\epsilon, t, X)$ (Eq. 16), eliminating the need for numerical simulation of the posterior SDE. This makes the training process simulation-free.

Questions:

1. Consistency Models (CM) as originally introduced in [2] learns a function $f_\theta(x_t, t)$ that approximates the solutions of a fixed marginal ODE starting at $x_t$ and ending at some $x_0 \sim p_{data}(x_0)$.

   In SDE Matching, we introduce a reparameterization function $F_\phi(\epsilon, t, X)$ (Eq. 16). Based on $F_\phi$, we derive a conditional ODE (Eq. 17) and then a conditional SDE (Eq. 19) that define the posterior process.

   While both approaches involve a pair of an ODE and its solver, the nature of these ODEs differs: in CM, the goal is to approximate solutions of an unconditional ODE, whereas in SDE Matching we derive an exact ODE as an intermediate step used to define a conditional SDE. Thus, while there is a connection, they are fundamentally different approaches.

2. SDE Matching does not incur additional learning cost by parameterizing the posterior process via the reparameterization function $F_\phi(\epsilon, t, X)$ (Eq. 16) — this is a central point of our approach. Unlike diffusion models, the posterior process in a Latent SDE is not parameter-free. It is a complicated process conventionally defined by the drift function $f_\phi(z_t, t, X)$ in the SDE (Eq. 2), and sampling from the posterior requires simulating this SDE.

   By reparameterizing the posterior process, we neither introduce new parameters nor increase the complexity of the dynamics. Instead, we reparameterise the same dynamics in a different way, enabling direct sampling of latent variables $z_t$ from the posterior marginals without numerically integrating the SDE.

   The key advantage of SDE Matching is that a single training iteration takes asymptotically less time (see rebuttal to JkpK). As discussed in Section 5.1 (lines 320–322), for the 3D Lorenz attractor dataset, one iteration of SDE Matching takes approximately $5$ times less time. Figure 2, while not the main result, additionally shows that SDE Matching not only has faster iterations but also converges more quickly.

3. To demonstrate the scalability of SDE Matching, we provide two additional experiments.

   First, to evaluate scalability with respect to time, we use the experimental setup described in Section 4.1 and compare the gradient norms of the objective function with respect to the model parameters for both SDE Matching and the adjoint sensitivity method. When integrating over time horizons $T \in {1, 2, 5, 10}$, the adjoint sensitivity method yields the following $\log_{10}$-based gradient norms: $6.26 \pm 0.14$, $6.95 \pm 0.20$, $7.91 \pm 0.23$, and $8.82 \pm 0.28$. This demonstrates that the gradient norms grow exponentially as the time horizon increases.

   In contrast, SDE Matching maintains a stable $\log_{10}$-based gradient norm of $4.92 \pm 0.24$ across all time horizons, indicating better stability for long time series modelling.

   Second, to demonstrate scalability with respect to dimensionality, we designed an experiment where the model learns a sequence of images (a video) depicting a moving pendulum. SDE Matching successfully captures the underlying dynamics and generates realistic samples.

   You can find examples of generated sequence at the anonymised link: https://imgur.com/a/aMfSuHB.

4. Please refer to rebuttal to Reviwer C8BY (section “Other Comments”), where we discuss the tightness of the variational bound in SDE Matching.

5. We hope that the clarifications highlighting the efficiency of our approach, along with additional discussion and experimental results demonstrating the scalability of SDE Matching, will strengthen your support for the acceptance of our paper.

We will include clarifications and additional results, along with detailed explanations of the experimental setups, in the camera-ready version if accepted.

[1] Li et al. “Scalable gradients for stochastic differential equations”

[2] Song et al. “Consistency models”