Improving Flow Matching by Aligning Flow Divergence
Authors: Yuhao Huang
Taos Transue
Shih-Hsin Wang
William M Feldman
Hong Zhang
Bao Wang
Venue: ICML 2025 poster (2025)
Keywords: Flow Matching
Divergence Matching
Total Variation Bound
Generative Modeling
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: We establish a new user-friendly total variation bound between the learned and exact probability paths in flow matching, informing efficient algorithms to improve flow matching.

Abstract: Conditional flow matching (CFM) stands out as an efficient, simulation-free approach for training flow-based generative models, achieving remarkable performance for data generation. However, CFM is insufficient to ensure accuracy in learning probability paths. In this paper, we introduce a new partial differential equation characterization for the error between the learned and exact probability paths, along with its solution. We show that the total variation gap between the two probability paths is bounded above by a combination of the CFM loss and an associated divergence loss. This theoretical insight leads to the design of a new objective function that simultaneously matches the flow and its divergence. Our new approach improves the performance of the flow-based generative model by a noticeable margin without sacrificing generation efficiency. We showcase the advantages of this enhanced training approach over CFM on several important benchmark tasks, including generative modeling for dynamical systems, DNA sequences, and videos. Code is available at \ref{https://github.com/Utah-Math-Data-Science/Flow_Div_Matching}.

Improving Flow Matching by Aligning Flow Divergence
Yuhao Huang 12
Taos Transue 12
Shih-Hsin Wang" 12
William Feldman
Hong Zhang Bao Wang"
Abstract
FM circumvents computationally expensive sample simu- lations to estimate gradients densities. The celebrated diffusion models (DMs) with variance preserving (VP) or variance exploding (VE) stochastic differential equations (SDEs) (Song et al., 2020) can be viewed as special cases of FGMs with diffusion paths (c.f. Section 2). Furthermore, FM excels in generative modeling on non- non-Euclidean spaces, broadening its scientific applications (Baker et al., 2024; Chen & Lipman, 2024; Jing et al., 2023; Bose al., 2024; Yim et al., 2024; Stark et al., 2024).
Conditional flow matching (CFM) stands out as an efficient, simulation-free approach for training flow-based generative models, achieving remark- able performance for data generation. However, CFM is insufficient to ensure accuracy in learning probability paths. In this paper, we introduce a new partial differential equation characterization for the error between the learned and exact proba- bility paths, along with its solution. We show that the total variation gap between the two probability paths is bounded above by a combination of the CFM loss and an associated divergence loss. This theoretical insight leads to the design of new objective function that simultaneously matches the flow and its divergence. Our new approach improves the performance of the flow-based gen- erative model by noticeable margin without sac- rificing generation efficiency. We showease the advantages of this enhanced training approach over CFM several important benchmark tasks, including generative modeling for dynamical sys- tems, DNA sequences, and videos. Code is avail- able at Uhab-Math-DDat--Siien
At the core of FM is the idea of regressing a vector field that interpolates between the prior noise distribution g(x:) typically the standard Gaussian and the data distribution p(z). Specifically, we aim to regress the vector field u, (z) that guides the probability flow pt(x)) interpolating between easy-to-sample noise distribution and the data distribu- tion, i.e., Po q and yı p. The relationship between u, and pt is formalized by the following continuity equation (Villani et al., 2009):
@pr(2))
+ + (pr(a)).(()) = 0.
Ōt
FM approximates u, using a neural network- pattork-aaramterizeeee vector field vi (x 0). seeking to minimize the FM loss:
CFM == (.p)
|| 0) - -,)
(1)
1. Introduction
where ~ U10 1] follows uniform distribution over the unit time interval [0, 1].
Flow matching (FM) - leveraging a neural network to learn a predefined vector field mapping between noise and data samples has emerged as an efficient simulation-free train- ing approach for flow-based generative models (FGMs), achieving remarkable stability, computational efficiency, and flexibility for generative modeling (Lipman et al., 2023; Albergo & Vanden-Eijnden, 2023; Liu et al., 2023). Com- pared to the classical likelihood- based approaches for train- ing FGMs, c.g., (Chen et al., 2018; Grathwohl et al., 2018),
However, equation (1) is intractable as ur (z) is unavail- able. To address this, an alternative simulation-free method, known as conditional flow matching (CFM) (Lipman et al., 2023: Albergo & Vanden-Eijnden 2023), is employed. In CFM, (2r, 0) is trained by regressing against predefined conditional vector field on per-sample basis, ensuring both computational efficiency and accuracy. Concretely, for any data sample ~ p(x), we can define a condi- tional probability path (=(2)/1) for E [0, 1] satisfying Po (a: (x,) = ((x) and pı (xr(z,)) Rs 8(x - æx),, and define the associated conditional vector field u, (ar(rr);; see Sec- tion 2 for review on several common designs of conditional probability paths. Once the conditional probability paths are defined, the marginal probability path pt is given by:
'Department of Mathematics. University of Utah, Salt Lake City, UT, USA SSientific Computing and Imaging (SCI) Institute, Salt Lake City, UT, USA 'Mathematics and Computer Science Division, 240 Argonne National Laboratory, Lemont, IL, USA. Correspondence to: Bao Wang <wangbaonj @ comi.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
(x = /
pr pr (peel)))(((()))))))))))

---

Improving Flow Matching by Aligning Flow Divergence
Similarly, the marginal vector field is defined as:
characterization describes how the error propagates over time, allowing us to derive a total variation (TVV)-baeed er- ror bound between the two probability paths; see Corollary 3.2 and Theorem 3.3. These theoretical results underscore the importance of controlling the divergence gap to en- hance the accuracy in learning P. (z)
u u,(x)= := / ue (z/ )
Pt (z) a (p(1)))
dx1..
Pr(x))
With these relations in mind, CFM regresses (z 0)
against uxr by minimizing the following CFM loss:
CCFM (f) == m. ) (1 0) (x( ) ||?
Informed by our established TV error bound, we develop new training objective by combining the CFM loss with the divergence gap. However, directly minimizing the divergence gap is intractable since the divergence of the marginal vector field is unavailable. To address this issue, we propose a conditional divergence gap an upper bound for the unconditional divergence gap. We refer to this new training objective as flow and divergence matching (FDM); see Section 4 for details.
(2)
It has been shown that the CFM loss is identical to the FM loss up to a constant that is independent of 0 (c.f. (Lipman
et al., [Theorm 2D. Therefore, minimizing (0) enables (2, 0) to be unbiased estimate for the marginal vector field (zr).
While CFM enables vx(2),) 0) to efficiently approximate (z()), we observe that their divergence gap' (z, 0)-
V V.u,(2) (z)] can be substantial, resulting in significant errors in learning the probability path and estimating sample likeli- hood. Figure highlights the challenges in learning Gaus- sian mixture distribution using CFM. The full experimental setup for this result is provided in Section 3. Additionally, we will prove the existence of an intrinsic bottleneck of FM. This underscores the importance of improving FM for generative modeling, especially for tasks requiring accurate sample likelihood estimation. Indeed, such tasks are ubiq- uitous in climate modeling (Finzi et al., 2023; Wan et al., 2023; Li al., 2024), molecular dynamics simulation (Pe- tersen et al.), cyber-physice systems (Delecki et al., 2024), and beyond (Hua et al., 2024).
We validate the performance of FDM across several bench- mark tasks, including synthetic density estimation, trajec- tory sampling for dynamical systems, video generation, and DNA sequence generation. Our numerical results, presented in Section 5, show that our proposed FDM can improve likelihood estimation and enhance sample genera- tion by remarkable margin over CFM.
1.2. Some Additional Related Works
The Kullback-Leibler (KL) divergence between the ex- act and learned distributions has been studied for DMs (c.f. (Song et al., 2021; Lu et al., 2022; Lai et al., 2023)) and FM (c.f. (Albergo et al., 2023)) with ODE flows, where it was observed that the FM loss in equation (1) alone is insufficient for minimizing the KI. KI. divergence between two probability paths, and the KL divergence bound depends on higher-order score functions.
1.5e-1
3
Pasta Pen
CCFM CCMM
1e-1
2e-1
1e-1
Several works have explored improving training DMs with higher-order score matching. For instance, Meng et al. (2021) have proposed high-order denoising score match- ing leveraging Tweedie's 's formula (Robbins, 1992; Efron, 2011) to provide a more accurate local approximation of the data density (e-g., its curvature). We notice that the trace of the second-order score matching proposed in (Meng et al., 2021) resonates with the idea of our proposed FDM in the context of DMs. Additionally, inspired by the KL divergence bound, high-order score matching matching up to third-order score - has been used to improve likelihood estimation for training DMs (Lu et al., 2022). Nevertheless, these higher-order score-matching methods are significantly more computationally expensive than our proposed FDM.
5e-2
1K
2K Epoch
3K
5
0
5
Figure 1. Experiments of training an FM model using CFM for sampling the ID Gaussian mixture distribution in equation (18). The left panel shows that the conditional divergence loss LCDM equation (14) much larger than the CFM loss LCrM, and the right panel shows the significant gap between the exact distribution (pData) and the distribution learned through FM (pex)).
1.1. Our Contribution
We summarize our key contributions as follows:
We characterize the error between the exact (p)(z) and learned (p: (x)) probability paths using a partial differen- tial equation (PDE); see Proposition 3.1. This new error
Enforcing the continuity equation for flow dynamics is an- other related work that has been studied in the context of DMs. In particular, Lai et al. (2023) shows that the score function satisfies a Fokker-Plande equation (FPE) and di- rectly penalizes the loss function with the error from plug-
'erre use the absolute value notation since the divergence of the vector field is a scalar.
2

---

