We thank all reviewers for constructive feedback for which we will include in revision. We are encouraged that reviewers recognise our work as addressing an important problem with a novel approach. Below we will first state our core contributions and address common concerns, and then answer reviewer specific questions.

Summary of contributions

We propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include

• Introduction of stochastic flows to the field of sequential deep generative models, enabling one-step computation of conditional distributions of an SDE at arbitrary time points,
• A conditional normalisation flow that preserves important properties of stochastic flows,
• The “flow loss” as a regulariser to enforce flow consistency.

Reviewers TqQA and YJDb endorsed the novelty of our contributions on one-step SDE conditional distribution evaluations, noting NSF’s mathematical soundness and substantial speedups. When compared with baselines, Reviewer aZhw commended on NSF being significantly faster while achieving comparable performance regarding accuracies.

Common questions

Scalability (YJDb, yfA9)

Regarding computation, our architecture is based on affine coupling-based flows, which is scalable, e.g., our stochastic moving MNIST experiment is conducted on videos with frame size 64 x 64, and RealNVP [12] and Glow [27] (based on affine coupling flows) has been applied to images. Future work will incorporate our proposed constraints into more expressive normalising flow architectures, e.g., TarFlow [51].

Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)

Nonlinear SDEs naturally capture real-world stochastic dynamics in human motion, where random perturbations interact nonlinearly with the system state. This generalises better than ODEs (deterministic) and linear SDEs (state‑independent noise).

Among nonlinear SDE approaches, Latent SDE models rely on controlled SDEs to represent the posterior, which involves learning complex drift function, that is difficult to optimise. SDE Matching addresses challenges by learning posteriors conditioned on the observations, $p(x_t|o_{t_1},o_{t_2},\ldots)$, but must learn complex, time-varying distributions. In contrast, our Latent NSF takes a simpler yet effective way: directly modeling transition laws $p(x_{t_j}|x_{t_i})$, using normalising flows. This avoids needing both control-based SDEs and complex time-dependent posteriors.

Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)

We conducted ablation studies on challenging partial observation tasks. We trained on Stochastic Lorenz Attractor data where only a random continuous segment of length 0.5 from the full trajectory [0,1] was observed.

Flow Loss Ablation on Missing Stochastic Lorenz

Effect of Auxiliary Optimisation Steps (λ=0.4)

* Simultaneous: auxiliary and main models updated jointly without inner loops

Directional KL Components (λ=0.4, K=5)

Without flow loss (λ = 0), the model fails to generalise to unobserved regions, with KL ≈ 23 at $t=1.0$. Adding flow loss dramatically improves performance; with λ = 0.2–0.8, KL stays around 1.2 across time steps.

Applying flow loss in only one direction already reduces KL divergence. However, since forward KL encourages broad coverage of the distribution, and reverse KL focuses on matching high-density regions, we decided to use both equally to balance these effects with symmetry.

A similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.

Main paper’s content organisation (TqQA, yfA9)

We will move key architectural details (NSF design, training algorithm) to the main text for clarity.

Specific questions from Reviewer yfA9

Thank you for your valuable feedback. We address your concerns regarding experiments below.

Stochastic Lorenz Attractor

1. The variable t represents the prediction time horizon from the initial state. Specifically, we evaluate the KL divergence between the true distribution $p(x_t|x_0)$ and our model's predicted distribution at times t ∈ {0.25, 0.5, 0.75, 1.0}. The non-monotonic KL divergence is natural for chaotic systems, since dynamics of chaotic attractors exhibit complex convergence/divergence patterns.
2. In this setting, NSF is trained to predict arbitrary time intervals within [0,1] in a one step. During inference, we can either make one-step predictions with $t_\mathrm{max}=1.0$, or use recursive predictions with smaller $t_\mathrm{max}$ values (e.g., 4 steps of 0.25 each to reach $t=1.0$). The optimal t_max is not trivial to determine a priori, as stated in our Limitations section. The model may have learned better representations at certain time scales during training, leading to this non-monotonic behaviour.
3. Figure 3(d) shows trajectories generated using one-step predictions from NSF. Each line in the visualisation shows an sample from the learned conditional distribution $p(x_t|x_0)$, where we predict from the initial state $x_0$ to time t in a one step, with the same random sample from Gaussian. Table 1 evaluates the same model but provides KL divergence. The $t_\mathrm{max}$ parameter indicates the maximum one-step prediction interval for NSF, e.g., $t_\mathrm{max}=0.5$ means that reductions up to $t=0.5$ are one-step; beyond $t=0.5$ requires two steps.

