Neural Stochastic Flows: Solver-Free Modelling and Inference for SDE Solutions
Authors: Naoki Kiyohara
Edward Johns
Yingzhen Li
Venue: NeurIPS 2025 poster ()
Keywords: Neural stochastic flows
stochastic differential equations
normalising flows
continuous-time models
latent dynamics
time series modelling
Primary area: probabilistic_methods
TL;DR: A solver-free approach for modelling stochastic differential equations using conditional normalising flows that enables direct sampling between arbitrary time points within trained horizon.

Abstract: Stochastic differential equations (SDEs) are well suited to modelling noisy and/or irregularly-sampled time series, which are omnipresent in finance, physics, and machine learning applications. Traditional approaches require costly simulation of numerical solvers when sampling between arbitrary time points. We introduce Neural Stochastic Flows (NSFs) and their latent dynamic versions, which learns (latent) SDE transition laws directly using conditional normalising flows, with architectural constraints that preserve properties inherited from stochastic flow. This enables sampling between arbitrary states in a single step, providing up to two orders of magnitude speedup for distant time points. Experiments on synthetic SDE simulations and real-world tracking and video data demonstrate that NSF maintains distributional accuracy comparable to numerical approaches while dramatically reducing computation for arbitrary time-point sampling, enabling applications where numerical solvers remain prohibitively expensive.

Neural Stochastic Flows: Solver-Free Modelling and Inference for SDE Solutions
Naoki Kiyohara 1.2+ Edward Johns'
Yingzhen Li!
'Imperial College London "Canon Inc.
(n. ln.kiyohar223, e. johos,, yingzhen yingzhen.11j@immerri. 1i)@imperial......m ac. uk
Abstract
Stochastic differential equations (SDEs) are well suited to modelling noisy and/or pregularly-samplee time series, which are omnipresent in finance, physics, and machine learning applications. Traditional approaches require costly simulation of numerical solvers when sampling between arbitrary time points. We introduce Neural Stochastic Flows (NSFs) and their latent dynamic versions, which learn (latent) SDE transition laws directly using conditional normalising flows, with architectural constraints that preserve properties inherited from stochastic flow. This enables sampling between arbitrary states in single step, providing up to two orders of magnitude speedup for distant time points. Experiments synthetic SDE simulations and real-world tracking and video data demonstrate that NSF maintains distributional accuracy comparable numerical approaches while dramatically reducing computation for arbitrary time-point sampling, enabling applications where numerical solvers remain prohibitively expensive.
1 Introduction
Stochastic differential equations (SDEs) underpin models in finance, physics, biology, and modern machine learning systems [39 50] they capture how state 2; E R² evolves following velocity field whilst being perturbed by random noise. In many real-time settings such as robots, trading
algorithms, digital twins, often requires the transition law p(æ; xy.) describing the probability
distribution of future states given earlier states over arbitrary time gaps f [30]- Conventional approaches handle this transition law by learning neural (latent) SDEs [33 from data and simulating numerical solvers with many small steps 119 300, which incurs high computational costs. In this regard, neural flow 4 techniques for ordinary differential equations (ODEs) bypass numerical integration by directly learning the flow map with architectural constraints and regularisation terms. However, these methods inherently cannot express stochastic dynamics required for SDE modelling. In the domain of diffusion models 12
20 47 line of work leverages the associated probability flow ODE (PF-ODE): earlier distillation approaches compress iterative samplers by matching teacher- student trajectories under the PF-ODE 44], while consistency-ste methods learn direct mappings between time points (48, with trajectory-level variants also proposed [26.. However, these techniques are tied to particular boundary conditions and diffusion processes.
An alternative approach for modelling stochastic dynamics data is through stochastic flows 32) which, under suitable conditions, describe families of strong solutions to SDEs via mappings that evolve initial states over time under shared stochasticity. These flows naturally define the transition law p(æ, wg), and when parameterised by neural networks, enable direct learning of efficient one-step sampling between arbitrary time points. However, it remains an unsolved challenge for such network design, due to the requirements of satisfying several properties, e.g., identity mapping when
Project page: https://nkkyyharrt httpyy/oarrr.ggguubb... wtrechb...//ss--nuuiippp./.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(a) Sample trajectories
(b) Conditional distribution
(c) One- step transition
p(aulaa)
State ar
(d) Two-step transition
Flow loss [p(oedleslplmmmaamddd p(anlan))
Time
Time
Figure 1: Comparison between (a) traditional neural SDE methods requiring numerical integration and (b) our NSF, where blue and red contours represent conditional distributions for one-step sampling and recursive application, respectively. Panels (c) and (d) illustrate our flow loss concept, ensuring distributional consistency between one-step and two-step transitions through intermediate states,
t= s, Markov property, and Chapman Chapman-Kolmogeen Kolmogorov flow property that ensures self-consistency of transition laws with different number of steps.
Contributions. We introduce Neural Stochastic Flows (NSFs) to address the network design challenge of parameterising stochastic flows. NSFs employ conditional normalising flows directly learn the SDE transition distributions p(z; x.) for any < t, where the normalising flow architecture design ensures the validity identity, Markov, and (for autonomous SDEs) stationarity properties. Specifically, given the present state and elapsed time At, NSF draws a single Gaussian noise vector and transforms it through specially designed affine coupling layers 13 28 that reduce to identity maps when At on (x, At, s) and scale their effect with the time interval. A bi-directional KL divergence based regularisation loss is designed to further encourage the network to satisfy Chapman-Kolmogon flow property. Since all transformations are bijective, the transition log-density is available in closed form, allowing both training and inference to proceed without SDE solvers. A side-by-side schematic of solver-based vs. NSF sampling is shown in Fig.
to
[4]. These transformations have parameters conditioned
To summarise, Neural Stochastic Flows can:
Learn SDE's weak solution, in the form of conditional distribution, directly for solver- free training and inference, with architectural constraints and loss function design to enforce desirable stochastic flow properties;
Enable efficient one-step sampling between arbitrary time points within the trained horizon (thus suitable for modelling irregularly sampled time series), with maximum gains for distant time points;
Be extended to noisy and/or partially observed data scenarios through latent dynamic modelling, in similar fashions variational state-space model and latent SDEs 31 33
Empirically, across diverse benchmarks such stochastie Lorenz attractor, CMU Motion Capture, and Stochastic Moving MNIST, our approach maintains distributional accuracy comparable to or better than numerical solver methods, while delivering up two orders of magnitude faster predictions over arbitrary time intervals, with the largest gains on long-interval forecasts.
Background
Stochastic Differential Equations. Stochastic differential equations (SDEs) model dynamical systems subject random perturbations:
dx; = pl(,,, t) dt a dW
(1)
where E R is the state, / R x R + Rd is the drift, B R4 Raxm is the diffusion, and W, is an m-dimensionaa Wiener process. Neural SDEs parameterise these terms with neural networks, enabling learning from data. In many applications, the conditional probability distribution

---

p(a; (x)) for 8 < t is crucial for uncertainty quantification and probabilistic forecasting, However, while numerical solvers can generate sample trajectories, they do not provide direct access this distribution and incur computational costs that scale with the time gap / 8.
Stochastic Flows of Diffeomorphisms. An alternative perspective SDEs comes from the theory of stochastic flows, which characterises the solution space through time-dependent mappings. For an Itô SDE in Eq. (T) with smooth coefficients, these solutions form a stochastic flow of diffeomorphisms B2. This is defined by measurable map Φ,,t R x S + R" where < & < +oo and 92 is the sample space, satisfying:
1. Diffeomorphism: For any s < t and w E Ω2 the map + &s (x, w) is almost surely diffeomorphism.
2. Independence: Maps Pia.t2..... Ptn-1.fn for any non-overlapping sequence s S S
S in are independent, which yields the Markov property of the flow.
3. Flow property: For any 0 S t S t2 S t3. any point æ E IR", and any w E Q::
P (r, w) Dig (Di1st (ar, w) w)
4. Identity property: For any 0, any point æ 12" and any f2: 22, w) = x.
Furthermore, when the SDE is autonomous, where the drift and diffusion terms are independent of time, an additional property holds:
5. Stationarity: For any S r.r 0, and w Ω2 the maps ( and Φ 2str.44 w)
have the same probability distribution.
Solver-based neural SDEs (also see Section 5 require costly numerical integration. Leveraging the flow properties formalised above, in the next section we introduce a solver-free alternative, NSFs, as efficient models for stochastic systems that provide direct access to transition probability distributions of SDE-governed processes.
3 Neural Stochastic Flows
While stochastic flows of diffcomorphisms provide a powerful theoretical framework for understand- ing strong solutions of SDEs, they require modelling infiite-dimmesiionll sample paths, which is computationally intractable. Instead, in this section we focus on developing neural network architec- tures for modelling the probability distributions that characterise weak solutions of SDEs, by deriving appropriate conditions and loss functions for the network using properties of stochastic flows.
For SDE, the relationship between strong solution (represented by stochastic flow a)) to its corresponding weak solution (represented by a conditional probability distribution) is expressed as:
plart Fe. 1, - 4.) [ (ari, - -D..e, (Itz wi) p(w) dw,
where 8(-) denotes the Dirac delta function and represents the dimensodii sample space containing complete Brownian motion trajectories. Instead of directly modelling the map &, approximate the resulting probability distribution po (or, it,i4, t2 - 4,) x p(æ,, Ze. tj 4)) using a Neural Stochastic Flow (NSF), parametric model fo that transforms Gaussian samples into the desired distributions via conditional normalising flows:
po (Ii, i4, -t,) = f
6(1z, =
fo (zt, t,,t, t, t), A( 0, de,
(3)
Therefore So serves as our parametric model to approximate the SDE's weak solution.
Conditions. Under the NSF framework, in below we reformulate the conditions of stochastic flows as strong solutions of SDEs to conditions of NSFs as weak solutions:
1. Independence: For any sequence 0 < t S 12 S t., the conditional probabilities
(x,, t1, - t1),... Po (t... xt,. 1. - -)) must be independent.
2. Flow property: For any < t, S to S tx. the joint distribution must satisfy the Chapman-
Kolmogorov equation:
po(x Izi; ti,th to tt)) = po( De,ity,te t))po(2r,, I baitt,,, t, tt) dxz,. (4)
3

---

3. Identity property: At the same time t. the distribution must reduce to delta function, indicating no change: Po(a;, = æ tt,it;,0) 8( xz,).
For autonomous SDEs, we restate the stationarity condition in terms of our parametric model:
4. Stationarity: The conditional distributions must exhibit stationarity such that for any &,, ty, x;; 1,, - (4) = po æ;; 4 + r, r, (2).
0, popoo
3.1 Conditional Normalising Flow Design
Let c (æt, A1,44) denote the conditioning parameters, where i is included only for non- autonomous systems and omitted otherwise, and let A == by 4. We instantiate the sampling procedure of the flow distribution (Eq. (3))
St. + At (L) + MLP O,, E~N(0,11)
(5)
(e)
a(c)
T1 fo(z, c) = fi(:e, 0z) o f1-1(;c.01-1) ..... f((z;,
(6)
Our architecture integrates a parametric Gaussian initialisation with sequence of bijective transfor- mations. The state-dependent Gaussian, centred at μıc) with noise scale (c). follows the similar form to the Euler-Maruyama discretisation with drift scaled by At and diffusion by V The subsequent bijective transformations f1 through f L are implemented
conditioned coupling flows 13 [8] whose parameters depend on c. Each layer splits the state z into two partitions (#A, ZB) and applies affine update one partition conditioned on the other and on
f((z; c, = Concat AAA @exp At MLP(1) scale (ZA, c; 0 ( ALMLP shift (zA, C; shin))
(6)
(7)
with alternating partitions across layers. The explicit A/ factor ensures fi c. 0:) z when At and, combined with the form of the base Gaussian (Eq. (5)). preserves the identity zero time gap. Stacking such layers yields an expressive diffcomorphism [49 while keeping the Jacobian log-determinant tractable for loss computation. Independence property ensured by the conditional sampling on the initial state Dt, without any overlap between the transitions, and stationarity is obtained by omitting f; from e when modelling autonomous SDEs.
3.2 A Regularisation Loss for Flow Property
To train an NSF, apart from using supervised learning loss functions based on e.g., maximum likeli- hood, we add an additional discrepancy term Laow that encourages the flow property as formalised Eq. We require a measure of mismatch between the one-step distribution po 2²;)) and the two-step marginal f ( ot,) PI | Fr;)) dæy,, Several candidates exist such as divergences because
Stein, adversarial, and kernel MMD. Among them, we choose KL minimise the upper-bounds of the KL divergences in a tractable way, and since KL asymmetric, we use both directions: the forward KL promotes coverage, while the reverse KL penalises unsupported regions. Direct KLs remain intractable due the marginalisation over in- termediate state tt, in two-step side. Therefore, we introduce a bridge distribution be (21, It,, xt,)) auxiliary variational distribution [I] 40 45 and optimise variational upper bounds for both directions (Appendix B
Specifically, for the forward KI. divergence:
DKL. po (Itr (ere) /
(Ier I (aee)) po (arey wx.) date,)
log po (20 t xe.)) be (20r, 1t, Dte
E mi ~Po (mn.))
E my Peg-be (1)
pe (z1z x2,,) po (xee, xz,)
=: Lnow, 2 (0, E; 4k).
(8)
And for the reverse KL divergence:
pe (xr xx.)) DKL (ʃ po (ara xi1)) po (xi, xc,) dæt,

---