Improving Flow Matching by Aligning Flow Divergence
ging the learned score function into the score FPE. To the best of our knowledge, developing a PDE characterization of the error between the exact and learned probability paths and bounding their TV gap using only the vector field and its divergence have not been considered in the literature.
by
A-1 (æ x1). -
Ur (x)
where oi denote the derivative of the function. Likewise, the VP SDE (Song et al., 2020) has the following conditional probability path:
1.3. Organization
(=(2))1)) = =N(a (1 (1-aj o)
We organize this paper as follows. We provide a brief review of FM in Section 2. In Section 3, we present our theoretical analysis of the gap between the exact and Iearned probabil- ity paths, accompanied by illustrative numerical evidence. We present FDM to improve training FGMs in Section 4. We verify the advantages of FDM over FM using few rep- resentative benchmark tasks in Section 5. Technical proof', additional experimental details, and experimental results are provided in the appendix.
where 02 = e TT()) and T(1) = So ((s)ds with ((s) being the noise scale function. The corresponding conditional vector field is
0i (a1-12 x1).
un (1(x)/1) = 1 - 2
1-a
Besides diffusion paths, the optimal transport (OT) path another remarkable choice (Lipman al., 2023). OT path uses the Gaussian conditional probability path with
2. Flow Matching
In this section, we provide a brief review of FM and preva-
Mi (,(a)) Iz and (a) = 1 - (1 0min)..
lent designs of conditional probability paths. A vector field u, [0, x Rd - Rd defines a flow th. : [0, 1] x Rd - Rd
The corresponding conditional vector field is given by
through the following ODE:
ue (x)(x))) = X1 (1 T )
d {1(2) (x)),
1 1 (1 Tmin
with the initial condition &ho (x) = æ. FGMs map a prior noise distribution Pa = to data distribution p1 R p via the following map:
3. Error Analysis for Probability Paths
In this section, we analyze the error between the two prob- ability paths associated with the exact and learned vector fields, respectively. Specifically, we show that this error sat- isfies a PDE similar the original continuity equation, but with additional forcing term. Using Duhamel's principle (Seis, 2017), we reveal that this forcing term directly gov- erns the magnitude of the error. The omitted proofs, along with the common assumptions employed by (Lu et al., 2022; Lipman et al., 2023; Albergo et al., 2023) and adopted in our theoretical results, are provided in Appendix A.
Ovp1 Ox (xx) Vx E Po-
(x()) (0 (ar) )det
For given sample x1 FM defines a conditional prob- ability path satisfying Po (z) |z)) = g(x) and pı (z|z)) x 8(x - (() the Dirac-delta distribution centered at æ1 and the corresponding conditional vector field u,(x) (x)xxx1) Then FM regresses neural network paraetezear unconditional vector field v. (2. 8) by minimizing CFM in equation (2).
A prevalent choice for (2 xx)) is the Gaussian conditional probability path given by
3.1. PDE for the Error Between Probability Flows
Recall that the marginal probability path pe(x) and the marginal vector field u, (2r satisfy the following continuity equation (Villani al., 2009):
with n4 (x 0 and (0( = 1. Moreover, i (x1)
xn and (x;) is small number so that pı (σ|x1) ) x ō(x ). Some celebrated DMs can be interpreted as FM models with Gaussian conditional probability paths. In particular, the generation process of the DM with VE SDE (Song et al., 2020) has the conditional probability path:
@p((x) Øl
+ V (pr(z)u((x) = 0.
(4)
We can rewrite the continuity equation into the following non-conservative form:
0p,(x) = -V ux()))p(()) - -u))) (p((x) (5) Øt
Pr (x) (n) = (a<az)..
where o₁ is an increasing function satisfying du = and 01 >> 1. The corresponding conditional vector field given
Similarly, consider the probability path (x associated with the neural network-paramerte vector field n, (z, 0).
3

---

Improving Flow Matching by Aligning Flow Divergence
This probability path satisfies the following continuity equa- tion, which has the same initial condition as the ground truth equation (5), i.e., po fo:
term as follows:
V- - -2) + (ur vv) V log Pr
(DM ==
(DM(O):=EE,pctt))
(p)(ar) -V (x, ()) ( v2 (()())) (6) Ot
(10)
The following theorem establishes an upper bound for the error term TV(p..pr pe) in terms of (D).
We now introduce the enror term ex(x) == p(())) (x))
Theorem 3.3. Under some common mild assumptions adopted in (Lu et al., 2022; Lipman et al., 2023; Al- bergo et al., 2023), the following inequality holds for any E 1]:
The following proposition shows that 6 satisfies a PDE similar to equation (4), but with an additional forcing term that reflects the discrepancy between the vector fields ut and v4.-
1 TV & (DDM(O).
Proposition 3.1. C == Pt- pe satisfies the following PDE:
(11)
=2"
0.4+ + V ((((v) = Lt: [<((x)= = =0,
Specifically, pr(x)) px( when LDM is zero.
(7)
4. Conditional Divergence Matching
where
In the previous section, we have highlighted the impor- tance of matching the divergence between u, and v, beyond matching the vector fields themselves. However, directly minimizing the divergence loss presents a computational challenge, as computing the divergence of the exact uncon- ditional vector field is intractable. To address this issue, we will leverage similar idea to the conditional flow matching to address the computational issue.
L -pt V (Ut= - -2)) + (w: vz) V logpt
(8)
3.2. Error Bound for Probability Paths
FM aims to minimize the discrepancy between pt and p. by reducing the difference between their associated vector fields, uz (x) and (2 0), through minimizing the CFM loss equation (2). However, Proposition 3.1 highlights that the error dynamics are not only influenced by u, or but also by V -u, - vz). as both terms contribute the forcing term in equation (7). To formalize this observation, we solve e, using Duhamell's 's formula (Seis, 2017). In particular, we have the following result:
We start by deriving the conditional version of LDM (f). We recall the following conditional form of the continuity equation from (Lipman et al., 2023):
(.p.(alx)) V
(xlla)))
(12)
Corollary 3.2. For any E [0, 11. the error e, satisfies
which relates the evolution of the conditional probability density Pe(z1(21) to the divergence of p: 21 21) ( 1)). By integrating over the conditioning variable 21 and ap- plying the continuity equation (4), we obtain the following connection between the conditional divergence and uncon- ditional divergence:
<+())())) det = (o(()(()) - det (z))dd,
where (x) is the flow induced by the vector field (x) in similar way as that in equation (3), det det (x) denotes the determinant of the Jacobian matrix Vol(m) and L,, is defined in Proposition 3.1.
V. (p(()))((z)) -pp((())
Corollary 3.2 suggests that minimizing the divergence gap is as important as reducing the vector field discrepancy in order to learn an accurate probability path.
/ (1p((()111)((())
(13)
/ v. (p dxy
To quantify the error ex, we consider the following TV dis- tance between pt and pr:
1 2 / (p((x))((())] pc(x)) dæ
Furthermore, we observe the following identity:
TV(p,,
(9)
= / |<((z)) dx.
Pt V (wt v.) + (m - wz)) V Vlogpt
=V (prue) - V (prvt).
Motivated by the error-related identity in Corollary 3.2 and the form of L, equation (8), we introduce an additional
This leads the following error estimation for the condi-
4

---

Improving Flow Matching by Aligning Flow Divergence
tional divergence loss:
the KL divergence, leave the exploration of developing computationally efficient method for controlling the KI. divergence in this direction for future work.
(CDM(0)
((lp))(()))))) ( (x(/2)))) V (.((x,)))
4.2. Synthetic Experiment.
+ +(u((a))))) 0)) 03) V (logp(()))))))
To solidify our theoretical results, present simple nu- merical example before moving to real-world applications. Specifically, we consider the problem of sampling from the following Gaussian mixture distribution
(14)
Now we are ready to establish the fact that the conditional di- vergence loss LCDM ( is an upper bound for the divergence loss LDM (0) and the TV gap TV(p,, We summary our results in the following theorem:
p(æ) =0 =0. -00
(18)
+ 0.15A(-1,0.) - 0.27.V 0.1),
Theorem 4.1. We have the following inequality:
(DM S (ODM(().
(15) using both standard FM and our proposed FDM defined in equation (17). We use a 3-layer MLP to approximate the VP diffusion path vector field by minimizing equation (17) with
Furthermore, we have:
A 1, A2 = 0 for FM and A = 1, A2 0.2 for FDM. We use 104 data points sampled from equation (18) for training.
1 TV p:) & (CDM(O),
(16)
for any t E 1].
Data
FM
FDM
4.1. Flow and Divergence Matching.
In practice, we observe that minimizing (CDM()) alone cannot yield appealing results, as the loss cannot go to exact zero in training. This nonzero loss comes from a balance between V u, (ar a1) - ve(2r, 0) and (ux(a) (x:) vx (z,0)) 0)) V log pr ( and both terms can be positive negative, resulting in cancellation. As such, there is guarantee that we can learn a vector field v. (z, 0) that is in proximity to un (x) by minimizing LCDM (f). In contrast, by using weighted sum of LCDM and LCFM as the training objective, we can directly control the gap between the vector fields and their divergences. Therefore, we propose the flow and divergence matching (FDM) loss:
Figure 2. Snapshots probability paths at = 0.6, 0.85, and 1
(left right). First/Second row: FM/FDM data distribution.
LFDM AICCFM + AeLCDM,
(17)
where dz > are hyperparameters; we choose them via hyperparameter search in this work. It is an interesting future direction to design a principle to choose As optimally.
Remark 4.2. It worth noting that minimizing the objective function LFDM offers a more computationally efficient ap- proach compared to higher-order control methods presented in (Lu et al., 2022; Lai et al., 2023), as it is computationally much cheaper than controlling differences higher-order quantities (c.g.. gradient of the divergence). Moreover, to introduce effi- cient squared conditional divergence- mettgenne-matching loss LCDM2 which adopts stop-gradien (Lu et al., 2022) and Hutchin- son trace estimation (Hutchinson, 1989) techniques. This adds only one extra backward pass compared with baseline flow-mn wow-matching training; see Appendix D for details. While a bounded TV distance does not necessarily imply a bound on
t.0
t=
t.5
t- 0.5
t=1.0
t=1.0
Figure 3, Comparison of probability paths over time learned by FM (left) vs. FDM (right).
further improve training efficiency,
Figures 2 and contrast the performance of our proposed FDM against the baseline FM. The numerical results con- firm that the probability path (at t = 1) learned by FM suf- fers from a substantial discrepancy from the exact Gaussian mixture distribution. In contrast, FDM learns the Gaussian
5

---

Improving Flow Matching by Aligning Flow Divergence
mixture much more accurately than FM. Specifically, the TV gaps between the learned and exact distributions are 0.0945 and 0.0587 for FM and FDM, respectively.
5. Experimental Results
(a) FM (OT)
(b) FDM (OT)
(c) Ground Truth
In this section, we validate the efficacy and efficiency of the proposed FDM in enhancing FM across various bench- mark tasks, including density estimation synthetic 2D data (Section 5,1.1) and image data (Section 5.1.2), DNA sequence generation (Section 5.2), and spatiotemporal data sampling tasks including trajectory sampling for dynamical systems (Section 5.3.1) and video prediction via latent FM (Section 5.3.2).
Figure 4. Generated samples from FM and FDM using the optimal transport (OT) path trained on the checkerboard dataset.
5.1. Density Estimation on Synthetic and Image Data
We train the models for density estimation on two datasets: synthetic 2D checkerboard and the image dataset CIFAR- 10 (Krizhevsky et al., 2009).
Our experiments confirm that our proposed FDM remark- ably improves FM with guidance, enhancing promoter DNA sequence design with class-conditional flow, well as refin- ing trajectory generation for dynamical systems and video predictions conditioning on the initial states over the first several time steps. In this section, we report the error be- tween the exact and learned distributions in terms of the TV distance, and the corresponding KL divergence results are further provided in Appendix C.
1.1.1. SYNTHETIC DENSITY ESTIMATION
In this experiment, we train models using FM and FDM for 2 x 10 iterations using batch size of 512. For each iteration,
numerically sample the data for the training set and use the same sampling method for validation and testing sets. We compare FM and FDM with the baselines for both OT and VP paths in the likelihood computed based on the test dataset. The results in Table I and Fig. 4 show that FDM consistently outperforms FM across different probability paths.
Software and Equipment. Our implementation utilizes PyTorc Lightning (Falcon, 2019) for synthetic density es- timation, DNA sequence generation, and video generation, while JAX (Bradbury et al., 2018) and TensorFlow (Abadi et al., 2016) are employed for dynamical systems-related ex- periments. Experiments are conducted on multiple NVIDIA RTX 3090 GPUs
5.1.2. DENSITY MODELING ON IMAGE DATASETS
In the experiment, we train models using both FM and FDM for image sampling on the CIFARIO dataset (Krizhevsky et al., 2009). We follow the experimental settings in the flow matching baseline paper (Lipman et al., 2023) and compare the performance in terms of the negative log- likeliold and FID scores of the sampled images as shown in Table 2.
Training Setup. See Appendix B.
Models and Datasets. We employ OT and VE/VP dif- fusion paths for the flow maps in most tasks except the Dirichlet flow for DNA generation. We follow the approach used in (Huang et al., 2024; Lu al., 2022) estimate the divergence, which employs Hutchinson's trace estimator (Hutchinson, 1989). Our experiments utilize numerical simulation-based dataset for density estimation and trajec- tory sampling, a dataset extracted from a database of human promoters (Hon et al., 2017) for DNA design, and the KTH human motion dataset (Schuldt et al., 2004) and the BAIR Robot Pushing dataset (Ebert et al., 2017) for video predic- tion.
Model
NLL(4)
FID(L)
FM(OT)
2.99
6.35
FDM(OT)
2.85
5.62
Table 2. Negative log-likelihood and sample quality (FID scores) estimation on CIFAR-10.
5.2. Sequential Data Sampling-DNA Sequence
In this experiment, we demonstrate that FDM enhances FM with the conditional OT path and Dirichlet path (Stark et al., 2024) on the probability simplex for DNA sequence genera- tion, both with and without guidance, following experiments conducted in (Stark et al., 2024). For this task, instead of directly parameterizing the vector field, the Dirichlet flow model constructs it by combining pre-designed Dirichlet probability path functions with a parameterized classifier Pr (m |a,, 0), where x sampled from the conditional proba-
Model
FM (OT)
FDM (OT)
FM (VP)
FDM (VP)
Likelihood (1)
2.38
2.53
2.34 02
2.461.02
Table Likelihood estimation of models on the checkerboard test set. Here, "OT" denotes the optimal transport path and "VP" de- notes the pariance-preserving path. Unit: x x10** -2
6

---

Improving Flow Matching by Aligning Flow Divergence
bility at time given the data point #1- Since #1 represents discrete categorical data with finite number of categories, it can be treated as class label. Since this approach only requires parameterizing the classifier (11 [z, f), we only need to penalize the norm of the gradient with respect to the input of the classifier, which is equivalent to minimiz- ing the divergence error; see Appendix B.2 for more de- tails. Additionally, we conduct experiments where the clas- sifier is parameterized with guidance, pa ar, y, 0) with representing the guiding information. We use the same experiment setup in (Stark et al., 2024) except the newly introduced hyperparameters X1 and 22; see Appendix B.2 for the detailed settings.
Method
MSE (4)
Bit Diffusion (One-hot Eacoding) (Albergo al., 2023)
3.951-2
DDSM (Albergo et al., 2023)
3.34E-2
Large Language Model (Sturk al., 2024)
3.331-2
Linear FM (Stark et al., 2024)
2 1.82 oz E-2
Linear FDM (ours)
2 2,78+0 +3,01..1 B-2
Dirichlet FM (Stark et al., 2024)
2.68 10.10..01 E-2
Dirichlet FDM (ours)
2.59 o.
Table 4. Evaluation of transcription profile guided promoter DNA sequence design of different models.
5.3. Spatiotemperal Data Generation
In this section, we evaluate our model on spatiotemporal data sampling tasks, both with and without guidance. Specif- ically, we consider two scenarios: trajectory sampling for dynamical systems and video generation.
5.2.1. SIMPLEX DIMENSION WITHOUT GUIDANCE
We first evaluate the performance of FM and FDM in non- guided simple generation task. The data is sampled from uniform Dirichlet distribution with a sequence length of 1 = 4 and K" = 40 categories. We compare the TV distance and KI. divergence between the generated distribution and the target distribution on the test dataset. The results in Table and Table 14 Appendix C.I demonstrate that FDM outperforms FM in generating the simple sequential categorical data.
5.3.1. TRAJECTORY SAMPLING FOR DYNAMICAL SYSTEMS
Sampling trajectories for dynamical systems under event guidance is crucial for understanding and predicting the climate and beyond (Perkins & Alexander, 2013; Mosavi et al., 2018; Hochman et al., 2019). Finzi et al. (2023) develop DM for sampling these events.
In this experiment, we compare FDM against FM and DM from (Finzi et al., 2023) on the Lorenz and FitzHugh- Nagumo dynamical systems (Farazmand & Sapsis, 2019); the details of these systems are provided in Appendix B.1. We test sampling trajectories from these systems with and without event guidance. A trajectory, either from a dataset sampled, is discrete time series of vectors concate- nated into 21 = (c((T,))] E Rddd where M is the number of time steps and d is the dimension of the sys- tem. Following (Finzi et al., 2023), an event E is set of trajectories characterized by some event constraint; for ex- ample, E = [x) C(x1) > 0}, where the event constraint function C' Md + R smooth. The challenge of this experiment is to sample trajectories in E when C is only known after the models have been trained. The detailed sampling procedure using DM can be found in (Finzi et al., 2023).
Method
TV Distance
Time (w/iter)
Linear FM
0.12 0o0s
0.10
Lincar FDM
0.16 0.08 1..00.000 .05
0,10+ 0044
Dirichlet FM
0.10
Dirichlet FDM
0.07 +00.04
0.16
Table TV distances between the generated and target distribu- tions.
5.2.2. PROMOTER DNA SEQUENCE DESIGN WITH GUIDANCE
We further evaluate the ability of FM and FDM in training generative models for designing DNA promoter sequences guided by a desired promoter profile. We train the models guided by profile by providing it as additional input to the vector field and evaluate generated sequences using mean- squared error (MSH) between their predicted and original regulatory activity, as determined by SEI (Chen et al., 2022). We include the discrete DM (Albergo al., 2023) and the language model (Stark al., 2024) for comparison in Ta- ble 4. For this task, we use dataset of 100,000 promoter sequences with 1024 base pairs extracted from a database of human promoters (Hon et al., 2017). See Appendix B.2 for more details about the dataset. The results confirm that FDM improves FM in training guided models for categorical data generation.
The event wenttgguiee sampling procedure from (Finzi et al., 2023) uses Tweedie's formula (Robbins, 1992; Efron, 2011 1), which requires the score function V log (ax). Since FM and FDM are not trained to approximate V log Pr P<(x)) di- rectly, we derive approximation formula using the learned vector field v, (ax, 0). Applying Lemma I of (Lipman et al., 2023) to the probability flow ODE (Song et al., 2020), the evolution of px(2:) satisfies:
1
(x) = -f 1-x +
2 g02(1 (1-t)-)Vooop logg (19)
where f the drift term and g the noise coefficient.
7

---

Improving Flow Matching by Aligning Flow Divergence
Rearranging equation (19), we express V p,(r)) in terms of (z)),, then approximate (x)))) by (2, 0).
Diffusion
FM
FDM
Dormand-
4
10a
We use the events defined in (Finzi et al., 2023) for our experiments. The event for the Lorenz system is when a trajectory stays on one arm of the chaotic attractor. This is characterized by C(x) = 0.6 - || - =||| > O, where F the Fourier transform over trajectory time T. || -I1 is the 1 I-norm summing over both over the frequency magnitudes and the three dimensions of 2T)), and I is the average of (T) over T. For the FitzHugh-Naguna system, the event is neuron spiking, which is characterized by C(x)== [ax + x2(r)]/2 2.5 > 0.
2
10-²2
0
4
100
2
10 2
0
4
10a
2
10-2
We compare the models' ability to generate trajectories ac- cording to p(z1) and p(æ1 (E) by computing test set of trajectories using the Dormand-Pring ODE solver (Dor- mand & Prince, 1980) and sampling trajectories using each model. Table 5 presents the TV distance between the model and the distributions. From the result, we observe that FDM achieves the lowest TV distance for every distribution. The TV distance of FDM is smaller than that of FM, which em- pirically demonstrates that the divergence mismatch has a significant effect on the error (2x:).
0
-1
-2
2
(a) Lorenz
(b) FizzHuuh-Naguma
Figure 5. Histograms of the constraint value C(æ1) where æn is an event trajectory computed by the Dormand PPincee ODE solver or sampled from the model with event guidance. The unguided sampling histograms are shown in Appendix C.2.
Lorenz
FirzHogh-Nagumo
Model
NLL(w) ) (4)
P(E)
NLL(z,) (U)
p(E)
Furthermore, this shows that the proposed loss CFDM effec- tively reduces the mismatch. This mismatch reduction also enables FDM to attain the lowest negative log-likellng ligeliiel (NLL) estimates. Table 6 shows the mean NLL over trajecto- ries and trajectory dimension with respect to to p( )) while Fig. 5 compares the histograms of the event constraint value of each event trajectory. Importantly, these improvements of
Donnand-Prince
0.197
0.035
Diffusion
-7.052
0.200
-7.3 37365
0.032
FM
-13.190
0.199
-13.942
0.034
FDM
-14.361
0.200
-14,408
0.033
Table 6. NLLs averaged over trajectories and trajectory dimension with respect to the trajectory distribution p(z:), ). and the likelihood of the user-defined event estimated by the proportion of trajectories contained in event E sampled from the model without guidance. Here, the Diffusion follows from (Finzi et al., 2023). FM and FDM are based on our own implementation.
FDM do not trade off with its accuracy in estimating p( E)) When p( is estimated based on the proportion of sampled trajectories that fall within E, all the models are comparable. Table 7 reports the KL. divergence between the histograms of the event constraint value ((x1) ) for event trajectories 21 from the dataset computed by an ODE solver and those sampled with event guidance from the models. The results show that our FDM consistently outperforms both the FM and Diffusion models.
Lorenz
Fitzllugh Nagiumo
Model
p(211)
p(=]|E)
p(z))
p(m]|
Diffusion
0.0056
0.2774
0.0260
0,3011
FM
0,0081
0,2560
0,0250
0.3468
HDMrours)
0.0049
0.2084
Table 7. KL divergence between the histograms of the event con- straint value C(æ1) for event trajectories æi in the dataset of trajectories computed by an ODE solver and event trajectories sampled with event guidance from the models.
Lovenz
Fitzlugh Nagumo
Model
p(w1) (4)
p( 0 (L)
(L)
p(#1 (
Diffusion
0.0314
0,1001
0.1192
FM
0.0348
0.0972
0,0314
0.2164
FDMcuu))
0.0306
0.0914
0,0266
0.1168
Table TV distances of the models from the trajectory distribution p(æ1) and from the distribution conditioned on event p(æ11/ | F.). Here, Diffusion results follow from (Finzi et al., 2023), while FM and FDM are based on our implementation, which builds the code provided by Finzi et al. (2023).
We train a latent FM (Davtyan et al., 2023) and a latent FDM for video prediction. We utilize a pre-trained VQGAN (Esser et al., 2021) to encode (resp. decode) each frame of the video to (resp. from) the latent space. We train the models using the latent state t and t r, where is randomly selected from [2,... .). by providing them as additional input guidance to the vector field at t > C, where C is a positive integer. At inference time, we use the frames at time = 0 to = C' of a video as the guidance and then utilize flow matching to predict the frames after C.
5.3.2. GENERATIVE MODELING FOR VIDEOS
We aim to show how FDM pushes the boundary of FM performance for sequential data generation in latent space.
8