MoCap

1. You are correct that Li et al. [32] reported MSE 4.03 for Setup 1, while ours shows 12.91. This difference stems from known reproducibility issues with the original Latent SDE implementation on the MoCap dataset. This issue is also raised by multiple researchers in public discussions (e.g., official GitHub repo issue trackers from 2022). To ensure fair comparison, we re-implemented the method using consistent data processing and evaluation procedures. Our reported values for Latent SDE (12.91 ± 2.90) are consistent with the accuracy-runtime curve in the Appendix and with other reproduction reports mentioned in these public discussions. We will add a footnote in the camera-ready version explicitly noting this reproducibility issue.
2. The increasing MSE with runtime for Latent NSF reflects a characteristic of our training approach. The model is trained across the full interval $[0, t_\mathrm{max}]$, it learns transitions for all time gaps, but it remains an open question about whether one-step or recursive predictions performs the best at inference. In MoCap, one-step prediction yields the best balance of accuracy and efficiency.
3. For explanation for the inductive bias, see common rebuttal. The key difference between Setup 1 and Setup 2 is in the temporal split. Setup 1 trains and tests within the same time range (interpolation), and Setup 2 trains only on steps 1–200 and tests on unseen future steps 201–300 (extrapolation), requiring generalisation beyond the training horizon.

Stochastic Moving MNIST

1. For computing these metrics, we used 256 generated samples for each of 32 different input conditions. We agree that comparing distributions in such high-dimensional spaces through finite samples is inherently challenging, which contributes to the high variance observed. This is a fundamental limitation when evaluating generative models for high-dimensional video data. However, as in Appendix F, our validation experiments show that these metrics do capture meaningful differences between static and dynamic aspects of the generated videos, despite the high variance, providing qualitative insights into the relative performance of the methods.
2. The trade-off between Static FD and Dynamics FD is controlled by the β parameter, which balances reconstruction error and the KL divergence term. We agree that analysing how this trade-off changes with β for both Latent SDE and Latent NSF would provide valuable insights. If accepted, we will include this analysis in the camera-ready version's appendix, showing how different β values affect the Static/Dynamics FD balance for both methods. The speed comparison is provided in Table 7 (Appendix). The difference between recursive and one-step NSF is tiny since the computational bottleneck is the decoder forward pass, not the NSF. The performance differences relate to which prediction horizons achieve highest accuracy for each model and dataset. As in our limitations, determining optimal prediction horizons a priori remains challenging. The flow loss is designed to maintain consistent accuracy across all prediction horizons, but dataset-specific characteristics still influence which sampling strategy performs best.

General questions

1. We followed the same dimensionality as baseline papers standard (Stochastic Lorenz Attractor: [32], MoCap: [2]). For Stochastic Moving MNIST, we followed the dataset's original paper [10]. We conducted systematic searches for key hyperparameters. For β (KL weight) and flow loss weight, as is standard for sequential VAE, we performed logarithmic grid search (…, 0.1, 0.3, 1.0, …) for both ours and baselines.
2. Noted in the common rebuttal “Scalability”.

References

For 1-50, please refer to the main text.

[51] Zhai et al., "Normalizing Flows are Capable Generative Models", ICML 2025.

We hope this addresses your concerns and would be happy to provide any additional clarifications.

I would like to thank the authors for their thorough answers to my questions. It cleared up most of my doubts about the paper, so I am increasing my score to 4. Nevertheless, if the paper is accepted, I insist that the authors add more details to the discussion about the difficulty of choosing the $t_{max}$ parameter. In particular, which factors might affect the optimal choice. I would also appreciate the addition of the number of parameters of each model right after its description in the appendix in order to convince the reader that all models in comparison have similar sizes.