Algorithm 1: Optimisation procedure for (Latent) Neural Stochastic Flows
Follow for NSF, for Latent NSF; lines without markers are run for both models.
Input: Transition dataset D; learning rates Tjo Te (and nes Opp for latent); inner steps K
Output: Trained parameters 0, & (and Φ, & for latent)
while not converged do
Sample minibatch from D
Encode sampled data into latent states using Eq. (5)
Sample time triplets (f,, ty, tx) with & < tj ty (see Appendix for details)
for k 1 to K do
Compute Laow (0,E) E) using Eq. (0
Update bridge parameters: 5 5
7/g Veenow
Compute total loss L(O, {) using Eq. 11
Compute total loss L(O, E. Φ v)) using Eq.
18
Update NSF parameters: 0 < 0 B) VoL
Update NSF/encoder//eecoder parameters: {0, Φ, wp + (0.0. - 110.0.01 V 130
Po (Iz, (t..) po (It, xt,))
E 21, ~pol (-12;;)
log be (It, Dt., Dt..) (I xt,.)
Laow, 2-1o-1 (0, E: ty, tyy
The flow loss Low combines the above two terms, balancing the consistency in both directions: + Lnow,
Roo (6, (0,E)
E [CRow, 1-002 (
(10)
The total loss for training NSFs combines the negative log-likelihood objective with the flow loss:
C(O
E (m., Fey (e.t,))))
llog pe (x:, Dt, t4, t - [:]] + À Lnow ( (0, E),
(11)
where D represents the dataset and X is hyperparameter that controls the strength of the flow consistency constraint. The model parameters 9 and bridge model parameters are optimised using gradient-based methods; procedure is summarised in Algorithm| (see Appendix C.I for details).
Latent Neural Stochastic Flows
Iregularly-sampler real-world sequences seldom expose the full system state; instead we observe noisy measurements o,, at times to t1 < ... tr. To model the hidden continuous time dynamics while remaining solver-free, we introduce the Latent Neural Stochastic Flows (Latent NSFs) as variational state-space models (VSSMs) whose transition kernels are governed by NSFs.
4.1
Recap: Variational State-Space Models (VSSMs)
VSSM
31 extends the variational autoeneoder framework 29 to sequential data by
endowing sequence of latent states To.7 with Markovian dynamics while explaining the observations Op. through a conditional emission process. The joint distribution factorises as
Po. (20nr; 0o:T) = Po (m)) II (a 2+ 1 ax:);
(12)
where each conditional is typically Gaussian with its distributional parameters produced by neural networks. Amortised inference is carried out using Gaussian encoder qo q x;-1, o;) for every time-step f, and the encoder and decoder are jointly trained using the negative evidence lower bound (NELBO) loss, also known the variational free energy (VFE):
[(,, 1b ④) E 9
log pe (Ko)
log pe (x; xx-1) log (oz (x.) log go (æ, |2,- 0:)]
(13)
5

---

At Ats
Af
Afg
Z1a
O
On
(a) Variational State-Space Model
(b) Latent Neural Stochastic Flow
(c) Skip KL Divergence
Figure 2: Graphical models. (a) Standard state-space model with discrete-time transitions between
states, (b) NSF model with continuous-úmme transitions parameterised by time intervals (At; == 1) (c) Example of skip KL divergence over two time steps. The upper row shows the two- step-ahead prediction from the posterior at to via NSF, while the lower row shows the posterior at t2 conditioned on all observations (ot,, Ot, 0t2) ). Across all models: solid arrows represent generative processes, dashed arrows represent inference processes, wavy arrows in (c) represent KI. divergences.
4.2 Neural Stochastic Flows as Latent Transition Kernels
To model
time series, we require the latent dynamics to evolve in continuous- time. We therefore implement the transition distribution pe(x; xt-1) (12) with an NSF pe (z Dt 14-1; At,) which parameterises the weak solution of an SDE follows. As illustrated in Fig. 2 a). a standard VSSM uses discrete-time transitions, whereas latent NSFs in Fig. 2 incorporates continuous-time dynamics through time intervals.
Generative model. Let Øt E IRd be the latent state and At; == 14 - 11-1- The joint distribution factorises as
Pe. Pioer Ota:r)
pe (2 II pe Dt, Ft. it--1 At.) Pap (o:. xx,))
(14)
where po (3 Dt; . is NSF, hence supports arbitrary A4; without numerical integration.
Variational posterior. For the inference model, we employ a recurrent neural network encoder that processes the observation sequence 77
(a(ator
(cstr)
II (ax, m., diag(si diag(s?.)) (mi;, 8t.) GRU Ali, Li, hi; .): t=0
(15)
with h O, and the absolute time (4-1 in Eq. and 4 in Eq. non-autonomous systems and omitted when modelling autonomous SDEs.
15 is only included for
Learning objective. Our goal is to train generative model Pe.v v (Fto Otoir and inference model
9 (Doorr Octr). A standard choice is the 3-weigtted negative ELBO (8-NELBO loss [18
Ca-NELBO Σ
log (o., [xx.)] + 3DL 3DK (ao, loss,) || po (æt, [=) (16)
E Q
i=0
which focuses only on adjacent time steps, potentially leading to the error accumulation for long-term dependencies. To address this, propose to add skip-ahead KI divergence loss (Fig. 2c)):
Cakip
E [DKL (g0(2,, loss,) po [x,,)]]-
(17)
=-0 (-2((++27))
Unlike traditional overshooting methods [66] that require recursive transitions, latent NSF enables direct sampling across arbitrary time gaps, enabling this objective to be computed efficiently.
Combining the above two losses and the flow loss (Eq. (10),, we get the total loss:
Lital NE-NLLO + X Laow Baip Laipp
(18)
where À and Baip are hyperparameters that control the strength of the flow loss and skip-ahead KL divergence loss, respectively.
training procedure for the latent model is summarised Algorithm (see Appendix C.2

---

Related Work
We categorise prior work in terms of whether learning and/or sampling avoid fine-grained numerical time-stepping, and the class of dynamics targeted, as summarised in Table T
Modelling general ODEs. Neural ODEs [6] learn parametric vector field that is integrated by a numerical solver at both training and test time. Their runtime there- fore scales with the number of function evaluations required by the solver. Neural flows 4 side-step this cost by learning the solution map directly.
Table 1: Solver requirement across continuous time differential equation models. 'Pre-defined diffusion SDEs/ODEs' refers to a boundary-conditioned diffusion process bridging fixed base distribution and the data distribution utilised in diffusion models.
(ethod(s)
Target
Solver-free training
Solver-free inference
dynamics
Modelling prescribed SDEs/PF-ODEs. Score-based generative models [47 and flow matching [34] are continuous-time generative models that learn reverse SDEs/PF-ODEs of noundary-conditions diffusion processes. Fast generation meth- ods have been developed through learn-
General ODEs
Neural ODES
Neural flows
ODEs
Score-based diffusion verse SDEsPF-ODES
Pro-defined
(77
diftusion
x
matching
SDEWODEE
Pre-defined
vector- field straightening [34]]]]]]] 53 However, they specifically designed for prescribed boundary-conditioned diffusion processes, and do not provide general transition den- sity for arbitrary SDEs.
ing direct mappings 26
48
Progressive distillation
SDESODES
diffusion
Pre-defined
Consistency models
45
diffusion
rectifled flows 35
SDESODES
Neural (latent) SDEs
25
53
General Itô
BX ARCTA
46 50
SDEs
9
SDE matching
General Tta
SDEs
Modelling general SDEs. Neural (latent) SDEs 25
Neural Stochastic Flows
General Itō
33
38 50 utilise neural net-
SDEs
works to model the drift and diffusion terms
of an SDE and approximate its trajectories
via stochastic solvers, whose computational cost grows linearly with the prediction horizon. Recent advances have aimed to mitigate this cost by improving sampling efficiency [46] or by introducing
solver-free approaches such as ARCTA [ and SDE matching (3) which learn latent SDEs by
modelling posterior marginals and aligning them to the prior dynamics. However, all these methods remain solver-dependent at inference time.
Position of the present work. NSF unifies the solver-free philosophy of neural flows with the expressive power of neural (latent) SDEs. It learns the Markov transition density as a conditional normalising flow, which satisfies the conditions of weak solutions SDEs by design and regularisa- tion objectives. Unlike diffusion model-specifi. accelerations, NSF handles arbitrary Itô SDEs and extends naturally to latent sequence models. The resulting method removes numerical integration while retaining closed-form likelihoods.
6 Experiments
We evaluate NSF and latent NSF on three diverse tasks: modelling a synthetic stochastic Lorenz
attractor 33 36), predicting real-world human motion (CMU Motion Capture [15), and generative video modelling (Stochastic Moving MNIST (TI. Our goal is to show that (latent) NSF matches or improves upon the accuracy -ask-speccfiic metrics) of solver-basee neural SDEs while drastically reducing computational cost in terms of FLOPs and runtime. Full details are in Appendix E
6.1 Stochastic Lorenz Attractor
We first test NSF on the stochastic Lorenz system 33 36 standard chaotic dynamics benchmark.
We use the setup from Li et al. 33 generating 1,024 training trajectories and comparing against
baselines including latent SDE 33 SDE matching 3 and Stable Neural SDE variants (Neural LSDE, Neural GSDE, Neural LNSDE) (38) as baselines. Performance is measured by KL divergence (estimated via kernel density estimation) between generated and true distributions, and computational cost via FLOPs and runtime (details in Appendix E.2)
7

---

Table 2: Comparison on KL divergence and FLOPs different time steps. Hpred indicates the maximum single-step prediction horizon for NSF; see text for details on recursive application. Average runtime per 100 samples for latent SDE (JAX): 124-148 ms; NSF (JAX): 0.3 ms (see
Appendix E.2.5
0.25
0.75
KI.
KFLOPs
K1.
KFLOPS
K1.
KFLOPs
KL.
kFLOPs
Latent SDE
26
2.1 0.9
959
1.8
0.1
1,917
0.9
0.3
2,839
1.5
0.5
3,760
Neural LSDE
1.3 0.4
1,712
7.2
1.4
3,416
74.5
21.6
5,057
53.1
29.3
6,699
Neural GSDE
[8
1.2
1,925
3.9
+ 0.3
3,848
20.2
7.6
5,898
14.1
8.4
7,548
Neural LNSDE
58
1.7 0.4
0.4
1,925
4.6 1:
0.7
3.848
57.3 1
16.4
5,698
44.6
23.4
7,548
SDE matching
0.0001
4.7
184,394
5.3 1.0
368,787
3.4 0.8
553,034
1.0
737.354
A 0.01
6.3
1,917
11.7 = 0.5
3,834
7.9 0.3
5,677
6.0 0.3
7,520
NSF (ours)
Musl 1.0
0.8
53
1.3
53
0.6
0.3
53
0.2
0.6
53
H mod
24
1.9
1.3
0.1
53
1.0 0.4
H 0.25
1.2
53
1.2
0.1
105
0.8 0.4
156
1.3 =
0.8
208
Fig. 3 shows paths from each method; NSF traces visually coincide with the true at- tractor whereas solver-based methods over- spread. Quantitatively in Table 2 NSF pred 1.0, single-step) achieves the low- est KL divergence and with significantly fewer FLOPs. This leads substantial runtime savings (approx. 0.3 ms vs. >> 100 ms per batch, Appendix (.2.5). No- tably, our flow-based training objective en- ables both single-step and recursive appli- cation variants to maintain relatively low KL divergence across different time hori- zons, with NSF remaining computationally cheaper in these configurations. This con- firms NSF's capability to accurately model complex SDEs efficiently.
Table 3: Test MSE and 95% confidence interval based t-statistic Motion Capture datasets. indicates results from Yildiz et al. from Course and Nair 9 and from Ansari al. [2]. All other results are our reproductions. Average run- time per 100 samples: latent SDE (JAX) 75ms; Latent NSF (JAX) 3.5ms (detailed in Appendix E.3.7).
(54) from Li et al. [33],
Methods
Setup
Setup
npODE 17
22.96'
Neural ODE
22.49 = 0.88
ODE2 KL
54
8.09 = 1.95t
Latent
5.98 0.28
31.62 1 0.05
Latent SDE
83
12.91 1 2.9
9.52 0.21
Latent Approx SDE (46
7.55 = 0.05
10.50
ARCTA 9
7.62 0.93*
+
NCDSSM
5.69 0.01
4.74 0.01
6.2 CMU Motion Capture
SDE matching
3
5.20 0.43
4.26 = 0.35
Latent NSF (ours)
8.62 + 0.32
3.41 + 0.27
We evaluate Latent NSF on the CMU Mo- tion Capture walking dataset [5]. The 23- sequence walking subset is down-sampled
300 time steps and split 16/3/4 for train/validattinneess, matching prior work 51 two established protocols.
54 We report
(a) Ground truth
(b) Latent SDE
(c) SDE matching
(d) NSF (onc-step, ours)
Figure 3: Comparison generated samples (64 samples per panel) on the stochastic Lorenz attractor. Baseline methods are simulated step-by-step. For NSF, each point is an independent sample from the learnt conditional distribution p(a; æ..) originating from the same initial state and seed, connected visually for comparison. This visualisation choice aids in assessing distributional accuracy over time.

---

(a)
(b)
30
40
Latent SDE
Latent SDE
Latent NSF
30
Latent NSF
20
MSE
MSE
20
10
10
10
10
10
10a
10
10"
10a
10a
10a
10*
Runtime (ms, log scale)
Runtime <- (ms, log scale)
Figure 4: Trade-off between prediction accuracy (MSE) and computational cost (runtime, using JAX) for latent SDE and Latent NSF on CMU Motion Capture dataset. (a) Setup Within-horizon forecasting, (b) Setup 2: Beyond-horize extrapolation.
Within-horizon forecasting (Setup 1) [33 54]. All 300 time steps are used during training; for tests, only the first three observations are revealed and the model must forecast the remaining 297 steps.
Beyond-horizon extrapolation (Setup 2) 2 The last third segment of every sequence (steps 200-299) is withheld from training. At test time the model receives the first 100 observations and must predict the next 200 steps that lie beyond the training horizon.
Table 3 shows the results. In Setup I, Latent NSF (single-step prediction) performs similarly to latent SDE models. Crucially, in the Setup 2 extrapolation task, latent NSF achieves state-of-the-ar MSE.
These extrapolation gains align with the relevance of state-dependent stochasticity human motion. Nonlinear SDE formulations often require learning complex, time-varying posterior structures (e.g.. controlled SDEs or complicated conditional marginals). By contrast, Latent NSF directly models the Markov transition with conditional normalising flows, sidesteps these optimisation challenges.
Fig. Iillstrates the trade-off between prediction accuracy and computational cost. While latent SDE adjusts the discretisation step size to balance accuracy and speed, Latent NSF controls this trade-off by varying the number of recursive applications. Latent NSF achieves better performance while being approximately more than two orders of magnitude faster than latent SDE in both setups, as detailed in Appendix E.3.7
6.3 Stochastic Moving MNIST
Finally, we test Latent NSF on high- ditheeiieniionll video using a Stochastic Moving MNIST |I] variant with physically plausible bouncing dynamics: digits undergo perfect specular reflection with small angular noise, avoiding the unphysical velocity resets of the original code. The task involves modelling two digits moving 64x64 frame. We train on 60k sequences and test on 10k held-out sequences. We evaluate using Fréchet distances (FD) on embeddings from a pre-trained SRVP model 14, measuring static content, dynamies, and frame similarity (validated in Appendix G We compare our Latent NSF model against the latent SDE 33 as the main baseline, representing the current representative solver-based neural SDE model.
Table and Fig. 5) show that Latent NSF (recursive and one-step) achieves competitive Static FD and Frame-wise FD compared to the latent SDE baseline. Dynamics FD for recursive NSF is slightly higher in this run, but visualisations suggest comparable dynamics capture (see Appendix E.4.5) The decline after step 35 across all models reflects the digits approaching a uniform spatial distribution, indicating a limitation of this metric when evaluating long-term predictions in feature spaces fitted to Gaussian distributions. Nevertheless, Latent NSF offers significant potential for computational speedup over the solver-based latent SDE (e.g., latent SDE: 751 ms vs. Latent NSF (one-step): 358
'Li et al. 33 report 4.03 + 0.20 for Setup 1; our re-implementation and other reproductions observe higher MSE. Details provided in Appendix E33.2
BBartosh et al. [ report 4.50 + 0.32 for Setup 1; however, at the time of writing, official code for CMU
Motion Capture is not yet available. We report our re- re-implementation here to maintain consistency across Setups I and 2. Details of our reproduction provided in Appendix E.3.3

---

Table 4: Fréchet distance comparison Stochastic Moving MNIST using pre- pre-trained SRVP embeddingg. The lower the better. Note that dynamics FD is ill- defined for one-step simulation of latent NSF, which only predicts the marginal distribution of the future frames conditioned on the last observation.
Freghet distance
1.5
(one-slop)
Latent
SDF
1.0
0.5
0.0
25
30
35
40
45
Methods
Static FD
Dynamics FD
Frame-wise FD
Time
Latent
2.66 0.87
+ 3.10
Figure 5: Comparison of frame-wise Fréchet distance across time-steps on Stochastic Moving MNIST.
Latent NSF (recusiive)
2.36 + 0.60
7.76 + 2.56
0.63 + 0.17
Latent NSF (ome-step)
1.67 +: 0.47
0.63 +: 0.15
ms for full-sequence prediction; see Appendix E.4.5). (..45. high-dimensional stochastic sequence modelling.
This demonstrates NSF's promise for scaling
7 Discussion
Neural Stochastic Flows (NSFs) represent a novel paradigm for continuous-time stochastic modelling by directly learning weak solutions to SDEs conditional distributions. Through carefully designed architectural constraints and regularisation objectives, NSF circumvents numerical integration while providing closed-fore transition densities via normalising flows. The latent NSF extension enables applications to partially observed or high-dimensional time series, preserving efficient one-step transitions within principled variational state-space modelling framework.
Our comprehensive evaluation across chaotic dynamical systems, human motion capture, and video sequences demonstrates that the proposed flow-based approach achieves comparable superior performance relative to solver-based neural SDEs while reducing computational requirements by two orders of magnitude. Moreover, it attains state-of-the-art long-horizon extrapolation performance in latent settings. These results collectively establish that learning distributional flows, rather than the underlying vector fields required by numerical solvers, constitutes an effective and computationally efficient paradigm for real-time stochastic modelling.
Limitations. Several limitations remain. First, Chapman Chapman-Kolmogenn Kolmogorov consistency is enforced only approximately via variational bounds, potentially yielding discrepancies beyond the training regime. We recommend systematic monitoring of flow losses on held-out data with appropriate adjustments. Second, the selection of maximum one-shot horizon necessitates careful balance between computational efficiency and predictive accuracy. Third, while our affine coupling architecture guarantees analytical Jacobian computation, it imposes constraints on the form of the architecture.
Future work. Extending NSF to action-condition settings could enable direct integration with control algorithms such as model predictive control and reinforcement learning. Additionally, bridging NSF with diffusion models presents opportunities for leveraging complementary strengths of both frameworks in continuous time modelling. Another natural direction is to explore stronger flow parameterisations. In particular, transformer-based flows such as TarFlow [55) may relax the structural constraints of affine coupling while retaining tractable Jacobians, potentially further improving expressivity. These extensions would significantly broaden the applicability of NSF to decision-making and generative tasks.
Application-wist the resulting analytical tractability, coupled with empirically demonstrated two- onder-of-magnituee computational speedups, renders NSF particularly suitable for diverse applica- tions, from large-scale digital twin implementations.
trading scenarios requiring real-time robotic control and fast simulations
Broader Impact Statement. This paper presents work whose goal is to advance machine learning research. There may exist potential societal consequences of our work, however, none of which feel must be specifically highlighted here at the moment of paper publication.
10

---

References
[1] F. V. Agakov and D. Barber. An auxiliary variational method. In Neural Information Pro- cessing, 11th International Conference, ICONIP 2004, Calcutta, India, November 22-25, 2004, Proceedings. volume 3316 of Lecture Notes in Computer Science, pages 561-566 Springer, 2004. doi: 10.1007/978- 5-0- 10400003499-00 86. URL https:///oi...... org 10. 1007/ 976-3-540-30499-9-58
[2] A. Ansari, A. Heng, A. Lim, and H. Soh. Neural continuous-discres state space models for ireeuularly sampled time series. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and Scarlett, editors. Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pages 926-951, PMLR, 23-29 Jul mmlr.rngs.........mmmmmm..m.mm www.prressvvv.e /v02/ ansarinaiii 23a. html
2023. URL
[3] G. Bartosh, D. Vetrov, and C. A. Naesseth. SDE matching: Scalable and simulation-free training of latent stochastic differential equations. In Proceedings of the 42nd International Conference on Machine Learning (ICML), 2025.
[4] M. Biloš, J. Sommer, S. S. Rangapuram, T. Januschowski, and S. Günnemann. Neural flows: Efficient alternative to neural odes. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[5] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. M.clluurin, G. Necula, A. Paszke, Waandeman-Milne and Q. Zhang. JAX: composable transformations of
J. VanderPlas, Python+NumPy programs, 2018. h/ghtttp::/ttteee
com/jax-ml/ja
[6] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, pages 6572-6583, 2018.
[7] K. Cho, B. van Merrienboer, Ç. Gülçehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In A. Moschitti, B. Pang, and W. Daelemans, editors, Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing, EMNLP 2014, October 25-29, 2014, Doha, Qatar, meeting of SIGDAT, a Special Interest Group of the ACL, pages 1724-1734. ACL, 2014. doi: 10.3115/V1/D14-111 d14-1179
URL https:tpp/.. pattpp://dor //doi. org 3115/v11
[8] Chung, K. Kastner, L. Dinh, K. Goel, A. C. Courville, and Y. Bengio. A recurrent latent variable model for sequential data. In Advances in Neural Information Processing Systems, 2015.
[9] K. Course and P. B. Nair. Amortized reparametrization: Efficient and scalable variational inference for latent SDEs. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URLlattps://openeevviee http://wwwweeeeeeeeooorrraaaaaaa
[10] R. Daems, M. Opper, G. Crevecoeur, and T. Birdal. Variational inference for SDEs driven by fractional noise. In The Twelfth International Conference on Learning Representations, 2024. URL
[11] E. Denton and R. Fergus. Stochastic video generation with learned prior. J. G. Dy and
A. Krause, editors, Proceedings of the 35th International Conference Machine Learn- ing, ICML 2018, Stockholmsmüssun, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1182-1191. PMLR, 2018. URL
http://proceedinng pagess..xx. mlr. http::twww..p..ee///////ee/// arress/v80/denton1155 html
[12] Dhariwal and A. Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 8780-8794, 2021.
[13] I.. Dinh, J. Sohl-Dickstein,, and Bengio. Density estimation using real NVP. In 5th Interna- tional Conference on Learning Representations (ICLR), 2017.
11

---

[14] J. Franceschi, E. Delasalles, M. Chen, S. Lamprier, and P. Gallinari. Stochastic latent residual video prediction. In Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 3233-3246. PMLR, 2020. httppttpp//prrceeenggsssss .ml. press/v119 franceschi 20a. html
[15] Z. Gan, Li, R. Henao, D. E. Carlson, and L. Carin. Deep temporal sigmoid belief networks for sequence modeling. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 28. Curran Associates. Inc., 2015. URL htt.hhttss:///stt //tpoeees///roooeennggggs newrsddddd.nnnnnnnn..e.e..e cc/ wwco/pppr__feee files/ 1aaper_f1eesspppe/22/////111 /appr//011///tlle 351110000644400rt000099fff6666
pdf
[16] D. Hafner, T. Lillicrap. I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2555 2565. PMLR, 09-15 Jun 2019. URL httpps// proceedings mlr..mm.... wwwppppssss////pee htwp:///aaaeerr....
[17]
M. Heinonen, C. Yildiz, H. Mannerström, J. Intosalmi, and H. Lähdesmäki. Learning unknown ODE models with Gaussian processes. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1959-1968. PMLR, 10-15 Jul 2018. URL httpttss/////ooceediii mlr. www.press/vvee www.ev//eeeeennneeee html
[18] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017. URL https
//openreview //pparrerier..etfffouufeeeeeeeee metf
[19] J. J. Higham. An algorithmic introduction to numerical simulation of stochastic differential equations. SIAM Review, 43(3):525-546, 2001. doi: 10.1 10.1137/5003614445007773 URL https://doi...
org/ 10. 1137 01157///00061144003778200
[20] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 6840-6851, 2020.
[21] M. J. Johnson, D. K. Duvenaud, A. B. Wiltschko, R. P. Adams, and S. R. Datta, Composing graphical models with neural networks for structured representations and fast inference. In Advances in Neural Information Processing Systems, 2016.
[22] M. Karl, M. Soelch, J. Bayer, and P. van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Conference Learning Representations, 2017. URL. httpps: net/ htt.//orummdd-HyTaaaSaass
[23] P. Kidger. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
[24] Kidger and C. Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
[25] P. Kidger, J. Foster, X. Li, H. Oberhauser, and T. Lyons. Neural sdes as -nfiniteddimensional gans. In Proceedings of the 38th International Conference Machine Learning (ICML), pages 5541-5553, 2021.
[26] D. Kim, C. Lai, W. Liao, N. Murata, Y. Takida, T. Uesaka, Y. He, Y. Mitsufuji, and S. Ermon. Consistency trajectory models: Learning probability flow ODE trajectory of diffusion. In International Conference Learning Representations (ICLR), 2024.
[27] D. P. Kingma and Ba. Adam: A method for stochastic optimization. In Y. Bengio and Y. LeCun, editors, 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Truck Proceedings, 2015. URL http
//arxiv. org/abs/////eeoooooeoo.gooeeeett /11./111...0 6980
12

---

[28] D. P. Kingma and P. Dhariwal. Glow: Generative flow with invertible 1xl convolutions. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, pages 10215-10224, 2018.
[29] D. P. Kingma and M. Welling. Auto-enn variational bayes. In Y. Bengio and Y. LeCun, editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL aattp://arxiv org/
abs/1312. 6114
[30] Kloeden and E. Platen. Numerical Solution of Stochastic Differential Equations. Stochastic Modelling and Applied Probability. Springer Berlin Heidelberg, 2. ISBN 9783540540625. URL httttt::///bkkkke http://booooo.ogll...ee cww.glee....ukooookksstt..t BiwpoookkkiidBBtsssssooooe
[31] R. G. Krishnan, U. Shalit, and D. Sontag. Structured inference networks for nonlinear state space models. In Proceedings of the Thirty-First AAAI Conference Artificial Intelligence, 1A page 2101-2109. AAAI Press, 2017.
[32] H. Kunita. Stochastic Flows and Stochastic Differential Equations. Cambridge Studies in Advanced Mathematics. Cambridge University Press, 1990. ISBN 9780521599252. URL https://books
google. co uk/books?ict -SRii--CSRRCCM
[33] X. Li, T. Wong, R. T. Q. Chen, and D. Duvenaud. Scalable gradients for stochastic differential equations. In Proceedings of the 23rd International Conference Artificial Intelligence and Statistics (AISTATS), 2020.
[34] Y. Lipman, R. T. Q. Chen, H. Ben- Ham,H M. Nickel, and M. Le. Flow matching for generative modeling. In International Conference on Learning Representations (ICLR), 2023.
[35] X. Liu, C. Gong, and Q. Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. In International Conference on Learning Representations (ICLR), 2023.
[36] E. N. Lorenz. Deterministic nonperiodic flow. Journal of Atmospheric Sei- ences, 20(2):130 141, 1963. doi: 10. 10L11755155000000000000000001110
2. URL https://journalss ametsoc wttp:aaa.....a....ooooooeeeeee/e http:////////////aaaa/////////// 1963 020 01300_030_0000 dnf ___2_0 2_0__o__2_ 2m
[37] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Con- ference on Learning Representations, 2019. URL. httpps://gpeeeeeveee.eeeeeeee//e Bkg6
[38] Y. Oh, D. Lim, and S. Kim. Stable neural stochastic differential equations in analyzing irregular time series data. In The Twelfth International Conference on Learning Representations, 2024.
URL SoLLittps///operreeee
[39] B. Øksendal. Stochastic Differential Equations: An Introduction with Applications (Universi-
text). Springer, 6th edition, 2003.
[40] R. Ranganath, D. Tran, and D. M. Blei. Hierarchical variational models. In Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 324-333 333. JMLR.org, 2016. URL pattp::/pproceddnggg mlr, press/v48/res /aaess/////raaggaaahh html
[41] D. Rezende and Mohamed. Variational inference with normalizing flows. In F. Bach and D. Blei, editors, Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1530- 1538, Lille, France, 07-09 Jul 2015. PMLR. StRLsttts:// htt:ttts/s//oceeeiinnee aeriinnnn.m........mmmmmmmr.mm.e www.s/rerss/// htt.3///eeeeeeeeee html
[42] D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference deep generative models. In Proceedings of the 31 International Conference on International Conference on Machine Learning- Volume 32, ICML'14 page (1-1278-1I-122 JMLR.org, 2014.
[43] Y. Rubanova, R. T. Q. Chen, and D. Duvenaud. Latent odes for iregularly-samppee time series. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
13

---

[44] T. Salimans and J. Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022. URL /ttpp::////poee net/forum? id=TIdIXIpzhol
[45] T. Salimans, D. P. Kingma, and M. Welling. Markov chain Monte Carlo and variational inference: Bridging the gap. In Proceedings of the 32nd International Conference on Machine Learning, ICML 2015. Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pages 1218-1226. JMLR.org, 2015. URL http://pprceeeeingg
mlr mlr. wtwppp/s/////aaaaaaaaaaa htmmll
[46] A. Solin, E. Tamir, and P. Verma. Scalable inference in sdes by direct matching of the Pokker-PPncckKKoloooorvy equation. In Advances in Neural Information Processing Systems (NeurlPSS, 2021.
[47] Y. Song, Sohl-Dicksteins generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021.
Kingma, Kumar, S. Ermon, and B. Poole. Score-based
[48] Y. Song, Dhariwal, M. Chen, and Sutskever. Consistency models. In Proceedings of the 40th International Conference on Machine Learning (ICML), pages 32211-32252, 2023.
[49] T. Teshima, I. Ishikawa, K. Tojo, K. Oono, M. Ikeda, and M. Sugiyama. Coupling- based invertible neural networks are universal diffeomorphism approximators. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 3362-3373, Curran Associates, Inc., 2020. URL http:ttt////occeedddaiee neurips. cc. soccpaper__fee files paaper_fileeapppper/22000000 /ptp////////////eeee/ 2290a 220077886e411c522222222222222222
.ddf
[50] B. Tzen and M. Raginsky. Neural stochastic differential equations: Deep latent gaussian models the diffusion limit. arXiv preprint arXiv: 009880, 2019.
[51] M. Wang, D. Fleet, and A. Hertzmann. Gaussian process dynamical models for human motion. IEEE Trans. Pattern Anal. Mach. Intell., 30(2):283-298, Feb. 2008. ISSN 0162-8828. doi: 00.1009/PAAM..00..1161
URL httthhtts///eo /doi org/10.g//.ooo/.oo 1119//1999PPII 2007. 1167
[52] C. Winkler, D. Worrall, E. Hoogeboom, and M. Welling. Learning likelihoods with conditional normalizing flows. arXiv preprint arXiv: 19/2,0002,, 2019.
[53] L. Yang, Z. Zhang. Z. Zhang, X. Consistency flow matching: Defining straight flows with velocity consistency, 2024. URL https://arxiv. org/ abb/24077 02398
Liu, M. Xu, W. Zhang, C. Meng, S. Ermon, and B. Cui.
[54] C. Yildiz, M. Heinonen, and Lahdesmaki. Ode2vae: Deep generative second order odes with bayesian neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL http:: http:tttts///oceeedinnee neeriip
neurips. cc/paper.pppp....ppppp fimeerrrfllss/__ pdf
[55]
S. Zhai, R. ZHANG, P. Nakkiran, D. Berthelot, J. Gu, H. Zheng, T. Chen, M. Á. Bautista, N. Jaitly, and M. Susskind. Normalizing flows are capable generative models. In Feroy- International Conference Machine Learning, 2025. URL httpp:://openreveewwneet forum?id=2uheUFoFFa
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction state the main contributions--vv our neural stochas- tie flow (NSF) model, latent NSF extension, and its empirical benefits-and and the experimental results Sections and Appendix directly support these claims without exaggeration.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The current manuscript contains an explicit "Limitations" paragraph in Sec- tion 7 discussing weaknesses.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
15

---

Answer: [NA]
Justification: The work presents an empirical method without formal theorems; no theoretical claims are made, so formal assumptions and proofs are not applicable.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
Implementation details and full hyper- hhper-paramene parameter are provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. .., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path reproducing verifying the results. 5. Open access to data and code
16

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: For now, we include the related references for the sources of the datasets. We will release the code and detailed running instructions upon publication.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: [ppenddx([eemmmeees every training and evaluation hyper-parameter (learning rate, optimizer, batch size, number of epochs, etc.) for each experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All reported metrics include standard deviation or 95% confidence inter- vals computed over five random seeds; methodology and seed information are detailed in Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: All experiments were conducted on a single NVIDIA RTX 3090 GPU (24 GB); typical training times under few days per dataset reported in Appendix E
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:///wwrrrpppcccccup//////aaa
Answer: [Yes]
Justification: In this paper, we introduce work designed to push the boundaries of machine learning. While our methods could carry ethical implications, we do not believe any require specific discussion at this point in the submission process.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: dedicated "Broader Impact Statement" section is provided near the end of the paper, outlining potential positive and negative societal impacts.
18

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out, For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI...
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Whenever we use any assets, we always cite the original asset source.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
19

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssippsssie thcode wwwccom/aaaoo//c datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We fully describe our contributions and any assets used in the paper. We will also releasing code after the publishing the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The study involves only publicly available simulation and motion-capture datasets and does not recruit human subjects or use crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human-subject experiments were conducted, so IRB approval is not required.
Guidelines:
20

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Generative LLM tools (e.g.. ChatGPT, GitHub Copilot, Claude) were used solely for auxiliary tasks such as code refactoring, generating data visualization scripts, and proofreading the manuscript. These tools did not contribute the development of the core methodology, experimental design, or interpretation of results.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to LLM policy http::///earripacccc//ooeerreees for what should or should not be described.
21

---

Appendix Contents
Architecture
24
A. Neural Stochastic Flow
24
A.2 Bridge Model
24
1.2.1 Conditioned Bijective Transformations
B Derivation of the Flow Loss
24
I One-Step to Two-Step KL. Divergence
25
B.2 Two-Step to One-Step KI. Divergence
26
Detailed Description of the Learning Process
26
Neural Stochastic Flow Optimisation)
26
Latent Neural Stochastic Flow Optimisation
27
D Sampling and Density Evaluation
27
Experimental Details
27
E.
General Setup
27
Stochastic Lorenz Attractor
28
................................
E.2.1
Data Generation
................................
28
E.2.2
Configurations for Neural Stochastic Flows
................................
28
E.2.3
Contigurations for Baselines
28
E22.4
Evaluation and Metries
................................
30
E,2.5
Detailed Results
................................
E.3
CMU Motion Capture
31
E.3.1
Data Preparation
31
E.3.2
Latent SDE Reproduction
E.3.3
SDE Matching Reproduction
31
E.3.4
Configurations for Latent NSH
31
E.3.5
Configurations for Baselines
34
3,6
Evaluation and Metries
36
E
E.3.7
Detailed Results
36
Stochastic Moving MNIST
36
Data Generation
36
E,4.2
Contigurations for Latent NSF
38
E.4.3
Confgurations for Baseline
38
2,44
Evaluation and Metrics
39
E.4.5
Detailed Results
39
Hyperparameter Sensitivity and Ablation Studies
40
Stochastic Lorenz Attracior with Missing Data
40
22

---

E.1.1
Experimental Setup
40
R.1.2
Results
43
CMU Motion Capture Dataset with Extrapolation Task (Setup 2)
F2.1
Experimental Sctup
43
Results
43
H3
Practical Guidance on Hyperparameters
43
F3.11
Time Horizon for Training Huain and Interence Hred
43
F.3.2
Flow Loss Weight
43
E3.3
Auxiliary Inner Steps Simultaneous Updates
43
F3..4
Directional Flow Loss Components
43
Validation of the Fréchet Image and Video Metrics
44
Protocol
44
G.I.I
Datasets.
44
G 1.2
Conditions.
44
G.1.3
Computation of distances,
44
2
Results
3
Discussion
46
23

---

A Architecture
A.1 Neural Stochastic Flow
We refer to the main text for the NSF architecture and notation (Section 3 Eqs. 5 and (6) in the main text). Here we only note implementation specifics. MI.P, and MLP, neural networks that share parameters and produce split outputs for the mean and variance components. Softplus activation is used for the variance component ensure positivity, For stability, the scale heads end with tanh multiplied by a learnt scalar to bound log-scales 13 Exact log-densities are computed via the thange-of-variable formula; see Section
A.2 Bridge Model
Here, we describe the architecture of the bridge model used as auxiliary variational distribution 40 45). The bridge model characterises the conditional distribution be(a, Øt, zt,) for any intermediate time point where t; < < tj. given the boundary states. Combined with flow loss (Eq. (T0) minimisation, this enables inference of intermediate states when both endpoints are known. Let Cbe = (It, Lt,, At, T, t,) denote the conditioning parameters, where r = (t - t2) At the normalised time and At = tj t, is the time interval, and the initial time t, is ommitted for autonomous systems.
Similarly to the main NSF, we first draw a sample from parametric Gaussian distribution:
Z1 = Et. + T (20e, - DX,)) + a(r) MLP, + V At Softplus (MLP, (Car)) Oe (19)
m(E
(t)
where ~ (O, I) is standard normal noise vector, a(T) := 7(1 r) is the standard Brownian- bridge factor, and MLP,, and MLP, are neural networks that share parameters and produce split outputs for the mean and variance components. The intermediate state is then obtained by applying flow transformation:
2t fg(Zz, =x) f Car, EL) Car,
o f Cbr, EL-1) ..... f f Cbr E1)
(20)
corresponding ensuring consistency with boundary conditions. The flow transformation fg is implemented using series of conditioned bijective transformations, as detailed below.
Note that aar) vanishes the endpoints (r O and
and /.
fz))
A.2.1 Conditioned Bijective Transformations
Each layer of the bridge flow follows the same coupling structure as in the main NSF (Section 3) Given input 2,, split it into two parts (z #B) Split(z,) with alternating partitioning across layers, and apply the transformation:
fę(2i; Cbr. E.)
(21)
= =Concat Æイi FB 3 exp(((() MLP
scale (ZA, Cbei (l + a(T) MLP(1) shill (2 Cbe <(a
Following the main model, we apply hyperbolic tangent function with learnt scale as the final activation of ( scale for training stability [13]. This construction preserves differentiability in and guarantees the consistency at the boundary conditions (r = 0, 1) where a(T) = 0.
B Derivation of the Flow Loss
In our research, we utilise a bridge distribution described in Appendix A.2 as auxiliary variational distribution [1 40 to constrain the upper bounds bidirectional KL divergences between the one-step and two-step distributions of the NSF By minimising these upper bounds, we guide the model towards satisfying the Chapman-KKlmogorov relation, as shown in Equation 4 thereby improving its consistency with the theoretical requirements of stochastic flows of diffeomorphisms combined with architecture design. Specifically, we define Laow as the expectation of a weighted sum of two components:
Laow (0.E)
I. p(t).tz.tu)
1n (0, Eit,, tj. tx)) + A2 Cnow, 2 (0. tt)).. (22)
24

---

Here, A1-1022 and A2-1o11 are weighting factors, which are set to in the main text, balancing the contribution of each component to the overall flow loss. The two components. Lnow, 1-to-2 and Lnow, 2.1,, 2. correspond to the upper bounds of the one-step to two-step and two-step to one-step KI. divergences, respectively.
In our experiment, which focuses on the autonomous case, the triplet (4/, Eyy (u)) is sampled according the following procedure:
The initial time i, sampled from the data D, representing the starting times the dataset. (For non-autonomous SDEs, it is recommended to sample i, uniformly to ensure the Chapman- Kolmogorov equation is satisfied across all time points.)
The final time ty is sampled from mixture distribution p(tx) 2/2 U(tt,, t, + Huain) 3 Pdnaaaty f4), where U(44, 44 Hmann) denotes the uniform distribution over the interval [i, 1i H and Pdata (lx (4) is the conditional data distribution of end times correspond- ing the sampled initial time 4. Here, Hiriin represents the maximum one-shot interval used during training.
The intermediate time tj uniformly sampled from the interval [t,, tkl, i.e.,, ty (ts, tx).
The inclusion Pana (fz t4) in the mixture distribution for sampling fu is motivated by our obser- vation that the model achieves higher accuracy in regions where data exists. By incorporating this data-driven component, we aim to enhance the efficiency of the learning process, leveraging the model'ss 's improved performance in data-rich areas.
In the following, we derive the upper bounds of the one-step to two-step and two-step to one-step KL divergences.
B.1 One-Step to Two-Step KL Divergence
We begin with the derivation of the upper bound of the KL divergence from one-step computation to marginalised two-step computation. To clarify the computation, subscripts pe and p distinguish between the distributions for 1-step sampling from T,, It, and two-step sampling from He, to Tt, via Tt and the time arguments are omitted from the probability distributions for clarity. The key steps are outlined as follows:
DKL polase xi..
pē (I
x21)) pā(20, x1.) dair,
(23)
= pàl (20e) xxe,)
dx
(24)
pg (2
D21)) p0 (Te;
xe:)) dx
= -f pore (x.)) log
p/ (a
(x,,) ) 110 (2r1,
dæ
(25)
P (x:
2
(21)) pă (xr,
-ʃ pb
pb (w
(ra.) gob (aie,
De xee.
0e, wrr) log (x.)) g6 (mn
p3 (De,, axa)) p (att, me)))
(26)
9 x2
log
dæ,,
dar d
2 Pg (z1,
Dee xite)
pb (xte
(x.)) 90 (wi,
It. 1e
E m m mm)
log
pbl Pb(
xt,,) pd (21,
xcc.)
(27)
-KLL (9 (22,,
xee || po (I:1,
pb (aF
Dt, 9o (x1
Ft, I.t..
E Pb(1t a., led]) (.))
log
(28)
x.i,) Po Iz x..)
Lnow, 1-t02 (0, Φi t, tt,))))))
(29)
25

---

The equation from the second line to the third line is obtained using Bayes' theorem and the Markov property:
I | (at, Teg) poo (mee, x4.)) dai,
18 (1 Fi1) 1/ ( av::)
(30)
pe (2, 0e; 1t2r.
B.2 Two-Step to One-Step KL Divergence
Similarly, explore the derivation of the upper bound of the KL divergence from marginalised two-step computation to a one-step computation. We denote this as Lnow, 2-10-1- The key steps are outlined as follows:
DKL. (f Po
(12,, (Ite, ) r (are) orc)) dar,
Po (a mre))
(31)
= // p& (1rta (ma,) ) ped ( mee)) dayt
pd (3
12; ) pd (2),; xe.))
log
dæe
(32)
16 (2it)
we. xeex p (orc xt..)
Iʃ pà (zt, xww) pp (are
p (ara,
mrr.) Po ( mr.))
xze)) log
(33)
Pd (20;; Fie De p/ (20 orce)
ro (12,
xx..) p/ (11 xity)
log
(34)
r3(w1, w., )p3( mxa))
be be(2,
0t,, vta) ph (20 (2..)
Lnoww 2-to-1 (0, E: ts tx)
(35)
From the first line to the second line, we use the relationship in Equation (30) Although the expression inside the log in the second line appears to depend æty- it does actually depend Lt, due to this relationship. Therefore, the integrand change from the second line to the third line is justified.
By minimising flow loss consists of these upper bounds, we encourage the model to satisfy the flow property in both directions simultaneously.
C Detailed Description of the Learning Process
C.1 Neural Stochastic Flow Optimisation
The optimisation process of NSF is summarised in Algorithm iin the main text. Here, we provide additional implementation details.
1. Batch and time triplet sampling: We sample batch {(x,, H from the dataset. For each element, we sample triplet of time points (ti, Ly (x) where t, < tú < tu, following the strategy described in Appendix
2. Bridge model optimisation: Before updating the main model parameters, we perform K inner optimisation steps on the bridge model parameters &. This ensures that the aux- iliary variational distribution be (zx, [xt,, xt,.) closely approximates the model's bridge distribution.
3. Main model optimisation: then sample transition pairs {( De, ti t ) from the dataset and compute the negative logllikel Combined with the flow loss, this forms the total loss used to update the main model parameters
This sequencing ensures that the main model is updated based on a bridge model that closely approximates the bridge distribution of the main model. Alternatively, the inner optimisation steps can be omitted and update the bridge model parameters simultaneously with the main model parameters when the main model parameters are updated.
26

---

C.2 Latent Neural Stochastic Flow Optimisation
The optimisation process of Latent NSF is summarised in Algorithm in the main text. Here, we provide additional implementation details.
1. Observation encoding: We sample batches of observation sequences {oter d bb=1 and encode them to latent states using the encoder network Q (b as defined in Eq. [15
Ota.r.b)
2. Time triplet tripletsammlinn sampling: Similar to NSF, we sample time triplets (ti, Egy (x) for each sequence,
with < tj < tg.
3. Bridge model optimisation: Before updating the main model parameters, we perform K inner optimisation steps on the bridge model parameters &.
4. Main model optimisation: We then compute the total loss using Eq. (18) in the main text and update the NSF parameters 0, encoder parameters Φ and decoder parameters tb.
The specific values of hyperparameters Saip. and À in the total loss (Eq. (8) in the main text) vary depending on the dataset and task, and are detailed in Appendix [ As well as NSF, the inner
optimisation steps be omitted and update the auxiliary model parameters simultaneously with the main model parameters when the main model parameters are updated.
D Sampling and Density Evaluation
To generate sample at time t, from known state It at time t, we compute fa t,, At) and o (z,, 1,, At) using the MLPs, draw (O, I) and compute 2 = μ + or Ô E. Then, apply the flow transformation: Σt, = f. o f o ... o fi (z; At,t;) This enables single-step prediction without numerical integration, following the standard approach used in normalising flows.
Using the change of variables formula, the log-density can be computed as:
0f
log logpdær, 04,14 At) logp(e) Σ log det
(36)
0f--1
where p(E) is the density of the standard normal distribution. When using the affine coupling formulation, this simplifies to:
d logpo( (2 (ryi Li, At) -- 2 log 2T
1
ΣΣ AL. ALP (1-1) scale (=) ,4, Al);, (37)
2m1 JEB.
where B; is the set of indices corresponding to the second partition in the i-th coupling layer, and (i-1) is the output of the - 1)-th layer. This density evaluation approach mirrors the standard technique in normalising flows, with appropriate conditioning the time parameters.
E Experimental Details
This section provides detailed information about the experimental setups, hyperparameters, archi- tectures, and computational analyses for the experiments presented in Section 6 in the main text. Implementations are available at wihhtps:/ pestts::/githuut wtcpaauu.ooommaaeaaameaaaaeeeeee wwt.coo./////.eee hamm//kiyooaara/// jax_nsf
E.1
General Setup
Unless otherwise specified, experiments were conducted with the following setup:
Frameworks:
Our (Latent) NSF models: JAX
5
with Equinox library
24
PyTorch baselines: torchsde 33
JAX baselines: Diffrax 23
Hardware: NVIDIA RTX 3090 GPU
Optimiser: AdamW 27 37
Hyperparameters: Specified below for each experiment
27

---

E.2
Stochastic Lorenz Attractor
E.2.1
Data Generation
Following ei al.
the data generation process is as follows:
SDE:
dr = o - df
dW1
dy (z(p 2) - y) dt ay dW2
dz (zy- - 32) dt +, dW3
Parameters: o = 10, p= 28, 3 = 8/3, a = (0.15, 0.15,
Initial states (Zo. No. 2n): Sampled from I),
Trajectories: 1,024 for training and 1,024 for testing
Time steps: E [0, 1] with time step size At 0.025 (41 time steps per trajectory)
E.2.2
Configurations for Neural Stochastic Flows
Architecture:
State dimension d_ l_state 3; conditioning dimension dond (state + time)
Gaussian parameter network: (aput(d_cond) + 2x[Linear (64) + SiLU(] 7
Linear -taddts splits into (mean, std) with std via Softplus ()
Scale-shift networks (in conditional flow): Input (d_ (dd_state/] + d_ cond)
2x [x[ineear (64) SiLUO] Linear (d_state); splits into (scale, shift in Eq. ( in the main text) (In this case, since d_state is odd, we split three dimensions of d stateee into and two dimensions)
Conditional flow (affine coupling; Eq. (5). (6)): 4 layers with alternating masking
Parameters:
NSF: 25,130
Bridge model: 26,410
Training:
Data conversion: Time series data is converted into pairs of states (R)) where
> t, to predict p(Te, xx.))
Epochs: 1000
Batch size: 256
Optimiser: AdamW W 27
Learning rate: 0.001
Weight decay: 10-5
Loss function: Combined negative log-likelihood and flow loss (Eq. 11
À 0.4 (0.2 for data component, 0.2 for sampled component)
* A1 = A2. = 1.0
Time triplets (t,, ty, tx) for Lnow Sampled as described in Appendix with Humin = 1
Auxiliary updates: bridge model be trained concurrently with K 5 inner optimisation steps per main model update
E.2.3 Configurations for Baselines
For our baseline comparisons, utilised the official implementations of latent SDE SDE matching [3]. SDE-GAN [25, and Stable Neural SDE series (38 with minimal modifications. We maintained the original model architectures, optimiser configurations, and hyperparameter settings as provided in their respective codebases, replacing only the input data to match our experimental setting. For SDE-GAN, we found that exhibited significant training instability when applied to the Stochastic I Lorenz Attractor dataset, preventing convergence despite multiple training attempts with varying hyperparameter configurations. Therefore, we excluded SDE-GAN results from our comparison in the paper.
28

---

Latent SDE.
The latent SDE model
33 has the following configuration:
Architecture:
Observation dimension d_obs 3; latent dimension d_latent 4; hidden size d hidden 128; sequence length T 41
Encoder: Input ((T, d_obs))
CRU(128) Linear (d_coxtxxt)
reverse Output (T, d_context))
restore time order
Posterior drift network: Input (d__latet + d_context) + 2 [Linear (128)
Softplus () (d_
Prior drift network: Input + 2x 2x[Linear (128) Softplus ())
Linear (d_latent)
Diffusion networks:
per-dimension Input (1) Linear (128) Softplus ()
Linear (1) + Sigmoid
Decoder: Input (d__latet) - Linear (d_obs)
Parameters:
Encoder: 59,328
Posterior initial state network: 520
Posterior drift network: 25,860
Prior drift network: 17,668
Diffusion networks: 1,540
Decoder: 15
Training:
Iterations: 5000
Optimiser: Adam
Initial learning rate: 0.01
Learning rate decay: Exponential (gamma=0.997)
KL divergence annealing: Linear annealing from 0 to over first 1000 iterations
Numerical integration: Euler-Maruyama method with step size At = 0.01
SDE Matching.
The SDE matching model has the following configuration:
Architecture:
Observation dimension d_obs 3; latent dimension d_latent = 4; hidden size d_ hidden 128
Encoder: (n, d_obs)) +
GRU(d_hidden)
Output ((T+1,
d_ (.hhddd)) (concatenate initial hidden with per-step outputs)
Prior initial distribution: learnable diagonal Gaussian; parameters m, log in d_latent dimensions each
Prior drift network: Input (d_latent) Linear (d_hidden) Softplus()
Linear (d__hiiden) - Softplus () - Linear (d_ (d_latent)
Diffusion networks: per-dimension Input (1) - Linear (d_hidden) + Softplus + Linear (1) - Sigmoid()
Posterior affine head: Input (d hidden + 1) + Linear (d hidden) - SiLU()
- Linear (d hidden) - SiLU() Linear (2+d latent); splits into (nt, log o.)
with o, = exp(log
Decoder: Input (d_latent) + Linear (d_obs); Gaussian likelihood with fixed std
0.01
Parameters:
Prior init:
Prior drift network: 17,668
Diffusion networks: 1,540
29

---

Decoder: 15
Encoder: 99,072
Posterior affine head: 66,560
Training:
Iterations: 10,000
Batch size: 1,024
Optimiser: Adam
Learning rate: 0.001
Stable Neural SDE Series. and Neural LNSDE.
The Stable Neural SDE series includes Neural LSDE, Neural GSDE,
Architecture (common to all variants unless specified):
Encoder: Input (3
Linear (64)
Embedding network: Input (64) + Linear (64)
Drift network: Input (64) 2x [Lineer (64) LipSwish Linear (64)
Diffusion network: Input (64) - 2x[Linear (64) - LipSwish()) Linear (64)
Decoder: Input (64) - Linear (3)
Parameters:
Encoder: 256
Embedding network: 4,160
Drift network: 12,480
Diffusion network: 12,480
Decoder: 192
Training:
Epochs: 100
Batch size: 16
Optimiser: Adam
Learning rate: 0.001
Numerical integration: Euler-Mon Marryaaa scheme with step-size At 0.01
E.2.4 Evaluation and Metrics
KL Divergence. For computing the KL divergence, we employed non- non-parametric Kernel Density Estimation (KDE) with Gaussian kernels. The bandwidth parameter was determined through 3-fold cross-valiattinn by optimising the log-likelihood on unseen data. To ensure statistical robustness, we computed the KL divergence and test log-likeliin log-likelihood across ten independent folds, excluding values beyond the 5th and 95th percentiles to mitigate outlier influence. The final reported metrics consist of mean and standard deviation of these filtered values.
FLOPs. We calculated FLOPs on a per-sample basis using different methods depending on the framework. For JAX-based models like Neural Stochastic Flows, analysed the JIT-compiled sampling function using jax sit: () .lower () compile () cost_ analysis () [ ['flops']] For byserch-based models such latent SDE and Stable Neural SDEs, torch. utils. _lop_counted Cowuterr.Fooooouue FlopCounterMode to count FLOPs during sample genera- tion. then divided by batch size. Our comparison table shows the FLOPs value for one sample.
utilised
Runtime. For all methods, we measured the wall-clock runtime for computing vectorised batches of 100 samples.
30

---

E.2.5 Detailed Results
Table 6
presents FLOPs comparisons across different prediction horizons. NSF requires constant FLOPs for single-step predictions regardless of time interval, while baseline methods scale linearly with the number of integration steps.
Table shows wall-clock runtime measurements for generating 100 samples. PyTorch implementa- tions exhibit higher runtime due to dynamic graph construction, while JAX implementations benefit from JIT compilation. NSF achieves consistent nub-millisecond runtime across all time horizons.
E.3
CMU Motion Capture
E.3.1
Data Preparation
Data was prepared follows:
Dataset: Preprocessed CMU Motion Capture (walking) from Yildiz et al.
[54
Features: 50-dimensional joint angle vectors
Sequence length: 300 time steps
Split:
16 training, validation, and 4 test sequences
Setups:
Setup
336
Train full 300 steps, test by predicting steps 4-300 given steps 1-3
Setup
Train on steps 1-200, test by predicting steps 101-300 given steps 1-100
E.3.2 Latent SDE Reproduction
0.20 for Setup 1. our re-implementation and other reproductions observe higher MSE values. These results are consistent with those reported in public discussions hhtpp:: //gi / wow.:o/ooooo/ooeeeeeeeeeeeeeeeee research/ forchsde/issues/ /112 Our implementation is available
As noted in footnote in the main text, whilst Li et al. 33 report 4.03
com/ nkiyohara/ latent Sstent-sde-a mocap, which is consistent with the results reported in the public discussion.
E.3.3 SDE Matching Reproduction
As noted the footnote the main text, at the time of writing, the official implementation of SDE matching [ for the CMU Motion Capture dataset has not yet been publicly released. To ensure consistent comparison across Setup and Setup 2, we re-implemented SDE matching to the best of our ability based on the description in their paper. While Bartosh et al. 3 report 4.50 0.32 for Setup 2, our re-implementation achieves 5.20 0.43. This discrepancy may stem from implementation details not specified in the original paper, such architectural choices, hyperparameter settings, or training procedures. We have made our best efforts obtain the smaller MSE by adjusting the hyperparameters. architecture, and training procedure. Our implementation is available at htttp::///// //github. www/ooimioioeeeeeccc htt///kiyoohrrra sde -ete-mmttctiininieeeeeee mocap
E.3.4
Configurations for Latent NSF
Latent NSF - Setup (within-horizon forecasting).
Architecture:
Latent dimension d_latent conditioning dimension d_cond
Encoder:
Input (150) + 2 2x[Linear (30) + Softplus ()]
Linear (2+d_latent): splits into (mean, std) with std via Softplus
Emission Po xx,): Input (d_ (d_latent) + 2x [Linear (30) Softplus ()]
- Linear (50); Gaussian likelihood with fixed std
NSF transition model (latent prior):
* Gaussian parameter network: Input (dndddd 2e(einieeaeee (30) SiLU()
+ Linear (2+d [ateent:: splits into (mean, std) with std via Softplus
31

---

Table FLOPs comparison (K) for Stochastic Lorenz Attractor experiments, showing computational costs per single sample prediction. EM refers to Euler-Maruyama method, and SRK refers to Stochastic Runge-Kure Kutta method.
Method
0.25
0.5
t= 0.75
t 1.0
Latent SDE 33
EM, At 0.003
3,133
6,193
9,253
12,349
EM, At = 0..01
959
1,917
2,839
3,760
EM, At 0.03
369
664
995
1,290
Milstein, At = 0.01
1,010
2,021
2,994
3,967
SRK, A 0.01
9,253
18,838
28,054
37,270
SDE matching 3
EM, A/ 0.0001
184,394
368,787
553,034
737,354
EM, A/ 0.001
18,506
37,011
55,443
73,875
EM, A/ = 0.01
1,917
3,834
5,677
7,520
Stable Neural SDEs 38
Neural LSDE (EM, At 0.01)
1,708
3,416
5,057
6,699
Neural LSDE (Milstein, At 0.01)
1,708
3,416
5,057
6,699
Neural LSDE (SRK, A/ =.00)
16,483
33,555
49,971
66,387
Neural GSDE (EM, Al =0.01)
1,925
3,848
5,698
7,548
Neural GSDE (Milstein,
0.01)
1,925
3,848
5,698
7,548
Neural GSDE (SRK, At 0.01)
18,571
37,807
56,303
74.799
Neural LNSDE (EM, At = 0.01)
1,925
3,848
5,698
7,548
Neural LNSDE (Milstein, At 0.01)
1,925
3,848
5,698
7,548
Neural LNSDE (SRK, At = 0.01)
18,571
37,807
56,303
74,799
NSF (ours)
Hirin 1.0, H pred = 0.25
53
105
156
208
Humin 1.0, H peed 0.5
53
53
105
105
Hunin 1.0, pred = 1.0
53
53
53
53
Table 6: FLOPs comparison (K) for Stochastic Lorenz Attractor experiments, showing computational costs per single sample prediction. EM refers Marurama method, and SRK refers to Stochastic Kutta method.
Method
t == 0.25
t= 0.5
t= 0.75
t = 1.0
Latent SDE (EM, At 0.01)
959
1,917
2,839
3,760
Latent SDE (Milstein, At = 0.01)
1,010
2,021
2,994
3,967
Latent SDE (SRK, At 0.01)
9,253
18,838
28.054
37,270
Neural LSDE (EM, At 0.01)
1,708
3,416
5,057
6,699
Neural LSDE (Milstein, At =4=0000
1,708
3,416
5,057
6,699
Neural LSDE (SRK, A/ =0.01)
16,483
33,555
49,971
66,387
Neural GSDE (EM, A/ 00.0)
1,925
3,848
5,698
7,548
Neural GSDE (Milstein, Al =
1,925
3,848
5,698
7,548
Neural GSDE (SRK, Al 0.01)
18,571
37,807
56,303
74,799
Neural LNSDE (EM, At 0.01)
1,925
3,848
5,698
7,548
Neural LNSDE (Milstein. At =t=000 0.01)
1,925
3,848
5,698
7,548
Neural LNSDE (SRK, At = 0.01)
18,571
37,807
56,303
74,799
NSF (Adtain 0.25, H pred 0.25)
53
105
156
208
NSF (Hmiin 0.5, poot = 0.5)
53
53
105
105
NSF (Huain 1.0, H pred 1.0)
53
53
53
53
32

---

* Scale-shift network (in conditional flow): 4 layers with alternating masking, each layer: Input (d_laten//2 d_nond) 2x[Linear (30) SiLU(]
- Linear (d_latent); splits into (scale, shift in Eq. 7 in the main text)
Bridge model:
Gaussian parameter network: Input (d_cond 1) [xiner (30) SiLUO] + Linear (2+d_latent) splits into (mean, std) with std via
Softplus( ()
Scale-shift network (in conditional flow): layers with alternating masking, each layer: Input (d_latent d_nond + 1) + 2x [Lineer (30) - SiLU() - Lineardd_ splits into (scale, shift in Eq. (20))
Parameters:
Stochastic flow: 8,454
Bridge model: 9,504
Decoder: 2,690
Posterior: 5,832
Training:
Objective: maximise Eq. (88
Steps: 100,000
Batch size: 128
Optimiser: AdamW W
Learning rate: 0.001
Hyperparameters: 0.1, 1.0, Bakip 0.1
Skip-ahead KI.: prediction horizons up to time steps; samples per skip
Latent NSF - Setup 2 (extrapolation).
Architecture:
Latent dimension d_latent conditioning dimension d_nond
Encoder: Input ( (T, 50)) + GRU(32) - 2 [Linear (32) - Softplus ()) +
Linear (2nd__aten)_ splits into (mean, std) with std via Softplus ()
Emission (o, xr.): Input (d_latent) 2x [Linear (30) Softplus ())
- LLinear (50); Gaussian likelihood with trainable std
NSF transition model (latent prior):
Gaussian parameter network: Input(d_
2-(Leeeerrere (64) SiLU()
+ Linear [E#d_latent)) splits into (mean, std) with std via Softplus ()
* Scale-shift network (in conditional flow): layers with alternating masking, each layer: Input (d latent/2 cond) 2x [X[Lnee (64) SiLU())
+ Linear (d_ (d_latent); splits into (scale, shift)
Bridge model:
* Gaussian parameter network: Input (d_ (oncond + 1) + 2x (32)
SiLU()] + Linear (2*d_ (atd_tatent;; splits into (mean, std) with std via
Softplus() ()
*
Scale-shift network (in conditional flow): layers with alternating masking, each layer: Input (d_ latent/2 + d_cond 1) - 2 [Linear (32) - SiLU()] Linear Lineardd_ (d_latttt) splits into (scale, shift)
Parameters:
Stochastic flow: 28,820
Bridge model: 10,452
Decoder: 4,240
Posterior: 10,604
Training:
Objective: maximise Eq.
18
33

---

Steps: 100,000
Batch size: 64
Optimiser: Adam
Learning rate: 0.001
Hyperparameters: = 0.3, 1.0, Saup = 0.3
Warm-up: / and flow loss weight are linearly increased over the first 2,000 steps
Time horizon for flow loss: Hmin
20
Auxiliary updates: K 3 inner steps per training iteration for the bridge model
Skip-ahead KL: prediction horizons up to 10 time steps; 10 samples per skip
E.3.5
Configurations for Baselines
Latent SDE - Setup 1 (within-horizon forecasting).
Architecture:
Observation dimension d_obs 50; latent dimension d_latent 6; context di- mension dontext = 3
Encoder: Input (150) 2x [Linear (30) Softplus() O) Linear (15); splits into (mean, std, context) with std via exp(), with 6, 6, and 3 dimensions respec- tively
Posterior drift network: Input (d_latent
d_context) - Linear (30)
Softplus () - Linear (d_latent)
Prior drift network: Input (d_ latent + 1) Linear (30) + Softplus ()
Linear (d_latent)
per-latent Input (2) + Linear (30) + Softplus () 7 Linear(1) Sigmoid()
Diffusion networks:
Decoder: (dpputd_lateeeee + Linear (30) + Softplus () Linear (30) 7
Softplus( () - Linear (2+d_obs)
Parameters:
Encoder: 5,925
Decoder: 4,240
Posterior drift network: 516
Prior drift network: 426
Diffusion networks: 726
Training:
Iterations: 5,000
Optimiser: Adam
Learning rate: 0.01 with exponential decay = 0.999 at every step
KL annealing: linear over the first 400 iterations
Numerical integration:
Buler-Maruyanm with step size At = 0.02
Latent SDE - Setup 2 (extrapolation).
Architecture:
Observation dimension d_obs = 50; latent dimension d_latent = 10; context di- mension d_context
Encoder: Input ((P, 50)) - ODE-GRU(30); linear heads: Input (30)
Linear (3) and Input (30) Linear (20); splits into (mean, std) with std via e)
Posterior drift network:
autonomous Input (d_latent + d_context) +
Linear (30) + Softplus Linear (d_latent)
Prior drift network: Input (d_ latent) Linear (30) Softplus ()
Linear (d_lattnt
34

---

Diffusion networks: per-latent Input (1) + Linear (30) + Softplus () +
Linear (1) + Sigmoid()
Decoder: Input (d_llaten) + Linear (30) - Softplus () - Linear (30) +
Softplus () > Linear (2*d_obs)
Parameters:
Encoder: 12,027
Decoder: 4,360
Posterior drift network: 730
Prior drift network: 640
Diffusion networks: 910
Training:
Iterations: 5,000
Optimiser: Adam
Learning rate: 0.01 with exponential decay = 0.9 every 500 iterations
KL annealing: Linear over the first 500 iterations
Numerical integration: Buler-Maruyanm with step size At 0.05
SDE - (within-horizon forecasting).
Architecture:
Observation dimension d_obs 50; latent dimension d_latent 6; posterior hid- den size d_ s_hidden 100
Encoder:
Input ((T, d_obs)) reverse
GRU CRU(d_hidden)
restore
Output ( ((T+1, d_ (_hidden) (concatenate initial hidden with per-step outputs)
Posterior affine head: Input (d_hiddeen 1) Linear (d_ hiddhdndn SiLU() Linear (d_ hidden) - SiLU() - Linear (2=d_latent] splits into (p,,or) with
0₁ Sof Softsuus()) time-weighted aggregation encoder outputs enabled
Prior initial distribution: learnable diagonal Gaussian; parameters m, log s Rddlntent
Prior drift network (autonomous): Input (d_latent) + Linear (30) +
Softplus () - Linear (d_latett)
Diffusion networks (diagonal, per-latent): Input(1) - Linear (30) Softplus ()
+ Linear (1) - Sigmoid
Decoder (observation model); Input (d_ latent) + [x[Lire (30)
Softplus ()] + Linear (d_obs); Gaussian likelihood with fixed std (op=1.0, floor 10-3)
Parameters:
Prior SDE: 942
Prior Observation: 2,740
Posterior Encoder: 45,400
Posterior Affine: 21,513
Training:
Iterations: 2,000
Optimiser: AdamW
Learning rate: 0.01 with exponential decay = 0.999 per step
KL weight: 1.0
Training sequence length: sampled from 24(3, 100}
SDE - (extrapolation).
Architecture:
Observation dimension d_obs 50; latent dimension d_ d_latent 10; posterior hidden size d hidden = 1000; prior SDE: autonomous
35

---

Encoder: Input d_obs))
GRU (d_iiddee))
time
time order
Output ((T+1, a_hidden))
Posterior affine head: Input ( hidden) - Linear (d hidden) - SiLU()
Linear (d_ hiddn)) ~ SiLU() + Linear laddet)) splits into (1r,0t) with
at Softplus())) time-weighted encoder aggregation enabled; no explicit time input
Prior initial distribution: learnable diagonal Gaussian in IRd_latent
Prior drift network (autonomous):
Input (d_latent)
Linear (30)
Softplus ~ Linear Linear (d_latent)
Diffusion networks (diagonal, per-latent): Imput (1) Linear (30) Softplus ()
+ Linear (1) + Sigmoid()
Decoder (observation model); Input (d_ lallent) + 2x (30)
+
Softplus ()] + Linear (d_obs); Gaussian likelihood with fixed std (0-=1.0, floor
10
Parameters:
Prior SDE: 1,550
Prior Observation: 2,860
Posterior Encoder: 3,154,000
Posterior Affine: 2,022,021
Training:
Iterations: 100,000
Optimiser: AdamW
Learning rate: 0.0001
KI. weight: 1.0
Sequence lengths: train=200, test=300;
Evaluation: Euler-Maruyama integration with fixed step At = 0.01
E.3.6 Evaluation and Metrics
We report the mean squared error (MSE) between the predicted mean trajectory and the ground truth joint angles over the forecast horizon (steps 4-300 for Setup 1, 101-300 for Setup 2). Results in Table 3 are averaged over 10 runs with different random seeds, reporting mean + 95% t-confidence intervals where available from original papers or our runs.
E.3.7 Detailed Results
Table 8 provides FLOPs and runtime comparisons. Despite comparable computational complexity (FLOPs), Latent NSF achieves significantly faster runtime due to its ability to predict states at arbitrary time points in parallel, unlike traditional SDE methods that require sequential time-stepping.
E,4
Stochastic Moving MNIST
E.4.1 Data Generation
Data generation for Stochastic Moving MNIST was as follows:
Frame size: 64x64 pixels
Sequence length: 25 steps
Digits: Two MNIST digits, chosen randomly
Initial conditions: Positions and velocities sampled uniformly at random
Bouncing mechanism: Modified from Denton and Fergus [II for perfect reflection off boundaries with small angular noise (0, o2 hounss) where Obounce 0.1 rad
Dataset size:
Training: 60,000 sequences (using original MNIST training set)
Test: 10,000 sequences (using MNIST test set)
36

---

Table Runtime comparison (ms) for Stochastic Lorenz Attractor experiments. All measurements mean times for generating batch of 100 vectorised samples (10 trials). EM refers to Euler- Maruyama method, and SRK refers to Stochastic Runge Kutte method.
Method
Framework
0.25
0.5
t
0.75
t= 1.0
Latent SDE (EM, At 0.01)
PyTorch
141.7
281.8
417.8
1,988.3
Latent SDE (EM, At 0.01)
JAX
123.5
130.7
150.6
147.8
Latent SDE (Milstein, At = 0.01)
Py PyTorch
493.4
951.2
1,449.9
3,294.2
Latent SDE (Milstein, At 0.01)
JAX
142.7
149.4
181.4
188.7
Latent SDE (SRK, At 0.01)
Py PyTorch
1,084.5
2,189.0
3,283.9
5,827.9
Neural LSDE (EM, At 0.01)
PyTorch
101.0
198.9
291.6
1,632.7
Neural LSDE (Milstein, At 0.01)
Py PyTorch
158.4
250.9
396.1
1,767.9
Neural LSDE (SRK, At 0.01)
Py PyTorch
886.5
1,799.9
2,681.0
4,807.4
Neural GSDE (EM, At = 0.01)
Py PyTorch
107.6
215.4
316.9
1,662.3
Neural GSDE (Milstein,
0.01)
Py PyTorch
137.7
274.9
431.9
1,852.9
Neural GSDE (SRK, Al 0.01)
Py PyTorch
964.2
1,962.1
2,925.3
5,155.3
Neural LNSDE (EM, Al 0.01))
Py PyTorch
104.7
209.7
310.9
1,644.0
Neural LNSDE (Milstein, =t=0.01)
Py ToTch
130.8
290.7
381.9
1,779.1
Neural LNSDE (SRK, At =001)
Py Torch
950.0
1,921.4
2,853.7
5,023.7
NSF pHrod = 0.25)
JAX
0.221
0.303
0.341
0.373
NSF (Herdd 0.5)
JAX
0.229
0,287
0.332
0.310
NSF (perdd = 1.0)
JAX
0.255
0,286
0.294
0.295
Table FLOPs, runtime, and MSE comparison for CMU Motion Capture experiments. FLOPSs are measured per sample (latent SDE uses Fuler-Maruyama integration), while runtime measurements represent the time required to generate batch of 100 vectorised samples.
FLOPs (M)
Runtime (ms)
MSE
Method
Framework
Setup 1
Setup
Setup
Setup 2
Setup
Setup 2
Latent SDE (At 0.005)
JAX
13.5
25.1
3,585.0
3,236.0
Latent SDE (Á 0.005)
PyTorch
13.5
25.1
10.254.9
17,039.7
13.0
10.2
Latent SDE (Âr 0.01)
JAX
17.8
1,894.0
1,664.0
Latent SDE (Ar 0.01)
PyTorch
8.0
17.8
5,338.7
8,594.0
12.9
10.1
Latent SDE (Ar 0.02)
JAX
5.2
14.1
854.2
909.6
Latent SDE (Ât 0.02)
PyTorch
5.2
14.1
2,663.2
4,455.8
13.0
10.1
Latent SDE (Ât 0.05)
JAX
3.6
11.9
374.1
494.2
Latent (Âr 0.05)
PyTorch
11.9
1,201.2
2,028.2
12.8
10.2
Latent SDE [A 0.1)
JAX
3.0
11.2
223.0
372.6
Latent SDE (At 0.1)
PyTorch
3.0
11.2
595.1
1,155.9
12.9
9.7
Latent SDE (Ár 0.2)
JAX
2.7
10.8
148.4
323.3
Latent SDE (ÅI 0.2)
PyTorch
2.7
10.8
279.6
764.6
12.9
10.6
Latent SDE (AI 0.5)
JAX
2.6
10.6
107.1
289.6
Latent SDE (A 0.5)
PyTorch
2.6
10.6
136.3
584.4
12.8
31.3
Latent SDE (At 1)
JAX
2.5
10.5
99.88
279.6
Latent SDE (At
PyTorch
2.5
10.5
95.4
522.9
13.3
281.3
Latent SDE (At 2)
JAX
2.5
10.5
92.19
266.1
Latent SDE (Ar = 2)
PyTorch
2.5
10.5
77.9
492.7
19.9
Latent SDE (Ar
JAX
2.5
10.5
89.13
269.7
Latent SDE (At
PyTorch
2.5
10.5
64.4
474.6
362.1
Latent NSF (ours)
JAX
6.0
17.7
0.44
1.8
8.7
3.4
37

---

E.4.2
Configurations for Latent NSF
Architecture:
Latent dimension d_latent 10; per-frame embedding dimension d_embed = 64;
scene dimension d_scene 64
Per-frame encoder (CNN, 64x64 + 4x4): 4 down blocks: Conv2d(3+3,
stride=1, padding=1) MaxPool2d (2,2) GroupNorm SiLU(); chan-
nels: 1-664-128- 2 spatial: Linear (anear(d__mttee (d_embed)
head: Flatten +
Posterior (inference network): Input (d_emddd) + (128); head: Input (128)
+ Linear (128) SiLU() (2nd splits into (mean, std) with
std via Softplus ()
Scene encoder: temporal median pooling over per-frame embeddings; (nput(dd_beed) Linear (128) + SiLU() Linear(d_scene)
Emission:
inputs [scene, x_t] with size d_scene + d_latent: MLP()
- Linear (4=4=4*d_
reshape (4*d_embed) embed, 4, 4); up blocks: Conv (3+3) - GroupNorm (8) - upsample(*2) + SiLU(), channels: 4+d_embed
4+d_ embed + a*d_ember - embed; spatial: + 8 - 16 - 32; then Conv (3+3) ++ SiLU() -> Conv (3+3) - Sigmoid() (); output size C=1, H=W=64; Gaussian likeli-
hood with uniform trainable std (shared across pixels)
NSF transition model (latent prior):
* Gaussian base:
Input () + 2x DX[iinear (64) + SiLU()]
Linear (2+d_latee (2+d_latent); splits into (mean, std) with std via Softplus()
Affine coupling flow: 4 layers with alternating masking (autonomous SDE)
*
Per-layer conditioner: Input (d_latent/2) - +2xx[inees (64) - SiLU()] -
Linear (d_ [d_eate:)); splits into (scale, shift)
Bridge model:
*
Gaussian base:
Input + 2 [Linear (64) 7 SiLU()] 7
Linear latentt:::::; splits into (mean, std) with std via Softplus
Affine coupling flow: layers with alternating masking
* Per-layer conditioner: Input (d_latent) - 2x 2x(Linear (64) + SiLU())
Linear (d_latent)
Parameters:
Encoder: 1 1,223,360
Decoder: 1,341,570
Posterior: 11,228
Scene encoder: 33,088
NSF prior: 33,740
Bridge model: 37,260
Training:
Epochs: 1,000
Batch size: 64
Learning rate: 0.001
Latent NSF objective parameters: 3 = 30.0, X = 0.1, Skip = 30.0, Hunin = 2.5
Warm-up: 8, Bakip, and flow loss weight linearly increased over the first 20 epochs Auxiliary updates: K 3 inner steps per iteration for the bridge model
Skip-ahead KL: up to 25-step horizons: 25 samples per skip
E.4.3
Configurations for Baseline
Latent SDE We used an implementation based on Daems et al. 10. The baseline latent SDE model characteristics are follows:
Architecture:
38

---

Latent dimension d_latent 10; per-frame embedding dimension d_embed 64;
content dimension d_contents 64
Prior drift network: Input (d_latent) + Linear (200)
+ Tanh ()
+
Linear (200) Tanh + Linear (d_latent)
Diffusion networks (per-latent): Input(1) Linear (200)
Tanh
Linear (200) + Tanh() + Linear (1) - Softplus()
Posterior control network: inputs [x, htt)]; Input (2*d_laten) latent d_embed) - Linear (200) - Tanh () + Linear (200) - Tanh () + Linear Aaneardd_ (d_latent;
last layer kernel initialised to zero
Encoder (2D conv, frames 64*64+1): four down blocks: Conv Conv(3+3) + MaxPool (2) + GroupNorm + SiLU(); channels 1 + d_mbed 4 2*d_embed - A*d_ember
4+d embed; spatial 64 32 16
4;
head: Flatten () - Linear (m_ed) per frame
Content extractor: temporal median pooling over per-frame embeddings; (n_ Linear (d_embed) SiLU() Linear (d_ contenttss) out- puts
Inference network for x_0: temporal Conv (3) + SiLU() Conv (3); then MLP() on concatenated features to output 2*d_latent (mean and log-variance)
Decoder: inputs [w, x_t] with size d_contents + d_latent; MLP() > Linear (4*4*4*d ; reshape to (4, 4, 4+d_embed); 4 up blocks: Conv (3+3) + GroupNorm (8) + upsample(*2) + SiLU(), with channels _*d_embed d_ embed; spatial + 32: then Conv (3+3) - SiLU() - Conv (3+3) + Sigmoid (); output size C=1, H=W=64
4+d_ embed + 4*d_embed
16
Parameters:
Encoder: 1,269,972
Scene encoder: 8,320
Decoder: 1,341,569
Prior drift network: 44,410
Posterior drift network: 59,210
Diffusion network: 408,010
Training:
Epochs: 100
Batch size: 32
Optimiser: Adam
Learning rate: 0.0003
KL weight: 0.1
Numerical integration: Stratonovich-Milscae with step size equal to one third of the observation interval
E.4.4 Evaluation and Metrics
Evaluation was performed using Fréchet distances with embeddings from the pre-trained SRVP (using publicly available weights). The computed metrics are:
model 14
Static FD: Based on time-averaged frame embeddings
Dynamics FD: Based on sequence-level dynamics embeddings
Frame-wise FD: Averaged per-frame embedding distances
Validation of these metrics is provided in Appendix G
E.4.5 Detailed Results
Table Table@provider computational comparisons for the Stochastic Moving MNIST experiments. Beyond the quantitative metrics reported in the main text, we provide qualitative visualisations to demonstrate
39

---

the reconstruction and prediction quality of our Latent NSF model compared to the latent SDE baseline.
a direct comparison of reconstruction quality between latent SDE and Latent NSF on input sequences. Both models tasked with encoding and then reconstructing the same 25-frame input sequences. The figure demonstrates that Latent NSF achieves comparable reconstruction quality while maintaining computational efficiency.
For predictive performance, Figure 7 iillustrates the models' ability to forecast future frames given the first 25 frames of a sequence. The ground truth trajectory (top row) shows the natural bouncing dynamics the two MNIST digits. The latent SDE predictions (middle row) and Latent NSF recursive predictions (bottom row) both capture the general trajectory and bouncing behaviour, though with different levels of visual fidelity and temporal consistency.
Finally, Figure X provides focused comparison at specific prediction horizon (20 steps into the future) by showing multiple predicted samples from each model. This visualisation allows for assessment of both the accuracy of individual predictions and the diversity of the stochastic predictions across different model variants.
Hyperparameter Sensitivity and Ablation Studies
To validate the sensitivity of the hyperparameters and effectiveness of the flow loss, we conduct an ablation study on the Stochastic Lorenz Attractor with missing data (Section and CMU Motion Capture dataset with extrapolation task (Setup 2) (Section
F.1 Stochastic Lorenz Attractor with Missing Data
We trained on Stochastic Lorenz Attractor data where only a random continuous segment of length 0.5 from the full trajectory t [0, 1] was observed.
F111
Experimental Setup
We modify the training protocol as follows:
Original data structure:
The complete dataset contains trajectories with 41 time points at
E (0, 0.025, 0.05. ... 0.975, 1.0), sampled at regular intervals of At = 0.025.
Missing data protocol: We only use random continuous segment of length 20 time steps from the training data.
Training pairs: From the 20 time points, we construct possible state transition pairs (Zi,, Bay) where 4 lj. This yields 20 19/2 = 190 unique pairs per trajectory, compared 820 pairs in the complete data setting. Crucially, no training pairs directly connect states across > 20 time steps.
Training conditions: All other hyperparameters remain identical to Section E.2.2:
Epochs: 1000
Optimiser: AdamW with learning rate 0.001 and weight decay 10-5
Batch size: 256
Inner steps for auxiliary model: 5 inner optimisation steps
Flow loss weights: = 0.4
Flow loss balancing: A140-2 = A240-1 = 1.0
Evaluation: We compute KL divergence between the model and true distributions at (0.5, 1.0} using kernel density estimation as described in Section E.2.4 Results are averaged over 5 random seeds.
Based on this baseline condition, we conducted experiments varying the inner steps K, flow loss weight X, and directional flow loss components X1-1-2 and A2-0-11
40

---

Table 9: FLOPs and runtime comparison for Stochastic Moving MNIST experiments, showing computational costs for encoding 25 input frames and predicting the subsequent 25 frames. Runtime measurements represent the time required to process a batch of 100 vectorised samples. SM: Stratonovich-Milstein solver.
Method
Framework
FL.OPs (M)
Runtime (ms)
Latent SDE (SM)
JAX
20,264
750
Latent NSF (recursive)
JAX
20,276
338
Latent NSF (one-step)
JAX
20,277
358
5
%
4F
4
4
%
4
5
1/4
%
4
4
4
55
*
4
4
4
4
Figure 6: Reconstruction quality comparison on Stochastic Moving MNIST. Top row: original input frames, middle row: latent SDE reconstructions, bottom row: Latent NSF (recursive) reconstructions. Each column shows consecutive time steps (frames 1-25). Both models successfully capture the appearance and positioning of the bouncing digits.
4
4F
45
4F
46
4
4
4
4
4
4
4F
4F
45
+6
45
+6
4F
4
4
4
o
+F
4F
75
4
4F
4F
45
46
4F
4F
4F
4
4
*
4
*
4f
4
4F
45
AP
4F
4F
44
4
+
4
4F
4F
*F
4
4F
4F
4F
+F
+F
+F
4
47
4
45
4F
4
4
4
4
4
4
4
4
6
5
44
34
4
4
4F
4
9
4
6
5
5
4
44
54
44
4F
+F
+
49
4
4
4
4
4
+
4
7
4
45
+
45
*
4
5
5
94
4
$
4
4
4
4
45
4
45
45
45
4/
%/4
4
4
+5
5
4
4
Figure 7: Future frame prediction comparison Stochastic Moving MNIST. Top: ground truth continuation, middle: latent SDE predictions, bottom: Latent NSF predictions. Both models predict 25 future frames given the first 25 frames. The comparison shows how well each model captures stochastic bouncing dynamics while maintaining visual quality and temporal coherence.
45
45
45
4
54
5
5
4
y
4
45
45
45
4
Ff
Gf
uy
4
45
45
45
4F
4/
45
4
45
44
4
(a) Ground truth
(b) Latent SDE
(c) Latent NSF (one-step)
(d) Latent NSF (recursive)
Figure 8: Stochastic prediction samples 20 steps into the future on Stochastic Moving MNIST, Each panel shows multiple independent samples generated after observing the first 25 frames and predicting 20 steps ahead.
41

---

(a) Ground truth
(b) 0.0
X 0.2
(d) X 0.8
Figure 9: Comparison of generated samples (64 samples per panel) on the stochastic Lorenz attractor.
Columns show ground truth and model samples with different flow loss weights (0.0, 0.2, 0.8).
Table 10: Flow loss ablation on missing Stochastic Lorenz. Reported values are mean + standard deviation over seeds.
Flow loss weight (A)
KL 0.5)
KL (t 1.0)
0.0 (no flow loss)
7.9 0.2
22.9 - 1.8
0,001
3.0
0.9
0.01
1.7
2.1
0.4 (default)
1.2
1.2
0.8
0.2
1.3
Table 11: Effect of auxiliary optimisation steps K on KL and training time (per epoch). A = 0.4. Simultancous: auxiliary and main models updated jointly without inner loops, described in Appendix
C.1
K (auxiliary inner steps)
KL (t = 0.5)
KL (t = 1.0)
Epoch time (s)
0 (no auxiliary steps)
7.9 - 0.2
22.9 + 1.8
66
(simultaneous)
1.1
1.8 1.2
106
1.1
1.7
125
(default)
1.2
1.4
201
10
1.1 1.1 0.2
1.2
295
Table 12: Effect of directional flow loss components on KL. X 0.4, K 5.
Flow loss direction
KL 0.5)
KL (t 1.0)
1-to-2 only 71-102 = 1, A2-10-1 =
1.1 + 0.1
1.4 0.6
2-to-1 only λ1-10- 0, A2-0-11 = =
1.2 0.1
1.4
Bidirectional A140-2 = A2t0-1 =
1.2 0.1
1.4
42

---

F.1.2 Results
Flow loss weight sensitivity is summarised in Table TO and visualisation of 64 trajectories for the ground truth and generated samples is shown in Figure which obtained in the same way as Figure in the main text with one-step prediction. The effect of auxiliary inner steps shown in Table [T] Directional components compared in Table [2] Overall, the sensitivity to the flow loss weight and the number of inner steps is low when the order of magnitude is around the baseline conditions.
F.2
CMU Motion Capture Dataset with Extrapolation Task (Setup
F.2.1
Experimental Setup
We the same experimental setup Setup 2 in Section E.3.4 and conducted the parametric study varying the inner steps K, flow loss weight X, and directional flow loss components X1.1022 and 12.0-11 on the CMU Motion Capture dataset with extrapolation task (Setup 2).
F.2.2 Results
Directional variants for flow loss at 1.0, K = 3 are summarised in Table 14 The number of inner steps for the auxiliary model (with 1.0, bidirectional) are shown in Table 15] In short, best MSE occurs near À 0.03, direction 2-to-1 is slightly better than 1-to-2 but bidirectional is comparable, and small K (e.g., 1) also works while K = fails, which means that, in this case, the sensitivity to the flow loss weight and the number of inner steps is low when the order of magnitude is around the baseline conditions.
Flow loss weight (bidirectional, 3) is reported in Table 13
F33 Practical Guidance on Hyperparameters
We provide practical guidance on hyperparameters specifically for the NSF or Latent NSF.
F.3.1 Time Horizon for Training Hirain and Inference prred
The time horizon for training should be the maximum one-shot interval used during training. In our experiments, we found that there no significant difference performance dependent on the Hped value. However we have not explored the effect the Htrain value. Since, in practice, large Hwain requires more expressive model capacity, we recommend balancing trade-off between model capacity and prediction runtime using H and using Hpred Harain in most cases, depending on the data complexity and prediction runtime requirements.
F.3.2 Flow Loss Weight
For any experiment we have conducted, found that sub-order of magnitude À values are sufficient to achieve good performance, assuming the data is normalised.
F.3.3 Auxiliary Inner Steps or Simultaneous Updates
The auxiliary inner steps K controls the number of inner steps for the auxiliary model. We recommend using simultaneous updates first, as described in Appendix C.2 and if there is training instability, recommend using K 0 and increasing the value of K until the flow loss is stable.
F.3.4 Directional Flow Loss Components
The directional flow loss components A1 and A control the strength of the flow loss in the 1-to- 2 and 2-to-1 directions respectively. In our experiments, though the performance is not significantly different, we recommend using bidirectional flow loss for symmetry and stability.
43

---

G
Validation of the Fréchet Image and Video Metrics
G.1
Protocol
To confirm that the three Fréchet distances introduced in $6.3 (static content, dynamics and frame- wise) behave intended Stochastic Moving MNIST benchmark, performed series of controlled experiments. Code for reproducing the experiments is available at https:/ //gtth// hom//kkiyoohar//rrpp-ff
G.1.1 Datasets.
For each trial, two datasets were instantiated as follows:
Content:
Each containing 256 sequences of 25 frames.
Production method: Same as in the main text with identical hyperparameters, varying only specific factors (digit classes, handwriting samples, dynamics).
Digit ordering: To remove ambiguity, a second copy of the first dataset was generated with digit indices swapped. The version yielding the smaller static Fréchet distance was retained
G.1.2 Conditions.
Four experimental conditions were repeated over 32 trials (each with fresh random seed);
Different dynamics: same digits and handwriting, different dynamics.
Different digits: different digit classes, identical dynamics.
Different handwriting: same digit classes, different handwriting, identical dynamics.
Different both: different digit classes and different dynamics.
G.1.3
Computation of distances.
Distance computation details:
Encoder: Public SRVP encoder of Franceschi al. [14
Frame-wise scores: Averaged the 25 frame distances to yield a single number per sequence pair.
G.2 Results
Table reports the mean and standard deviation (over the 32 trials) of each metric. The metrics reflect the intended factors:
Varying only the dynamics (different dynamics) drives the dynamics distance up by order magnitude while leaving the static distance low.
Altering appearance while keeping motion fixed (different digits/handwriting) raises the static distance, with digits > handwriting as expected, and leaves the dynamics distance low.
When both factors change (different both), both metrics are simultaneously high.
Frame-wise distances stay near zero unless both content and motion differ, indicating that they act as a weak sanity check.
'Without this post-hoc swap the static metric can be artificially inflated when the same digits appear in reverse order.
44

---

Table 13: Effect of flow loss weight À on test MSE (mean + 95% confidence interval t-statistic) for CMU Motion Capture (Setup 2). Auxiliary inner steps fixed to K 3. Flow loss weight 1.0, 2 = A240-1 1.0), and KL weight 3 = Sakip = 0.3.
Flow loss weight (A)
MSE
0.001
4.08 + 0.46
0,003
3.77 + 0.37
0.01
3.73 0.48
0,03
:.26
0.1
3.43 + 0.29
0.3
3.60 + 0.33
1.0
3.41 + 0.27
Table 14: Effect of directional flow loss components on test MSE (mean + 95% confidence interval in t-statistic) for CMU Motion Capture (Setup 2). Flow loss weight 1.0, inner loop steps for bridge model K 3, and KL weight 3 = Sskip 0.3.
Flow loss direction
MSE
1-to-2 only A1402 = 1, A2 = 0
3.66 + 0.50
2-to-1 only Ai-t022 = 0, A2-1--1 3.48 + 0.46
Bidirectional A1-to-2 A240-1 =1 3.41 + 0.27
Table 15: Effect of auxiliary optimisation steps K on MSE (mean + 95% confidence interval in
I-statistic) for CMU Motion Capture (Setup 2). Flow loss weight X = 1.0, bidirectional =
A2-w-1 1.0), and KL weight 3 Sakip 0.3. "Simultaneous" refers to updating auxiliary and main models jointly without inner loops for bridge model, as described in Appendix C.2
K (auxiliary inner steps)
MSE
0 (no auxiliary training) (simultaneous)
261.82 + 124.57
3.59 0.61
3.35 + 0.39
3 (default)
3.41 + 0.27
5
3.68 + 0.40
10
3.59 0.38
Table 16: Mean + standard deviation of the three Fréchet distances across 32 trials.
Condition
Static
Dynamics
Frame mean
Different dynamics
2.57 1.40
14.85 + 5.48
0.80 0.14
Different digits
5.66 + 1.82
1.75 + 1.83
0.13 + 0.06
Different handwriting
3.74 + 1.65
1.79 + 1.66
0.09 0.05
Different both
6.78 + 1.95
14.74 4 5.79
0.87 + 0.13
45

---

G.3
Discussion
The clear separation between conditions, the low variance within each condition and the alignment with human intuition together demonstrate that our Fréchet metrics are both reliable and interpretable. They therefore provide a sound basis for quantitative evaluation of generative video models on Stochastic Moving MNIST, capturing complementary aspects of visual content and motion without leakage between two.
46


Reviews:
[{"review_id": "6Hceg6TSWl", "reviewer": "Reviewer_TqQA", "summary": "The paper focuses on using neural stochastic flows to model stochastic dynamic data without solving SDEs. To do this, it is proposed to use conditional normalising flows to learn the SDE transition distribution. It also allows efficient one-step sampling to facilitate the computational efficiency. Experiments were conducted with Stochastic Lorenz Attractor, CMU Motion Capture, and Stochastic Moving MNIST.", "questions": "Algorithm in Appendix C can be written in the main text, since it helps to better follow the process. Minor: (1) line 130: typo in \"regularisation\"; (2) Line 170: in the description of the total objective, \"the flow loss\" might refer to the Eq (25) in the appendix instead of Eq (24) as in the current version.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "* The work focused on important problem. To avoid learning the transition law via simulations with a large amount of steps, the paper follows a solver-free approach. The paper proposed a specific design (architecture) for parameterising the conditional normalising flows, and the loss to train this. To capture the noisy scenarios, the framework is also extended to Latent-NFS.\n\n* The paper is written in self-contained way. Figure 1 can well-illustrate the concept behind the proposed approach (neural stochastic flows-NSF). Section 2 briefly introduce to the SDEs and Stochastic Flows to follow the next sections of the paper. \n\n* The approach (NSF and Latent NSF) has strong motivations and technically sound. \n   - To avoid costly computation of characterising strong solutions of SDEs, the approach derives the alternative weak solutions. A parametric model is then proposed to approximate it (Eq (2) and (3), with design in section 3.1). The conditions for the converted weak solutions of SDEs are described provided in section 3.\n   - To enforce the flow property (especially the property 2 in Section 3) in NSFs, the approach utilised variational technique to construct upper-bounds on KL divergence for $\\mathcal{L}_{\\text{flow}}$ beside the normal supervised loss. \n   - The connection made for the choice of the structure of Latent NSF is described in lines 157-164 with support in the Figure 2. \n\n* The experiments is showed that the computational cost is significantly reduced, which is an important contribution in term of the efficiency of this approach. Compared to other Latent- and Neural- SDE existing approaches, NSF obtained low KL divergence (high performance) with much lower computational cost (Table 1). Analyses about the trade-off between computation and performance is also conducted in Figure 4 (with CMU Motion Capture dataset) and Table 3 (for Stochastic Moving MNIST).", "was_revised": true, "final_justification": "Dear Area Chairs,\n\nIn the rebuttal phase, authors have discussed and adequately addressed specific questions and suggestions in my reviews of the paper. For the final justification, I maintain my overall score of 5: Accept.\n\nKind regards", "rebuttal": "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.\n# Summary of contributions\nWe propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include\n\n- 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,\n- A conditional normalisation flow that preserves important properties of stochastic flows,\n- The “flow loss” as a regulariser to enforce flow consistency.\n\nReviewers 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. \n# Common questions\n## Scalability (YJDb, yfA9)\nRegarding 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].\n## Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)\nNonlinear 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).\n\nAmong 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.\n## Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)\nWe 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.\n\n**Flow Loss Ablation on Missing Stochastic Lorenz**\n|Flow Loss Weight ($\\lambda$)|KL ($t=0.5$)|KL ($t=1.0$)|\n|-|-|-|\n|0.0 (no flow loss)|7.9 ± 0.2|22.9 ± 1.8|\n|0.001|3.0 ± 0.2|4.5 ± 0.9|\n|0.01|1.7 ± 0.1|2.1 ± 0.6|\n|0.4 (default)|1.2 ± 0.1|1.2 ± 1.0|\n|0.8|1.3 ± 0.2|1.3 ± 0.6|\n\n**Effect of Auxiliary Optimisation Steps ($\\lambda=0.4$)**\n|K (auxiliary steps)|KL ($t=0.5$)|KL ($t=1.0$)|Epoch Time (s)|\n|-|-|-|-|\n|0 (no auxiliary)|7.9 ± 0.2|22.9 ± 1.8|66|\n|0 (simultaneous)*|1.1 ± 0.1|1.8 ± 1.2|106|\n|1 |1.1 ± 0.1|1.7 ± 0.6|125|\n|5 (default)|1.2 ± 0.1|1.4 ± 0.6|201|\n|10|1.1 ± 0.2|1.2 ± 0.7|295|\n\n\\* Simultaneous: auxiliary and main models updated jointly without inner loops\n\n**Directional KL Components ($\\lambda=0.4, K=5$)**\n\n|Flow Loss Direction|KL ($t=0.5$)|KL ($t=1.0$)|\n|-|-|-|\n|Forward only (1→2)|1.1 ± 0.1|1.4 ± 0.6|\n|Reverse only (2→1)|1.2 ± 0.1|1.4 ± 0.7|\n|Bidirectional|1.2 ± 0.1|1.4 ± 0.6|\n\nWithout 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.\n\nApplying 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.\n\nA similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.\n## Main paper’s content organisation (TqQA, yfA9)\nWe will move key architectural details (NSF design, training algorithm) to the main text for clarity.\n\n---\n# Specific questions from Reviewer TqQA\n\nWe 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:\n\n## Algorithm Placement\n\nWe 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.\n\n## Minor Corrections\n\nThank you for catching typo and wrong reference. We will fix them in the camera-ready version.\n\n---\n## References\nFor 1-50, please refer to the main text.\n\n[51] Zhai et al., \"Normalizing Flows are Capable Generative Models\", ICML 2025.\n\n------\nThank you again, and we hope this addresses your concerns and would be happy to provide any additional clarifications."}, {"review_id": "hoQr5jY3DT", "reviewer": "Reviewer_YJDb", "summary": "The paper proposes neural stochastic flows, which learn transition laws via conditional normalizing flows. The normalizing flows are constructed in such a way that the identity and Markov properties are respected and trained with an additional regularization loss for the flow property/Chapman Kolmogorov equation. The model is tested on three experiments, improving over baselines in terms of speed and performance.", "questions": "- The way $x_{t_i}$ is sampled in the expectations in eq. (8) and eq. (10) is a bit unclear. Is it just obtained as a random sample from the dataset at time $t_i$?\n- l238: \"this is due to the inductive bias of SDEs for better generalization\" What is meant by that?", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper is well written and the method seems mathematically sound\n- Neural stochastic flows show strong empirical performance while being much faster than non solver-free methods. \n\nWeaknesses:\n- A lot of the desired properties of the flows are incentivized through additional loss terms leading to multiple hyperparameters (K, $\\lambda$, $\\lambda_{\\text{1to2}}$, $\\lambda_{\\text{2to1}}$, $\\beta$, $\\beta_{\\text{skip}}$). It is not clear how to best select all these parameters making extensive hyperparameter tuning necessary. How sensitive is the method to the choice of hyperparameters in practice? \n- The dependence on affine coupling layers makes the method not very scalable. Is there a way to lift this dependence?\n\nMinor: \n- l26: before consistency models, general distillation approaches for diffusion models also tried to achieve this using the probability flow ODE to define the mapping [1]\n- l130: \"regularisaion\" -> regularisation\n \nI am not familiar enough with the recent related work to properly assess the novelty.\n\n[1] Progressive Distillation for Fast Sampling of Diffusion Models https://openreview.net/pdf?id=TIdIXIpzhoI", "was_revised": true, "final_justification": "The authors have addressed my concerns and questions in the rebuttal. I maintain my positive rating.", "rebuttal": "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.\n# Summary of contributions\nWe propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include\n\n- 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,\n- A conditional normalisation flow that preserves important properties of stochastic flows,\n- The “flow loss” as a regulariser to enforce flow consistency.\n\nReviewers 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. \n# Common questions\n## Scalability (YJDb, yfA9)\nRegarding 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].\n## Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)\nNonlinear 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).\n\nAmong 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.\n## Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)\nWe 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.\n\n**Flow Loss Ablation on Missing Stochastic Lorenz**\n|Flow Loss Weight ($\\lambda$)|KL ($t=0.5$)|KL ($t=1.0$)|\n|-|-|-|\n|0.0 (no flow loss)|7.9 ± 0.2|22.9 ± 1.8|\n|0.001|3.0 ± 0.2|4.5 ± 0.9|\n|0.01|1.7 ± 0.1|2.1 ± 0.6|\n|0.4 (default)|1.2 ± 0.1|1.2 ± 1.0|\n|0.8|1.3 ± 0.2|1.3 ± 0.6|\n\n**Effect of Auxiliary Optimisation Steps ($\\lambda=0.4$)**\n|K (auxiliary steps)|KL ($t=0.5$)|KL ($t=1.0$)|Epoch Time (s)|\n|-|-|-|-|\n|0 (no auxiliary)|7.9 ± 0.2|22.9 ± 1.8|66|\n|0 (simultaneous)*|1.1 ± 0.1|1.8 ± 1.2|106|\n|1 |1.1 ± 0.1|1.7 ± 0.6|125|\n|5 (default)|1.2 ± 0.1|1.4 ± 0.6|201|\n|10|1.1 ± 0.2|1.2 ± 0.7|295|\n\n\\* Simultaneous: auxiliary and main models updated jointly without inner loops\n\n**Directional KL Components ($\\lambda=0.4, K=5$)**\n\n|Flow Loss Direction|KL ($t=0.5$)|KL ($t=1.0$)|\n|-|-|-|\n|Forward only (1→2)|1.1 ± 0.1|1.4 ± 0.6|\n|Reverse only (2→1)|1.2 ± 0.1|1.4 ± 0.7|\n|Bidirectional|1.2 ± 0.1|1.4 ± 0.6|\n\nWithout 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.\n\nApplying 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.\n\nA similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.\n## Main paper’s content organisation (TqQA, yfA9)\nWe will move key architectural details (NSF design, training algorithm) to the main text for clarity.\n\n---\n# Specific questions from Reviewer YJDb\n\nThank 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.\n\n## Hyperparameter sensitivity\n\nWe 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.\n\n## Scalability and Architecture\n\nPlease see “Scalability” above in common response, where we discuss the limitation on the scalability and architecture.\n\n## Sampling procedure for expectation\n\nThank you for your careful reading and specific questions.\nThe 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]$.\nWe 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.\n\nSimilarly, 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]$.\nWe 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.\n\nWe will clarify these expectation in camera-ready version.\n\n## Inductive bias of SDEs\n\nPlease see “Clarifying inductive bias of SDEs in MoCap dataset” above in common response.\n\n## Minor Points\n\nThank you for the reference to Progressive Distillation. We'll add this important connection and will also fix the typo on l130.\n\n---\n## References\nFor 1-50, please refer to the main text.\n\n[51] Zhai et al., \"Normalizing Flows are Capable Generative Models\", ICML 2025.\n\n---\n\nWe 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."}, {"review_id": "4Mc3zNFuHc", "reviewer": "Reviewer_aZhw", "summary": "The authors propose to use conditional normalizing flows to approximate transition densities in SDE. This yields a solver-free method to sample between arbitrary timepoints in a single step. \"Neural stochastic flows (NSF)\" is a conditional coupling flow aiming to estimate transition densities $p(x_t|x_s, s, t-s)$ of SDEs (for all s,t). The authors enforce that some known properties of them are architecturally embedded, and develop a regularization term to enforce the \"flow\" property. It's shown in experiments on synthetic SDEs and real-world tracking and video data that NSF stays distributionally accurate while strongly reducing computational cost.", "questions": "See weaknesses i.e.:\n- How much does the regularization term help (compared to no regularization) ?\n- Why forward and backward KL for regularization? This also introduces the \"weighting\" hyperparameters. As reverse and forward KL are often on quite different \"scales,\" this might be important to \"tune\". Did the authors investigate just using the forward term (or just the backward term)?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- Well written, structured, and motivated, i.e., first pointing out the target's properties, then incorporating them in the model.\n- Overall, the empirical results show that the proposed method is accurate and comparable to other approaches, but much more computationally efficient.\n\nWeaknesses:\n- One of the main methodological novelties is the regularization term for the Flow property. However, it only occupies ~0.5 page in the method section and some design decisions about it remain unclear (i.e., \"A natural design for such a regularization loss is to combine both forward and reverse KL ...\" Why is this the natural choice? There are several other possibilities to achieve the same. \n- Missing ablations e.g. given that this regularization term adds quite a lot of complexity to the approach, I miss an ablation that this term does empirically improve upon the standard maximum likelihood loss (which should enforce the Flow Property implicitly in a data-driven manner). How much slower is it to train with regularization? I expect quite a bit.\n- Table 2 raises questions: SDE Matching performs best in Setup 1 (by quite a big margin compared to the proposed method), but results for this approach on Setup 2 are not reported (why?). Given that the authors claim \"state-of-the-art long-horizon extrapolation performance\" this result should be included.", "was_revised": true, "final_justification": "The authors adressed my main concerns in the rebutal:\n- \"Flow Loss Ablation\" in the rebutal did clearly demonstrate that the auxiliary loss is important.\n- The authors did to what is possible to include a fair comparission with SDE matching.", "rebuttal": "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.\n# Summary of contributions\nWe propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include\n\n- 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,\n- A conditional normalisation flow that preserves important properties of stochastic flows,\n- The “flow loss” as a regulariser to enforce flow consistency.\n\nReviewers 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. \n# Common questions\n## Scalability (YJDb, yfA9)\nRegarding 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].\n## Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)\nNonlinear 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).\n\nAmong 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.\n## Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)\nWe 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.\n\n**Flow Loss Ablation on Missing Stochastic Lorenz**\n|Flow Loss Weight ($\\lambda$)|KL ($t=0.5$)|KL ($t=1.0$)|\n|-|-|-|\n|0.0 (no flow loss)|7.9 ± 0.2|22.9 ± 1.8|\n|0.001|3.0 ± 0.2|4.5 ± 0.9|\n|0.01|1.7 ± 0.1|2.1 ± 0.6|\n|0.4 (default)|1.2 ± 0.1|1.2 ± 1.0|\n|0.8|1.3 ± 0.2|1.3 ± 0.6|\n\n**Effect of Auxiliary Optimisation Steps ($\\lambda=0.4$)**\n|K (auxiliary steps)|KL ($t=0.5$)|KL ($t=1.0$)|Epoch Time (s)|\n|-|-|-|-|\n|0 (no auxiliary)|7.9 ± 0.2|22.9 ± 1.8|66|\n|0 (simultaneous)*|1.1 ± 0.1|1.8 ± 1.2|106|\n|1 |1.1 ± 0.1|1.7 ± 0.6|125|\n|5 (default)|1.2 ± 0.1|1.4 ± 0.6|201|\n|10|1.1 ± 0.2|1.2 ± 0.7|295|\n\n\\* Simultaneous: auxiliary and main models updated jointly without inner loops\n\n**Directional KL Components ($\\lambda=0.4, K=5$)**\n\n|Flow Loss Direction|KL ($t=0.5$)|KL ($t=1.0$)|\n|-|-|-|\n|Forward only (1→2)|1.1 ± 0.1|1.4 ± 0.6|\n|Reverse only (2→1)|1.2 ± 0.1|1.4 ± 0.7|\n|Bidirectional|1.2 ± 0.1|1.4 ± 0.6|\n\nWithout 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.\n\nApplying 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.\n\nA similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.\n## Main paper’s content organisation (TqQA, yfA9)\nWe will move key architectural details (NSF design, training algorithm) to the main text for clarity.\n\n---\n# Specific questions from Reviewer aZhw\nThank you for your valuable feedback. We address your specific concerns below.\n## Motivation for Bidirectional KL Divergence in Flow Loss\nWe appreciate this opportunity to clarify our design choice.\n\nIn 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.\n\n- 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.\n- 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.\n## Ablation Studies on Regularization Term\nPlease see “Hyperparameter Sensitivity and Ablation Studies” in common response.\n## SDE Matching Comparison (Table 2)\nSDE 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.\n\nWe 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.\n|**Method**|**$t=0.25$**||**$t=0.5$**||**$t=0.75$**||**$t=1.0$**||\n|-|-|-|-|-|-|-|-|-|\n||KL|kFLOPs|KL|kFLOPs|KL|kFLOPs|KL|kFLOPs|\n|**SDE Matching**|||||||||\n|$\\Delta t=0.0001$|4.3 ± 0.7|184338|5.3 ± 1.0|368677|3.4 ± 0.8|553015|3.8 ± 1.0|737354|\n|$\\Delta t=0.01$|6.3 ± 0.4|1880|11.7 ± 0.5|3760|7.9 ± 0.3|5640|6.0 ± 0.3|7520|\n|**NSF (ours)**|||||||||\n|$t_\\mathrm{max}=1.0$|0.8 ± 0.7|53|1.3 ± 0.1|53|0.6 ± 0.3|53|0.2 ± 0.6|53|\n|$t_\\mathrm{max}=0.5$|2.4 ± 1.9|53|1.3 ± 0.1|53|1.0 ± 0.4|105|1.7 ± 1.1|105|\n|$t_\\mathrm{max}=0.25$|1.2 ± 0.7|53|1.2 ± 0.1|105|0.8 ± 0.4|156|1.3 ± 0.8|208|\n\nWe 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.\n\n---\n## References\nFor 1-50, please refer to the main text.\n\n[51] Zhai et al., \"Normalizing Flows are Capable Generative Models\", ICML 2025.\n\n------\nThank you again, and we hope this addresses your concerns and would be happy to provide any additional clarifications."}, {"review_id": "vDfXEbz7RQ", "reviewer": "Reviewer_yfA9", "summary": "The paper presents neural stochastic flows, a novel framework for modelling trajectories of stochastic differential equations for time series data. The method utilizes normalising flows to get a direct access to the probability distribution of SDE transitions, which allows it to obtain a sample from an arbitrary time step without requiring a numerical solver. Experimental results show that the proposed approach significantly speeds up the sampling process while maintaining high sample quality.", "questions": "All my questions are related to the Experiments section.\n\n**Stochastic Lorenz Attractor**   \n1. In Table 1, the meaning of the variable $t$ is unclear. What does it represent, and why does the KL divergence vary non-monotonically with changes in $t$?\n2. Why does NSF with $t_{\\text{max}} = 0.5$ perform worse than NSF with $t_{\\text{max}} = 1$ and $t_{\\text{max}} = 0.25$?\n3. How does Figure 3 relate to Table 1? Were the depicted trajectories generated in a one-step fashion, or sequentially by conditioning each point on the previous one?\n\n**CMU Motion Capture**   \n1. In the Latent SDE paper [1] I see that the reported MSE for Setup 1 is 4.03, which is notably better than the result of your reproduction. Please explain this discrepancy.\n2. Figure 4 (a) clearly indicates that MSE grows with increasing runtime for Latent NSF. How can you explain this behavior?\n3. You attribute the better performance of Latent NSF in Setup 2 to the SDE's inductive bias. Could you please clarify the idea of inductive bias here and elaborate on why Setup 2 constitutes extrapolation while Setup 1 does not for this particular dataset?\n\n**Stochastic Moving MNIST**   \n1. The variance of the Fréchet Distance in Table 3 and Figure 5 appears high relative to the mean differences between methods. This raises concerns about the statistical significance of the results. How many samples were used to compute these metrics, and what is the level of significance?\n2. There is no analysis of the results. How can you explain that recursive NSF is much better then SDE in terms of Static FD and much worse in terms of Dynamics FD? How does the speed of recursive NSF compare with one-step NSF? Are the results aligned with expectations, and how do they compare with findings on the other datasets? Why does recursive NSF perform better for Moving MNIST, while appearing to underperform on the Motion Capture dataset?\n\n**General questions**   \n1. I noticed that you use different dimensionality of latent space for different methods (Appendix E.3 and E.4), which might have an impact on quality. Could you please explain how you chose these values and all other hyperparameters, and how different methods compare in terms of the number of parameters?\n2. All experiments are conducted on relatively small-scale problems (latent space dimensionality ≤ 10). Is the proposed method applicable to larger-scale settings, such as in the video domain?\n\n[1] X. Li, T. Wong, R. T. Q. Chen, and D. Duvenaud. Scalable gradients for stochastic differential\nequations. In Proceedings of the 23rd International Conference on Artificial Intelligence and\nStatistics (AISTATS), 2020.", "limitations": "The limitations are discussed.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**   \nThe paper structure is good, the proposed method is well-motivated and its theoretical justification is strong.\n\n**Weaknesses**   \nThe experimental evaluation lacks clarity and depth in analysis (see Questions section). Some of the important implementation details are not presented in the main text of the paper. For example, the architecture of the proposed normalizing flow is a crucial part of the approach and it should not be moved to the appendix.", "was_revised": true, "final_justification": "My main concerns were addressed by the authors in their rebuttal:\n\n* The difference between Latent SDE results in the original paper and the authors' reproduction is explained by the well-known issue with the original Latent SDE implementation.\n* Experiment setup became clear for the Stochastic Lorenz Attractor.\n* The difference between Static FD and Dynamics FD is explained for the Stochastic Moving MNIST and the statistical significance of the results is proven.", "rebuttal": "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.\n# Summary of contributions\nWe propose Neural Stochastic Flows (NSF) as a solver-free method to directly learns the weak solution of an SDE from data. Key components include\n\n- 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,\n- A conditional normalisation flow that preserves important properties of stochastic flows,\n- The “flow loss” as a regulariser to enforce flow consistency.\n\nReviewers 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. \n# Common questions\n## Scalability (YJDb, yfA9)\nRegarding 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].\n## Clarifying inductive bias of SDEs in MoCap dataset (YJDb, yfA9)\nNonlinear 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).\n\nAmong 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.\n## Hyperparameter Sensitivity and Ablation Studies (YJDb, aZhw)\nWe 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.\n\n**Flow Loss Ablation on Missing Stochastic Lorenz**\n|Flow Loss Weight (λ)|KL (t=0.5)|KL (t=1.0)|\n|-|-|-|\n|0.0 (no flow loss)|7.9 ± 0.2|22.9 ± 1.8|\n|0.001|3.0 ± 0.2|4.5 ± 0.9|\n|0.01|1.7 ± 0.1|2.1 ± 0.6|\n|0.4 (default)|1.2 ± 0.1|1.2 ± 1.0|\n|0.8|1.3 ± 0.2|1.3 ± 0.6|\n\n**Effect of Auxiliary Optimisation Steps (λ=0.4)**\n|K (auxiliary steps)|KL (t=0.5)|KL (t=1.0)|Epoch Time (s)|\n|-|-|-|-|\n|0 (no auxiliary)|7.9 ± 0.2|22.9 ± 1.8|66|\n|0 (simultaneous)*|1.1 ± 0.1|1.8 ± 1.2|106|\n|1 |1.1 ± 0.1|1.7 ± 0.6|125|\n|5 (default)|1.2 ± 0.1|1.4 ± 0.6|201|\n|10|1.1 ± 0.2|1.2 ± 0.7|295|\n\n\\* Simultaneous: auxiliary and main models updated jointly without inner loops\n\n**Directional KL Components (λ=0.4, K=5)**\n\n|Flow Loss Direction|KL (t=0.5)|KL (t=1.0)|\n|-|-|-|\n|Forward only (1→2)|1.1 ± 0.1|1.4 ± 0.6|\n|Reverse only (2→1)|1.2 ± 0.1|1.4 ± 0.7|\n|Bidirectional|1.2 ± 0.1|1.4 ± 0.6|\n\nWithout 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.\n\nApplying 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.\n\nA similar trend was observed on MoCap setup 2 and we will include these ablations in the revision.\n## Main paper’s content organisation (TqQA, yfA9)\nWe will move key architectural details (NSF design, training algorithm) to the main text for clarity.\n\n---\n# Specific questions from Reviewer yfA9\nThank you for your valuable feedback. We address your concerns regarding experiments below.\n## Stochastic Lorenz Attractor\n1. 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.\n2. 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$).\nThe 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.\n3. 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.\n## MoCap\n1. 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.\n2. 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.\n3. 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.\n## Stochastic Moving MNIST\n1. 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.\n2. 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.\nThe 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.\nThe 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.\n## General questions\n1. 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].\nWe 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.\n2. Noted in the common rebuttal “Scalability”.\n\n---\n## References\nFor 1-50, please refer to the main text.\n\n[51] Zhai et al., \"Normalizing Flows are Capable Generative Models\", ICML 2025.\n\n---\nWe hope this addresses your concerns and would be happy to provide any additional clarifications."}]

Author rebuttal:

Meta-review: This submission introduces a new view of continuous-time dynamics modeling based on approximating a stochastic flow, the transition probabilities of a stochastic process. The approach relies on coupling-style normalizing flows that by definition satisfies 3 of the 4 standard properties of a stochastic flow. The flow property is satisfied approximately by combining a cross-entropy/likelihood objective with a regularisation loss. Neural tochasic flow can be thought of as a stochastic analogue to the deterministic distillation-like techniques such as consistency trajectory models, flow map matching, etc.

The reviewers reached a consensus that this is an interesting approach for time-series modeling that is theoretically sound and potentially practically useful and thus I recommend acceptance. The main concerns were regarding experiments and significant amount of hyperparameters which the authors mostly resolved during the rebuttal phase. When revising the paper, please pay attention to the promised improvements and extra results. I also strongly urge the authors to include both the original result of the Li et al. (2020) paper as well as the failed replication.
Decision: Accept (poster)