---

Improving Flow Matching by Aligning Flow Divergence
We consider the human motion dataset- KTH (Schuldt et al.,
2004) and BAIR Robot Pushing dataset (Ebert al., 2017). We follow the experimental setup of (Davtyan et al., 2023); see Appendix B.3 for details. To evaluate the generated samples, we compute the Fréchet video distance (FVD) (Unterthiner et al., 2018) and peak signal-to-noise ratio (PSNR) (Huynh-Thu Ghanbari, 2008).
Walking; a-FM
KTH Dataset: For KTH, we use the first 10 frames guidance and predict the next 30 frames. The results in Table 8 indicate that FDM enhances latent FM for temporal data generation. Furthermore, Fig. 6 presents illustrative cases showing that our FDM consistently maintains high visual quality throughout the video, whereas the FM model exhibits noticeable degradation in later frames, including loss of fine motion details, missing body parts, and motion failure.
Walking; a-FDM
Boxing; b-FM
BAIR Dataset: For BAIR, we predict 15 future frames based on a single initial frame, with each frame having a resolution of 64 x 64 pixels. Because the highly stochas- tic motion in the BAIR dataset, following (Davtyan et al., 2023), we generate 100 samples per test video each condi- tioned the same initial frame and compute metrics 100 x 256 generated samples against 256 randomly selected test videos. To highlight the effectiveness of FDM, we omit the frame refinement step used in (Davtyan et al., 2023). As mentioned in (Davtyan et al., 2023), many models for the BAIR task are computationally expensive, whereas latent FM achieves a favorable trade-off between FVD and compu- tational cost. Our approach further improves latent FM with acceptable additional computational overhead, shown in Table 9.
Boxing: b-FDM
Hand Waving; c-FM
Hand Waving; c-FDM
Figure 6. Samples on KTH human motion dataset frame 13, 27, 40 from left right generated by latent FM (a-FM, b-FM, c-FM) and latent FDM (a-FDM, b-FDM, c-FDM).
We notice that the experiments in Chen et al. (2024) achieve very impressive results for video generation, and it is an interesting future direction to integrate our approach into their framework.
6. Concluding Remarks
In this paper, we have developed a new upper bound for the gap between learned and ground-trute probability paths using FM. Our new error bound shows that FM can be im- proved by ensuring the divergences of the vector fields are in proximity. To achieve this, we derive a new conditional divergence loss with computational efficiency. Our new training approach flow and divergence matching signiff cantly improves FM various challenging tasks. There are several avenues for future work. A particularly intriguing di- rection to develop a computationally efficient method for controlling the KL divergence, for example by integrating deep equilibrium models (Bai et al., 2019) into our frame- work similar how prior works have incorporated them into diffusion bscore-based) models (Huang al., 2024; Bai & Melas- Kariazi,i,zi, 2024). This remains open problem and an important avenue for future research. Moreover, explor- ing our approach in the Schrödinger bridge setting (Tong et al., 2024) is also an interesting problem.
Method
FVD())
PSNR(1)
Time(s/iter)
SRVP (Franceschi et al. 2020)
222
29.7
SLAMM (Akan al., 2021)
228
29.4
Latent FM (Davtyan et al., 2023)
180
30.4
0.18
Latent FDM (ouns)
158.5
4-5
31.2
0.27
Table KTH dataset evaluation. The evaluation protocol is to predict the next 30 frames given the first frames.
Method
FVD(L)
MEMIGB)
Timechoum))
THYDD-AN-FF 2020)
Video Tranffrmmrr /Weissenbon al., 2019)
512
336
LVT (Raihimov al., 2000)
126
128
RaMViD (Uiftusion) (Hlippe
Latent FM [Davtyan 2023)
24.2
(oums)
4.5
Table BAIR datasct evaluation. We adopt the standard evalua- tion setup, where the model predicts 15 future frames conditioned single initial frame. MEM stands for peak memory footprint.

---

Improving Flow Matching by Aligning Flow Divergence
Acknowledgement
Bai, S., Kolter, J. Z.,, and Koltun V. Deep equilibrium mod-
els. Advances in neural information processing systems, 32, 2019.
This material is based on research sponsored by NSF grants DMS-2152762, DMS-2208361, DMS-2219956, and DMS- 2436344, and DOE grants DE-SC0023490, DE-SC0025589, and DE-SC0025801. HZ acknowledges the support from the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research research grant DOE-FOA-2493 "Data-intensive scientific machine learn- ing". under contract DE-ACO2-06CHI 1357 at Argonne Na- tional Laboratory.
Bai, X. and Melas-Kyriazi L. Fixed point diffusion models. Proceedings of the IEEE/CVE Conference on Com- puter Vision and Pattern Recognition, pp. 9430-9440, 2024.
Baker, J. Huang, Y., Wang, S.-H., Pasini, M. L., Bertozzi, A. L.,, and Wang, B. Stabilized e(n)- equivariant graph neural networks-assisted generative models, 2024. URL https: https::/openreviiee openrevienrreeeeeeeeeeeeee .eet:/ forum? Siaummid=Neiin iF FKLBB.
Impact Statement
This work presents a new theoretical bound on the gap be- tween the exact and learned probability paths using flow matching. The new theoretical bound informs the design of a new efficient training objective to improve flow match- ing. Our work directly contributes to advancing doweb generative modeling. Flow-based models have achieved remarkable results in climate modeling and molecular mod- eling. By developing new theoretical understandings and fundamental algorithms with performance guarantees, we expect our work will advance climate modeling and molecu- lar sciences using generative models. Our work contributes to basic research, and we do not see potential ethical con- cerns or negative societal impact beyond the current AI.
Bao, F,, Nie, S., Xue, K., Cao, Li, Su, H., and Zhu, J. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 22669-226799 2023.
Bose, J., Akhound-Sadegle T., Huguet, G., FATRAS, K., crotorBrrooks, J., Liu, C.-H., Nica, A. C., Korablyov, M., Bronstein, M. M., and Tong, A. SE(3)-stochastic flow matching for protein backbone generation. In The Twelfth International Conference on Learning Representations, 2024. URL https: htap:::://ppereeee ppwp/eveeeeeeeeee//////ee net forum? id=kJFTH23hXb.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Moritz, P., Paszke, A., VanderPlas, J., Wanderman- Milne, S., and Zhang, Q. Jax: Autograd and xla. GitHlub repository, 2018.
References
Abadi, M., Barham, P.. Chen, J.. Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., al. TensorFlow: a system for large-scale machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16), pp. 265-283, 2016.
Chen, K. M., Wong, A. K., Troyanskaya, O. G., and Zhou, J. A sequence-based global map of regulatory activity for deciphering human genetics. Nature genetics, 54(7): 940-949,10 2022.
Akan, A. K., Erdem, E., Erdem, A., and Güney, F. Slamp: Stochastic latent appearance and motion prediction. In Proceedings of the IEEE/CVF International Conference on Computer Vision pp. 14728-14737 2021.
Chen, R. T., Rubanova, Y.. Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Akiba, T., Sano, S., Yanase, T., Ohta, T., and Koyama, M. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2019.
Chen, R. T. Q. and Lipman, Y. Flow matching on general geometries. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview net woo.mooddrgooooiii
Chen, Y., Goldstein, M., Hua, M., Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Probabilistic forecast- ing with stochastic interpolants and föllmer processes. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A.. Oliver, N., Scarlett, J., and Berkenkamp. F. (eds.), Pro- ceedings of the 41st International Conference on Mu- chine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 6728 -7556 PMLR, 21-27 Jul 2024. URLhttps: httpsss///ccceeeeeeeeee mlr press/ v235/ chen24n html.
Albergo, M. S. and Vanden- E. Building normaliz- ing flows with stochastic interpolants. The Eleventh International Conference on Learning Representations, 2023. URL https:// http:///ppeeeeew.ee.../// forum? id=1i7geBbCRIt.
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08777 2023.
10