We thank all reviewers for constructive feedback for which we will include in revision. We are encouraged that reviewers recognise our work as addressing an important problem with a novel approach. Below we will first state our core contributions and address common concerns, and then answer reviewer specific questions.

Summary of contributions

We propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include

• Introduction of stochastic flows to the field of sequential deep generative models, enabling one-step computation of conditional distributions of an SDE at arbitrary time points,
• A conditional normalisation flow that preserves important properties of stochastic flows,
• The “flow loss” as a regulariser to enforce flow consistency.

Reviewers TqQA and YJDb endorsed the novelty of our contributions on one-step SDE conditional distribution evaluations, noting NSF’s mathematical soundness and substantial speedups. When compared with baselines, Reviewer aZhw commended on NSF being significantly faster while achieving comparable performance regarding accuracies.

Common questions

Scalability (YJDb, yfA9)

Regarding computation, our architecture is based on affine coupling-based flows, which is scalable, e.g., our stochastic moving MNIST experiment is conducted on videos with frame size 64 x 64, and RealNVP [12] and Glow [27] (based on affine coupling flows) has been applied to images. Future work will incorporate our proposed constraints into more expressive normalising flow architectures, e.g., TarFlow [51].

Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)

Nonlinear SDEs naturally capture real-world stochastic dynamics in human motion, where random perturbations interact nonlinearly with the system state. This generalises better than ODEs (deterministic) and linear SDEs (state‑independent noise).

Among nonlinear SDE approaches, Latent SDE models rely on controlled SDEs to represent the posterior, which involves learning complex drift function, that is difficult to optimise. SDE Matching addresses challenges by learning posteriors conditioned on the observations, $p(x_t \mid o_{t_1}, o_{t_2}, \ldots)$, but must learn complex, time-varying distributions. In contrast, our Latent NSF takes a simpler yet effective way: directly modeling transition laws $p(x_{t_j}\mid x_{t_i})$, using normalising flows. This avoids needing both control-based SDEs and complex time-dependent posteriors.

Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)

We conducted ablation studies on challenging partial observation tasks. We trained on Stochastic Lorenz Attractor data where only a random continuous segment of length 0.5 from the full trajectory $[0,1]$ was observed.

Flow Loss Ablation on Missing Stochastic Lorenz

Effect of Auxiliary Optimisation Steps ($\lambda=0.4$)

* Simultaneous: auxiliary and main models updated jointly without inner loops

Directional KL Components ($\lambda=0.4, K=5$)

Without flow loss ($\lambda =0$), the model fails to generalise to unobserved regions, with KL $\simeq$ 23 at $t=1.0$. Adding flow loss dramatically improves performance; with $\lambda = 0.2\text{-}0.8$, KL stays around 1.2 across time steps.

Applying flow loss in only one direction already reduces KL divergence. However, since forward KL encourages broad coverage of the distribution, and reverse KL focuses on matching high-density regions, we decided to use both equally to balance these effects with symmetry.

A similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.

Main paper’s content organisation (TqQA, yfA9)

We will move key architectural details (NSF design, training algorithm) to the main text for clarity.

Specific questions from Reviewer aZhw

Thank you for your valuable feedback. We address your specific concerns below.

Motivation for Bidirectional KL Divergence in Flow Loss

We appreciate this opportunity to clarify our design choice.

In principle, a flow consistency loss should make sure the agreement between the one-step transitions $p_\theta\left(x_{t_k} \mid x_{t_i}\right)$ and the two-step composition $\int p_\theta\left(x_{t_k} \mid x_{t_j}\right) p_\theta\left(x_{t_j} \mid x_{t_i}\right) d x_{t_j}$ by minimising a statistical divergence/distance between the two. However, the two-step composition involves integration over $x_{t_j}$, which is not tractable. This makes the choice of statistical divergence tricky and motivates our development of the flow loss.

