Elucidating the Design Choice of Probability Paths in Flow Matching for Forecasting
Authors: Soon Hoe Lim
Yijin Wang
Annan Yu
Emma Hart
Michael W. Mahoney
Sherry Li
N. Benjamin Erichson
Venue: Accepted by TMLR ()

Abstract: Flow matching has recently emerged as a powerful paradigm for generative modeling and has been extended to probabilistic time series forecasting. However, the impact of the specific choice of probability path model on forecasting performance, particularly for high-dimensional spatio-temporal dynamics, remains under-explored. In this work, we demonstrate that forecasting spatio-temporal data with flow matching is highly sensitive to the selection of the probability path model. Motivated by this insight, we propose a novel probability path model designed to improve forecasting performance. Our empirical results across various dynamical system benchmarks show that our model achieves faster convergence during training and improved predictive performance compared to existing probability path models. Importantly, our approach is efficient during inference, requiring only a few sampling steps. This makes our proposed model practical for real-world applications and opens new avenues for probabilistic forecasting.

Annan Yu Published in Transactions on Machine Learning Research (07/2025)
Elucidating the Design Choice of Probability Paths in Flow Matching for Forecasting
Soon Hoe Lim
shlim@kth.cs
Deportment of Mathematics, KTH Royal Institute of Technology Nordia, KTH Royel Institute of Technology and Stockholm Universily
Yijin Wang International Computer Science Institate
yijín. wjjn..wang@heerkeeeee edu
ay2620cormell.edv
Center for Applied Mathemalicss Cornell Universily
Emma Hart
emmahar1@lbl.gg
Department of Mathematics, Emory University
Michael W. Mahoney
mmahoney@stat. berkeley. edu
Deportment of Statistics, University of California at Berkeley International Computter Science Institute Lauvence Berkeley National Laboralory
Xiaoye S. Li Laurence Berkeley National Laboratory
rali@lbl gov
N. Benjamin Erichson
erichson berkeley, edu
International Computer Science Institute
Laurence Berkeley National Laboratory
Reviewed on OpenReview: https: // openreview. net/ forum? id- JApMDWLbU
Abstract
Flow matching has recently emerged as a powerful paradigm for generative modeling and has been extended probabilistic time series forecasting. However, the impact of the specific choice of probability path model on forecasting performance, particularly for high- dimensional spatio-temporal dynamics, remains under-cxplored. In this work, demonstrate that forecasting spatio-temporal data with flow matching is highly sensitive to the selection of the probability path model. Motivated by this insight, we propose novel probability path model designed to improve forecasting performance. Our empirical results across various dynamical system benchmarks show that our model achieves faster convergence during training and improved predictive performance compared to existing probability path models. Importantly, our approach is efficient during inference, requiring only a few sampling steps. This makes our proposed model practical for real-world applications and opens new avenues for probabilistic forecasting.
1 Introduction
Generative modeling has achieved remarkable success in recent years, especially for generating high- high-dimensional objects by learning mappings from simple, casily-sampled reference distributions, To. to complex target distributions, T.. In particular, diffusion models have pushed the capabilities of generating realistic samples across various data modalities, including images (Ho et al., 2020; Song et al., 2020b; Karras et al., 2022),

---

Published in Transactions on Machine Learning Research (07/2025)
videos (Ho et al., 2022; Blattmann et al., 2023; Gupta et al., 2023), and spatio-temporal scientific data like climate and weather patterns (Pathak et al., 2024; Kohl et al., 2024). Despite their impressive performance, diffusion models often come with high computational costs during training and inference. Additionally, they typically assume a Gaussian reference distribution, which may not be optimal for all data types and can limit modeling flexibility.
One promising alternative is flow matching, where the mappings are learned via a stochastic process that transforms To into म through ordinary differential equation (ODE), approximating its marginal vector flow (Lipman et al., 2022; Albergo et al., 2023; Liu et al., 2022; Tong et al., 2023b; Pooladian et al., 2023; Lipman et al., 2024). While score-based models (Song & Ermon, 2019; Song et al., 2020a;b; Ho et al., 2020) are specific instances of flow matching, with Gaussian transition densities, the general framework allows for a broader class of interpolating paths. This flexibility can lead to deterministic sampling schemes that are faster and require fewer steps (Zhang & Chen, 2022). Recent work has demonstrated the remarkable capabilities of flow matching models for generating high-dimensional images (Esser et al., 2024) and discrete data (Gat et al., 2024).
Building on this, flow matching in latent space has recently been applied to forecasting spatio-tempont data (Davtyan et al., 2023) (predicting future frames in videos). This approach leverages latent representations to capture the complex dynamics inherent in temporal data. However, spatio-temporal forecasting, especially for videos and dynamical systems data, presents unique challenges. A video prediction model capable of generalizing to new, unseen scenarios must implicitly "understand" the scene: detecting and classifying objects, learning how they move and interact, estimating their 3D shapes and positions, and modeling the physical laws governing the environment (Battaglia et al., 2016). Similarly, accurate weather forecasting requires capturing intricate physical processes and interactions across multiple scales (Dueben & Bauer, 2018; Schultz et al., 2021).
We observe that, in the context of spatio-temporal forecasting, the performance of flow matching is highly sensitive to the choice of the probability path model, an important topic which has not been widely explored within a unified framework. Different. probability paths can significantly impact the accuracy and convergence of forecasting models, particularly when dealing with complex dynamical systems characterized by partial differential equations (PDEs) and chaotic bchaviors. Motivated by this, we propose novel probability path model tailored for probabilistic forecasting of dynamical systems, with a particular focus on PDE datasets relevant to scientific applications. Our model leverages the continuous dynamics intrinsic to temporal data by interpolating between consecutive sequential samples. This approach ensures better alignment with the constructed flow, leading to improved predictive performance, more stable training, and greater inference efficiency. Existing probability path models often fail to fully capture the continuous nature of spatio-temporal data, resulting in a misalignment with flow-based methods and suboptimal outcomes. Our proposed model addresses these limitations directly.
Building on previous approaches, we provide theoretical framework and efficient algorithms tailored to probabilistic forecasting of high-dimensional spatio-temporal dynamics using flow matching in latent space. Within this framework, we demonstrate that our probability path model outperforms existing models across several forecasting tasks involving PDEs, achieving faster convergence during training and requiring fewer sampling steps during inference. These advances enhance the practicality of flow matching approaches for real-world applications, particularly in scenarios where compulational resources and time are critical constraints. Our main contributions are the following.
Theoretical Framework and Efficient Algorithms: We present theoretical framework and efficient algorithms for applying flow matching in latent space to the probabilistic forecasting of dynamical systems (see Algorithms 1-2), extending the approach of (Lipman et al., 2022) and (Davtyan et al., 2023). Our approach is specifically tailored for time series data, enabling effective modeling of complex temporal dependencies inherent dynamical systems.
Novel Probability Path Model: Using dynamical optimal transport and the Schrödinger's bridge perspoctive (sce Theorem 3.2), we motivate and propose new probability path model (see Scction 3.3) specifically designed for modeling dynamical systems data. We provide theoretical insights show that the variance of the vector field (VF) generating our proposed path can be lower than
2

---