---

Improving Flow Matching by Aligning Flow Divergence
Davryan, A., Sameni, S., and Favaro, P. Efficient video prediction via sparsely conditioned flow matching. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 23263 222744, 2023.
Hochman, A., Alpert, P., Harpaz, T., Saaroni, H., and Mes- sori, G. A new dynamical systems perspective at- mospheric predictability: Eastern mediterranean weather regimes as a case study, Science advances, 5(6):eaau0936, 2019.
Delecki, H., Schlichting, M. R., Arief, M., Corso, M., and Kochenderfer, M. Diffusion-based failure sampling for cyber-physical sys- tems. arXiv preprint arXiv:2406. 14761. 2024.
A.,
Hon, C.-C., Ramilowski, J. A., Harshbarger, J.. Bertin, N., Rackham, O. J., Gough, J., Denisenko, E., Schmeier, S., Poulsen, T. M., Severin, J., al. An atlas of human long non-coding mas with accurate 5' ends. Nature, 543 (7644):199-204, 2017.
Dormand, J. R. and Prince, P. J. A family of embedded runge-kutta formulae. Journal of computational and ap- plied mathematics, ((1):19-26, 1980.
Höppe, T., Mehrjou, A., Bauer, S., Nielsen, D., and Dittadi, A. Diffusion models for video prediction and infilling. arXiv preprint arXiv:2206.076690 2022.
Ebert, F. Finn, C., Lee, A. X., and Levine, S. Self- supervised visual planning with temporal skip connec- tions. CoRL, 12(16):23, 2017.
Hua, X., Ahmad, R., Blanchet, J., and Cai, W. Acceler- ated sampling of rare events using neural network bias potential. arXiv preprint vrXiv:2401.06928 2024.
Efron, B. Tweedie's s formula and selection bias. Journal of the American Statistical Association, 106(496):1602- 1614, 2011.
Huang, Y., Wang, Q., Onwunta, A., and Wang, B. Efficient score matching with deep equilibrium layers. In The Twelfth International Conference on Learning Represen- tations, 2024. URL https: oppnrevien. net/ Forum?id-Jre -inn-ddddJJaLAAAMM.
Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Com- munications in 183):1059-1076, 1989.
Falcon, W. Pytorch lightning. GitHub. Note: wwtp.///iiiiinn///iiiiiiiiiiia/e
2019.
and Computation,
Farazmand, M. and Sapsis, T. P. Extreme events: Mecha- nisms and prediction. Applied Mechanics Reviews, 71(5): 050801, 2019.
Huynh- Thu,, Q. and Ghanbari, M. Scope of validity of psnr in image/video quality assessment. Electronics letters, 44 (13):800-801, 2008
Finzi, M. A., Boral, A., Wilson, A. G., Sha, F,, and Zepeda- Núñez, L. User-defined event sampling and uncertainty quantification in diffusion models for physical dynam- ical systems. In International Conference on Machine Learning, pp. 10136-10152. PMLR, 2023.
Jing, B., Berger, B., and Jaakkola, T. Alphafold meets flow matching for generating protein ensembles. In NeurIPS 2023 Generative Al and Biology (GenBio) Workshop, 2023. URL https: /tpp///ppeeeee.ee//////e net/ net/forunt forum? i id=yDcebRgOfF.
FitzHugh, R. Impulses and physiological states in theoret- ical models of nerve membrane. Biophysical journal, (6):445 466, 1961.
Krizhevsky, A., Hinton, G., al. Learning multiple layers of features from tiny images. 2009.
Franceschi, J.-Y., Delasalles, E., Chen, M., Lamprier, S., and Gallinari, P. Stochastic latent residual video prediction. In International Conference on Machine Learning, pp. 3233-3246. PMLR, 2020.
Lai, C.-H., Takida, Y., Murata, N., Uesaka, T., Mitsufuji, Y., and Ermon, S. Fp-diffusion: Improving score-based dif- fusion models by enforcing the underlying score fokker- planck equation. In International Conference on Machine Learning, pp. 18365-18399 PMLR, 2023.
Frans, K., Hafner, D., Levine, S., and Abbeel, P. One step diffusion via shortcut models. arXiv preprint arXiv:2410. 12557, 2024.
Li, L., Carver, R., Lopez-Gomezz I., Sha, F., and Anderson, J. Generative emulation of weather forecast ensembles with diffusion models. Science Advances, 10(13)eaak4489, 2024.
Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D. Ffjord: Free-form continuous dy- namics for scalable reversible generative models. arXiv preprint / 2018.
Lipman, Y., Chen, R. T. Q.. Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The

---