• Why KL divergence: Among many divergence choices, KL divergence is appealing, because by introducing a variational auxiliary distribution $q_\phi\left(x_{t_j} \mid x_{t_i}, x_{t_k}\right)$, we can minimise the upper-bounds of both the forward and reverse KL in a tractable way, and this is the core idea of our flow loss. These bounds only require sampling from $q_\phi$ and evaluating log-densities from the one-step flow, making training stable and low-cost. Other options like Wasserstein distance are computationally expensive, Stein discrepancies require (intractable) score functions and careful tuning, and adversarial methods often lead to unstable training.
• Why bi-directional KL: Although our ablation study shows that using forward- or reverse-only KL can already improve performance, we argue that combining both directions with equal weight is a more principled choice. The forward KL encourages the one-step transition to cover the support of the two-step composition, and vice versa. Importantly, both losses share the same auxiliary distribution $q_\phi\left(x_{t_j} \mid x_{t_i}, x_{t_k}\right)$, so optimising either direction helps refine $q_\phi$, and improvements in $q_\phi$ in turn tighten the bound in both directions. This can stabilise training and encourages $q_\phi$ to approximate the true bridge distribution more efficiently. Moreover, the symmetric formulation avoids introducing directional bias into the consistency regularisation.

Ablation Studies on Regularization Term

Please see “Hyperparameter Sensitivity and Ablation Studies” in common response.

SDE Matching Comparison (Table 2)

SDE Matching's code was not publicly available at NeurIPS submission time, and their paper didn’t have sufficient implementation details for reproduction. After submission deadline, SDE Matching paper authors have since released their Stochastic Lorenz attractor implementation. Still SDE Matching’s code for other experiments is not available.

We have now conducted the comparison on Stochastic Lorenz attractor. Overall our method achieves better KL divergence with more than 2 orders of magnitude fewer FLOPs.

We also tried to extend their Stochastic Lorenz implementation to Motion Capture, but it did not reproduce their reported scores. We have contacted the authors, they confirmed that an official implementation for Motion Capture experiment will be released soon, but not before the NeurIPS author feedback period ends. Hopefully we can include comprehensive SDE Matching comparisons for all tasks in our camera-ready version. If their Motion Capture implementation is still unavailable by camera-ready deadline, we will provide results using our best-effort reproduction with appropriate caveats about potential implementation differences.

References

For 1-50, please refer to the main text.

[51] Zhai et al., "Normalizing Flows are Capable Generative Models", ICML 2025.

Thank you again, and we hope this addresses your concerns and would be happy to provide any additional clarifications.

We thank all reviewers for constructive feedback for which we will include in revision. We are encouraged that reviewers recognise our work as addressing an important problem with a novel approach. Below we will first state our core contributions and address common concerns, and then answer reviewer specific questions.

Summary of contributions

We propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include

• Introduction of stochastic flows to the field of sequential deep generative models, enabling one-step computation of conditional distributions of an SDE at arbitrary time points,
• A conditional normalisation flow that preserves important properties of stochastic flows,
• The “flow loss” as a regulariser to enforce flow consistency.

Reviewers TqQA and YJDb endorsed the novelty of our contributions on one-step SDE conditional distribution evaluations, noting NSF’s mathematical soundness and substantial speedups. When compared with baselines, Reviewer aZhw commended on NSF being significantly faster while achieving comparable performance regarding accuracies.

Common questions

Scalability (YJDb, yfA9)

Regarding computation, our architecture is based on affine coupling-based flows, which is scalable, e.g., our stochastic moving MNIST experiment is conducted on videos with frame size 64 x 64, and RealNVP [12] and Glow [27] (based on affine coupling flows) has been applied to images. Future work will incorporate our proposed constraints into more expressive normalising flow architectures, e.g., TarFlow [51].

Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)

Nonlinear SDEs naturally capture real-world stochastic dynamics in human motion, where random perturbations interact nonlinearly with the system state. This generalises better than ODEs (deterministic) and linear SDEs (state‑independent noise).