Published in Transactions on Machine Learning Research (07/2025)
that of the optimal transport (OT) VF proposed by (Lipman et al., 2022) for sufficiently correlated spatio-temporal samples (see Theorem 3.3). We further provide intuitions to understand why our model leads to smoother training loss curve and faster convergence compared to other models.
Empirical Validation: We provide extensive empirical results demonstrate that our proposed probability path model can outperform other models on several forecasting tasks involving PDEs (see Section 5). Our results demonstrate that the proposed probability path model outperforms existing flow matching models, achieving faster convergence during training and improved predictive performance across several evaluation metrics, while requiring fewer sampling steps.
2 Flow Matching for Probabilistic Forecasting
In this section, we first introduce the objective of probabilistic forecasting, and then we discuss how flow matching can be used for learning conditional distributions latent space.
Probabilistic forecasting framework. Suppose that we are given training set of trajectories, with each trajectory of length m, S, = (z))()) E RR), coming from an underlying contimuous-tine dynamical system. For simplicity, we denote the trajectories as x x = (x',...,,x") unless there is need to specify the corresponding n. The trajectories are observed at arbitrary time points fimm- = (4,..... (m) such that x² == xx44) E Rd and (xr()) are the observed states of the ground truth system. In practice, we may have access to only few trajectories, i.e, n is small or even n 1, and the trajectories themselves may be observed at different time stamps.
where (x!=m)()))
(x")),)))))))) (with the
The goal of probabilistic forecasting is to predict the distribution of the upcoming I elements given the first k elements, where m = [ + k, ie.: propose to model each one-step predictive conditional distribution via probability density path. Instead of using score- based diffusion models to specify the path, we choose flow matching, a simpler method to train generative models. With flow matching, we directly work with probability paths, and we can simply avoid reasoning about (forward) diffusion processes altogether. Moreover, we shall work in a latent space.
"'""""",,,,,)) = I [[==o(=+++ g(25+t/=")....... ....,, x (=,"",..,,,,,t))) We
=== Flow matching in latent space. Let = = =(x") for T= I...., m, where & denotes a pre-trained encoder that maps from the data space to lower dimensional latent space. Working in the latent space, our goal is to approximate the ground truth distribution q(= .o........-.) by the parametric distribution p(==|2 which can then be decoded as x- = D(=T). The latent dynamics can be modeled by an ODF:
2 = u. (Z)).
(1)
where ue is the (velocity) VF describing the instantenous rate of change of the state at time t. Learning the dynamics of the system is oquivalent to approximating the VF u4 This can be done by regressing a neural network using the mean squared error (MSE) loss.
Following the idea of flow matching, we infer the dynamics of the system generating from the collection of latent observables by learning time-dependent VF 0 [0, 1] x Rd + Rd, t E [0, 1], such that the ODE
d.(2) = (2)), de (2) 2,
defines a time-dependent diffeomorphic map (called a flow), (Z) : [0, 1] x Rd - R4, that pushes a reference distribution Po (Z) towards the distribution P1 (Z) ~ (()) along some probability density path p. (Z) and the corresponding VF un (Z). In other words, Pt- = [øt]-Po. where [.]. denotes the push-forward operation. Here, q the ground truth distribution, p denotes a probability density path, i.e., p [0, 1] Rd R>0, and f (p(())dZ== 1. 1. We also write Z, o(Z); and thus the ODE can be written as Z, 02 (z)). Zo Z. Typically the reference distribution Po is chosen to be the standard Gaussian (Lipman et al., 2022; Liu et al., 2022).
In other words, the main goal of flow matching is to learn a deterministic coupling between p and by learning a VF i such that the solution to the ODE (2) satisfies Zo ~ p and Z, ~ g. When Z = (Z)) solves Eq. (2) for a given function 0t, we say that Z is a flow with the VF vt. If we have such a VF, then (Zo. 71) ) is a coupling of (p. q)- If we can sample from p, then we can generate approximate samples from
3

---

Published in Transactions on Machine Learning Research (07 2205))
the coupling by sampling Zu p and numerically integrating Eq. (2). This can be viewed continuous normalizing flow (Chen et al., 2018).
If one were given a pre-defined probability path (Z( and the corresponding VF u. that generates the path, then one could parametrize (2( with a neural network (Z), with f the learnable parameter, and solve the least square regression by minimizing the flow matching loss:
min fm = (2) (t)]lo) - (0---u(()) (z)),
(3)
0
where & E 24 1], Z ~ (z) and w(t) > O is weighting function. The weighting function weights the importance of the 2 loss at different times & (noise level), balancing the importance of high frequency and low frequency components. Following the standard framework (Lipman et al., 2022), we take w(t) = 1. Note that we choose to include w(t) in Eq. (3) for generality as different choice of weighting function corresponds to different paranctrization of the network output (sce (Anonymous, 2025) for more details).
However, we do not have prior knowledge for choosing Pt and ut, and there are many choices that can satisfy Pi R q. Moreover, we do not have access to closed form We that generates the desired Pt. We shall follow the approach of (Lipman et al., 2022) and construct target probability path by mixing simpler conditional probability paths. This probability path is the marginal probability path:
Pr(Z) f
f pr pz(Z)2)g(2)d2
obtained by marginalizing the conditional probability density paths p. (Z over observed latent trajectorios
Z, with po(Z)Z) p(Z) and pi (Z)2) = -V(2)2, e2 for small probability Pı which is a mixture distribution that closely approximates q. Then, assuming that pe(Z) > 0 for all Z and t, we can also define marginal VF as:
Doing so gives us a marginal
ue (e = f ue (Z)2 Pz (212))) d7,
(5)
(2)
where uz(z)z is a conditional VF (conditioned on the latent trajectory 2). It turns out that this way of mixing the conditional VFs leads to the correct. VF for generating the marginal probability path (4). We can then break down the intractable marginal VF into simpler conditional VFs which depends on single sample.
To deal with the intractable integrals in Eq. (4)-(5) which complicates computation of an unbiased estimator of fimm; shall minimize the conditional loss proposed by (Lipman et al., 2022):
min Cefm (0) = B. (((( 21
-
(6)
where t E 4/00 1], 2 ~~~00 12 q(2), Z ~ (2)2), and un (2 is the VF defined per sample 2 that generates the conaltional probability path pr (2)2). Importantly, one can show that the solution of (6) is guaranteed to converge to the same result in (3); see Theorem D.1 in App. D. Therefore, the conditional flow matching loss can match the pre-defined target probability path, constructing the flow that pushes po towards pl. Since both the probability path and VF are defined per sample, we can sample unbiased estimates of the conditional loss efficiently, particularly with suitable choices of conditional probability paths and VFs.
3
Probability Path Models for Probabilistic Forecasting
In this section, we describe the family of probability paths that we consider for flow matching, and we propose improved model tailored for probabilistic forecasting of spatio-temporal data.
3.1 Common Probability Path Models
The family of Gaussian conditional probability paths gives us tractable choices to work with since the relevant quantities in Eq. (6) and thus the conditional flow can be defined explicitly. Therefore, we will work with
4

---

Published in Transactions on Machine Learning Research (07/2025)
Table 1: Choices probability density paths that we study in this paper. For VE-diffusion, oy is increasing in i, oo = For VP-diffusion, noise scale.
Model
0t
by
VE-diff. (Song et al., 2020b; Lipman et al., 2022)
1
0
of
VP-diff. (Song et al., 2020b; Lipman et al., 2022)
e (2(--)))
0
2(1-1) ,,,+) T(t)) = Sa 3(8)ds
OT-VF (Lin et al., 2022)
t
0
(1 (1 - Fiin )t)² Emiin 2
SI (Chen al., 2024)
1 t or t2
1-t
27111 1)²
Ours
t
min 02 t), Omin, 0 2
Gaussian probability paths. Moreover, we are going to solve (6) over the dataset of all transition pairs Ppair {(=-=, (00=,,2=))0 , and use a pair of points for Z, setting 2 (Zo, Zi) E D Pyair. In particular, we consider the following class of models for the probability paths:
Pt(Z = (Zo. Z1)) = N(Zlaazo + b, b,
(7)
where at, by and & are differentiable time-dependent functions on [0, I], and denotes the identity.
Table provides five different choices of probability paths, including our proposed choice, that we study here. The optimal transport (OT) VF model was initially proposed by (Lipman et al., 2022), and setting Emin = 0 gives the rectified flow model of (Liu et al., 2022), which proposed connecting data and noise straight line. The stochastic interpolant (ST) model in Table 1 is the one considered by (Chen et al., 2024). The VE- and VP-diffusion conditional VFs (derived with Theorem D.2) coincide with the VFs governing the Probability Flow ODE for the VE and VP diffusion paths proposed in (Song et al., 2020b). It has been shown that combining diffusion conditional VFs with the flow matching objective leads to a training alternative that is more stable and robust when compared to cxisting score matching approaches (Lipman et al., 2022).
As remarked in (Lipman et al., 2022), there are many choices of VFs that generate given probability path. We shall use the simplest VF that generates flow whose map is affine linear. Let pt (z(2)) be the Gaussian probability path defined in Eq. (7) and consider the flow map ahe defined as (Z) == a Zo + b, + c, with & > 0. Then the unique VF that defines the is (see Theorem D.2 and the proof in App. D):
2 (Z (a,Zo b((1)) + dZo 6,Z1. C2
u (Z)2))
(8)
where prime denotes derivative with respect to t, and uz 2z( generates the Gaussian path pa (Z(2).
In view of this, minimizing the conditional loss becomes:
min Lafm (0) = Br,, in 1, 1011 11 (a1z" + b,2--1)) 0,2" -,,5-11
(9)
a
where ~~ (0, 1], 2 ~ pr(2 1 and g(z). We refer to this as the Flow Matching loss parametrization and shall work with this parametrization. There are other parametrizations: most popular ones are the Score Matching loss, Score Flow loss and DDPM loss. See App. C for a comparison of different loss parametrizations and App. for connections SDE based generative models.
In practice, it may be beneficial to learn the vector field of using additional context information, in which case (9) becomes:
min Lefm (0) E,,,y, 2 w(1) o (2()))
1z (2-x7 t + b,2--1)) 4
di=" -12
(10)
where C represents the context information. The choice of C is lask-dependent. For spatio-lemporal lasks, we adopt the sparse conditioning scheme of Davtyan et al. (2023), choosing C' to be past frames as follows. Given a sampled frame 2", we sample another index uniformly from {1..... we call context frame, together with the previous frame ---1, as the two conditioning frames. Thus, this case C (=7-1, T c) (see Algorithm 1). While conditioning on many past frames as possible
2) and use 29, which
5

---

Published in Transactions on Machine Learning Research (07/2025)
is desirable to improve prediction accuracy, sparse conditioning is sufficient to achieve favorable trade-off between computational efficiency and accuracy (Davtyan et al., 2023). Datasets with more challenging dynamics may require the use of more context frames.
3.2 Optimal Probability Paths
We consider the problem of selecting the optimal probability paths within the class of the Gaussian probability paths. We shall make use of the Schrödinger bridge (Léonard, 2013; Chen et al., 2021b) perspective, and seek to find the optimal stochastic processes that evolve a given measure into another, subject to marginal constraints and based on a prior belief. This optimal process describes a novel probability path model that we propose for probabilistic forecasting of spatio-temporal data.
To be more precise, let v/ be two given probability measures and let Q be the path measure of an arbitrary stochastic process. The Schrödinger bridge (SB) is the solution to the following constrained minimization problem over all path measures P (which are absolutely continuous with respect to Q) of stochastic processes on the finite time interval [0, 1]:
Po=*, Pi=u! (PII))
min
(11)
where Dx denotes the Kullback Leibler divergence and P., denotes the time marginal at time t. Typically v and v' are the (empirical) marginal distributions of an unknown continuous-tin time dynamics observed at the initial and terminal times, and Q is the path measure of prior (or reference) process that represents our belief of the dynamics before observing the data. The solution F to the problem can then be interpreted as the optimal dynamics that conforms to the prior belief Q while respecting the data marginals P7 v, Pt = =v. In other words, the SB is the (path measure of the) unite-time diffusion which admits as initial and terminal distributions the two distributions of interest and is the closest in KT. divergence to (the path measure of) a reference diffusion. Recent work (Shi et al., 2024; Pooladian & Niles- Weed, 2024) focuses on developing improved algorithms to solve SB problems for general classes of reference diffusions. SBs have also been used formulate generative models by interpolating distributions on a finite time interval (Wang et al., 2021; De Bortoli et al., 2021; Chen et al., 2021a; Peluchetti, 2023; Tong et al., 2023a; Liu al., 2023; Goltwald & Reich, 2024).
We will consider Gaussian prior processes for Q. The Gaussian probability paths that we consider can be formulated via the differential equation: dZ, = Zo + b + &t)dtt where & ~ I). Note that B Zo, Z1] = Zo + b,7n =: at- We shall take Q to be the path measure of the linear SDE with drift oz:
dY. a,dt + wdW,, t [0,1],
(12)
where w > a regularity parameter and We is the standard Wiener process. While many choices for Q exists, the rationale behind our choice is that the Gaussian process Y, admits a minimal representation Y = Zo, Z1]+ that incorporates Brownian motion as a reference noise to characterize the stochasticity surrounding the conditional estimation of Z, Importantly, such choice allows interpretation of the SBs as generalized dynamical optimal transport (DOT) (Chewi et al., 2024) between two (not necessarily Gaussian) measures. In particular, adapting Theorem in (Bunne et al., 2023) to our setting, we have:
Proposition 3.1. Consider the SB problem with Ye as the reference process:
min
DKI((P))))
(13)
Po=u,
where Y is the path measure induced by (Yi) Then, (13) is equivalent to inf E [To C(p,, vx) at with
|| wa
C(pe, vz)
||V pa (ac. V Vlog Pt),
(14)
2w² 8
where the infiruum is taken over all pairs v1) such that pa = v, p = v/, Pt absolutely continuous, and satisfies O.p.- = -V ((v, + (x))..
The DOT problem (14) can be seen as generalized version of DOT with quadratic cost (minimizing a 'kinetic energy"); see Eq. (4.41))((4.41c in (Lipman et al., 2024). The GDOT (14) not only minimizes the
6

---

Published in Transactions on Machine Learning Research (07/2025)
quadratic cost, but also minimizes the Fisher information and maximizes the expected alignment of the score function with the drift
In our case, given a pair (Zo, Z1) of data points, we are interested in constructing SBs where the marginal constraints are Gaussians centered around Zo and Zz, which we denote as Eo == Ni minn) and { :=
M(Z1, o mmnn 1) respectively, for some given Umin 2 O. These are the Gaussian SBs:
min Po=Eo P==
(PI()))
(15)
It turns out that the solution of these Gaussian SBs admits a closed-form expression (Bunne et al., 2023).
Theorem 3.2. The solution 1 to the Gaussian SB (15) is the path measure of a Markov Gaussian process with the marginal variable Xt- ~ NN I) where
[ = (1 - + ar 40- - t(ai - ao)) Zo + (1+ b, - bo - -((b - -bo))
(16)
o, = min + 02((1 1), (),
(17)
with 0² 40 min + w
20 min > 0.
See App. D.3 for the proof. In particular, imposing the boundary constraints du = = 1,, by 0 on the functions a and b,, we have u 0r Zo + Zi in Theorem 3.2. Note that of is independent of a, and b,
3.3 A Novel Probability Path Model
Our proposed probability path model is a stochastic extension of the straight-line trajectory connecting consecutive latent samples (Zo. Z1), i.e., Me (1 - 2o + tZz. This choice of a and b, can be motivated by the following optimality principle. Given a pair of data points (Zo, Z1) E Rd x Rd, the linear interpolation path Pr arises naturally as the solution to the following variational problem:
min S/' |/ subject to μo = Zo, 111 = Z1. ==(0.1)-R4
(18)
This is a classical energy minimization problem, where the goal is to find the smoothest curve connecting the endpoints Zo and Zi. It can be solved using the Euler-Lagrange equations under the given boundary conditions. This formulation corresponds to minimizing the kinetic energy of the curve among all smooth interpolations between Zo and Z1. The solution, which is constant-velocity motion along the straight line, has a natural physical interpretation: it is the least effort way to move from Zo to Zi in Euclidean space.
More generally, when the boundary conditions are probability distributions (e.g., Gaussians centered at Zo and 7,,, this linear interpolation corresponds to the displacement interpolation in the 2-aasserstein space. In the case where the marginals are Gaussian and the cost function is quadratic, the interpolating path of distributions is also Gaussian with linearly interpolated means, minimizing an upper bound on the expected kinetic energy among all stochastic paths with prescribed marginals at t= 0 and t = 1; see Section 4.7 in (Lipman et al., 2024) for more details. Alternative interpolants are possible, but the lincar form admits analytical tractability and natural interpretation in terms of velocity fields induced by diffusion processes. Importantly, this leads to trajectories that are generally easier to sample with ODE solvers, leading to ODE simulations with smaller errors (Liu et al., 2022).
For the stochastic extension, we build on the deterministic backbone and relax the boundary constraints Gaussians centered around Zo and Z1. We shall look for the "most likely" probability paths connecting the boundary distributions as follows. First, we consider the noise perturbed paths Z, Mr + wW, where W, is the standard Wiener process, as our reference process. Then, we take the solution to the Schrödinger bridge (15) with respect to the reference process as our proposed probability path model.
More precisely, in view of Theorem 3.2, we propose to choose a; = 1 - 1, b, =1, o = -mm + +02(11 - 1), in which case we have the probability path described by:
N + (1 1)7o: 2 mim + 821(1 ())
(19)
7

---

Published in Transactions on Machine Learning Research (07 (07/2025)
which transports a Gaussian distribution centered around Zu with variance owin t = 0 to a Gaussian distribution centered around Z1 at t= 1 with variance minn Here mmin; 2 0 are tunable parameters. In the case when min = 0, it describes a Brownian bridge that interpolates between Zo and Z (Gasbarra et al., 2007). To ensure numerical stability when sampling t ~ [4 1], it is beneficial to use a small min > 0. Note that o² a scale factor determining the magnitude of fluctuations around the path interpolating wetween Zo and Z1. The variance of minimum with the value of mni at the endpoints t = and 1, and the maximum variance is ommn o²/4 /4 which occurs in the middle of the path at = 1/2.
The variance schedule is designed to balance exploration and stability. Low variance at the start ensures stable initialization, preventing the trajectory from deviating too far from the initial distribution. High variance in the middle allows the model to explore diverse paths in the latent space, avoiding mode collapse and enhancing diversity in the generated trajectories. Low variance at the end sharpens the trajectory, ensuring accurate reconstruction of the desired output. This strategy is inspired by findings in diffusion models that utilize a forward noising process and backward denoising process (Ho al., 2020; Song et al., 2020b), where such variance patterns have been shown to effectively manage the trade-off between exploration and refinement.
The corresponding VF that defines the flow is then given by (applying Theorem D.2):
o² 1 21 2 o + 02 (1(1-)) -)
us (z(Z Z1 Zo +
(20)
where e : 7 (1Z) + (1-1) Under reasonable assumptions, the variance of this VF is lower than the variance of the OT-VF of (Lipman et al., 2022) with Emin := 0 (rectified flow). Here the variance is taken with respect to the randomness in the samples z" and the Gaussian samples drawn during gradient descont updates.
More precisely, denote the VF that corresponds to our model as u, and the VF that corresponds to the rectified flow model as ue For a given 7, they generate the probability path Z, = t2" + (1 f)2--1+ + Ge and Z, t= + (1 t)n respectively, where E. ~ NNO0 I), a v Ymin + o²t(1 - t) and t E 10. Applying Eq. (8), we have:
ue 4z,2112-11 ,2) - -- dec de Ae(Zol2== -- = =---1_ n7.
(21)
(22)
Theorem 3.3. Suppose that (=")r=1,...,m with the 2" Rd,, is a discrele-time stochastic process with nonzero correlation in time and let t E [0, 1] be given. If Coo(=--1, 1,2) (10 min
((xg" ((10"-11)
- -1)
I + Vor(st)) then
Var(ūr(Z))2--11) 2 Var(us(Zalz-
See App. D.4 for a proof and related discussions. Theorem 3.3 implies that if the consecutive latent variables 2", --1 are sufficiently correlated and a is chosen small enough relative to mmin, then the variance of the VF that corresponds to our probability path model is lower than that corresponds to the rectified flow model.
4 An Efficient Probabilistic Forecasting Algorithm
In this section, we present efficient algorithms for training and inferencing the flow matching model.
Recently (Davtyan et al., 2023) proposed an efficient algorithm for latent flow matching for the task of video prediction, using the probability path generated by the OT-VF of (Lipman et al., 2022). To enable efficient training, we shall follow (Davtyan et al., 2023) and leverage the iterative nature of sampling from the learned flow and use single random conditioning element from the past at each iteration. However, our method differs from (Davtyan et al., 2023) as we use different probability paths and target VFs.
Training. Recall that Ž (Zo Zi) denote training samples from Dpair {(---1,27 )T) r=....m In other words, we set Z1 to be the target element and Zo to be the reference element chosen to be the previous element before the target element. In this way, our target probability path model maps distribution centered
8

---

Published in Transactions on Machine Learning Research (07 (07/2025)
around previous state to the distribution of the current states, which is more natural from the point of view of probabilistic forecasting whose goal is to obtain an ensemble of forecasts. Note that this differs from the approach of Davtyan et al. (2023), where = Z1 (i.e., they do not use reference element to define their probability path, whereas we use a pair of elements (Zu, Z1)). Algorithm summarizes the training procedure of our method.
Both the autoencoder and the VF neural network can also be jointly trained in an end-to-end manner, but our empirical results show that separating the training lead to improved performance. Moreover, doing allows us to better assess the impact of using different probability paths.
Algorithm Flow matching for spatio-temporal data
Input: Dataset of sequences D, number of iterations M for i in range(1, M) do
Sample sequence from the dalasel D Encode it with a pre-trained encoder to obtain z
Choose a random target element x", E (3.... |z]]: from z Sample a step t ~ U[
Sample a noisy observation Z ~ Pe (Z =", ,- where Pt given by Eq. (7)
Compute ue |=(2|=",²* 1-= according to (20)
Sample condition frame =0 {l,... 2}
Updale the parameters 0 via gradient descent.
(Z ,2, -- c) u16 2==11=1
end for Return: A learned VF, vg
Inference. We use an ODE sampler during inference to generate forecasts. The ODE sampler is described as follows. Let (YT))= denote the generation process, where N - is the number of integration steps and the superscript denotes the time index for which the generation caatatioonnooreeee fforecast is intended for. Given the previous elements (x), ,,.,"" 2"1 of time series sample, in order to generate the next element (i.e., the T-th element), we start with sampling the initial condition Yς from V[=T~1,022mm where 2" -T-1_ E(x"-1). This is in contrast to the procedure of (Davtyan et al., 2023), which simply uses mean-zero Gaussian sample instcad. We then use an ODE solver to integrate the learned VF along the time interval [0, 1] to obtain Y We use D(Y as an est of x² and forecasting is done autoregressively.
for some small 0aam 0,
Algorithm summarizes this procedure when the sampling is done using the forward Euler scheme. Note that we can also use computationally more expensive numerical schemes such as the Rungo-Kutta (RK) schemes.
Algorithm 2 One-step ahead forecasting with forward Euler
Input: A sequence (x ....,. 11) containing the previous elements, number of integration steps 1, grid 80 0 81 <... < - 1, a learnt VF of for [0, 1] Set As, 8n41 8n for Sample YiT ~ -N(ExxTT1
N..., -
T- ),0aam I), Jeam 2 0
for n in range(O, N - do
Sample c ~ 11
TT E(xT-9
YTr YT + As, 00 sn
(Y Yo ,y" .y T- -T
end for
Return: An estimate of of of D(FF=1)
9

---

Published in Transactions on Machine Learning Research (07 (07/2025)
5 Empirical Results
In this section, present our main empirical results to elucidate the design choice of probability paths within the flow matching framework (comparison with other generative modeling frameworks is not our focus here). We focus on PDE dynamics forecasting tasks here (additional details and results can be found in App. E-F). We test the performance of our probability path model, i.c., Eq. (19) with d = t, b, = t and G = Vonun + +021(1- - 1) on these tasks. We pick mwin = 0.001, and treat a and Saam as tunable parameters. We compare our proposed model with four other models of probability paths:
RIVER (Davtyan et al., 2023): RIVER uses the OT-VF model in Table 1, i.e. ag 0, b t, 4 = (1 (1- Cmintt, choosing Cmin 10-?.
2t
VE-diffusion in Table 1: We use oe = min V
- with Mmin = 0.01, Faa = 0.1, and sample 10-5, following (Song et al., 2020b).
Amaus omin
uniformly from [0, 1 <| with
VP-diffusion in Table 1: We use B(s) = Amin + s(3maz - Bmin)) where Amin 0.1, Amaz = 20 and t is sampled from 4[0, 1 - e] with = 10-5, following (Song et al., 2020b). Thus, T(s) = s 8Bmin + 2 g pmin).
The stochastic interpolant (SI) path in Table We consider the path proposed by (Chen et al., 2024) and use the suggested choice of a =1 t, f2 and c = c(1 e)v- (see Eq. (2) in (Chen et al., 2024) and note that Var((1 f)W, (1 t)2t for the standard Wiener process W.). We also consider the choice with 6 = instead. We choose
0.01 for both choices. This is a path that is similar to ours, but with the variance of chosen such that the maximum occurs at t = 1/v3 instead of at the middle of the path at t = 1/2. We shall see that different forms of variance can lead to vastly different performance in the considered tasks.
Evaluation metrics. We evaluate the models using the following metrics. First, we use the standard mean squared error (MSE) and the relative Frobenius norm error (RFNE) measure the difference between predicted and true snapshots. Second, we compute metries such as the peak signal-to-noise ratio (PSNR), and the structural similarity index measure (SSIM) to further quantify the quality and similarity of the generated snapshots (Wang et al., 2004). Third, we use the Pearson correlation coefficient to assess the correlation between predicted and true snapshots.
Training details. We use an autoencoder (AE) to embed the training data into a low-dimensional latent space, which enables the model to capture the most relevant features of the data while reducing dimensionality (Azencot et al., 2020); see App. F for further discussion of the motivation. We then train a flow matching model in this latent space. Training generative models in latent space has also been considered by (Vahdat et al., 2021) for score matching models and by (Dao al., 2023) for flow matching models. To train the AE, we minimize a loss function that consists of reconstruction error, in terms of MSE, between the input data and its reconstructed version from the latent space. The choice of the AE architecture is tailored to the complexity of the dataset (see App. for details).
5.1 Probabilistic Forecasting of Dynamical Systems
We evaluate the performance of our proposed probability path model on challenging dynamical systems to demonstrate its effectiveness in forecasting complex continuous dynamics. Specifically, we consider the following lasks (for details see App. F.2):
Fluid Flow Past a Cylinder (FPC): This task involves forecasting the vorticity of a fluid flowing past a cylinder. The model conditions on the first 5 frames and predicts the subsequent 20 frames at resolution of 64 x 64 with 1 channel representing vorticity.
Shallow-Wator Equation (SWE): This datasct models the dynamics of shallow-water equa- tions (Takamoto al., 2022), capturing essential aspects of geophysical fluid flows. We use the first. 5 frames for conditioning and predict the next 15 frames at a resolution of 128 x 128 with 1 channel representing horizontal flow velocity.
10

---

Published in Transactions on Machine Learning Research (07 (07/2025)
This dataset models the dynamics of 2D diffusion-reaction equation (Takamoto et al., 2022). We use the first 5 frames for conditioning and predict 15 future frames at 128 x 128 with 2 channels representing velocity in the I and directions.
Diffusion-Reaction Equation (DRE):
Incompressible Navier-Stokes Equation (NSE): As a more challenging benchmark, consider forecasting the dynamics of a 2D incompressible Navier-Stokes equation (Takamoto et al., 2022). We use the first frames for conditioning and prodict the next 20 frames at a resolution of 512 x 512 with 2 channels representing velocity the x and y directions.
Table 2 sumumarizes the performance of our model compared to other probability path models across all tasks. It can be seen that our probability path model achieves the lowest test MSE and RFNE across all tasks, indicating more accurate forecasts. Moreover, the higher PSNR and SSIM scores indicate that our model better preserves spatial structures in the predictions. Despite the similarity of our proposed model with the stochastic interpolant of (Chen et al., 2024), in that both models use consecutive samples to define the path, our model outperforms the stochastic interpolant model, suggesting that choosing the maximum variance to occur at the middle the path is a better choice. Importantly, our model is highly efficient during inference time since it requires very few sampling sleps; this is significantly lower compared to other models.
Figure 1 shows the Pearson correlation coefficients of the predicted snapshots over time for all models. Our model's predictions shows slower decay of correlation coefficients compared to other models, indicating better temporal consistency and long-tern term predictive capability. Correlation coefficients about 95% indicate performance on par with physics-based numerical simulators.
Table 2: Results for forecasting dynamical systems using different probability path models for flow matching. Results are averaged over 5 generations obtained with sampling steps (N = 10) using RK4. For our model, we use min 0.001 and sam for all the considered tasks.
Task
Model
Test MSE (1)
Test RFNE (1)
PSNR (t)
SSIM (1)
RIVER
1.86e-03
4.48e-02
44.30
0.99
VE-diffusion VP-diffusion (h, = (2) SI (b, f)
2.29e-01
4.74e-01
27.36
0.55
Flow past Cylinder
3.58e-03
6.09e-02
42.37
0.98
3.40e-03
6.10e-02
41.81
0.98
1.19e-02
9.96e-02
39.64
0.99
Ours (ar 0.01)
3.79e-04
2.30e-02
48.88
1.00
RIVER
9.18c-04
1.49c-01
34.92
0.92
VE-diffusion VP-diffusion
1.32c-02
5.66o-01
28.10
0.55
Shallow- Water
1.39c-03
1.81c-01
34.07
0.87
(b, (2) SI (he f)
1.05c-03
1.53c-01
35.59
0.80
6.74e-01
1.29e-01
36.08
0.93
Ours (a 0.1)
6.60e-04
1.28e-01
36.10
0.93
RIVER
2.87e-03
2.28c-01
38.12
0.82
VE-diffusion
1.04c-01
1.71
32.98
0.36
Diffusion-Reactinn
VP-difuusion
2.02e-02
7.03e-01
34.98
0.52
SI (h (2)
6.17e-02
8.62e-01
45.68
0.76
SI (he ()
3.72e-04
1.18e-01
34.24
0.89
Ours (a = 0)
3.56e-04
1.16e-01
34.34
0.90
RIVER
2.84e-02
8.67e-01
30.75
0.63
VE-diffusio
1.58e-01
2.31
26.90
0.33
Navicr-Stokes
VP-diffusion
2.09<-01
2.48
27.
0.30
SI (b: (2)
1.27c-03
2.66c-01
30.73
0.90
SI (b,
1.136-03
2.54o-01
30.66
0.93
Ours (a 0.1)
1.13c-03
2.53c-01
30.66
0.93
Figure 2 compares the training loss curves of our model with others for the FPC and the SWE task. Our method lcads to faster convergence during training and smoother loss curves. This suggests that our model requires fewer iterations to generate high-quality samples when compared to other models. We find that there is significant difference in training time between the models when trained for the same number of epochs.
11

---

Published in Transactions on Machine Learning Research (07 (07/202) /025)
Simple Fluid Flow
Shallow-Ware Equation
0.998
correlation
0.998
0.96
0.994
0.94
VE
0.992
VP
0.92
Interp
Ours
5
20
Draccion-Diffusiin Equation
Ravier-Stokes Equation
correlation
0.96
0.96
0.94
0.94
0.92
10
15
10
15
20
Prediction Steps
Prediction Steps
Figure 1: Pearson correlation coefficient to assess the correlation between predicted and true snapshots at various prediction steps for different probability path models. Our probability path model shows the best. performance on all three tasks. Note that the first 5 snapshots are the conditioning snapshots.
Simple Fluid Flow
Shallow- Water Equation
10a
10
Sooch Dhrs
10
10
Loss
Loss
10"
200
400
600
800
1300
1200
1400
1900
1800
2000
100
200
300
400
500
600
700
800
900
1000
Epoch
Epoch
Figure 2: Training loss for different models of probability path for the fluid flow past a cylinder task. Our model leads to fastest convergence and smoothest loss curve among all models.
5.2 Discussion
Our empirical results showed that our proposed model consistently outperformed other models across several forecasting tasks involving different PDEs. Our model shows improved training efficiency, with faster convergence reducing the computational resources and time required for model training. Morcover, our model efficient during inference time since only requires a few sampling steps, making practical for real-world applications where computational efficiency is crucial. Additionally, the model maintained better temporal consistency, as indicated by slower decay of Pearson correlation coefficients over longer prediction horizons. These findings validate the effectiveness of our approach in modeling complex dynamical systems. Our focus in this paper is on near- to mid-term forecasting (15 20 steps), which is a standard set-up in peatio-temporal video and PDE benchmarks. While our framework compatible with longer rollouts in principle, scaling to significantly longer horizons remains a long-standing challenge, as for most forccasting tasks, due compounding errors and distributional drift, especially high-dimensional dynamics. This is important direction for future work.
12

---

Published in Transactions on Machine Learning Research (07/2025)
Our ablation study (see App. E.2) further validated the advantages of our proposed probability path model. We found that larger a values not only contributed to smoother training loss curves but also enhanced the overall stability and efficiency of the model. Moreover, we saw that Faam = 0 can be fixed without compromising accuracy. We also provide studics to understand the effect of using context frames and the choice of min on lest performance in App. E.2. In particular, in Table 5 we see noticeable improvement in test performance using min > O, showing the advantages of going beyond the commonly used Brownian bridge model (see (Tong et al., 2023b;a)). Importantly, the fact that our model achieved improved performance even with the simplest sampler (Euler scheme) and minimal number of sampling steps (as few one for the FPC task; see Figure 4) validates its practical applicability, especially in scenarios where computational resources and time are limited.
Lastly, we provide an expanded discussion to position the proposed probability path relative to the stochastic interpolant (ST) path. While the ST path offers simplicity and analytical tractability, it lacks principled guidance for choosing by. which an important design choice. Our proposed probability path is grounded in the dynamical SB framework, which, in our formulation, models the most. likely stochastic evolution between two Gaussian distributions centered around time-adjacent samples (with a minimal variance o under prior dynamics. From this perspective, the path is not simply a hcuristic interpolation and the choice of b, = is justified, whereas the inclusion of a non-zero oii. is important in the context of probabilistic forecasting. Note that when miin = 0, our proposed model simplifies to the SI model with b, = A key implication of this viewpoint is that the variance structure of the path plays a crucial role: it reflects the uncertainty in matching the two endpoints under the prior dynamics. Our proposed path introduces a variance profile that is symmetric and peaks at. the midpoint i = 1/2, better reflecting the intrinsic uncertainty in interpolating between endpoints. In contrast, the preferred SI path suggested in (Chen et al., 2024) (with b, = has asymmetric variance peak at / t= 1/v3, which may misaligu with the data geomelry. Empirically, see that our probability path model mostly outperforms the SI models across the considered tasks (see Table 2). As expected, when min is very close to zero, the test performance of our model is only marginally better than that of the SI model with by = f for most tasks. As discussed earlier (see also App. E.2), using bigger values of miin > can not only improve test performance but also help improve training stability.
6 Conclusion
In this work, we investigated the use of flow matching latent space for probabilistic forecasting of spatio- temporal dynamics, providing a thcoretical framework and efficient algorithms. We demystified the critical role of the probability path design in this setting and proposed an improved probability path model. Our model is theoretically motivated via the SB and dynamical optimal transport perspective. It leverages the inherent continuity and correlation in the spatio-temporal data, lcading to more stable training and faster convergence. Our empirical evaluations on several PDE forecasting tasks demonstrated that our model performs favorably when compared to existing models. While we focus on the flow matching approach, we leave comprehensive comparisons with more computationally demanding approaches, such as score matching (Song et al., 2020b) and bridge matching (De Bortoli et al., 2021), to future work.
Our findings position flow matching as a promising alternative to diffusion-based generative models in PDE forecasting, While diffusion models have shown strong performance, they are typically not simulation-free and require many iterative sampling steps. They rely on simulating reverse SDEs or ODEs, often using 25-100+ steps for generation. Therefore, they can be computationally costly and difficult to tune for high- dimensional spatio-temporal systems. In contrast, flow matching enables training wihout simulating stochastic processes, by directly learning a continuous vector field aligned with a designed probability path and few-step sampling, offering improved inference efficiency. Our results show that careful design of the probability path, especially those tailored to contimuous-time dynamics, can significantly improve training stability and forecast performance. These insights suggest that flow- methods not only offer theoretical elegance via connections to optimal transport and Schrödinger bridges, but also hold practical advantages for scalable and controllable forecasting in complex dynamical systems.
13

---

Published in Transactions on Machine Learning Research (07 (025)
Acknowledgments
The computations were enabled by resources provided by the National Academic Infrastructure for Super- computing in Sweden (NATSS), partially funded by the Swedish Research Council through grant agreement. no. 2022-06725 (NAISS 2024/5-269).
SHL would like to acknowledge support from the Wallenberg Initiative Networks and Quantum Information (WTNQ) and the Swedish Research Council VR22021-03648).. AY was supported by the SciAT Center, funded by the Office of Naval Rescarch under Grant Number N00014-23-1-2729, EH was supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC00243S6, NBE would like acknowledge LBL's LDRD initiative for providing partial support. XSL was supported in part by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research's 's Applied Mathematics SciML program under Contract No. DE-AC02-05CH1122) al Lawrence Berkeley National Laboratory.
References
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprinl arXiv:2303.08797 2023.
Juan Miguel Lopez Alcaraz and Nils Strodthoff. Diflusion-based time series imputation and forecasting with structured state space models, 2023.
Anonymous. Diffusion models and Gaussian flow matching: Two sides of the same coin. ICLR Blogposts 2025, 2025.
Omri Azencot, N Benjamin Erichson, Vanessa Lin, and Michael W Mahoney. Forecasting sequential data using consistent Koopman autoencoders. In International Conference on Machine Learning, pp. 475-485. PMLRR 2020.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in Neurol Information Processing Systems, 29, 2016.
Yoshua Bengio, Nicolas Boulange-Lewaaddosst, and Razvan Pascamu. Advances in optimizing recurrent networks. In 2013 IEEE International Conference Acoustics, Speech and Signal Processing, pp. 8624-8623 IEEE, 2013.
Marin Biloš, Kashif Rasul, Anderson Schneider, Yuriy Nevmyvaka, and Stephan Günnemann. Modeling temporal data as continuous functions with stochastic process diffusion, 2023.
Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominík Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffusion: Scaling latent video diffusion models to large datascts. arXiv preprint arXiv:2311.15122 2023.
George EP Box, Gwilym Jenkins, Gregory C Reinsel, and Greta M Ljung. Time Series Analysis: Forecasting and Control. John Wiley Sons, 2015.
Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. The Schrödinger bridge between Gaussian measures has closed form. In International Conference Artificial Intelligence and Statistics, pp. 5802 5833. PMLR, 2023.
Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. Dyffussio: A dynamics-informen diffusion model for spatiotemporal forecasting, 2023.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018.
14

---

Published in Transactions on Machine Learning Research (07/2025)
Tianrong Chen, Guan-Horng Liu, and Evangelos A Thcodorou. Likelihood training of Schrödinger bridge using forward-backward sdes theory. arXiv preprint arXiu:210011129) 2021a.
Yifan Chen, Mark Goldstein, Mengjian Hua, Michael S Albergo, Nicholas M Bolli, and Eric Vanden- Eijnden. Probabilistic forecasting with stochastic interpolants and Föllmer processes. arXiv preprint arXiv:2403. 19724, 2024.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge on Schrodinger bridge. SIAM Review, 63(2):249-31) 313, 2021b.
Sinho Chewi, Jonathan Niles-Woor Weed, and Philippe Rigollet. Statistical optimal transport. ar preprint arXiu:2407.18163, 2024.
Tim Colonius and Kunihiko Taira. A fast immersed boundary method using a mullspace approach and multi-domain far-field boundary conditions. Computer Methods in Applied Mechanics and Engineering, 197 (25-28):213122)) 2146, 2008.
Quan Dao, Hao Phung, Binh Nguyen, and Anh Tran. Flow matching in latent space. arXiv preprint arXiv:2307.08606 2023.
Aram Daviyan, Sepehr Sameni, and Paolo Favaro. Efficient video prediction via sparsely conditioned flow matching. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 93263-23274, 2023.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to scoro-based generative modeling. Advances in Neural Information Processing Systems, 34: 17695 17709, 2021.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-damped Langevin diffusion. arXiv preprint orXiv:2112, 12.00688, 2021.
Peter D Dueben and Peter Bauer. Challenges and design choices for global weather and climate models based on machine learning. Geoscientific Model Development, 11 (1(1)):39994400 -0099, 2018.
N Benjamin Erichson, Lionel Mathelin, J Nathan Kutz, and Steven L Brunton. Randomized dynamic mode decomposition. SIAM Journal on Applied Dynamical Systems, 18(4):1867-158) 1891, 2019.
N Benjamin Erichson, Lionel Mathelin, Zhewei Yao, Steven L Brunton, Michael W Mahoney, and J Nathan Kutz. Shallow neural networks for fluid flow reconstruction with limited sensors. Proceedings of the Royal Society A, 476(2233):20200005 2020.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference Machine Learning, 2024.
Lawrence C Evans. Partial Differential Equations, volume 19. American Mathematical Society, 2022.
Dario Gasbarra, Tommi Sottinen, and Esko Valkeila. Gaussian bridges. In Stochastic Analysis and Applications: The Abel Symposium 2005, pp. 361 382. Springer, 2007.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnacve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXio preprint arXiv:2407.15598 2024.
Georg A Gottwald and Sebastian Reich. arXiv:2409.07900 07968, 2024.
Localized Schrödinger bridge sampler. arXie preprint
Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. arXiv preprint wxXi:2312.066633 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-666 6851, 2020.
15

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advonces in Neural Information Processing Systems, 35:8633 8646, 2022.
Yang Hu, Xiao Wang, Lirong Wu, Huatian Zhang, Stan Z Li, Sheng Wang, and Tianlong Chen. FM-IS: Flow matching for time serics generation. arXiv preprint arXiv:2411007500 2024.
Tero Karras, Miik Aittala, Timo Aila, and Sammli Laine. Elucidating the design space diffusion-based generalive models. Advances in Neural Information Processing Systems, 35:26565 26577, 2022.
Diederik P Kingma and Ruiqi Gao. Understanding the diffusion objective as a weighted integral of ELBOs. arXiv preprint Xr:::290330000 00848, 2023.
Georg Kohl, Liwei Chen, and Nils Thucrcy. Benchmarking autorogressive conditional diffusion models for turbulent flow simulation. In ICML 2024 Al for Science Workshop, 2024.
Aditi S Krishnapriyan, Alejandro F Queiruga, N Benjamin Erichson, and Michael W Mahoney. Learning continuous models for continuous physics. Communications Physics, 6(1):319, 2023.
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. arXiv preprint arXi::1388. Naec:338..211 0215, 2013.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, orXiv preprint 2.Xi::2200002/43 2022.
Yaron Lipman, Marton Havasí, Peter Holderricth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:2412.0626,, 2024.
Gnan-Hormg Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A Theodorou, and Ricky TQ Chen. Generalized Schrödinger bridge matching. arXiv preprint xrXiv22310002200 2023.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arX:::2209.000 030003, 2022.
Kai Lv, Liang Yuan, and Xiaoyu Ni. Learning autoencoder diffusion models of pedestrian group relationships for multimodal trajectory prediction. IEFE Transactions Instrumentation and Measurement, 73:1-12, 2024. doi: D0.109/TTM..002.333777
Dongwei Lyu, Rie Nakata, Pu Ren, Michael W Mahoney, Arben Pitarka, Nori Nakata, and N Benjamin Erichson. WaveCastNet: An Al-enabled wavefield forecasting framework for earthquake early warning. arXiu preprint 2r5 20516, 2024.
Caspar Meijer and Lydia Y Chen. The rise of diffusion models in time-series forecasting. arXiv preprint w:Xiv:441.03006 2024.
Jaideep Pathak, Yair Cohen, Piyush Garg, Peter Harrington, Noah Brenowitz, Dale Durran, Morteza Mardani, Arash Vahdat, Shaoming Xu, Karthik Kashinath, et al. Kilometer-scale convection allowing model emulation using generative diffusion modeling, arXiv preprint orXiv:2408.10955 2024.
Stefano Peluchetti. Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023.
Aram-Alexandre Pooladian and Jonathan Niles-Weed Plug-in estimation of Schrödinger bridges. arXiv preprint arXiv:2408. 11686, 2024.
Aran-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enricn Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint 2rXio:2304.14777 14772, 2023.
16

---

Published in Transactions on Machine Learning Research (07/)025)
Hao Qu, Yongshun Gong, Meng Chen, Junbo Zhang, Yu Zheng, and Yilong Yin. Forecasting fine-grained urban flows via spatio-temporal contrastive sell-supervision. IEFE Transactions on Knowledge and Dala Engineering, 35(8):8008 8023, 2023. doi: D..1199/TKDE:2222.200000
Alejandro Queiruga, Benjamin Erichson, Liam Hodgkinson, and Michael W Mahoney. Stateful ODE-nets using basis function expansions. Advances in Neural Information Processing Systems, 14-21770-27887 2021.
Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In International Conference on Machine Learning, pp. 8857 8868. PMLR, 2021.
Pu Ren, Rie Nakata, Maxime Lacour, llan Naiman, Nori Nakata, Jialin Song, Zhengfa Bi, Osman Asif Malik, Dmitriy Morozov, Omri Azencot, et al. Learning physics for unveiling hidden earthquake ground motions via conditional generative modeling. arXiv preprint (xxx:2447.15088. 2024.
Simo Särkkä and Arno Solin. Applied Stochastic Differential Equations, volume 10. Cambridge University Press, 2019.
Martin G Schultz, Clara Betancourt, Bing Gong, Felix Kleinert, Michael Langguth, Lukas Hubert Leufen, Amirpasha Mozaffari, and Scarlet Stadtler. Can deep learning beat numerical weather prediction? Philo- sophical Transactions of the Royal Sociely A, 379((194):0020007, 2021.
Xingjian Shi, Zhourong Chen, Hao Wang, Dan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional LSTM network: A machine learning approach for precipitation nowcasting. Advances in Neural Information Processing Systems, 28, 2015.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
Gaurav Shrivastava and Abhinav Shrivastava. Video prediction by modeling videos as contimuous multi- dimensional processes. In Proceedings of the IEEE/CVF Conference on Computer Vision and Paltern Recognition, pp. 7236 7236-7245, 2024.
Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, Devi Parikh, Sonal Gupta, and Yaniv Taigman. Make-a-vides Text-to-video generation without text-vidoo data. In The Eleventh International Conference on Learning Representations, 2023.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprinl arXiv:2010.02550 2020a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modcling through stochastic differential equations. arXiv preprint ar Xiw:2011.134500 2020b.
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34:1415- 11:1415-428,, 2021.
Makolo
Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. PDEBench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 35:1596-1611, 2022.
Ella Tamir, Najwa Laabid, Markus Heinonen, Vikas Garg, and Arno Solin. Conditional flow matching for time series modelling. In ICML 2024 Workshop on Structured Probabúlistic Inference & Generative Modeling, 2024.
17

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: Conditional score- diffusion models for probabilistic time series imputation. 2021.
Jakub M Tomezak. Latent variable models. In Deep Generative Modeling, pp. 57-127. Springer, 2021.
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free Schrödinger bridges via score and flow matching. arXiv preprinl arXiv:2307.0367. 2023a.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Xivi preprint ar Xi::2302.00482, 2023b.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score- generative modeling in latent space. Advances in Neural Information Processing Systems, 34:11287-11300 2021.
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via Schrödinger bridge. In International Conference on Machine Learning, pp. 10794 10804. PMLR, 2021.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Fero P Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE Transactions on Image Processing, 13(4):600 062-6620011120 2004.
Haomin Wen, Youfang Lin, Yutong Xia, Huaiyu Wan, Qingsong Wen, Roger Zimmermann, and Yuxuan Liang. DiffST:: Probabilistic spatio-temporal graph forecasting with denoising diffusion models, 2024.
Tijin Yan, Hongwei Zhang, Tong Zhou, Yufeng Zhan, and Yuanqing Xia. ScoreGrad: Multivariate probabilistic time series forecasting with continuous energy-baste generative models, 2021.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint ar 220: 13902, 2022.
18

---

Published in Transactions on Machine Learning Research (07/ 2025)
Appendix
This appendix is organized follows. In App. A, we provide detailed discussion of related work and connect them to our studies. In App. B, we provide some remarks on the connection of flow matching models to other generative models. In App. C, we provide commonly used loss parametrizations and compare them with our flow matching loss. In App. D, we provide theorems and proofs to justify the discussions in Section 2-3 (see Theorem D.1-D.2 and Theorem 3.2-3.3), as well as additional discussions. In App. E, we provide additional empirical results. In App. F, we provide the missing experimental dctails.
A Related Work
Modeling and forecasting dynamical systems. Traditionally dynamical systems arising in scientific applications have been studied through mathematical models based on physical laws, such ODEs and PDEs (Evans, 2022). These methods often require significant domain knowledge and strong assumptions, limiting their flexibility in data-driven scenarios where governing equations are unknown. Statistical and machine learning (ML) methods have emerged as powerful alternatives to study these systems. Classical approaches include autoregressive integrated móving average (ARMA) and state-space models such as Kalman filters (Box et al., 2015), which have long been used for time series forecasting but struggle with highly nonlinear dynamics. ML. methods such as recurrent neural networks (RNNs) (Bengio et al., 2013) and variants for temporl data (Shi et al., 2015; Lyu et al., 2024) are capable of learning complex temporal dependencies but they are challenging to train. Neural ODFs (Chen et al., 2018; Queiruga et al., 2021) offer a structured approach to continuous-time modeling by integrating neural networks with ODE solvers. However, these models still face challenges with gencralizing to new systems, and lcarning continuous physical dynamics (Krishnapriyan et al., 2023).
Generative modeling for time series forecasting. Generative models, particularly score-based diffussion models and flow-based approaches, have gained significant attention for learning complex data distributions in tigh-dimensional spaces such as images (Song et al., 2020b; Karras et al., 2022; Esser et al., 2024), videos (Davtyan et al., 2023; Shrivastava & Shrivastava, 2024), and dynamical systems (Pathak et al., 2024; Kohl et al., 2024; Ren et al., 2024). In the context of time series forecasting, diffusion models like TimeGrad (Rasul et al., 2021; Meijer & Chen, 2024) incorporate conditioning on previous time steps into both the forecasting process and the loss function. Building upon TimeGrad, CSDI (Tashiro et al., 2021) enhances performance in imputation tasks by replacing the RNN encoder with a transformer. ScoreGrad (Yan et al., 2021) adapts this framework to a score-based diffusion model for multivariate probabilistic forecasting. Methods such as DSPD and CSPD (Bilos et al., 2023), based on DDPM and SDE respectively, model dynamics as continuous functions and introduce time-corrclated noise functions. Another relevant work, SSSD (Alcaraz & Strodthoff, 2023), uses state-space models (S4) to encode time series and performs diffusion over the temporal domain instead of across multivariate components. While these models demonstrate strong performance, they often require many sampling steps during generation. Related to our approach, (Chen et al., 2024) proposes an SDE framework utilizing stochastic interpolants (Albergo et al., 2023) for probabilistic forecasting, but their sampler also requires many steps during generation. Another recent work is (Tamir et al., 2024), which introduces a flow matching framework for time scries modeling within the data space; however, it concentrates on small ODE datasets and does not address forecasting tasks. Meanwhile, (Hu al., 2024) proposes a model based on rectified flow with ODE-based straight line transport for efficient time series generation.
Applications of diffusion models to spatio-temporal data have primarily focused on video generation (Singer et al., 2023). For spatio-temporal forecasting on sensor grids, graph neural networks (GNNs) are widely employed in tasks such as traffic prediction (Qu et al., 2023; Ly et al., 2024; Wen et al., 2024) and air quality forccasting (Wen et al., 2024). Another recent work is DYffusson (Cachay et al., 2023), which utilizes a non-Markovian diffusion process to propagate images temporally, similar to DDIM. This method iteratively refines initial predictions at larger time steps by conditioning them on predictions at smaller time steps.
19

---

Published in Transactions on Machine Learning Research (07/2025)
B Connection to SDE Based Generative Model Frameworks
In this section, we provide some remarks on the connection between flow matching models and SDE based generative models (Song et al., 2020b).
Consider the following continuous- time Gaussian latent variable model (Tomczak, 2021): Z, E(X,),
2= = A,7o + B,Z1 + Lec, t E [), X, = D(z))
(23)
(24)
where is the contimuous variable, Xo E I2 represent data samples, E IRP the latent variable, e E~ NO, /) is independent of the random variables Xo. Zo. Z1. Here A., B, and L, 2 0 are pre-specified coefficients which are possibly matrix-valued and time-dependent, 'D and & denote the decoder and encoder map respectively, and D 3 I. Note that Zu and Z1 are initial and terminal point of the path (Z)))c(o.1) in the latent space.
The above latent variable model can be identified (up to equivalence in law for each t) with the linear SDE of the form:
d² F,Zrdt + H, + G,dW,, 2o = Zo: t tE (0,
(25)
where (Wi)ee(o.1) is the standard Wiener process. By matching the moments, obtain
F. = A,A He = B, - A B₁,
(26)
(27)
G.G = L + L.I. - A,A, LLT -
(28)
where the overdot denotes derivative with respect to t and AT denotes the transpose of A.
Under the above formulation, various existing generalive models such as DDPM (Ho et al., 2020), VP-SDE and VE-SDE of (Song et al., 2020b;a), the critically damped SDE of (Dockhorn et al., 2021), the flow matching models (Lipman et al., 2022; Tong et al., 2023b; Liu et al., 2022) and the stochastic interpolants of (Albergo et al., 2023) can be recovered, and new models can be derived.
The following proposition establishes the connection between flow matching using proposed probability path model, the Gaussian latent variable model (23) and the lincar SDE model (25).
Proposition B.1. For every t E 1], the Ze defined Eq. (19) can be identified, up to equivalence in law, with the Ze generated by the latent variable model (23) with A = (1 (1-t)I, B, tl, L = V + o²t 1) For E [0, 1), it can also be identified with the solution 2, of the linear SDE (25) with F, = 11/(1 - t),
H, = (1 + I and G. - -2= 20 I. Moreover, lime-11 2e _d Zi + mintt, where ~ ~N0 1) and =d I
denotes equivalence in distribution.
Proof. The identification follows from matching the moments of Z, and 2,, i.., applying Eq. (26)-
To prove that lim 2, -d Z1 + Ominc, we use the explicit solution of the SDE:
2, = o(t + S( s(t,s) [' (ttts)
where ltt, s) is the fundamental solution of the homogeneous equation d s) = F, s)dt with ((s = I. For our F, I/(1 t), we have (t s) exp(- 1 1 du) (1=)). Substituting this and the formula
for H, into the solution, obtain 2, (1 Zo + 1or (1 + + Z + S6 ( 0² 20 IdW.. 2.
Now, let us examine each term as t + 1. First, (1 Z) 0 as t - and 56 (111-) 1- 1,ds 1Z + Z
+ 1. Tt remains to deal with the stochastic integral term M, = f (1=1) 0² T-s dW., Note that
20
20

---

Published in Transactions on Machine Learning Research (07 (07/2025)
M, is an Itō integral that has zero mean, i.. E.M 0, and using Itô's formula (Särkkä & Solin, 2019),
1 ((1 - -1) s))² 0²
20
EM,2
- ds
(29)
/"
1
ds + 20 min f"
1
(1 1)2
a² 0 (1
(1 8)3 ds
(30)
8)2
1
(1 t)²
0²
o
1133 -1
(31)
(1 1)2
(1 t)t + o t(2 t),
(32)
which tends to Minn t + 1. Combining the above results, lim 2, ~ ~N o2 ominn I)-
C On Different Loss Parametrizations
In this section, we list popular choices loss parametrization considered in the literature and connect them our flow matching loss. We refer to (Kingma & Gao, 2023) for more comprehensive discussion. Recall that the Gaussian path that we consider is: Z, a.Zo + 6,21 + ofE, where & NV(,,)) In general, these loss parametrizations take the form of:
[( = R, (212)
m2 (m) 2)12;
(33)
where w(t) > 0 is a weighting function, m, m, 2) is the object (conditioned on 2) to be learnt and m is a neural network model used to learn the object of interest. Depending on which object one would like to learn/match, we have different loss parametrizations.
Flow matching. The flow matching loss that focus in this paper is:
Cefm (0) == E, (z) (z)2).((2) (())lo)( - un (z(222)11
(34)
where one aims to learn the flow generating vector field:
( (z- (a, + br?1)) + 7,7o 8,7.1 Ct
u (7l2)
(35)
Score matching. The score matching loss is:
Lem E ...2(22)) (Z)) - V log p. (2
(36)
where A(t) > 0 is weighting function and one aims to learn the score function:
V log (z(2)) 0,Z0 + b. - -Z
(37)
of
If X(t) of then this roduces the original score matching loss (Song & Ermon, 2019), whereas if X(I) 3(1 - (), this becomes the score flow loss (Song el al., 2021).
Noise matching. The noise matching loss is:
Lum (0) = E,
-
(38)
where one aims to learn the noise:
((()2)= 7 - (a+70 + b,7,)
(39)
Ct
21

---

Published in Transactions on Machine Learning Research (07 (07/2025)
D Theoretical Results and Proofs
In this section, provide theorems and proofs to justify the discussions in Section 2 and Section 3.
D.1 Connecting Flow Matching with Conditional Flow Matching
The following theorem justifies the claim that minimizing Lim is equivalent to minimizing Lefm-
Theorem D.1. If the conditional vector field (Z) generates the conditional probability path pe pe(Z))), then the marginal veclor field u, in Eq. (5) generales the marginal probabilily palh Ps in Eq. (4). Moreover, iJ (Z) > for all i, Z, then fm and feem are equal up to constant independent of 0.
Proof. The proof is a straightforward extension of the proofs of Theorem 1-2 in (Lipman et al., 2022) from conditioning on data samples to conditioning on latent samples and allowing an arbitrary weighting function w()).
Suppose that the conditional vector field (z) generates the conditional probability path pe pr (z(Z)), we would like show that the marginal vector field u, in Eq. (5) generates the marginal probability path pt in Eq. (4). To show this, it suffices to verify that Pt and ue satisfy the continuity equation:
d 2i(2) diolpi =0,
(40)
where the divergence operator, div, is defined with respect to the latent variable Z (Z'...,,Z%) i.., div Σd 87
We begin by taking derivative of (2) with respect to time:
d p2(2)2)((2ddz di p((2)= dt
(41)
/ pilz (((2)12)(2))
(42)
/ (bnuuu(222px(222)))((()()
(43)
-dvv (ʃ ( (fm(222)px(22)2(()) -diicuu(Z))rt2
--div
(44)
(45)
In the third line, we use the fact that 22(-12) generates 22). In the last line, we use Eq. (5). In the second and forth line above, the exchange of integration and differentiation can be justified by assuming that the integrands satisfy the regularity conditions the Leibniz rule
Next, we would like to show that if (Z) > for all t, Z, then Lym and Lefm are equal up a constant independent of 0. We follow (Lipman et al., 2022) and assume that g(Z) and pt (z(2)) are decreasing to zero sufficiently fast as || + O, that ue, vt, Vav, are bounded, that all integrals exist and exchange of integration order is justified via Fubini's theorem.
Using the bilinearity of the 2-norm, we decompose the squared losses as:
2(00( u (Z)) [of(2)) - 22 - [[])(2) (2(2) 12)1 -
+
(2)11, + 1lu, (212)11
(46)
(47)
22

---

Published in Transactions on Machine Learning Research (07 (025)
Now,
E Ep((2))lo((2)11 / 10, (1)
(48)
/ (Z (2)
(49)
= (2) (2)11 (1))
(50)
where we use Eq. (4) in the second equality above and exchange the order of integration in the third equality. Next, we compute:
= / (18(2). ( (zl2)2)m
E
Pr Pz(2)dZ
(51)
Pz(Z)
/ (8(22) / (u1211))x/2122)11244 (( )22
(52)
/ (v((2), (.11))))2212)))
(53)
E (0 (0 (Z), us (
(54)
where we first plug in Eq. (5) and then exchange the order the integration in order to arrive at the last equality.
Finally, noting that un are (t) independent (and are thus irrelevant for computing the loss gradients). have proved the desired result.
D.2 Identifying the Vector Field that Generates the Gaussian Paths
Similar to Theorem (Lipman et al., 2022), we have the following result, which identifies the unique vector field that generates the Gaussian probability path.
Theorem D.2. Let (Z) be the Gaussian probobility path defined in Eq. (7) and consider the flow map v defined (Z) a,Zo + b + aZ with ez > O. Then the unique vector field that defines the is:
u, (Z)z) = 1 [ z - (a,Zo b,21)) + dZz + 6,21,
(55)
where prime denotes derivative with respect to t, and ue (Z)²) generates the Gaussian path Pr (Z)
Proof. Let (z(z) be the Gaussian probability path defined in Eq. (7) and consider the flow map th defined as (2) = a, Zo + b. + a. We would like to show that the unique vector field that defines the is:
2 (Z (a,Zo b(Zz)) + diz b,Zz.
u (Z)2))
(56)
C₁
and (Z(zz) generates the Gaussian path pr (7)
We denote w, (Z(2) for notational simplicity. Then,
d
d vr(Z) = wr (vh (Z)).
(57)
Since 1i is invertible (as c > 0), we let - = 2 (Y) and obtain
- = wY(Y),
(58)
where the prime denotes derivative with respect to and we have used the apostrosphe notation for the derivative to indicate that oi is evaluated at (Y).
23

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
Inverting (Z) gives:
(x) Y - (m(
(59)
ct
where (z(2) Zo + b,Z1.
Differentiating 2 with respect to gives 22(2) = e, + (2).
Plugging the last two equations into Eq. (58), we obtain:
1 ( 4 "0rr (x()) + 12
(60)
which is the result that wanted to show.
D.3 Solution the Gaussian Schrodinger Bridge Problem
We prove Theorem 3.2 in this subsection. First, we recall Theorem from (Bunne al., 2023). then follows from an application of the theorem.
Theorem 3.2
Let to ~ (po, Σo) and & ~ ((11, Σ1) be two given Gaussian random variables. From now on, by abusing the notation we will also denote the measures of these Gaussians by Ea and & respectively.
Let 2 and set:
D, = (4E) 1 Ey 12/2 + o + 1/2
(61)
C -= (x) Dr 2o 1/2 o²1).
(62)
Consider the following Schrodinger bridges with Gaussian marginal constraints:
min Po=Eo
DKL [P]))),
(63)
where Q is the path measure of the linear SDE:
dY, (c,Y, + at)dt + g,dW, == fidt + g:dW,
(64)
Here, c: ::: + R, ae R Rd and gr R R are smooth functions, and W, is the standard Wiener process.
The following theorem is simplified version of Theorem 3 in (Bunne et al., 2023).
Thcorem D.3 (Theorem 3 in (Bunne et al., 2023)). The solution P to the Gaussian Schrodinger bridges is (the path measure of) Markow Gaussian process whose marginal variable X, ~ NW,, Ex), where
# = (Te (eT))oo + ru/1 + ((() - r.C(1), Io + E + +re( - rr) + C..) m(t, (]( P.)
(65)
Σ, (Te ret))
(66)
In the above,
x((,1') TTT (-" T 2
7o = exp
exp
ri= = n(1, 1)/x(1,1),
a² = voiin(1111 ((I) = Ti (' Taaad
fo "g
Pt. So To 9 ds
Now we prove Theorem 3.2.
24

---

Published in Transactions on Machine Learning Research (07 2025))
Proof of Theorem 3.2. We shall apply Theorem D.3 to the Gaussian SB (15) by setting Mo == Zo. 11 Z1,
Do E1 = 0 min 1, Ot == â, 70 + b, 7,, = 0, and It = w. Then, we have = 1, k(t,t') = w Fr = 4
o- = -w, ((t) = So a,ds = ( d Zo + ( bs Z1 = (at - oo) + (b. b1) and P = t. Also,
Do. (40 min + w and Co. C 1 (440 min + w W
Thus,
Nt = (1-1)20 + 121 + (a+ - an) 20 + (b. 20 - u((a1 an) Zo + (b) bp) 21)
(67)
and
2 = [[((1 -t² + +10+ Poomin +t(1-10 (cl--)(ee ()((0 + w wP) + w w2t(1 t)]
(68)
-VAotte 10 + wt 20 min t((
= o min +
(69)
Theorem 3.2 then follows from Theorem D.3 with the above formula for be and E..
0
D.4 Comparing the Variance of Different Vector Field Models
We begin with providing intuitions for understanding our proposed probability path model. We expect that our model improves upon the other considered models, as it takes advantage of the inherent continuity and correlation in the spatio-temporal data. Intuitively, for time series samples whose underlying dynamics are continuous and obey a physical law, a Gaussian sample is lypically further from the time series samples, so the distance between a Gaussian sample and a time series sample should generally be larger than the distance between consecutive time series samples (which can be highly correlated). Therefore, using probability path that connects consecutive time series samples could lead to faster convergence and more stable training, when compared to using path that simply connects time series sample to a Gaussian sample, since the resulting flow model in the former case can better capture the true system dynamics with less effort.
Moreover, if the consecutive samples are sufficiently correlated, then the variance of the VF corresponding to our proposed probability path model can be lower than the variance of the VF corresponding to the other choices of probability paths. Precise statements capturing this are the contents of Theorem 3.3, which focuses on comparison between our proposed model and the optimal transport VF model of (Lipman et al., 2022).
Before proving Theorem 3.3, we start with the following lemma.
Lemma D.4. Let A, B, C, be random vectors where C' and D are independent, both A and B are independent of C and D (but A and B could be dependent). If Cov(A, B) 2 (Var(C) - Var(D) + V Var(B))/2 then
Var(A 2 Var(A B + -B)
(70)
where A 2 B means that A - B is positive semidefinite.
Proof. We compute:
Var(A + D) -Vrr(A-B+C+B++
(71)
= (OM D + - DD
-2Ce(AABB+CC)- 2Cou D.C)
(72)
(MM Var(C) 2Cor( CCo CBB))
+2Co((ABBCC.)
(73)
= -Var(A- Can(A-B+C)) +) B + D) - Var(C) + 2Coo A, B) - 2V 2Var(B) B)
(74)
= Var(A- - B + C) - Var(C) + Var(D) + 2Cov( A, B) - Var(B),
(75)
where we have simply rearranged the terms in the first equality, used the formula Var (A + B + C) = (ar(A) + Var(B)+ Var(C) + 2Cou(A, B) + ()) + 2Co C), bilincarity of covariance, the facts that Cou(A, A) Var(A) and (ar(cA) = e² for scalar c, as well as the assumption that both A and B are independent of C, D in the last four equalities.
25

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Therefore, if -Var(C)) (ar())))Var))) Var(B) 2 0, then have Var(A+ D) Var(A- - C).
Now we prove Theorem 3.3.
Proof. Note that of = a2 21)
2vo2.... +02x(1-1) Var(c)) (c))² and Var(-n) = Therefore, using these and
applying Lemma D.4 with A == ====-11, B = 2", C c/ and D == -7, allow us to establish the claim that
(ar("((Z)l==") 2 Var(u (2(1==11,=7) if
- 4(0 min + t - t))
Cou(2" -- ;-,=,"²)) 2 2
+ Var(=")-1)
(76)
Since the function f(t) := K(o +()(())) maximized at the endpoints 0, with the maximum value of o 44 mins the desired result stated in the theorem follows.
o((1-())
Lastly, we provide some discussions following Theorem 3.3 (using the notations there).
Discussions. Let us denote (Z) (2T, 2- .)) to show the explicit dependence of the vector field neural net on the random samples t, x", =-1 and & ~~ (O, I) drawn during each update of gradient descent during training. During each gradient descent update, our model involves computation of
TaLofm(0:t,,,22)
o2 2- 12)² ( 1,-)) u, ( 21-11)
(77)
=: n ;--= ,g) (=
(78)
with ~ - and the latent samples 27, -- drawn randomly.
Similarly, for the rectified flow model let denote if (Z) = of (=--1, n7) to show the explicit dependence of the vector field neural net on the random samples =T- and r7 (O, drawn during each update of
gradient descent during training. Each update of gradient descent using the rectified flow model involves
computation of
(et,em(0:t,,.222 2V
2V .))= (6f(2" 1) =: 2Vo ( ,n)"" (1 ,)).
un (Zcl2=
(79)
(80)
with 1], 7- I(0, and the latent sample -T- drawn randomly.
Lower gradient variance results in smoother training loss curve and potentially faster convergence, so it is useful to compare the variances of the loss gradient for the two models. However, the variances are highly dependent on Voug, Voof and their covariance with the other random vectors appearing in Eq. (77) and Eq. (79), making such comparison challenging without strong assumptions. Heuristically, the difference in the variances of the loss gradient during each update for the two models is primarily determined by the difference between Var(2 (27-1, 20,0,,,,_00000 .)) and Var(A) n)) Vorg and VoËg are relatively stable. In this case, we have Var(A)(z" ,2 S (2
if we suppose the assumptions in Theorem 3.3,
(ar( 2 Var(o?) and Cou(oe, ux) 2 Cov(Tf, T.).
The implications of Theorem 3.3 together with the heuristics above could partially explain why our probability path model leads to smoother loss curve and faster convergence (see Figure 2) compared to the RIVER method of (Davtyan et al., 2023). On the other hand, the dependence of the lower bound in the theorem on a and Jmin suggests that using values of that relatively large enough might not keep the variance of the vector field low, which could partially explain the phenomenon displayed in Figure 2, where using o = 0.1 and miin 0.001 leads to large loss fluctuations.
E Additional Empirical Results
In this section, we provide additional experimental results.
26

---

Published in Transactions on Machine Learning Research (07 (025)
E.1
Visualization of Flow Patterns and Dynamics
Figure 3 provides visual results of the predicted snapshots by our model for each task. The visualizations highlight our model's ability to capture complex flow patterns and dynamics.
last con. frame
time
GT
pred.
GT
pred.
GT
pred.
pred.
Figure Visualization of predicted frames using our model of probability path for the considered (asks. From top to bottom: fluid flow past cylinder, shallow-wate water equation, diffusion-reaction equation, and incompressible Navier-Stokes equation. In each case, GT indicates the ground truth frames and pred. indicates the predicted frames.
E.2 Ablation Studies
To further assess model, conducted detailed ablation study focusing the impact of various hyperparameters for the fluid flow (FPC) task. Specifically, study the impact of the values of a, the choice of sampler, and the number of sampling steps during inference. For small cam, we find that this parameter has negligible impact on test performance, so we fixed deam = for all experiments in this section.
Impact of a on training stability. Figure illustrates the effect of different a values on the training loss curve for our method on the fluid flow past cylinder task. We observed that larger values of (c.g., = 0.1) resulted in smoother loss curves and more stable convergence during training.
27

---

Published in Transactions on Machine Learning Research (07 (075))
Effect of o, sampler choice, and sampling steps test performance. Table 6 investigates how different values of a, the choice of sampler (Fuler or RK4), and the number of sampling steps affect test. performance. It can be seen, that even with as few as one sampling step (N = 2) using the Euler scheme, our model perform reasonably well. However, increasing the number of sampling stops or employing the more computationally intensive RK4 sampler can help to lead to better results. From Figure 4, we see that our model leads the smallest test MSE for both samplers at all sampling steps. In particular, lowest MSEs can be achieved using as few as one sampling step for both sampling schemes. Moreover, while using the RK4 sampler lead to lower test MSE for all models, the performance gap is much smaller for our model, showing robustness of our model to sampling choice.
Culer
Test MSE (log scale)
10
(log scale)
N (log scale)
Figure Test MSE vs. number of sampling steps N - 1 (we consider N 2,5, 10, 20, 50, 100) for the five models with the Euler sampler (left) and the RK4 sampler (right)..
Effect of context frames on test performance. Table 3 shows the test performance in terms of the considered evaluation metries in the case when the neural networks are trained without using the context. frames, in which case C' mull in (10) and we are minimizing the loss (9). We see that without using the proposed context frames, the test performance of models degrades across the evaluation metrics, showing the benefits of using the additional context information. The degradation particularly significant for RIVER, VE-diffusion and VP-diffusion.
Table 3: Ablation study to assess the impact of context frames on test performance of the considered models, given that the same pre-trained autoencoder is used. Results are averaged over 5 generations.
Model
Test MSE (L)
Test RFNE (4)
PSNR (t)
SSIM (t)
RIVER
1.21
1.13
20.84
0.25
VF-diffusion
1.01
1.07
23.21
0.23
VP-diffusion
1.04
1.06
21.18
0.24
Stochastic interpolant
5.37e-03
7.04e-02
41.36
0.97
Ours (a = 0.01, .am 0)
8.42e-04
8.18c-02
46.75
0.99
We further assess the impact of the random conditioning frames 2", considering the case when the neural networks are trained with the reference frames but without using random conditioning frames, in which case C 2-1 in (10). Table shows the results. We see that in this case RIVER and VE-diffusion achieve lower Lest. performance across all evaluation metrics, which shows that these models indeed benefit from the use of conditional frames. Interestingly, our model, the stochastic interpolant model and the VP-diffusion model remain relatively robust, with comparable or even improved test performance. Despite this, our model outperforms the other models across all evaluation metrics regardless whether the conditioning frames are involved. For consistency and fair comparison with RIVER, retain the conditioning scheme by default.
28

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Table 4: Ablation study to assess the impact of the random conditioning frames on test performance of the considered models, given that the same pre-trained autoencoder is used. Results are averaged over generations.
Model
Test MSE (4)
Test RFNE (1)
PSNR (f)
SSIM (1)
RIVER
4.68e-03
6.19e-02
43.26
0.97
VE-diffusion
3.96e-01
6.04e-01
26.40
0.53
VP-diffusion
1.25e-03
3.67e-02
45.73
0.99
Stochastic interpolant
2.33e-03
4.90e-02
43.99
0.98
Ours (a 0.01, saa 0)
3.21e-00
2.16e-02
49.23
1.00
Euler Discretization
RK4 Discretization
10
0.01
0.01
10a
Loss
10
L.oss
10'
10
10
10
10-2
10
10
2000
1000
Epoch
Epoch
Figure 5: Training loss for different values of using probability path model. The left subplot shows results for the Euler sampler, and the right subplot for the RK4 sampler. We see that the loss curve is sensitive to the choice of a, with larger values of a giving smoother loss curves.
Sensitivity to Omin- Table shows that using min > 0 (while fixing the other hyperparameters) leads to noticeable improvement in lest performance of our model across the evaluation metrics. This is in line with the observation that using min > O improves training stability, since this alleviates the singularity of the target VF (20) at = 0, 1. In fact, we see that using Jmin 0.001 (the default value that we use) leads to better test results for the FPC task.
Table 5: Sensitivity analysis of our model to Omin- Results are averaged over 5 generations.
Umin
Test MSE (1)
Test RENE (L)
PSNR (t)
SSIM (t)
0.0
1.98e-08
4.41e-02
44.74
0.99
0.001
3.79e-04
2.30e-02
48.88
1.00
0.01
3.78e-04
2.26e-02
49.11
1.00
0.1
3.17e-04
2.12e-02
49.52
1.00
1.0
3.46o-04
2.17e-02
49.44
1.00
F Experimental Details
In this section, we provide the experimental details for the tasks considered in Section
F.1 On the Choice of Gaussian Reference Processes
Our choice of a Gaussian reference process not directly derived from the characteristics of a particular
spatio-temporal dataset, but rather follows a modeling convention that is widely adopted, particularly when physical or statistical smoothness assumptions are in play. Gaussian processes are standard modeling choice in these domains for several reasons:
29

---

Published in Transactions on Machine Learning Research (07 (02))
Table Ablation study for the fluid flow past cylinder task. Results are averaged over generations.
sampler
N
Test MSE (1)
Tost RFNE (L)
PSNR (1)
SSIM (1)
0.0
Euler
5
7.02e-04
2.96e-02
47.15
0.99
0.01
Euler
5
3.63e-04
4.14-03
48.89
1.00
Euler
5
2.97e-03
4.11e-02
43.99
0.98
0.0
Euler
10
7.31c-04
3.01c-02
47.02
0.99
0.01
Euler
10
3.76e-04
2.31c-02
48.76
1.00
0.1
Fuler
10
2.85e-03
5.04e-02
44.12
0.98
0.0
RK4
10
3.89c-04
2,26e-02
49.19
1.00
0.01
RK4
10
3.79e-04
2.30e-02
0.1
RK4
10
6.56e-03
7.83e-02
40.68
0.97
Euler
20
7.40c-04
3.03e-02
46.98
0.99
0.01
Euler
20
3.82c-04
2.33c-02
48.71
1.00
0.1
Buler
20
2.78e-03
4.98e-02
44.19
0.98
0.0
RK4
20
3.88c-04
2.26e-02
49.19
1.00
0.01
RK4
20
6.53e-04
2.80c-02
47.63
0.99
0.1
RK4
20
6.51e-04
7.80c-02
40.70
0.97
They arise solutions to linear stochastic partial differential equations (e.g., the heat equation with additive white noise), which are common in spatio-temporal physical systems.
They offer desirable analytical properties such as closed-form marginals, smooth sample paths, and tractable likelihoods.
In many practical settings (e.g., climate modeling, diffusion, spatial statistics), Gaussian priors are used because they act as regularizing priors on functions or fields that evolve over space and time.
F.2 Details on the Datasets
Fluid flow past cylinder (FPC). We the fluid flow past stationary cylinder at Reynolds number of 100 simple test problem. This fluid flow is canonical problem in fluid dynamics characterized by periodically shedding wake structure (Erichson et al., 2020; 2019). The flow dynamics are governed by the two-dimensional incompressible Bavier-Stoker equations:
du + (u v.-V)u- -
I -vp+vvu
V u-u=
where = (u,v) is the velocity field, p is the pressure, p the fluid densily, and is the kinematic viscosity. The vorticity field w is obtained from the velocity field via:
V u,
providing insights into the rotational characteristics of the flow.
For simulating the data, the Immersed Boundary Projection Method (IBPM) has been used (Colonius Taira, 2008). The flow tensor has dimensions 199 x 449 x 151, representing 151 temporal snapshots on a 449 199 spatial grid. We crop and spatially subsample the data which results in a x 64 spatial field.
Shallow-water equation (SWE). The shallow-water equations, derived from the compressible Navier- Stokes equations, be used for modeling free-surface flow problems. We consider the 2D equation used in (Takamoto et al., 2022), which is the following system of hyperbolic PDEs:
1 112h + - = -g.hVb,
0 + Vhu = 0, 0,hu V
(81)
30

---

Published in Transactions on Machine Learning Research (07/2025)
where u u, being the velocities in the horizontal and vertical direction respectively, h describes the water depth; and b describes a spatially varying bathymetry. hu can be interpreted as the directional momentum components and describes the gravitational acceleration. The mass and momentum conservation properties can hold across shocks in the solution and thus challenging datasets can be generated. This equation finds application in modeling tsunamis and flooding events.
We use the dataset generated and provided by PDEBBech (Takamoto et al., 2022). The data file (2D_ wwtteeeeee/eeeeeeeeeeeeeeeeeeeee
rdb_ MAANAA.NA h5) can be downloaded from https://gichub. dataadownlee The data sample is series of 101 frames at 128 128 pixel resolution and come with channel. The simulation considered in (Takamoto et al., 2022) is a 2D radial dam break scenario. On a square domain Ω = [-2.5, 2.5]², ; the water height is initialized as a circular bump in the center of the domain:
2. for r < +
01 = 0,
1, for r 2 V +2
with the radius randomly sampled from ((0.3,0.7). The dataset is simulated with finite volume solver using the PyClaw package. We apply standardization and then normalization to the range of [-1, 1] to preprocess the simulated data.
Incompressible Navier-Stokes equation (NSE). The Navier-Stokes equation is the incompressible version of the compressible Navier-Stokes equation, and it can be used to model hydromechanical systems, turbulent dynamics and weather. We use the inhomogeenous version of the equation (which includes a vector field forcing term u) considered by (Takamooo al., 2022):
v-v=0, plO Vv) -- nA u,
(82)
where p is the mass density, v is the fluid velocity, is the gas pressure and is shear viscosity. The initial conditions and inhomogeneous forcing parameters u are each drawn from isotropic Gaussian random fields with truncated power-law decay 7 of the power spectral density and scale a, where Trp = -3, do = 0.15, Tu = -1, o2 = 0.4. The domain taken to be the unit square Ω = [o, 1J2 and the viscosity 1j = 0.01. The equation is numerically simulated using Phiflow. Boundary conditions are taken to be Dirichlet to clamp the field velocity to zero at the perimeter.
We use the datasct generated and provided by PDEBench (Takamoto et al., 2022). The data file ins_incom, inhom_ 2d_ raa_512-0.30 h5) can be downloaded from https: //gp::///ttte tigghhu...... com Batpeeeccch/DDBecc//eeeeeee hain/pddbench/datte httpa/daaaaaadlaeeeeeeaaaaaaa download. The data sample is a series of 1000 frames at 512 x 512 pixel resolution and come with channels. We do not apply any data preprocessing procedure here.
Diffusion-reaction equation (DRE). We use the 2D extension of diffusion-reaction equation of (Takamoto et al., 2022) which describes two non-linearly coupled variables, namely the activator u u(t,x,y) and the inhibitor u= vtt, x, y). The equation is given by:
0,u D + Duōgn + Rus
(83)
020 D,O Dr + Rv,
(84)
where Dn and D, are the diffusion coefficient for the activator and inhibitor respectively. R. = Ru, v) and Rs = Ro (u, v) are the activator and inhibitor reaction function respectively. The domain of the simulation includes E (-1,1), y E (-1,1), (0, 5). This equation can be used for modeling biological pattern formation.
The reaction functions for the activator and inhibitor are defined by the Fitzhugh-Nagumo equation Ra(u, v) = u - 3 - k- - v R2 (u, v) = u where h = 5 x 10-3, and the diffusion coefficients for the activator and inhibitor are D, = 1 x 10-3 and D. = 5 x 10-3 respectively. The initial condition is generated as standard Gaussian noise w(0, g) ~ ~N(00, 1.0) for x E (-1,1) and y E 1). We take a no-flow Neumann boundary condition: D,O,.u = 0, D.O 0, D.O,u u = 0, and D, = 0 for x, E (-1,1)?.
We use a downsampled version of the dataset provided by PDEBench (Takamoto al., 2022). The data file (2D diff-react NA NA. h5) can be downloaded from https:// //gi thub. com httm://ddbeenc////eeccc////
31

---

Published in Transactions on Machine Learning Research (07/2025)
mii/pp/ebbenc///ttt_ddoomni daw.ddaaaaddolddd.oooooo The data sample is series of 101 frames at a 128 x 128 pixel resolution and come with 2 channels. The sample frames are generated using the finite volume method for spatial discretization, and the time integration is performed using the built-in fourth order Runge-Kutta method in the scipy package. We do not apply any data preprocessing procedure here.
F.3 Details on Pre- Training the Autoencoder
We provide details on pro-training the autoencoder here. The choice of first pre-training an autocncoder is motivated by the computational challenges of working directly with the high-dimensional spatial resolution of PDE datasets. Training directly in the ambient space requires substantial GPU memory and computational resources, making it impractical for large-scale or high-resolution datasets. By leveraging a latent-space representation, we achieve significant dimensionality reduction while preserving the essential structure of the data, enabling efficient training and inference with standard hard ware configurations. For these datasets, latent-space modeling provides a critical balance between computational efficiency, scalability, and performance.
We use the same architecture for the encoder and decoder for all the tasks, with the architecture parameters chosen based on the complexity of the task.
The encoder. The encoder first applies 2D convolution (conv_in) to the input frame, which reduces the number of channels from on_nhenelss to mid sud_chhnneels and processes the spatial dimensions. Then, series of ResidualBlock layers, which progressively process and downsample the feature map, making it smaller in spatial dimensions but more enriched in terms of features, are applied. After the residual blocks, the feature map undergoes an attention process via a multi-head attention layer. This layer helps the encoder focus on important parts of the input, learning relationships between spatial positions the imago. For the post-allention step, the feature map is further processed by residual blocks and normalized, preparing it for the final convolution. The output of the encoder is obtained by applying final 2D convolution (out_conv),, which maps the processed feature map to the desired number of output channels cout_channels).
The decoder. The decoder takes the encoded feature map and transforms it back into an output with similar spatial dimensions as the input. Similar to the encoder, the decoder starts with a convolution that adjusts the number of channels from in channels to aad_channels. Then, an attention mechanism (similar to the encoder) is applied to focus on important aspects of the encoded features. Next, series of ResidualBlock layers, combined with UpBlock layers, are used progressively increase the spatial dimensions of the feature map (upsampling), undoing the compression applied by the encoder. After the upsampling, the output is conv), mapping the internal feature representation
normalized and passed through final convolution to the desired number of output channels prut_channels).
Table summarizes the architecture parameters used for the considered tasks.
Table 7: Parameters chosen for the encoder (decoder) architecture.
Task
Fluid flow
Shallow-water eq.
Navier-Stokes eq.
Diffusion-reaction eq.
in_ in_channels
(1)
1 (1)
2 (2)
2 (2)
out_ chtanesss
1 (1)
1 (1)
2 (2)
2 (2)
mid_ channelss
64 (128)
128 (256)
128 (256)
128 (256)
Training details. We train the auloencoder using AdamW with batch size of 32, no weight decay and B = (0.9, 0.999). We use the cosine learning rate scheduler with warmup. For the fluid flow past a cylinder task, we train for 2000 epochs and use learning rate of 0.001. For the Navier-Stokes task, we train for 500 epochs and use learning rate of 0.0001. For the other two tasks we train for 5000 epochs and use learning rate of 0.0005. Our implementation is in PyTorch, and all experiments are run on an NVIDIA A100-SXM4 GPU with 40 GB VRAM belonging to an internal SLURM cluster.
Impact of the autoencoder (AE) size on test performance. We consider using pre-trained AEs with increasing number of middle channels (mid_ channels) in the encoder and decoder, resulting in AFs with
32

---

Published in Transactions on Machine Learning Research (07 (07/2025)
374,605, 1,485,589, 5,916,709 and 23,615,557 trainable parameters respectively. Table 8 shows the effect of the size of AE on the test performance for all models. We see that there is an optimal size of AE that leads to the best test performance for RIVER, the stochastic interpolant model and our model. For VE-diffusion, the best test performance is achieved with the smallest AE, suggesting that the model might overfit when AE increases size or sulfer from optimization issues bigger lalent spaces. For VP-diffusion, the test results are relatively insensitive the size of the AE used.
Table 8: Ablation study to assess the impact of the size of the pre-trained autoencoder (AE) on test performance using the fluid flow (FPC) task. The number of channels in the decoder is in parenthesis. Results are averaged over 5 generalions.
Model
mid_ ca_anneell
Test MSE (4)
Test RFNE (4)
PSNR (t)
SSIM (t)
RIVER
(32)
2.54e-03
5.43o-02
42.32
0.98
32 (64)
2.56e-03
4.89c-02
44.55
0.99
64 (128)
1.86e-03
4.48e-02
44.30
0.99
128 (256)
2.69e-03
5.52e-02
42.48
0.98
VE-diffusion
16 (32)
1.91e-01
3.99e-01
29.45
0.68
32 (64)
4.28c-01
6.19c-01
27.20
0.49
64 (128)
2.29e-01
4.74e-01
27.36
0.55
128 (256)
5.12e-01
7.22e-01
26.65
0.42
VP-diffusion
(32)
2.02e-03
5.11e-02
42.32
0.98
32 (64)
2.84--03
5.11e-02
44.14
0.98
64 (128)
3.58e-03
6.09e-02
42.37
0.98
128 (256)
2.03e-03
4.97e-02
42.97
0.98
Stoch. interp.
16 (32)
1.44e-02
1.22c-01
36.35
0.92
32 (64)
1.09c-02
9.01c-02
41.15
0.95
64 (128)
3.40e-03
10e-02
41.81
0.98
128 (256)
2.41e-02
1.36e-01
38.02
0.90
Ours
(32)
2.49e-03
5.70e-02
41.30
0.98
(a 0.01, Jaam 0)
32 (64)
2.87e-04
1.93e-02
50.62
1.00
64 (128)
3.79e-04
2.30e-02
48.88
1.00
128 (256)
3.53e-04
2.24e-02
48.96
1.00
F.4 Details on Training the Flow Matching Models
Architecture. The vector field regressor is transformer-basel model designed process latent vector fields and predict refined outputs with spatial and temporal dependencies. It uses key parameters like depth and mid_depth, which control the number of transformer encoder layers in the input, middle, and output stages. The state_ ssizesieee and state res parameters define the number of channels and spatial resolution of the input data, while tiner_dime sets the embedding dimension for processing. The model uses learned positional encodings and a time projection to incorporate spatial and temporal context into the input, which can include input latents, reference Iateetteeellaettss,, and conditioning, latengts The input is projected into the inner dimension and passed through series of transformer layers, with intermediate outputs from the input blocks concatenated with the output layers to refine predictions. Finally, the model projects the processod data back to the original spatial resolution and channel size using BatchNorm, producing the final vector field output.
Table 9 summarizes the architecture parameters used for the considered tasks.
Training details. For all the considered tasks, we train the regressor using AdamW with batch size of 32, learning rate of 0.00005, no weight decay and 8 = (0.9, 0.999). We use the cosine learning rale scheduler with warmup. For the fluid flow past cylinder, we train for 2000 epochs, for the shallow-water equation and diffusion-rraaction task we train for 1000 epochs, and for the Navier-Stoke task we train for 100 epochs. Our implementation is in PyTorch, and all experiments are run on a single NVIDIA A100-SXM4 GPU with 40 GB VRAM belonging to internal SLURM cluster. For the Navier-Stokes task, due to GPU memory
33

---

Published in Transactions on Machine Learning Research (07/2025)
Table 9: Parameters chosen for the vector field neural network.
Parameter
Fluid flow
Shallow-water eq-
Navier-Stokes eq.
Diffusion-reaction eq.
state_stt size
4
4
state_res res
[8,8]
[16,
16]
[64, 64]
[16, 16]
inner_di dim
512
512
512
512
depth
4
4
4
4
5
5
constraints in the experiments, we use gradient accumulation to simulate the desired batch size while training with a smaller per-step batch.
F.5 Details on the Evaluation Metrics
In addition to the standard mean squared error and relative Frobenius norm error (RFNE), we use the Pearson correlation coefficient to measure the linear relationship between the forecasted frames and the target frames. The range this coefficient is [-1, I], with zero implying no correlation. Correlations of -1 or +1 imply an exact linear relationship. Positive correlations imply that as increases, does y. Negative correlations imply that increases, y decreases. In addition,
peak signal-to-noise ratio (PSNR) to evaluate the quality of signal representation against corrupting noise, and structural similarity index measure (SSIM) (Wang et al., 2004) to assess perceptual results. The presented results are computed by averaging over batch size and number of sample generations.
F.6 Standard Deviations for the Presented Results
Table 10-14 provide the standard deviation of the results presented in the main paper.
Table 10: Standard deviation results for the fluid flow past cylinder task using different choices of probability paths for flow matching, Results are averaged over 5 generations.
Model
Test MSE (1)
Test RFNE (4)
PSNR (t)
SSIM (t)
RIVER
1.04e-03
2.70e-02
1.46
4.92e-03
VE diffusion
5,22e-02
3.34e-01
6.02e-01
3.14e-02
VP- Ii-diffusion
7,09e-03
3.96e-02
4.26e-01
3.30e-03
SI (b, t2)
6.35e-05
3.53e-02
5.39e-02
4.03e-04
SI (be = t)
1.48e-04
8.70e-02
2.97e-02
5.30e-04
Ours (a 0.01, Jam 0, RK4)
4.26e-06
5.29e-03
3.61e-02
3.63e-05
34

---

Published in Transactions on Machine Learning Research (07/2025)
Table 11: Standard deviation for the results ablation study for the fluid flow past a cylinder task. Results are averaged over 5 generations.
sampler
Test MSE (0)
Test RENE (L)
PSNR (1)
SSIM (1)
0.0
Euler
1.00e-05
1.25e-02
4.37e-02
9.76e-05
0.01
Euler
6.71e-06
1.14e-03
5.17e-02
4.81e-05
0.1
Euler
4.88c-05
4.11e-02
5.39e-02
3.91e-04
0.0
Euler
10
1.24e-05
1.29e-02
4.28e-02
8.27e-05
0.01
Euler
8, 8.10e-06
4.37e-03
6.42e-02
6.21e-05
0.1
Euler
10
4.32e-05
4.02e-02
5.13e-02
3.06e-04
0.0
RK4
10
3.45e-06
7.56e-03
2.57e-02
2.46e-05
RK4
4.26e-06
5.29c-03
3.61e-02
3.63e-05
0.1
RK4
10
8.62e-06
5.88c-02
2.69e-02
2.61e-04
0.0
Euler
20
8. 1,44006
1.31c-02
3.11e-02
5.02u-05
Euler
4.25c-06
4.51e-03
3.66e-02
2.77u-05
0.1
Euler
20
2.12e-05
3.96e-02
2.21e-02
1.03--44
0.0
RK4
20
1.54e-06
7.49e-03
1.84--02
1.26e-05
0.01
RK4
3.99e-06
1.29e-02
1.74e-02
2.95e-05
0.1
RK4
20
2.84e-05
5.85e-02
8.10e-03
8.33e-05
Table 12: Standard deviation results for the shallow-water equation task using different choices of probability paths for flow matching. Results are averaged over 5 generations.
Model
Test MSE (1)
Test RFNE (4)
PSNR (t)
SSIM (t)
RIVER
2.28e-05
8.67e-02
8.39e-02
1.23e-03
VE-diffusion
7.29e-04
3.29e-01
2.17e-01
1.68e-02
VP-diffusion
2.55e-04
1.10e-01
3.62e-01
1.56e-02
SI (b, = f)) ST (b, t)
2.97e-06
1.03e-01
1.86e-02
2.03e-04
1.37e-06
7.18e-02
4.11e-03
2.12e-04
Ours (o 0.1, sem 0, RK4)
1.55e-06
7.06e-02
5.98e-03
2.11e-04
Table 13: Standard deviation results for the diffusion-reaction equation task using different choices of probability paths for flow matching. Results are averaged over generations.
Model RIVER
Test MSE (1)
Test RFNE (1)
PSNR (t)
SSIM (t)
5.56e-01
2.51e-11
7.88e-01
1.37e-02
VE-diffusion VP-diffusion
1.51e-02
1.11
5.89e-01
1.22x-02
1.55e-03
6.15e-01
1.06
7.08e-03
SI (b, 12) SI (br
6.44e-04
1.15
8.60e-02
6.98e-04
8.15e-07
5.01e-02
5.67e-03
8.25e-05
Ours (or O, aam O, RK4)
8.47e-07
4.77e-02
6.22e-03
9.53e-05
Table 14: Standard deviation results for the Navier-Stokes equation task using different choices of probability paths for flow matching. Rosults are averaged over genorations.
Model
Test MSE (4)
Test RFNE (4)
PSNR (t)
SSTM (t)
RIVER
7.52c-04
8.31e-01
9.34e-02
3.80e-03
VE-diffusion VP-diffusion
6,53e-04
1.21
1.19e-01
1.61e-03
3.88c-03
1.62
1.64e-01
5.74e-03
SI (b, f2) SI (b.
5.11e-06
3.27e-02
2.37e-02
3.34e-04
4.87e-07
1.09e-02
1.98e-03
6.22e-05
Ours (a 0.1, Jsam RK1)
1.81e-07
1.06e-02
2.30e-03
6.01e-05
35


Reviews:
[{"review_id": "i3EZ0GwMIP", "reviewer": "Reviewer_zAaW", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n1. The paper addresses an underexplored aspect in flow matching for probabilistic forecasting, i.e., design of probability paths.\n\n2. The connection to Schrödinger bridges and dynamical optimal transport adds a principled foundation to the proposed method.\n\n### Weaknesses\n\n1. The discussion is centered around general flow matching, with limited connection to spatio-temporal PDEs. For example, it is unclear why Gaussian prior processes are chosen as the reference process, whether this choice is specifically motivated by characteristics of spatio-temporal data.\n\n2. Some key steps in the paper are insufficiently justified. For instance, although the derivation from Schrödinger bridges can lead to $\\mu_t = a_tZ_0 + b_tZ_1$, the rationale for selecting $a_t = 1 - t, b_t = t$ rather than other forms is not clearly explained.\n\n3. The experiments compare only different probability paths within the flow matching framework, omitting comparisons with other state-of-the-art probabilistic forecasting methods (e.g., transformer-based models).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the role of probability path design in flow matching models for probabilistic forecasting, particularly of high-dimensional spatio-temporal data governed by PDEs. Authors consider the class of folows with probability paths $p_t(Z|Z_0, Z_1)=\\mathcal{N}(Z|a_tZ_0+b_tZ_1,c_t^2I)$ and flow maps $\\psi_t(Z)=a_tZ_0 + b_tZ_1 +c_tZ$. Restricting this class to be close to a Gaussian prior process with Schrödinger bridge framework, the optimal probability path model is derived. The authors theoretically show that, under certain assumptions, this path leads to a lower vector field variance than the rectified flow model. Experiments are conducted on four PDE-based spatio-temporal forecasting tasks to demonstrate the effectiveness of the proposed path over other paths.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Make the relevance of the proposed framework to spatio-temporal PDEs more explicit, especially in the modeling choices and theoretical assumptions. Highlight any PDE-specific insights that influence the design of the probability path or flow structure.\n\n2. Add baselines from alternative state-of-the-art probabilistic forecasting models, such as transformer-based models or classical deep learning approaches. This would better contextualize the advantages of the proposed flow matching framework.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0VZYwtzDai", "reviewer": "Reviewer_1pPU", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n1.  The paper investigates the role of probability path selection in flow matching for spatio-temporal forecasting, an area with aspects yet to be fully explored.\n2.  The paper proposes a probability path model with connections to dynamical optimal transport and Schrödinger bridges. Theorem 3.3 offers a theoretical perspective on the potential for lower vector field variance under specific conditions.\n3.  Experiments are conducted on multiple PDE datasets (Fluid Flow Past a Cylinder, Shallow-Water Equation, Diffusion-Reaction Equation, and Incompressible Navier-Stokes Equation). The proposed model is compared against other probability path models like OT-VF (RIVER), VE-diffusion, VP-diffusion, and Stochastic Interpolants, using several standard metrics.\n4.  The proposed model showed differences in performance metrics (accuracy, training speed, inference efficiency) when compared to the baselines in the conducted experiments. For instance, it was reported to use 10 sampling steps with RK4 for inference.\n5. The paper is generally organized, presenting the methodology and algorithms for training and inference. The appendices provide supplementary details.\n6. The reported inference efficiency suggests potential for applications where computational resources may be a consideration.\n\n**Weaknesses:**\n\n1.  The experimental setup focuses on predicting a limited number of future frames (e.g., condition on 5, predict 15 or 20). While common in video prediction, \"forecasting\" in other time-series contexts can imply longer horizons. A discussion on scalability to much longer prediction horizons could be informative.\n2.  Flow matching is performed in a latent space derived from a pre-trained autoencoder (AE). The overall system's performance may be influenced by the AE's characteristics. The paper notes that joint training of the AE and Vector Field (VF) was empirically less effective, but further discussion on sensitivity to AE choices could be beneficial.\n3.  Ablation studies indicate that performance and training stability (loss curves) can be sensitive to the hyperparameter $\\sigma$ in the proposed probability path's variance term $c_t^2 = \\sigma_{min}^2 + \\sigma^2 t(1-t)$. While a range is explored, this points to a tuning aspect. The value of $\\sigma_{min}$ is set to 0.001, and clarification on the specific $\\sigma$ values used for the main results in Table 2 would improve reproducibility.\n4.  The paper focuses on comparing probability paths *within* the flow matching framework, explicitly stating that comparison with other generative modeling frameworks is not its primary focus. Providing some discussion on how these findings position flow matching relative to other generative approaches (e.g., different types of diffusion models) in PDE forecasting could offer useful context.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper examines the selection of probability paths within flow matching models for spatio-temporal forecasting. The authors report that the choice of the probability path influences forecasting results. They introduce a probability path model, with theoretical motivation from dynamical optimal transport and the Schrödinger bridge perspective. This model is formulated to align with the continuous dynamics of spatio-temporal data by interpolating between consecutive sequential samples. Empirical results on several PDE-based dynamical system benchmarks indicate that this model leads to faster convergence during training and different predictive performance metrics (such as MSE/RFNE, PSNR/SSIM, and temporal consistency) when compared to some existing probability path models. The paper also notes the model's inference efficiency, requiring a limited number of sampling steps.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.  In Section 5.1 and Table 2, please explicitly state the value(s) of $\\sigma$ used for the \"Ours\" model to achieve the reported results. Information in ablation studies or appendices (e.g., Table 6 suggests $\\sigma=0.01$ for FPC) should be clearly referenced or reiterated in the main text for the primary results.\n2.  The paper sets $\\sigma_{min}=0.001$. A brief justification or sensitivity analysis concerning this choice would be useful, particularly its role in numerical stability versus impact on performance.\n3.  A brief discussion on potential challenges or adaptations if the model were applied to significantly longer forecasting horizons in PDE settings.\n4. A comment on how the latent space dimensionality (from the AE) might affect flow matching performance or the probability path parameters.\n5.  While inference efficiency is noted, a brief summary of training time observations (e.g., based on Figure 3 showing epochs to convergence for different probability paths in the main text.\n6.  A short paragraph discussing the implications of these improved flow matching results in the wider context of generative forecasting models for spatio-temporal data.", "broader_impact_concerns": "N.A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "VxxrApUjL9", "reviewer": "Reviewer_vmBN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- Addresses an important application setting of learning FMs for time-series data\n- Motivates a few method improvements for FMs in this application setting\n- Theoretical justification for why the proposed model might empirically perform better, in the form of a lower-variance training objective.\n- Demonstrates empirical performance gains\n\n### Weaknesses\n- (W1) It's not clear how the loss function used in Alg 1. is obtained. The \"Flow Matching loss parameterization\" introduced in Eq. 9 uses $v_t^\\theta(Z)$, whereas in Alg 1., the $v_t^\\theta$ is given as $v_t^\\theta(Z|z^{\\tau-1}, z^c, \\tau - c)$. The paper needs to have a clearer, precise, and easy-to-follow motivation of the loss function used in the training algorithm.\n- (W2) Related to the concern about the missing derivation of the loss function used in Alg. 1, the importance of conditioning the vector flow network on $z^c$ is not obvious. The paper's results would be improved with empirical proof that $z^c$ is important conditioning information. This could be done by ablating the $z^c$ input to the model.\n- (W3) The results are evaluated at a single number of noise-free evaluations, and this number is somewhat small (10). The paper would be improved with extended experiments that vary the number of sampling steps used by each method. It is possible that the performance gaps would be smaller for a larger number of sampling steps. For example, plots of the MSEs as a function of the number of sampling steps used by the sampler would be illuminating (e.g. N={1,2,5,10, 20, 50, 100}) \n- It's not clear why the weighting function is added to Eq. (3) and Eq. (6), which are not present in the corresponding objectives in (Lipman et al., 2022)\n- Table 1 would be improved by including references to the prior works where the methods were proposed\n- The symbol $\\mathcal D$ is overloaded: it's used both for the decoder $x=\\mathcal D(z)$ and the dataset. Fix this throughout by changing one of them.\n- This paper is related yet it is undiscussed in the paper: https://arxiv.org/pdf/2411.07506v1", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper addresses an important application setting of learning FMs for time-series data. The main idea is to define a new probability path model that results in a lower-variance training objective for time-series data, by using smoothness in the temporal targets. The paper demonstrates empirical performance gains using several time-series forecasting datasets.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Address (W1), (W2), and (W3) with the requested solutions. The biggest issue is W1 -- there's a gap between the presented theory and the implemented training algorithm. This may be simply an issue of presentation (some missing or overlooked steps in the derivation), but as it stands, it undermines the claimed contribution of a theoretical framework with efficient algorithms, as the connection between the framework and the algorithm is too weak.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)