Improving Flow Matching by Aligning Flow Divergence
Eleventh International Conference on Learning Represen- fations, 2023. URL https::////c http//pppnrreeweeeeeeeeeee net./ Forum?id= FidddPP P...d-Pgg.D..C..
Schuldt, C., Laptev, I., and Caputo, B. Recognizing human actions: local approach. In Proceedings of the 17th International Conference Pattern Recognition, 2004. ICPR 2004., volume 3, pp. 32-36. IEEE, 2004.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023. URL https: ppppeeeeee..eeeeee net/
Seis, C. A quantitative theory for the continuity equation. In Annales de l'Institut Henri Poincaré C, Analyse non linéaire, volume 34, pp. 1837-1850, Elsevier, 2017.
Shiraki, T., Kondo, S., Katayama, S., Waki, K., Kasukawa, T., Kawaji, H., Kodzius, R., Watahiki, A., Nakamura, M., Arakawa, T., et al. Cap analysis gene expression for eigh-throughput analysis of transcriptional starting point and identification of promoter usage. Proceedings the National Academy of Sciences, 100(26):1776-117781 2003.
Lorenz, E. N, Deterministic nonperiodic flow. Journal of atmospheric sciences, 20(2):130-144 1963.
Lu, C., Zheng, K., Bao, F., Chen, J., Li, C., and Zhu, J. Maximum likelihood training for score-basee bared diffusion odes by high order denoising score matching. In Inter- national Conference on Machine Learning, pp. 14429- 14460. PMLR, 2022.
Song, Y. and Dhariwal, P. Improved techniques for training consistency models. arXiv preprint arXiv:2310. 14189, 2023.
Luc, P., Clark, A., Dieleman, S., Casas, D. d. L., Doron, Y., Cassirer, A., and Simonyan, K. Transformation-ben adversarial video prediction on large-scale data. arXiv preprint 20t1::0000000035, 2020.
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 13456, 2020.
Meng, C., Song, Y., Li, W. and Ermon, S. Estimating high order gradients of the data distribution by denoising. Advances in Neural Information Processing Systems, 34: 25359-25369, 2021.
Song, Y., Durkan, C., Murray, I., and Ermon, S. Maxi- mum likelihood training of score-based diffusion models. Advances neural information processing systems, 34: 1415-1428, 2021.
Mosavi, A., Ozturk, P., and Chau, K.-w. Flood prediction using machine learning models: Literature review. Water, 10(1 ((lll):5530 2018.
Stark, H., Jing, B., Wang, C., Corso, G., Berger, B., Barzi- lay, R., and Jaakkola, T. Dirichlet flow matching with applications to dna sequence design. arXiv preprint 2.XX::2402.0841, 2024.
Nagumo, J., Arimoto, S., and Yoshizawa, S. An active pulse transmission line simulating nerve axon. Proceedings of the IRE, 20(10)2001120700 1962.
Paszke, A., Gross, S., Chintala, S., Chanan, G.. Yang, E., Veiito, Z., Lin, Z., Desmaison, A., Antiga, I., and Lerer, A. Automatic differentiation in pytorch. NIPS- W, 2017. URL https: https://openreviii net/pdE? id=BJJsrmfCa
Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L.., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulation- free schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics, pp. 1279-1287 PMLR, 2024.
Perkins, S. E. and Alexander, L. V. On the measurement of heat waves. Journal climate, 26(13) (6613))4500-4517, 2013.
Unterthiner, T., Van Steenkiste, S., Kurach, K., Marinier, R., Michalski, M., and Gelly, S. Towards accurate generative models of video: A new metric & challenges. arXiv preprint
Petersen, M., Roig, G., and Covino. Dynamicsdiffusion: Generating and rare event sampling of molecular dynamic trajectories using diffusion models. In NeurlPS 2023 Al for Science Workshop.
2018.
Villani, C. al. Optimal transport: old and new, volume 338. Springer, 2009.
Rakhimov, R., Volkhonskiy, D., Artemov, A., Zorin, D., and Burnaev, E. Latent video transformer. arXiv preprint arXiv:2006.1077 2020.
Wan, Z. Y., Baptista, R., Boral, A., Chen, Y.-F., Anderson. J., Sha, F., and Zepeda L. Debias coarsely, sample conditionally: Statistical downscaling through optimal transport and probabilistic diffusion models. Advances in Neural Information Processing Systems, 36:47749 47763, 2023.
Robbins, H. E. An empirical bayes approach to statistics. In Breakthroughs in Statistics: Foundations and basic theory, pp. 394. Springer, 1992.
12

---

Improving Flow Matching by Aligning Flow Divergence
Weissenborn, D., Täckström, O., and Uszkoreit, J. Scal- ing autoregressive video models. arXiv preprint arXiv: 02634, 2019.
Yim, J., Stärk, H., Corso, G., Jing, B., Barzilay, R., and Jaakkola, S. Diffusion models in protein structure and docking. Wiley Interdisciplinary Reviews: Computational Molecular Science, 14(2) 2024.
13

---

Improving Flow Matching by Aligning Flow Divergence
Appendix for Improving Flow Matching by Aligning Flow Divergence
A. Missing Proofs
Proposition 3.1. Pt Pr satisfies the following PDE:
O3E: + V ((rvs) Lr, <o(2)) =
where
L = -Pt V (Ue vz) + (ut - vz) V log
(8)
Proof of Proposition 3.1. For simplicity, denote 0 är by 0j. From the continuity equations and 6, we have:
d2c; = O.p.- 0.p.
V w6 U Vp: - (. (e((vou vv)) - Vpr
(20)
Ph V uu) fo (V . mz
Vpt + Vp.
Pa V (u: vz)) (V v.),e (uu v.) Vpt V4
Rewriting it, find:
8ee + V (E(22) -pc(V (wt vt)) Po(u vz) V log Pt
(21)
Let us define L2 = -p ( (
(u 01)) - pi(u: - vz) V log Pt- This gives the following PDE for 4 with the initial
condition Eo = Pu Pu = 0:
0,et V. (sso.) L,, <o(2)) = 0.
(22)
Corollary 3.2. For any E [0,1], 11, the enror satisfies
c(()((x)) (z) = Lo(o)())) (z)
where (e(xx) the flow induced by the vector field vx(x) in a similar way as that equation (3), detVo. (x) denotes the determinant of the Jacobian matrix (d)( and L.g is defined in Proposition 3.1.
Proof of Corollary 3.2. Let or denote the flow of the vector field w.. i.e.
0, v. (1) Φb(x)) = x.
(23)
Using Duhamel's formula (refer to (Seis, 2017)), we have the following formula for cz:
ec ())) det (o)( = =e(()) + / (z))(())) det (z) do = 1/' (o(())))) det (x)
(24)
Theorem 3.3. Under some common mild assumptions adopted in (Lu et al., 2022; Lipman al., 2023; Albergo et al., 2023), the following inequality holds for any [0,1]: 1]:
TV pe) < LDM(
(11)
Specifically, (x)) = (xx) when LDM zero.
14

---