Among nonlinear SDE approaches, Latent SDE models rely on controlled SDEs to represent the posterior, which involves learning complex drift function, that is difficult to optimise. SDE Matching addresses challenges by learning posteriors conditioned on the observations, $p(x_t \mid o_{t_1}, o_{t_2}, \ldots)$, but must learn complex, time-varying distributions. In contrast, our Latent NSF takes a simpler yet effective way: directly modeling transition laws $p(x_{t_j}\mid x_{t_i})$, using normalising flows. This avoids needing both control-based SDEs and complex time-dependent posteriors.

Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)

We conducted ablation studies on challenging partial observation tasks. We trained on Stochastic Lorenz Attractor data where only a random continuous segment of length 0.5 from the full trajectory $[0,1]$ was observed.

Flow Loss Ablation on Missing Stochastic Lorenz

Effect of Auxiliary Optimisation Steps ($\lambda=0.4$)

* Simultaneous: auxiliary and main models updated jointly without inner loops

Directional KL Components ($\lambda=0.4, K=5$)

Without flow loss ($\lambda =0$), the model fails to generalise to unobserved regions, with KL $\simeq$ 23 at $t=1.0$. Adding flow loss dramatically improves performance; with $\lambda = 0.2\text{-}0.8$, KL stays around 1.2 across time steps.

Applying flow loss in only one direction already reduces KL divergence. However, since forward KL encourages broad coverage of the distribution, and reverse KL focuses on matching high-density regions, we decided to use both equally to balance these effects with symmetry.

A similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.

Main paper’s content organisation (TqQA, yfA9)

We will move key architectural details (NSF design, training algorithm) to the main text for clarity.

Specific questions from Reviewer YJDb

Thank you for your thoughtful review and constructive feedback. We appreciate your recognition of our work's mathematical soundness and strong empirical performance. Below we address your concerns.

Hyperparameter sensitivity

We understand your concern about multiple hyperparameters. Please see "Hyperparameter Sensitivity and Ablation Studies" above, where we present comprehensive ablation results demonstrating the method's robustness.

Scalability and Architecture

Please see “Scalability” above in common response, where we discuss the limitation on the scalability and architecture.

Sampling procedure for expectation

Thank you for your careful reading and specific questions. The sampling procedure in Eq. 8 corresponds to the nested expectation $\mathbb{E}\_{ p\_\theta(x_{t_k} | x_{t_i}) } \left[ \mathbb{E}\_{q_\phi (x_{t_j} |x_{t_i}, x_{t_k}) }[ \cdot ]\right]$. We first sample the initial state $x_{t_i}$ from the dataset. Given a time triplet $(t_i, t_j, t_k)$ sampled according to the strategy described in Appendix l.868, we sample $x_{t_k} \sim p_\theta\left(x_{t_k} \mid x_{t_i}\right)$ using the neural stochastic flow. Then, conditioned on $x_{t_i}$ and $x_{t_k}$, we sample $x_{t_j} \sim q_\phi\left(x_{t_j} \mid x_{t_i}, x_{t_k}\right)$ from the auxiliary bridge distribution.

Similarly, this expectation in Eq. 10 corresponds to the nested expectation $\mathbb{E}\_{p\_\theta(x_{t_j} \mid x_{t_i})}\left[\mathbb{E}\_{p\_\theta(x_{t_k} \mid x_{t_j})}[\cdot]\right]$. We first sample the initial state $x_{t_i}$ from the dataset. Then, using our neural stochastic flow (NSF) model, we sample $x_{t_j} \sim p_\theta\left(x_{t_j} \mid x_{t_i}\right)$, and subsequently sample $x_{t_k} \sim p_\theta\left(x_{t_k} \mid x_{t_j}\right)$ using the same model.

We will clarify these expectation in camera-ready version.

Inductive bias of SDEs

Please see “Clarifying inductive bias of SDEs in MoCap dataset” above in common response.

Minor Points

Thank you for the reference to Progressive Distillation. We'll add this important connection and will also fix the typo on l130.

References

For 1-50, please refer to the main text.

[51] Zhai et al., "Normalizing Flows are Capable Generative Models", ICML 2025.