Improving Flow Matching by Aligning Flow Divergence
Proof of Theorem 3.3. Note that the total variation distance is defined as:
1 (V Po) - _ |p((x)) (x())) dæ = _// |e(((t) dæ
(25)
Using the change of variables twice and applying the formula in Corollary 3.2, obtain:
1
/ [a((r)] dx
(p), = 2
2 (((0(())) (() d
1
1
/ (((((())) det Vor(z) do
2
(26)
1
/ I (z) det (2.(x) da
2
1
/ / (.( det (x) ds dx
2
1 [ / |Lo(co)] doc ds
2
Substituting the expression for Lg we see that:
(-' / pe( (V. (ug + pp v2) V log Pe
(p -
dx ds
Jo
5 Ep.
V. (us - Us) + (u. - v.) Vlog ps ds
(27)
T
Ep.
V. (ue - vz) +wu - v6) V Vlog Pt dt
(0).
This completes the proof.
Theorem 4.1. We have the following inequality:
(D) s (CDM(0).
(15)
Furthermore, we have:
TV(p.,))) < (CDM(())
(16)
for any t 10.
Proof of Theorem 4.1. From equation (13), show that:
Pe V Wz + V po)
V (prur)
= / v (p=(a)))r) dx
(28)
f (pi(tolo))) - (x]221) (p((2))))) p(æ)) dx
(x(
15

---

Improving Flow Matching by Aligning Flow Divergence
On the other hand, we have
Pr V oz + Vs loggt) Pr
2s + 0z Vpt
(/ Pt p=(alz1))((z)) dxy v 2s + 7 (ʃ pal (2)(z)z)))1 )p(x) dx
/ (p((a)m)) m))())) daa / (m Va( da
(29)
/ (p( 1. p(( dx1.
Combining equation (28) with equation (29), we deduce that:
Pe (V u1 + u, Vlogp:) Pa (V vi + vs Voggp:)
(pl(ol))) (1(a))))) + (x//a)) (pp(ar) (( p(m)) dmi
(f (p.(zlz,) u v Vp.(= day
(30)
/ (V (x/xx) (1)p((2 /
(u](z)z1)) - vx(x)) (=(())))) (x) )p(x)) dæ
/ Vv (,(x)x))- v v,) + (x|[, - vx(x)) . (1l - ) pr x pp(x)) dæ
Now from the definitions of (f) and (f), we deduce that
T / pt V. (ue ve) (ue - v.) V VlogPt Pe dx dt
(D)
1-²0 /
1
pe th + us Voggi: - Pc (V 22 + v Pl) dx dt
1.") / / (V - vz) + (x|x1)) -
log (x)xx)) (((l))))p(a dæ dx dt
("' (v (((x)x))- - V v.) + (x(xx) - vi(x)) V log (x/ (x<x)))))((( dæ dx dt
(CDM(O).
(31)
B. Experiments Details
B.1. Trajectory Sampling for Dynamical Systems
For this experiment, repeatedly the Dormand-Prin Prince ODE solver with an absolute tolerance 1.4 x 10 and relative tolerance 1 x 10-6
Lorenz
The Lorenz system (Lorenz, 1963) is a chaotic dynamical system given by
o - x11)
12
F(x)
x1 (p 333) 32
T132 ,333
Following (Finzi et al., 2023), we set or = 10, p = 28 and 3= 8/3, and we used a scaled version of the Lorenz system to bound the system components 2; to [-3, 3] for i E [1,2,3] while preserving the original dynamics. The scaled system is given by F(x) F(20x)/20)
16

---

Improving Flow Matching by Aligning Flow Divergence
FitzHuugh-Nagumo The HizZHugh-Nagume system (FitzHugh, 1961; Nagumo et al., 1962) a dynamical system modeling an excitable neuron and is given by
4; z: - -x;)(x;; Ui
k A15a xz)
i = b,30 Cili
for i E (1,2). Following (Farazmand & Sapsis, 2019; Finzi et al., 2023), the parameters are set as follows: @1 02 -0.025794, b1 = 0.0065, b2 = 0.0135, c = = 0.2, k = 0.128, and Asj dij where ó the Kronecker delta.
Trajectory Dataset Construction Trajectories for the dataset are computed using the ODE solver. The trajectories' initial conditions sampled from Gaussian distributions (O -Noioonsss-(O.... I) for Lorenz, and (0, (0.2)²I) I) for FitzHugh-Nagumo. Each trajectory has 60 consecutive and evenly spaced time steps, where the first time step occurs after some trajectory "burn-in" time to allow the system to reach its stationary trajectory distribution. The first 30 and 250 time steps computed by the ODE solver are "burn-in" for Lorenz and FitzHugh-Nagumo, respectively. The time step sizes are 0.1 and 6.0, respectively.
Model Hyperparameters and Training All the models used the same UNet architecture as in (Finzi et al., 2023), and we used a variance exploding schedule (Song et al., 2020). We train the models on a training set of 32,000 trajectories computed by the ODE solver using Adam for 2,000 epochs with a batch size of 500. For FM and FDM, we also used an exponential decay learning rate scheduler with decay rate of 0.995. The initial learning rate for the diffusion model and FM was 10-4 The learning rate and regularization coefficients for FDM were tuned using Optuna (Akiba et al., 2019) for the lowest CFM loss produced by the EMA parameters and are given in Table 10. We sampled the times for the diffusion and CFM loss on a shifted grid following (Finzi et al., 2023).
Dynamical system
Learning rate A1
X2
Lorenz
0.000796
1
0.000385
FitzHugh-Nagumo
0.000245
1
0.00552
Table 10. Learning rate and regularization coefficient used to train FDM for Lorenz and FitzHugh-Nagumo dynamical systems.
We evaluated the models with the exponential moving average (EMA) of the parameters with a 2,000 epoch period.
Loss Weighting Functions The loss of the diffusion model equation (7) of (Song et al., 2020) where we used A(4) = of as the weighting. For both FM and FDM, the term LCFM in their loss was weighted by 1/(o(_;)². The term CCDM the loss of FDM weighted by o where M 60 is the number of trajectory time steps and d is the dimension of the dynamical system.
We estimated the divergence of FDM with respect to its trajectory input using the Hutchinson tracer estimator (Hutchinson, 1989; Grathwohl et al., 2018) where the noise vector is sampled from (0, I).
Estimating the Divergence
Likelihood Estimation We computed a test set of 32,000 trajectories using the ODE solver and evaluated their log- likelihood using the continuous change-of-variables formula from (Grathwohl al., 2018) with the ODE solver. Table was produced by computing the mean ligeeiioldood over the trajectories and their dimension.
B.2. DNA Sequence Generation
In this task, the model approximates classifier
p(an ar, 0) Pr(æ ( )p(x))
(32)
(xc)
instead of directly approximating the vector field (x 0) Rv (x)). Then, it constructs a vector field based on the classifier
as follows:
(x, 0) = Σ u = =c)p(a; =,,2,,
(33)
17

---

p(æ1 ei(z) is not continuous in D. Improving Flow Matching by Aligning Flow Divergence
where K is the number of categories and the divergence term is given by
V (z 0) Σ (Vp(z) - e, 0), ur (2 - =ei) + p(2) e,(2, u1(20/2 = ei)
(34)
we directly learn V 0 0), it requires computing V [ue( x = ) = =, 0) for = 1, K which can be very expensive memory footprint and time consumption. Furthermore, notice that u (z/ = e;) is pre- vector field that is independent of parameters 0 and is V. (e( a e;). Thus, there is need learn it.
For p( = e, æ, 0), Appendix A of (Stark al., 2024) states that 0r (x, 0) approximates the vector field if p(x1 x 0) ideally approximates the classifier p( [1 |z).. Consider an ideal classifier p(æ1 = e ez) for class œ1 = c₁ then p(z) e, [x) =
belongs to class 21 else O. Let æ E D, where D is the domain of this classifier, then we have
Suppose Dj is the union of all the differentiable sub-domains of D, then Vap(2: c,(a)) 0 for x E D1.
Therefore, the remaining thing is to include [V>p((r; = ex(2)]] for x Dj the training objective. In practice, we train the classifier by empirically estimating the cross entropy based on the perturbed points with its corresponding initial data æ the class label. We can just assume any point in sufficiently small ball around such perturbed data point belongs to the same class 11 the classifier is differentiable inside this ball, then we penalize || VV ax, 0)l in training the model.
Promoter Data We use a dataset of 100,000 promoter sequences with 1,024 base pairs extracted from a database of human promoters (Hon et al., 2017). Each sequence has a CAGE signal (Shiraki et al., 2003) annotation available from the FANTOM5 promoter atlas, which indicates the likelihood of transcription initiation at each base pair. Sequences from chromosomes 8 and 9 used as a test set, and the rest for training.
Model Hyperparameters and Training We just follow the experimental setup of (Stark et al., 2024). For the simplex dimension toy experiment, we train all models for 450,000 steps with batch size of 512 to ensure that they have all converged and then evaluate the KL of the final step. For promoter design, we train for 200 epochs with a learning rate of x 10 and early stopping on the MSE on the validation set. We use 100 inference steps for generation. Table show how we set A₃ and A2 for divergence loss.
Tasks
Learning rate
A1
A2
Simplex Dimension
5 10
0.5
0.05
Promoter Design
5 10-4
0.01
Table 11. Learning rate and regularization coefficient used to train FDM for DNA sequence.
B.3. Generative Modeling for Videos
We follow the experimental setting and models used in (Davtyan al., 2023).
Architechture We use U-VIT (Bao et al., 2023) model the flow matching vector field and use VQGAN (Esser et al., 2021) to encode (resp. decode) each frame of the video to (resp. from) the latent space with the following configurations
Model Hyperparameters See Table 13.
C. Additional numerical results
C.1. Dirichlet Flow Matching
Table 14 shows the test KI. divergence of models for the simplex dimension toy experiment of DNA sequence generation.
18

---

Improving Flow Matching by Aligning Flow Divergence
Parameter
KTH
BAIR
embed_dim
4
n_embed
16384
16384
double z
False
False
z.channels
4
4
resolution
64
64
in_channels
3
3
out.ch
3
3
ch
128
128
ch.mult
[1,222.4][1.2222.4
num res blocks
2
2
attn_resolutions
[16]
[16]
dropout
0.0
0.0
disc conditional
False
False
disc wtr_iinchhannee
3
3
disc_start
20k
20k
disc_wet disc_weight
0.8
0.8
codebook wwiggtt
1.0
1.0
Table Parameters VQGAN the KTH dataset.
Hyperparameter
Values/Search Space
Iterations
300000
Batch size
[16. 32, 64]
Learning rate
[2e-4, 2e-5]
Learning rate scheduler
polynomial
Learning rate decay power
0.5
Weight decay rate
lc-12
\1, A2
[[0.5, 1e-2], [1, 1e-2]]
Table 13. Training hyperparameters of video prediction.
Method
KL Divergence
Linear FM
2.5 E-2
Lincar FDM
2.1 B-2
Dirichlet FM
1.8 Dirichlet FDM
1.2
Table 14, KL divergence the generated distribution to the target distribution.
C.2. Flow Matching for User-defined Events
Unguided Sampling Histograms: The histograms of the event constraint values for the trajectories sampled without guidance by each model are shown in Fig. 7.
KL Divergence Table 7 shows the KL divergence between the histogram distributions. For Lorenz, FDM's unguided sampling has the lowest KI. divergence, with the divergence of the diffusion model and FM being 0.0007 and 0.0032 larger. In guided sampling, the FM has a lower KL divergence than the diffusion model and FDM by about 0.02 and 0.05, respectively. For FitzHugh-Nagumo, the diffusion model has a lower KL divergence than FM and FDM by 0.002. In guided sampling, FDM attains the largest performance gap with a KL divergence of about 0.1 and 0.14 lower than the diffusion model and FM, respectively.
19

---

Improving Flow Matching by Aligning Flow Divergence
Diffusion
FM
FDM
Dormand-Prince
100
1
10-2
0
2
100
1
10
100
1
10-2
0
-2
0
2
(a) Lorenz
(b) FirzHugh-Nagumo
Figure Histograms event constraint C evaluated on the data and trajectories generated from models.
Lorenz
FitzHugh-Nagumo
Model
p(x1)
p(z, E))
p(x1)
p(x,
Diffusion
0.0056
0.2774
0.0260
0,3011
FM
0.0081
0.2560
0.0280
0.3468
FDM
0.0049
0.3045
0.0280
0.2084
Table 15. KL divergence between the histograms the event constraint value C(x1) for event trajectories æ1 in the dataset of trajectories computed by ODE solver and event trajectories sampled with event guidance from the models.
D. Efficient Squared Loss
We observe that the conditional divergence loss in equation (14) is an absolute-value objective, whose non-differentiaabiity at the origin removes smoothness and can be less efficient than squared-error alternatives. In practice, we replace with the following squared loss:
LCDM-2 (0) (mml))))) (V (x(w;) vz(x,,)) 0)) + (ut(w))1)) - vx(x) 0)) log)(( { (35)
Theorem D.1 (Upper Bound for LCDM). The squared loss (CD provides an upper bound:
(CDM(() S V LCDM)2(0)
(36)
Proof. Let
((x = (V (x(x)/1)) (x 0)) + (uu(())))) - (x(x)))) (zgg)
20

---

Improving Flow Matching by Aligning Flow Divergence
have
/ / (f ttl.i,ttppppaa (t)p:p/pppeepppplidddddd
CCDM(0)=
= / [ (0 (p()ddp()))ddda dt
(37)
[
(2c, æn, t)p((t)) (ttp.//pppppppppxddxxx
Cauchy-Schward Schwarz Ineq.
/ 2()
where we define p=(2) = p( æ1, t)p(t.
D.I. Efficient Squared Loss for High-dimensional Data
D.1.1. ESTIMATED OBJECTIVE VIA HUTCHINSON TRACE ESTIMATOR
Let d be the data dimension. The second-order term in equation (35) requires computing the trace of the full Jacobian of the vector regressor model 2y (at f), which typically incurs a computational time complexity of O(d2) and becomes impractical high-dimensional data. Following (Lu al., 2022; Lai al., 2023), this cost can reduced to ((d) by employing Hutchinson's 's trace estimator (Hutchinson, 1989) and automatic differentiation (Paszke et al., 2017) provided by general deep learning frameworks, requiring only single backpropagation pass.
For d-by-d matrix A, its trace can be unbiasedly estimated by
tr(A) A) Ep(e){e Ac] By(e) [E (A.e)]
where p(s) d-dimensional standard Gaussian. Then we proposed the following estimation:
Cust CDM-2 (0) (al. | (Vur(x)a))) e (2,
(38)
(u((r))))) (x, 0) E) (V log pog e) e)
where (x(xa) and log (x)laa) are already given pre-detined matrix and vector. The term (v, 0) can be efficiently computed by the jvp interface, such as torch. func. jvp in PyTorch or jax. jvp in JAX.
Theorem D.2 (Upper Bound for LCDM-2). The estimated squared loss C (f) provides upper bound:
LCDM22(0) I Lest (CDM2(()
(39)
Proof.
[[V ux(x)x))) V vx(x, 0)) + (u/(z)))) v1(2,0) ( (logg((<)))))
LCDM-2(0) =
p. [Bpces{ T (Vu(( (Valaa (o, 0))
+ (1 (2, 0)) & (logm((aaan)) |
T (Vu,( - (0, 0))
Cauchy- Schway looq.
(40)
+ + ((u((1/21) ve 0)) @ V [logp(()(z1)) 12
[|e (Vu((x))1)) (VVu.(()))) e - Vui(2, ()
= :ee
(u((z)=1) v2(2, 0) -) (V (Vlog)((z)z.))
cest (0D
21

---

Improving Flow Matching by Aligning Flow Divergence
1.1.2. STOP GRADIENT
In practice, a stop-gradient operation is applied on the ve( 0) in C (CDM.2(0) following common practice (Frans al., 2024; Lu et al., 2022; Song & Dhariwal, 2023). In our case, train the model by combining LCFM(9) and the conditional divergence matching loss discussed in Section 4 so the stop-gradient operation climinates the need for "double backprop- agation" through 2y (z, 0), making the training more efficient. So we define the squared efficient conditional divergence matching loss:
Ceff (0) = (en( | (Vu,(x) E (vve2 0) E)
(41)
(x/l))) - & sg( 0)) (e) (Vlogm(())a)) e) e)
where sg denotes stop-gradient operator, which prevents gradients from propagating to through term (, 0) in Coff CDD-22 Thus, optimizing the flow and divergence matching loss equation (17) requires only one extra backward pass compared to the baseline LCFM-
22


Reviews:
[{"review_id": "k1HvBF4vUK", "reviewer": "Reviewer_uLYh", "summary": "The paper seeks to use PDEs to construct a theoretical bound on flow matching, and improve upon it using said insight by adding a divergence mismatch to the loss term which improves upon the probability path. They construct experiments on simple generative examples, along with DNA sequence generation and video prediction and show that it improves upon existing methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.\n\n---- \n\n**Q1. I would like to see the performance of their method on some more diverse video datasets or standard image benchmarks for generative modeling (CIFAR, Imagenet).**\n\n**Response:** During the rebuttal period, we conducted experiments to showcase the efficacy of our proposed FDM in improving flow matching (FM) for these reviewer-mentioned tasks. We first compare the performance of FDM against the baseline conditional FM (CFM) for generative modeling of CIFAR10 here, and we will show the advantages of FDM for other video generation tasks in **Q2**.\n\nWe follow the experimental settings in Appendix E of the flow matching baseline paper [1] to train the model for CIFAR10 generation. Due to the time constraint, we only compare our FDM against CFM for CIFAR10 generation using the better-performing optimal transport (OT) path in the baseline paper [1]. The following results confirm the advantages of FDM.\n\n|  Method  |  NLL ↓   | FID ↓   |\n|-------------|------------|-----------|\n| CFM       |   2.99    |   6.35   |\n| **FDM (ours)**       |   **2.85**    | **5.62**    |\n\n[1] https://arxiv.org/pdf/2210.02747\n\n--- \n\n**Q2. The KTH dataset isn't used usually for generative modeling benchmarks due to its limited diversity in terms of samples. It works as a simple example but some more diverse motion datasets such as BAIR (as done in Davtyan et. al (2023)) could be more interesting.**\n\n**Response:** We appreciate your feedback. For the BAIR dataset, we predict 15 future frames based on a single initial frame, with each frame having a resolution of $64 \\times 64$ pixels. Due to the highly stochastic motion in the videos of the BAIR dataset, we evaluate the model as follows (following [2]): We randomly select 256 test videos, and generate 100 samples per test video where each model is conditioned on the same initial frame of the video. Finally, we compute the FVD for each model comparing the $256\\times 100$ generated samples against the 256 test videos. Due to time constraints, we omit training the frame refinement network, which operates independently of the main model but could potentially enhance sample quality. The result is as follows:\n\n| Method                  | FVD ↓           | Memory (GB) | Time (hours) |\n|-------------------------|------------------|-----------|----------------|\n| TriVD-GAN-FP            | 103              | 1024      | 280            |\n| Video Transformer       | 94               | 512       | 336            |\n| LVT                     | 126              | 128       | 48             |\n| RaMViD (Diffusion)      | 84               | 320       | 72             |\n| Latent FM           | 146              | 24.2      | 25             |\n| **Latent FDM (ours)**   | **123 ± 4.5**    | 35        | 36             |\n\nAs mentioned in [2], many models for the BAIR task are computationally expensive, whereas latent FM achieves a favorable trade-off between FVD and computational cost. Our approach further improves latent FM with acceptable additional computational overhead.\n\n[2] https://arxiv.org/abs/2211.14575\n\n---- \n\n**Q3. Minor typos (does not effect valuation) L_{FDM} is used in the graph of Figure 1, but caption reads L_{DM}**\n\n**Response:** Thank you. We have fixed this in the revised manuscript.\n\n\n------\n\nThank you for considering our rebuttal.", "claims_and_evidence": "Their claims seem to be sufficiently supported mathematically and experimentally. I would like to see the performance of their method on some more diverse video datasets or standard image benchmarks for generative modeling (CIFAR, Imagenet).", "methods_and_evaluation_criteria": "The non-toy experiments seem to be the DNA sequence generation and the KTH dataset. The KTH dataset isn't used usually for generative modeling benchmarks due to its limited diversity in terms of samples. It works as a simple example but some more diverse motion datasets such as BAIR (as done in Davtyan et. al (2023)) could be more interesting. Evaluation criteria seem to be the de-facto norm for all the experiments.", "theoretical_claims": "I checked section 3 for correctness, and the claims seem to be valid.", "experimental_designs_or_analyses": "Specific experiment soundness seem valid and include error bars for some of the experiments, which seems sufficient.", "relation_to_broader_scientific_literature": "The paper offers some key insights for flow-based generative modeling by refining the flow matching approach, which also offer enhanced likelihood estimation. They also demonstrate their strengths with applications in structured data such as dynamical systems, dna sequences, and videos.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "Minor typos (does not effect valuation)\nL_{FDM} is used in the graph of Figure 1, but caption reads L_{DM}", "questions_for_authors": "None.", "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": "", "extra_scores": {"ethical_review_concerns": "None"}}, {"review_id": "g9Ef2CzPQL", "reviewer": "Reviewer_247t", "summary": "The paper proposes a modification to the flow matching / stochastic interpolant loss so as to better control the total variation distance between the model and the target at the final time of sampling, motivated by the fact that the standard loss is not sufficient to control the KL divergence (based on some assumptions on the target). \n\nThe modification they propose is a sort of \"divergence matching\" loss whereby the divergence of the model vector field is trained to match the divergence of the ground truth vector field defined by the interpolant. \n\nThey test the method on synthetic datasets used in previous works, trajectory sampling, dna sequencing, and video forecasting.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.\n\n--- \n\n**Q1. At the end of page 6, … it is not clear that the equation at the bottom of page 6 is an exact relation, and would introduce extra errors.**\n\n**Response:** Thank you for your comment, which helped us identify areas where we can improve our presentation, and we have revised the paper accordingly.\n\nTo address your concern, we stress that our theoretical results do not depend on this relation for the learned vector field $\\mathbf{v}_t(x,\\theta)$. As the reviewer noted, this relation holds only for the ground-truth vector field $\\mathbf{u}\\_t$ and the score function of its corresponding probability flow $p_t$ (due to the uniqueness of solutions to the continuity equation). Our derivations rely solely on the continuity equation and initial conditions, not on this specific relationship, which aligns with the reviewer’s observation.\n\nOur original intent in introducing this relation was to define a flow-matching model with the vector field $\\mathbf{u}\\_t$ corresponding to the score-based diffusion model with drift term $\\bf f$ and noise coefficient $g$. Notice that approximating $\\mathbf{u}\\_t$ with a parameterized vector field $\\mathbf{v}\\_t(x,\\theta)$ is enough to generate data effectively. However, for the numerical results (e.g., conditional probability computations in Table 2), we require an approximation of the score function $\\nabla \\log p_t$. Instead of computing $\\nabla \\log \\hat p\\_t$, we substitute the learned vector field $\\mathbf{v}\\_t(x,\\theta)$ into the following relation to estimate the score:\n$$\\nabla \\log p\\_{1-t}(x) = 2\\frac{\\mathbf{u}\\_t(x)+{\\bf f}\\_{1-t}(x) }{g^2_{1-t}}.$$\nWhile this introduces an approximation error, it is a practical trade-off for computational feasibility.\n\n---\n\n**Q2. The bounds rely on an analytic form for the score $\\nabla\\log p_t(x)$ arising from solving the continuity equation up to time $t$. This quantity is not equivalent to the score arising from the model i.e. $\\hat s_t\\neq \\nabla \\log p_t$, nor is it something that can be written down by relating a learned velocity field to a score.**\n\n**Response:** Indeed, computing $L_{DM}$ requires $\\nabla\\log p_t(x)$, which is derived by solving the continuity equation. To address this challenge, we introduce its conditional variant $L_{CDM}$. For $L_{CDM}$, we only need access to $\\nabla\\log p_t(x|x_1)$. In practice, we design $p_t(x|x_1)$ such that this quantity is readily available by plugging in values, bypassing solving an ODE explicitly. This approach ensures computational feasibility while maintaining the integrity of our bounds; see Theorem 4.2.\n\n---\n\n**Q3. The method seems to perform better, but only marginally. Given the fact that loss is more expensive …, how does one measure the advantage of the technique properly?**\n\n**Response:** The improvements for all tasks except trajectory sampling for dynamical systems are way over twice the standard deviation. For dynamical systems, we further perform a t-test with the null hypothesis that the mean negative log-likelihoods (NLLs) computed by two distinct models are equal. We use 32,000 test trajectories to compute each mean NLL. The t-test confirms that the improvement in the likelihood estimation of FDM over FM is significant.\n\nOur approach introduces additional computational cost, but it gives better generation results and better likelihood estimation. For tasks requiring accurate likelihood estimation, our approach is especially more valuable compared to the baseline flow matching.\n\n---\n\n**Q4. How does the method compare for forecasting videos to [1]?**\n\n**Response:** Thank you. We have cited this paper in the revision.\n\nIndeed, the results in [1] are impressive. We only found the codes for Navier-Stokes and CIFAR at github.com/interpolants/forecasting, which is the reason we did not include this approach as a benchmark in our study.\n\n[1] https://arxiv.org/pdf/2403.13724\n\n---\n\n**Q5. The authors should probably also cite Liu et al, Flow Straight and Fast (2023) for its contributions to the flow matching literature as well.**\n\n**Response:** We have cited this paper in the revision. \n\n---\n\n**Q6. Memory footprint in computing divergence. Does the method work when replaced with a Hutchinson estimator?**\n\n**Response:** In our implementation, we use the Hutchinson estimator to approximate the divergence term; see Lines 307-309 in our submission. It works well and does not significantly raise the memory footprint, which consistently stays within 1.5×.\n\nFor the 2D tasks, we test both computing the Hutchinson divergence estimator and the exact one; both result in nearly identical results. We use the Hutchinson estimator in all the other experiments due to memory constraints, following [2].\n\n[2] https://arxiv.org/pdf/2206.08265\n\n---\n\nThank you for considering our rebuttal.", "claims_and_evidence": "The motivation for the project is clear. They seek to control a divergence metric for their loss for a deterministic, ODE-based flow model. Other work has only done this using higher-order derivatives than the ones they consider here. However, the related work can get control on the KL, while this work is only on the TV. The authors acknowledge this fact.\n\nOne issue with the application of the method in the experiments that I'd like the authors to clarify: At the end of page 6, the authors state \"we need to express the score function $\\nabla \\log p_t(x)$ learned by diffusion models in terms of the learned vector field $v_t (x, θ)$. This limits us to choosing a conditional probability path corresponding to an SDE with a known drift term $f$ and noise coefficient $g$.\" Note that this relation between the velocity field and the score *is only true if you have actually found the velocity field associated to the interpolant/probability path you specified*. If you have not learned this, then the equation relating $v_t$ to $s_t$ is not valid. *Moreover*, this score only corresponds to the actual $\\nabla \\log p_t$ if the velocity is exactly learned. So it is not clear that the equation at the bottom of page 6 is an exact relation, and would introduce extra errors.", "methods_and_evaluation_criteria": "The proposed benchmarks make sense. They are basically doing a set of controlled trials comparing the standard method to their modification. It would also be nice to see how this performs versus the broader literature, though.", "theoretical_claims": "I checked the correctness of the proofs. They seem like decently straightforward applications of Jensen and Young's inequality. \n\nOne thing I was looking for in particular but did not necessarily find: The bounds rely on an analytic form for the score $\\nabla \\log p_t(x)$ arising from solving the continuity equation up to time $t$. This quantity is *not* equivalent to the score arising from the model i.e. $\\hat s_t \\neq \\nabla \\log p_t$, nor is it something that can be written down by relating a learned velocity field to a score.", "experimental_designs_or_analyses": "There is a wide breadth of experiments based on existing applications of generative models that people have tried.  The method seems to preform better, but only marginally. A question, then, is: given the fact that loss is more expensive (i.e you have to compute and match divergences), how does one measure the advantage of the technique properly? Can\n\nIn addition, how does the method compare for forecasting videos to: https://arxiv.org/pdf/2403.13724 ? It seems they report much lower FIDs, but also claim that Latent FM gets lower FIDs as well. It is probably worth citing this relevant literature, which seems to excel at this task.", "relation_to_broader_scientific_literature": "The contribution is to describe an avenue for improving a very commonly used generative modeling paradigm. There is not existing work as far as I know on divergence matching this way.", "essential_references_not_discussed": "The authors should probably also cite Liu et al, Flow Straight and Fast (2023) for its contributions to the flow matching literature as well.", "other_strengths_and_weaknesses": "In general, it would be nice to have a better understanding of the memory footprint given the divergence you have to compute. Does the method work when replaced with a hutchinson estimator, or is it too noisy?", "other_comments_or_suggestions": "n/a", "questions_for_authors": "n/a", "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": "", "extra_scores": {}}, {"review_id": "lUnlBnKnYX", "reviewer": "Reviewer_rKjb", "summary": "The paper proposes a very simple KL loss combined with CFM loss to improve the training of flow-based models. The optimizing results are very general across different tasks with a basic improvement.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.\n\n---- \n\n**Q1. Lack of implementation details which may indicate unfair comparisons.**\n\n**Response:** In our submission, we have provided experimental and implementation details in Appendix B. To ensure fair comparisons with benchmarks, we use the same neural network architectures and model hyperparameters as the baseline methods. In particular, we use the same dataset and splitting, and employ the same search space for the batch size with the same training iterations.\n\n---\n\n**Q2. The organization of the contribution part should be improved.**\n\n**Response:** Thank you for your suggestion. We have made the contribution section more concise and bulleted the key points in the revised paper.\n\n---\n\n**Q3. Quantify the individual impact of the divergence loss term $L_{CDM}$ on performance.**\n\n**Response:** Our training loss is a weighted sum of $L_{CFM}$ and $L_{CDM}$ as shown in Equation (19) of our submission. Only using $L_{CFM}$ is the same as classic conditional flow matching, while, as we state in Lines 268-270 of our manuscript, directly minimizing $L_{CDM}$ cannot yield appealing results. Detailed explanations are given after those lines. In particular, we notice empirically that training the model by minimizing $L_{CDM}$ alone can be quite noisy.\n\nUsing a weighted sum of $L_{CFM}$ and $L_{CDM}$ can improve the performance of the classic conditional flow matching, solidifying the benefits of our proposed approach.\n\n---\n\n**Q4. How does FDM’s training time and memory footprint scale with data dimensionality, especially for high-resolution videos?**\n\n**Response:** Compared to the standard conditional flow matching (CFM), our proposed loss function involves an additional divergence term and essentially has a similar scalability as CFM. In practice, instead of explicitly computing the Jacobian matrix, we follow the approach used in [1] and [2], which employs the Hutchinson estimator to approximate the divergence term (see Lines 307-309 in our submission). Then, we only need to compute the derivative of a scalar-valued output with respect to the input to construct the loss, which involves a single call to torch.autograd.grad() and adds just one additional backward pass. Since we're only doing one extra backward pass, the additional memory footprint and computational time of computing the divergence loss scales similarly to that of the original CFM loss with respect to the data dimensionality. We compared the training times and peak memory usage, and found that incorporating the divergence loss does not significantly hinder the model's usability, with computational time and memory footprint remaining mostly within 1.5 times that of the original CFM.\n\nRegarding the video experiment mentioned by the reviewer, we first use a pretrained model to map the video data into a latent space. As a result, the computational cost of the flow matching model depends only on the dimensionality of the latent variables, not the raw video data.\n\n[1] Lu, Cheng, et al. \"Maximum likelihood training for score-based diffusion odes by high order denoising score matching.\" International Conference on Machine Learning. ICML, 2022. https://arxiv.org/pdf/2206.08265\n[2] Lai, Chieh-Hsin, et al. \"Fp-diffusion: Improving score-based diffusion models by enforcing the underlying score fokker-planck equation.\" International Conference on Machine Learning. ICML, 2023. https://arxiv.org/pdf/2210.04296\n\n---\n\n**Q5. Can FDM be applied to non-Gaussian or non-OT probability paths?**\n\n**Response:** Yes, FDM can be applied to non-Gaussian paths. In our DNA sequence generation experiment, we use a Dirichlet probability path, which was designed in [3].\n\n[3] Stark et al. Dirichlet Flow Matching with Applications to DNA Sequence Design, ICML, 2024.\n\n------\n\nThank you for considering our rebuttal.", "claims_and_evidence": "NA", "methods_and_evaluation_criteria": "method", "theoretical_claims": "NA", "experimental_designs_or_analyses": "NA", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "S:\n1\\ Demonstrates consistent improvements over CFM across diverse tasks (synthetic data, dynamical systems, DNA, videos) with minimal computational overhead.\n2\\ Avoids costly higher-order score matching, making divergence alignment computationally practical.\nW:\n1\\ Lack of implementation details which may indicate unfair comparisons.", "other_comments_or_suggestions": "1\\ The organization of the contribution part should be improved.\n2\\ Quantify the individual impact of the divergence loss term L_cdm on performance.", "questions_for_authors": "1\\ How does FDM’s training time and memory footprint scale with data dimensionality especially for high-resolution videos?\n2\\ Can FDM be applied to non-Gaussian or non-OT probability paths", "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": "", "extra_scores": {}}, {"review_id": "Yk9MQn3Aul", "reviewer": "Reviewer_WCvg", "summary": "The paper makes a very keen insight that the \"TRUE\" goal of flow matching is to \napproximate the probabilty time series $(t \\mapsto p_t)$ with the approximate probability time series \n$(t \\mapsto \\hat{p}_t)$.  In doing so, they begin from the  difference \n$\\hat{p}_t - p_t$  and that, in order to make this small in terms of total variation, one must not only \nalign the corresponding vector field v that satisfies the CE wrt  p_t, but also the divergence. \nIn particular if $(p_t, v_t)$ and  $(\\hat{p}_t, \\hat{v}_t)$ are both the CE satisfying solutions, then both \n- $| v- \\hat v|$ as wellas \n- $| \\nabla \\cdot v- \\nabla \\cdot \\hat{v} |$\nmust become small. \n \nJust as there is a conditional-localized counterpart ($L_{CFM}$) for  the loss of $|v- \\hat v|$  ($L_{FM}$)through the trick resembling Denoising score matching,  they construct the localized counterpart $L_{CDM}$ for the original target about $| \\nabla \\cdot v- \\nabla \\cdot \\hat{v} |$ ($L_{DM}$) .  \nTheir proposal is to balance both of them and simultaneoulsy optimize them at the same time.  \nUnlike $L_{CFM}$ that is known to be just a scaler different from $L_{FM}$,  they were able to show only $L_{DM} < L_{CDM}$.\nBut this inequality is in the friendly direction, because LHS can be made smaller by making RHS smaller. \n\n\nThey first show that this is obviously effective on synthetic dataset in terms of both final accuracy and the speed of learning. \nThey showcase their result on categorical domain (DNA-Sequence) as well as PDE, and even Video Prediction with latent modeling.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. We have revised the paper according to all reviewers’ feedback. In what follows, we provide point-by-point responses to your comments.\n\n---- \n\n**Q1. It will be great if some knowledge can be shared regarding the relation between the choice of conditional vector field and the suggested choice of $\\lambda$s.**\n\n**Response:** It is an interesting comment. Our design of the new loss function - Equation (19) of our submission - is based on the empirical observation and discussion in Lines 268-274 (left column) and 220-229 (right column) of our submission. During our submission, we did not consider the meaning of $\\lambda$s in terms of data, but it is an interesting future direction to design a principle to choose $\\lambda$s optimally. We have acknowledged this point in the revised manuscript.\n\n---- \n\n**Q2. While I find the localization of DM with CDM very inspiring, … It feels that, when the method is combined with the works like (Simulation-free Schrodinger bridges via score and flow matching), things would also align really well. Has this work been considered to be applied to Schrodinger Bridge?**\n\n**Response:** Thank you for the very insightful comments. We have not considered the Schrodinger bridge yet, while we believe it is indeed an interesting research problem. We have briefly discussed this and cited the related reference that the reviewer pointed out in the revised manuscript.\n\nIn the flow matching setting, the conditional probability path and the associated vector field are often designed on a per-sample basis. Therefore, our proposed CDM is tractable, similar to the baseline conditional flow matching.\n\n\n---- \n\n**Q3. How does the computational overhead scale by the addition of $L_{CDM}$? The gradient of $\\nabla\\cdot \\hat v(x,\\theta)$ would, in the process of backpropagation, require second derivatives, and that sounds a bit scary from the engineering perspectives when we think of scaling this up to the level of very large networks. Would you please comment on that?**\n\n**Response:** In practice, the divergence term $\\nabla\\cdot \\hat v(x,\\theta)$ is computed efficiently using the Hutchinson estimator (see Lines 307-309 in our submission).  We follow the approach used in [1] and [2], which employs the Hutchinson estimator to approximate the divergence term. Then, we only need to compute the derivative of a scalar-valued output with respect to the input to construct the loss, which involves a single call to torch.autograd.grad() and adds just one additional backward pass. Since we're only doing one extra backward pass, the additional memory footprint and computational time of computing the divergence loss scales similarly to that of the original CFM loss with respect to the data dimensionality. We compared the training times and peak memory usage, and found that incorporating the divergence loss does not significantly hinder the model's usability, with computational time and memory footprint remaining mostly within 1.5 times that of the original CFM.\n\n[1] Lu, Cheng, et al. \"Maximum likelihood training for score-based diffusion odes by high order denoising score matching.\" International Conference on Machine Learning. ICML, 2022. https://arxiv.org/pdf/2206.08265 \n[2] Lai, Chieh-Hsin, et al. \"Fp-diffusion: Improving score-based diffusion models by enforcing the underlying score fokker-planck equation.\" International Conference on Machine Learning. ICML, 2023. https://arxiv.org/pdf/2210.04296\n\n----\n\nThank you for considering our rebuttal.", "claims_and_evidence": "Their claims are mathematically proven. The efficacy of the model is proven both on synthetic dataset, \ndataset of continuous domain, dataset of discrete domain, and dataset of practically an infinite dimensional domain.", "methods_and_evaluation_criteria": "The proposed method is general in nature, and I believe that it shall be evaluated not on the basis of the SOTA result on each domain but on the basis that it can make improvement on the competitive method based on FM, which is done properly with reasonable metric.", "theoretical_claims": "Theoretical claims seems mostly valid on the appendix.", "experimental_designs_or_analyses": "Experiment on the Dynamical system and Video prediction seems particulaly well designed and their evaluations seems to be in alignment with the conventional metric.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "Nothing in particular that came to my mind.", "other_strengths_and_weaknesses": "## Strengths: \n- The paper is well motivated\n- The proposition is convincing, and it is theoretically supported. \n- When the exact same technique as CFM did not work for the DM, the paper proposed a work-around solution of upper bound, which is very noble \n-  experiments are well designed, and the results are promising \n\n## Weakness: \n-  I will ask a question or two in the questions section\n-  While the paper's proposition is convincing, the proposition introduces an hyperparameter(the balancing of CDM and CFM)  whose meaning is a little hard to interpret in terms of data (it is clear in terms of definition, but its choice seems hard to make based on the nature of the data itself).    This might however be somehow inferrable from the choice of the path distributions from which to create the supervisory signal (using OT, for example, make a difference about the optimal choice of $\\lambda$.  It will be great if some knowledge can be shared regarding the relation between the choice of conditional vector field and the suggested choice of $\\lambda$s.", "other_comments_or_suggestions": "Please see the section above.", "questions_for_authors": "### Q1\nWhile I find the localization of DM with CDM very inspiring, as is noted in the paper this requires the computation of $p(x|x_1)$, which in the pure FM setting is a delta function, which would cause a problem. Indeed, this is not an essential problem of the philosophy of the work itself, beause DM only looks at $\\nabla \\log p_t(x)$,  in which $p_t$ is guaranteed to be absolutely continuous whenever $p_0$ is. \n\nAs is mentioned in the early part of the manuscript,  the paper is resolving this problem by assuming the path to be taking the form of \n$$ \\psi(t | x_1, x_source, \\epsilon) = (1-t) x_{source}  + t x_1 +  \\sigma_t \\epsilon$$ \nwhere $\\epsilon$ is a Gauss, and in experiments $\\sigma_t$ is modeled in VP or VE way, and derive the corresponding vector field. \nBut this $\\sigma_t$ looks a bit artificial, especially when $x_{source}$ is not Gaussian, for example--- and one of the strength of FM \nlies in its application of mapping an arbitrary distribution to another. \n\nIt feels that, when the method is combined with the works like (Simulation-free Schrödinger bridges via score and flow matching), things would also align really well.  Has this work been considered to be applied to Schrodinger Bridge?\n\n### Q2 \nHowq does the computational overhead scale by the addition of $L_{CDM}$? I very much believe that the philosophy of this work deserves much attention, but the gradient of $\\nabla \\cdot \\hat{v | \\theta}$ would, in the process of backpropagation, require second derivatives, and that sounds a bit scary from the engineering perspectives when we think of scaling this up to the level of very large networks. Would you please comment on that?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The proposed method is well motivated and it is theoretically supported. The empirical evaluation indicates limited improvements over the baselines.
Decision: Accept (poster)