We hope this addresses your concerns. We're committed to making these clarifications prominent in the final version to help future practitioners effectively use our method.

We thank all reviewers for constructive feedback for which we will include in revision. We are encouraged that reviewers recognise our work as addressing an important problem with a novel approach. Below we will first state our core contributions and address common concerns, and then answer reviewer specific questions.

Summary of contributions

We propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include

• Introduction of stochastic flows to the field of sequential deep generative models, enabling one-step computation of conditional distributions of an SDE at arbitrary time points,
• A conditional normalisation flow that preserves important properties of stochastic flows,
• The “flow loss” as a regulariser to enforce flow consistency.

Reviewers TqQA and YJDb endorsed the novelty of our contributions on one-step SDE conditional distribution evaluations, noting NSF’s mathematical soundness and substantial speedups. When compared with baselines, Reviewer aZhw commended on NSF being significantly faster while achieving comparable performance regarding accuracies.

Common questions

Scalability (YJDb, yfA9)

Regarding computation, our architecture is based on affine coupling-based flows, which is scalable, e.g., our stochastic moving MNIST experiment is conducted on videos with frame size 64 x 64, and RealNVP [12] and Glow [27] (based on affine coupling flows) has been applied to images. Future work will incorporate our proposed constraints into more expressive normalising flow architectures, e.g., TarFlow [51].

Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)

Nonlinear SDEs naturally capture real-world stochastic dynamics in human motion, where random perturbations interact nonlinearly with the system state. This generalises better than ODEs (deterministic) and linear SDEs (state‑independent noise).

Among nonlinear SDE approaches, Latent SDE models rely on controlled SDEs to represent the posterior, which involves learning complex drift function, that is difficult to optimise. SDE Matching addresses challenges by learning posteriors conditioned on the observations, $p(x_t \mid o_{t_1}, o_{t_2}, \ldots)$, but must learn complex, time-varying distributions. In contrast, our Latent NSF takes a simpler yet effective way: directly modeling transition laws $p(x_{t_j}\mid x_{t_i})$, using normalising flows. This avoids needing both control-based SDEs and complex time-dependent posteriors.

Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)

We conducted ablation studies on challenging partial observation tasks. We trained on Stochastic Lorenz Attractor data where only a random continuous segment of length 0.5 from the full trajectory $[0,1]$ was observed.

Flow Loss Ablation on Missing Stochastic Lorenz

Effect of Auxiliary Optimisation Steps ($\lambda=0.4$)

* Simultaneous: auxiliary and main models updated jointly without inner loops

Directional KL Components ($\lambda=0.4, K=5$)

Without flow loss ($\lambda =0$), the model fails to generalise to unobserved regions, with KL $\simeq$ 23 at $t=1.0$. Adding flow loss dramatically improves performance; with $\lambda = 0.2\text{-}0.8$, KL stays around 1.2 across time steps.

Applying flow loss in only one direction already reduces KL divergence. However, since forward KL encourages broad coverage of the distribution, and reverse KL focuses on matching high-density regions, we decided to use both equally to balance these effects with symmetry.

A similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.

Main paper’s content organisation (TqQA, yfA9)

We will move key architectural details (NSF design, training algorithm) to the main text for clarity.

Specific questions from Reviewer TqQA

We sincerely thank you for your thorough review and positive assessment of our work. We are delighted that you found our approach technically sound and appreciate your recognition of our contributions to computational efficiency in SDE modeling. Below we address your specific questions:

Algorithm Placement

We appreciate your suggestion to move the training algorithm from Appendix C to the main text, as we wrote in “Main paper’s content organisation”. We agree this would improve readability and help readers better follow the training process.

Minor Corrections

Thank you for catching typo and wrong reference. We will fix them in the camera-ready version.

References

For 1-50, please refer to the main text.

[51] Zhai et al., "Normalizing Flows are Capable Generative Models", ICML 2025.

Thank you again, and we hope this addresses your concerns and would be happy to provide any additional clarifications.

Dear Authors,

Thank you for your great effort to respond and have further detail discussions about the approach. The responses adequately address my questions and comments in the review.

Kind regards