Variational Rectified Flow Matching
Authors: Pengsheng Guo
Alex Schwing
Venue: ICML 2025 poster (2025)
Keywords: Flow Matching
Diffusion Model
Generative Model
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: We study Variational Rectified Flow Matching, a framework that improves classic rectified flow matching by modeling multi-modal velocity vector fields, and demonstrate its compelling results on synthetic data, MNIST, CIFAR-10, and ImageNet.

Abstract: We study Variational Rectified Flow Matching, a framework that enhances classic rectified flow matching by modeling multi-modal velocity vector-fields. At inference time, classic rectified flow matching 'moves' samples from a source distribution to the target distribution by solving an ordinary differential equation via integration along a velocity vector-field. At training time, the velocity vector-field is learnt by linearly interpolating between coupled samples one drawn from the source and one drawn from the target distribution randomly. This leads to ''ground-truth'' velocity vector-fields that point in different directions at the same location, i.e., the velocity vector-fields are multi-modal/ambiguous. However, since training uses a standard mean-squared-error loss, the learnt velocity vector-field averages ''ground-truth'' directions and isn't multi-modal. In contrast,  variational rectified flow matching learns and samples from multi-modal flow directions. We show on synthetic data, MNIST, CIFAR-10, and ImageNet that variational rectified flow matching leads to compelling results.

Variational Rectified Flow Matching
Pengsheng Guo! Alexander G. Schwing
Abstract
et al., 2023; Ma al., 2024). Those in turn can be viewed as variant of hierarchical variational auto-encoder (Luo, 2022). At inference time, flow matching 'moves' a sample from a source distribution to the target distribution by solv- ing an ordinary differential equation via integration along a velocity vector-field. To learn this velocity vector-field at training time, flow matching regresses to a constructed vector- field/ffield/llowww connecting any sample from the source distribution think of the data-domain positioned at time zero to any sample from the target distribution attained at time one. Notably, in a 'rectified flow,' the samples from the source and target distribution are connected via a straight line as shown in Figure I(a). Inevitably, this leads to multi-moddlityyaabiiiigg i.e., flows pointing in different directions the same location in the data-time space, as illustrated for one-dimensional data-domain in Figure (a)) Since classic rectified flow matching employs standard squared-norm loss to compare the predicted ve- locity vector-field to the constructed velocity vector-field it does not capture this multi- modality. Hence, rectified flow matching aims to match the source and target distribution in alternative ways. This is illustrated in Figure 1(b).
We study Variational Rectified Flow Match- ing, a framework that enhances classic rectified flow matching by modeling multi-mode velocity vector-fields. At inference time, classic rectified flow matching 'moves' samples from source distribution to the target distribution by solving an ordinary differential equation via integration along a velocity vector-field. At training time, the velocity vector-field is learnt by linearly in- terpolating between coupled samples one drawn from the source and one drawn from the target dis- tribution randomly. This leads to "ground-truth" velocity vector-field that point in different di- rections at the same location, i.e., the velocity vector- fields are multi-modal/ambiiuuon ever, since training uses a standard mean-squared- error loss, the learnt velocity vector-field averages "ground-truth" directions and isn't multi-modal. In contrast, variational rectified flow matching learns and samples from multi-modal flow di- rections. We show on synthetic data, MNIST, 10, and ImageNet that variational recti- fied flow matching leads to compelling results.
How-
To enable rectified flow matching to capture this multi- modality in the data-time space, we study variational rec- rified flow matching. Intuitively, variational rectified flow matching introduces a latent variable that permits to dis- entangle multi-modia /aboiuou flow directions at cach location the data-time space. This approach follows the classic variational inference paradigm underlying expec- tation maximization variational auto-encoders. Indeed, as shown in Figure I(c), variational rectified flow match- ing permits to model flow trajectories that intersect. This leads to learned trajectories that more closely align with the "ground-truth" flow. The latent variable can also be used to disentangle different directions.
1. Introduction
Diffusion models (Ho et al., 2020; Song et al., 2021a:b)
and flow matching (Liu et al., 2023; Lipman et al., 2023;
Albergo & Vanden-Eijnden, 2023; Albergo et al., 2023) have been remarkably successful in recent years. These techniques have been applied across domains from computer vision (Ho et al., 2020) and robotics (Kapelyukh et al., 2023) to computational biology (Guo et al., 2024) and medical imaging (Song et al., 2022).
Note that flow matching, diffusion models, and variational auto-encoders are all able capture multi-modality in the data-domain, as one expects from generative model. Im- portantly, variational rectified flow matching differs in that it also models multi-modality in the data-time space. This enables different flow directions at the same data-time space point, allowing the resulting flows to intersect at that loca- tion.
Flow matching (Lipman et al., 2023; Liu et al., 2023; Al- bergo & Vanden--ijjnden, 2023) can be viewed as continu- ous time generalization of classic diffusion models (Albergo
Apple. Correspondence to: Pengsheng Guo <peng- sheng.guo @pple..ppp.........
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Variational Rectified Flow Matching
1.00
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.00
0.00
0.0
-1.
0.0
-1.5
0.0
1.5
(a) Ground Truth
(b) Rectified FM (Baseline) (c) Variational Rectified FM (Ours)
Figure 1: Intuition and motivation: Rectified flow matching randomly couples source data and target data samples, as illustrated in panel (a). This leads to velocity vector-fields with ambiguous directions. Panel (b) shows that the classic rectified flow matching averages ambiguous targets, which leads to curved flows. In contrast, the proposed variational rectified flow matching is able to successfully model ambiguity which leads to less curved flows as depicted in panel (c).
We demonstrate the benefits of variational rectified flow matching across various datasets and model architectures. On synthetic data, our method more accurately models data distributions and better captures velocity ambiguity. On MNIST, it enables controllable image generation with im- proved quality. On CIFAR-10, our approach outperforms classic rectified flow matching across different integration steps. Lastly, on ImageNet, our method consistently im- proves the FID score of SiT-XL (Ma et al., 2024).
is then used as input in the distribution p0(1112) character- ized by the decoder. The loss encourages a high probability of the output space samples while favoring an approximate posterior distribution qo (2)2)a1, e) that similar the prior distribution p(z). To achieve this, formally, VAEs maximize lower-bound on the log-likelhoood, i.e..
Eg ER log p(z1)
> ExynD [Ezng, [logo(<rl22) - Dxr. (xo_(go--x1))lp
In summary, our contribution is as follows: we study the properties of variational rectified flow matching, and, along the way, offer an alternative way interpret the flow match- ing procedure.
2.2. Rectified Flow Matching
For flow matching, inference time, source distribution (o(o)) is queried to obtain a sample zo. This is akin to sam- pling of latent variable from the prior in VAEs. Different from VAEs which perform a single forward pass through the decoder, in flow matching, the source distribution sample Wi is used as the boundary condition for an ordinary differ- ential equation (ODE). This ODE is 'solved' by pushing the sample z forward from time zero time one via in- tegration along trajectory specified via learned velocity vector-field (o f) defined at time i and location Tt, and commonly parameterized by deep net weights 0. Note, the velocity vector-fiedd is queried many times during integra- tion. The likelihood of a data point Ij can be assessed via the instantaneous change of variables formula (Chen et al., 2018: 8; Song al., 2021b; Lipman et al., 2023),
2. Preliminaries
Given a dataset D = {xx1)} consisting of data samples J1, e.g., an image, generative models learn a distribution p(z1). often by maximizing the likelihood. In the following discuss how this distribution is learnt with variational auto-encoders and rectified flow matching, and why the corresponding modeled data distribution is multi-modal,
2.1. Variational Auto-Encoders (VAEs)
Variational inference generally and variational auto- encoders (VAEs) (Kingma & Welling, 2014) specifically have been shown to learn multi-modal distributions. This is achieved by introducing a latent variable 2. At inference time, a latent is obtained by sampling from the prior distri- bution p(z). typically a zero mean unit covariance Gaussian. A decoder which characterizes distribution p(ra 2) over the output space is then used to obtain output space sam- ple .1
log Pi (x,) = log (2o) + J," div vo t)dt, (I)
which is commonly (Grathwohl et al., 2018) approximated via the Hutchinson, 1990). Here, div denotes the divergence vector operator.
trace estimator (Skilling, 1989;
At training time, variational auto-encoders use an encoder to compute an approximate posterior distribution qo (2(22) 21) over the latent space. As the approximate posterior distri- bution is only needed at training time, the data Ii can be leveraged. Note, for the approximate posterior distribution Gaussian with parameterized mean and covariance is often used. A sample from this approximate posterior distribution
Intuitively, by pushing forward samples zo, randomly drawn from the source distribution Po(x)), ambiguity in the data domain is captured as one expects from a generative model.
At training time the parametric velocity vector-field (z t) needs to be learnt. For this, coupled sample pairs (#o, i)) are constructed by randomly drawing from the source and the target distribution, often independently from each other.
2

---

Variational Rectified Flow Matching
A coupled sample (1) and time L E [0, 1] is then used to compute time-dependent location 24 time via function (I,,,11 = Zz. Recall, rectified flow matching uses 2t = 6( t). = (1 t)ao + tr1. In- terpreting z2 as a location, intuitively, the "ground-truth" velocity vector-field v(zo, t) is readily available via v(zo, æ1, I) = 00( 31, 1)/0t, and can be used as the target to learn the parametric velocity vector-field (z t). Concretely, flow matching learns the parametric velocity vector field (z 1) by matching the target via an &g loss, i.e., by minimizing w.r.t. trainable parameters 0 the objective
or its integral form given in Equation (1).
Solving the partial differential equation given in Equa- tion (2) in general analytically is challenging, even when assuming availability of the probability density functions, i.e., when addressing a classic boundary value problem.
However, if we assume the probability density functions to be Gaussians and if we restrict the velocity vector-field be constant, i.e., of the simple parametric form to (2 t) 0, we can obtain an analytic solution. This is expressed in the following claim:
2r [2) t) v( I1, ()1]
Claim 1. Consider two Gaussian probability density func- tions po = -Nooo; To, 1) and in [f 7) with mean To and 2 respectively. Assume a constant velocity vector- field volEi, = 0. Then 0 = 21 - 2o solves the partial differential equation given in Equation (2) and its integral
Consider two different couplings that lead to different "ground-truth" velocity vectors at the same data-time space point (x,, (). The parametric velocity vector-field t) is then asked to match/regress to different target given the same input tt) This leads av- eraging and the optimal functional velocity vector-field v" (z,, t) = (((m).ary.t) (too:,)).)::((0002)))))==== [e(2o: t)]. Hence, Moltiimmddality in the data-time space is not captured. In the following we discuss and study a method that is able to model this multi-modality.
form given in Equation (1) and Ft (1 -1- t)xg + trı
Proof: Given the constant velocity vector-field vy (E(, t) 0, we have JT_ div vo(Ee, t)dt = 0. Plugging this and both probability density functions into Equation (1) yields
(Eo - =x - (5) - x²1² = 0 V E1. Using E = Eo- + So "olE,, f)dt = t0+ leads to (Eu -o)= == (f0-00+11 = 0 V.o which is equivalent to (an - xp 0) xo - x1 + 0) = 0 VEg. This can only be satisfied V if = 21 - To-
3. Variational Rectified Flow Matching
leading to F = zu + to (1 t) + tz,, which proves the claim.
Our goal is to capture the multi-modality inherent in "ground-truth" " velocity vector-fields obtained from typically used couplings (Io, 21) that connect source distribution samples ze ~ Po with target data samples 21 E D. Here, po a known source distribution and D is a considered dataset. This differs from classic rectified flow matching which does not capture this multi-modality even for simple distributions as shown in Figure I and discussed in Section 2. The struggle to capture multi-modalle tor fields that may be more curve and consequently more difficult to integrate at inference time. In turn, this leads to distributions that may not fit the data as well. We will show evidence for both, more difficult integration and less accurately captured data distributions in Section 4.
The arguably very simple setup in Claim I provides intuition for the objective of classic rectified flow matching and offers an alternative way to interpret the flow matching procedure. Specifically, instead of two Gaussian probability density functions po and pl, we assume the real probability den- sity functions for the source and target data are composed of Gaussians centered at given data points To and x; re- spectively, e.g., Po(Eu) Exses [N[Eu; Zo T)/[S}. More- over, importantly, let us assume that the velocity victor-field vs (æt, f) at a data-time space location (Iz, t) is character- ized by a uni-modal standard Gaussian
leads to velocity vec-
p(u| t) = (o,e,,))
To achieve our goal we combine rectified flow matching and variational auto outo-encoders. In the following we first discuss the objective before detailing training and inference.
with a parametric mean uo(2z, t). Maximizing the log- likelihood of the empirical "velocity data" is equivalent to the following objective
3.1. Objective
Eg Eg [lo p(g zol t)]
The goal of flow matching is to learn a velocity vector- field vo( 1) that transports samples from a known source distribution po at time = 0 to samples from a commonly unknown probability density function pi (21) at time = The probability densities Po, Pi and the velocity vector-field Ug are related to each other via the transport problem
(3)
x E, Er,00 [1Iv0 t) 21 + xo
Note that this objective is identical to classic rectified flow matching. Moreover, note our use of the standard rectified flow velocity vector-field, also derived in Claim 1.
i log pe (x.) = div 2a (I t),
This derivation highlights a key point: because the vec- tor field is parameterized via a Gaussian at each data-time
(2)
Øt
3

---

Variational Rectified Flow Matching
space location, multi-modality cannot be captured: the Gaus- sian distribution is uni-modal. Hence, classic rectified flow matching averages the "ground-truth" velocities.
Algorithm Variational Rectified Flow Matching Training
Data: source distribution Po and target sample dataset DD while stopping conditions not satisfied do sample xo po. 21 E DD [we use a mini-batch] sample L ~ (00, (different t for each mini-batch sample}
As mentioned before, this can be sub-optimal. To capture multi-modality, we study the use of a mixture model over velocities at each data-time space location. For this, we assume an unobserved continuous random variable 2, drawn from a prior distribution p(z), governs the mean of the conditional distribution of the velocity vector-field, i.e.
3t = (1 t)=o + tx
get latent = - Me 21, 1) + eag ( 21, Tea t)
with (0, 1) (reparameterization trick]
compute loss following Equation (5)
p(t|a,, 6, 2) = Nu; (0 i, 1),
perform gradient update on 8, Φ
end while
Note, this model captures multi-modalin as p(olaa,t) (p ()) is a Gaussian mixture.
A Agoorihmm Variational Rectified Flow Matching Inference
Data: source distribution Po
We now derive the variational flow matching objective. Since the random variable 2 is not observed, at training time, we introduce a recognition model (= (no, 21, tt, I) a.k.a. an encoder. It is parameterized by & and approximates the intractable true posterior.
sample xo ~ Po
get latent p(2) ODE integrate Zo from == 0 to = 1 using velocity vector-field ug (22, 2)
Using this setup, the marginal likelihood of individual data point can be lower-bounded by
Note that the mean of the approximate posterior is obtained from the deep net Mo(3o, x1, (4,4) and the standard deviation is obtained from (z(2 21, 21, t). Further, use the re-parameterization trick to enable optimization of the objective w.r.t. the trainable parameters A and Φ. Moreover, we use a single-sample estimate for the expectation over the unobserved variable z. We summarize the training procedure in Algorithm 1. Note, it's more effective to work with a single data point, which was merely used for readability in Algorithm 1.
log p(u(xe.)) t) 2 Ezaga [log p(v xe, t, 2)] - (go no 31, , ())p(.))
(4)
Replacing the log-probability of the Gaussian in the deriva- tion of Equation (3) with the lower bound given in Equa- tion (4) immediately leads to the variational rectified flow matching objective Er,,ro.a [logp(x) - sol/2,, Q)] 2
of samples rather than a
Ex, Er -EEzag (11v0((2, 11 + æo
(5)
Note that variational rectified flow matching training differs from training of classic rectified flow matching in only a single step: computation of latent sample z in Algorithm I. From a computational point of view we add a deep net forward pass obtain the mean Ps and standard deviation de of the approximate posterior, and backward pass to obtain the gradient w.r.t. Φ. Also note that the velocity vector-field architecture vs (z t, might be more complex as the latent variable 2 needs to be considered. However, the additional amount of computation is likely not prohibitive.
(a (-120, #1, ())((-
We note that this objective could extended in a number of ways: for instance, the prior p(z) could be trainable deep net conditioned on 2o and/or f. Note however that this leads to a more complex optimization problem with moving target. We leave a study of extensions to future work.
In Appendix A, we provide a theoretical proof demonstrat- ing that the distribution learned by the variational objective preserves the marginal data distribution, as previously estab- lished for classic rectified flow matching (Liu et al., 2023).
We provide implementation details for the deep nets va (It, L, 2). p (200 21 Te, and (2 21, 2z, t) in Sec- tion 4, as their architecture depends on the data.
In the following we first discuss optimization of this objec- tive before detailing the inference procedure.
3.3. Inference
3.2. Training
We summarize the inference procedure in Algorithm 2. Note that sample latent variable only once prior to classic ODE integration of random sample xi ~ Po drawn from
To optimize the objective given in Equation (5), we follow the classic VAE setup. Specifically, we let the prior p(z) = 0, and we let the approximate posterior 4s(2|20, 21, t) = (22: (2 21, Ft. t), as I1, te. This enables analytic computation of the divvevereee in Equation (5).
the source distribution po. To obtain the latent z we sample from the prior 2 ~ p(z) = NV(2; ODE integrate the velucity field uo( t, 2) from time t = 0 to time = 1 starting from 2g.
0,; T). Subsequently, we
4

---

Variational Rectified Flow Matching
25
0.8
20
15
0.4
10
05
14
Ma
0.0
0.0
X
0.5
2.0
2.0
2.0
1: (-1.0. 0.95)
0.0
0.0
0.0
0.0
0.0
0.0
0.5
0.2
0.2
0.2
0.0
0.0
0.0
0.0
0.0
0.0
2.0
3: 10.0, 0.5)
2.0
0.2
0.0
0.0
0.0
0.0
0.0
4: (0.0.0.0) 0.5
0.5
25
0.5
0.0
0.0
0.0
0.0
0.0
0.0
(a)
(b)
(c)
(d)
(e)
(f)
Figure 2: ID velocity ambiguity analysis with various conditioning options and sampling strategies. (a) Ground Truth (GT), (b) Baseline (Rectified Flow), (c) Ours (zg + 21 24), (d) Ours (zg). (e) Ours (æ1). (f) Ours (2²)). The heatmap illustrates the velocity standard deviation for sampled bins in www.-doman--im--dmmaae along with histograms of the velocity four sampled locations. Our method effectively models velocity ambiguity, while the baseline produces deterministic outputs.
4. Experiments
Letthood 1.
Perzen Wirdow
LOS spciioood
We evaluate the efficacy variational rectified flow match- ing and compare to the classic rectified flow (Lipman al., 2023; Liu et al., 2023; Albergo & Vanden-Eijndee 2023) across multiple datasets and model architectures. Our experi- ments show that variational rectified flow matching is able to capture the multi-mode velocity in the data-time space, lead- ing to compelling evaluation results. Moreover, we demon- strate that explicitly modeling multi-modalits through conditional latent can enhance the interpretability of flow matching models, leading to controllability. Implementation details for all experiments are provided in Appendix D.
Baseline
Base
Ours
Ou'
18
20
16
20
Evaluation Steps
Evoluction Steps
Figure 3: Quantitative evaluation synthetic 1D data for varying evaluation steps. Metrics are averaged over three runs. For True and Parzen Window Log-Likelihood, higher values are better.
of MLP layers. We follow this structure in our variational rectified flow matching, but add an encoding layer for the latent variable 2. The posterior model de follows a similar design as ve, outputting ps and a At inference time, 9s isn't used. Instead, we sample directly from the prior distribution p(z) A(2 0, 7). The KL loss weight is 1.0.
4.1. Synthetic 1D Data
For synthetic ID experiments, the source distribution is zero-mean, unit-variance Gaussian, while the target distri- bution is bimodal, with modes centered at 1.00 and 1.0.
We assess the performance using the Fuler ODE solver and vary the evaluation steps. Results are presented in Figure 3. Across both metrics, i.e., True Log-Likelihood and Parzen Window Log-Likelihood and most evaluation steps, our method outperforms the baseline. Notably, as
For the rectified flow baseline, we use a multi-layer MLP network uo to model the velocity. The network operates on inputs 24 and and predicts the velocity through a series
5

---

Variational Rectified Flow Matching
0.5
0.5
0.5
-0.5
-0.5
-0.5
-1.0
-1.0
-1.0 -0.5
0.0 0.5 1.0
-1.0 -0.5
0.0 0.5 1.0
-1.0 -0.5
0.0
0.5
1.0
(a) Ground Truth
(b) Rectified FM
(c) Variational Rectified FM (Ours)
Figure 4: Flow visualization for synthetic 2D data using the Fuler solver with 20 function evaluations. Sampled points from the source distribution are shown in red, and points from the target distribution in purple. Different from Rectified FM, which predicts flow trajectories with sharp curvature and U-turns to avoid crossings, our model captures velocity ambiguity and predicts flows that intersect.
The Logulellccd
0.0
WWestm Uieltood
0.0
mean curvature max curvature
-2.5
Rectified FM Variational Rectified FM
171.35 21.03
-5.0
-5.0
0.98
4.23
-75
Baseline
Baseiine Ours
Ours
-10.0
45
10
16
20
-10.0
10
16
20
Table 1: We calculated the curvature for 2D data results and find significantly lower curvature for our method.
Eva uation Steps
Evoluction Steps
Figure 5: Quantitative evaluation on synthetic 2D data for varying evaluation steps. Metrics are averaged over three runs with different random seeds.
æ conditioning (c), the model fails to capture the ground-
truth distribution at later timesteps (z = - 1100, = 0.95)
as the influence of Φ diminishes. With v, conditioning (f), the ambiguity plot follows the baseline as no extra data is provided to the posterior.
the model handles multi-modality in the data-time space, it produces reasonable results even for 2 or 5 evaluation steps. Qualitative visualizations of flow trajectories are provided in Appendix C.2.
4.2. Synthetic 2D Data
To better understand the multi multi-moddlltty of the velocity and to assess the efficacy of our model in handling it, we ran- domly sample different trajectories and plot the velocity range standard deviation across predefined bins in the data- time space, as shown in Figure 2. The ground-truth flow in Figure 2(a) shows that the standard deviation increases with time, peaking at (r 0.0,f f- = 0.75). The veloc- ity distribution transitions from a bi-modal distribution at early times f to a uni-modal distribution at later times t. Fig- ure 2(b) shows that the rectified flow baseline, which uses an MSE loss, fails to model the velocity distribution faithfully, collapsing to Dirac-delta distribution as expected. In con- trast, Figure 2(c) demonstrates that our model captures the distribution with higher velocity standard deviation range, matching the ground-truth reasonably, albeit not perfectly.
We further test efficacy using synthetic 2D data. Following Liu et al. (2023), we model the source distribution as mix- ture of Gaussian components positioned six equidistant points along circle with a radius of 1/3, shown in red in Figure 4(a). The target distribution follows similar struc- ture, but with components arranged along a larger circle with a radius of 1, shown in purple.
For the architecture we follow Section 4.1 and condition the posterior on [To, X1, æt]. We report the True Log-Likelihood and the Parzen Window Log- Likelihoodd for various evalu- ation steps of the Fuler ODE solver, shown in Figure 5. Compared to the ID data, our model shows more signif- icant performance boost here. We hypothesize that this is due to the increased complexity of the task: explicitly mod- eling multi modii-moddllty avoids the need for curved trajectories, making it easier to fit the target distribution. The qualitative flow visualization in Figure 4 supports this hypothesis: the rectified flow requires a U-turn avoid collisions, while our model, aided by the variational training objective, moves in trajectories that intersect and aren't 't curved. To verify, we calculated the curvature of flow trajectories for both the base- line and our variational rectified flow matching method, and report the results in Table 1. Our method yields significantly lower curvature, indicating smoother flow trajectories.
As discussed in Section 3.2, the posterior do can be condi- tioned in different ways. To understand the implications, we performed ablation studies and visualized the velocity distribution maps in Figure 2 (f)-(f) For Zu conditioning (d), the model struggles to predict the bi-modal distribution at carly timesteps (I,- = 0.0, t = 0.0) due to the absence of x information. However, when f is sufficiently large, the model can infer 21 from Fy, enabling it to predict a bimodal distribution again (r = 0.0, 0.5). Conversely, with
6

---

Variational Rectified Flow Matching
Our method achieves better FID scores than the baseline, even with 2-dimensional conditional latent space. Further details and results can be found in Appendix C.3.
4.4. CIFAR-10
Next, evaluate on CIFAR- 10, widely used benchmark in prior work (Lipman et al., 2023; Tong et al., 2024), For a fair comparison, we use the architecture and train- ing paradigm of Tong et al. (2024), but train the UNet model with the variational rectified flow loss detailed in Equation (5). The UNet consists of downsampling and upsampling residual blocks with skip connections, and a self-attention block added after the residual block at 16 16 resolution and in the middle bottleneck layer. The model takes both x, and t as input, with the time embedding used to regress learnable scale and shift parameters and 3 for adaptive group norm layers.
0 (a) 20
(b) 2
Figure 6: Visualization of learned MNIST manifold with
different random noise Tn-
NFE
2
5
10
20
50
100
200
Baseline
98.8
26.9
13.0
8.8
7.0
6.4
6.0
Ours
61.0
13.4
6.8
5.1
4.7
4.7
4.6
Table 2: FID score evaluation for the MNIST experiment conducted across three trials with different random seeds. Our model with a latent dimension of 2 outperforms the baseline. Note, the latent dimension of chosen for controllability analysis rather than being optimized for FID score improvement.
The posterior model ge shares a similar encoder structure as Us: image space inputs are chosen from [Zo, æ1, 2x] and concatenated along the channel dimension, while time t is conditioned using adaptive group normalization. The network predicts He and dg with dimensions 1 x 1 x 768. During training, the conditional latent & is sampled from the predicted posterior, and at test time, from a standard Gaussian prior. The latent is processed through two MLP layers and serves as conditional signal for the velocity network Un. We identify two effective approaches as con- ditioning mechanisms: adaptive normalization, where is added to the time embedding before computing shift and offset parameters, and bottleneck sum, which fuses the la- tent with intermediate activations at the lowest resolution using a weighted sum before upsampling.
4.3. MNIST
Modeling multi-modalic also enables more explicit control without additional conditioning signals. To show this we use variational rectified flow matching to train a vanilla convolutional net with residual blocks (He et al., 2015) on MNIST data (LeCun et al., 1998). We use (ro, (1, as input to 46 and set the KL loss weight to le
Following Kingma& & Welling (2014), we set the latent vari- able to be 2-dimensional. During inference, we sample linearly spaced coordinates on the unit square, transforming them through the inverse CDF of the Gaussian to generate latents z. Using these latents, we integrate the samples with an ODE solver and plot the generated samples in Figure 6. To show the effects of the source distribution sample 2o and the latent z, we visualize the learned MNIST manifold for two randomly sampled xo values in Figure 6(a.b). The results demonstrate that the latent space enables smooth in- terpolation between different digits within the 2D manifold, providing control over the generated images. By adjusting z, we can transition between various shapes and styles. The initial noise zo enhances the generation process by intro- ducing additional variations in character styles, allowing the model to better capture the target data distribution. We also evaluate the FID scores of our method using this 2- dimensional conditional latent space and report the results in Table 2. Despite the small latent dimension, it still en- ables the velocity model Up to achieve better FID scores than the rectified flow matching baseline.
We evaluate results using FID scores computed for varying numbers of function evaluations, as shown in Table 3. Four model variants were tested, differing in fusion mechanisms, posterior model de inputs, and KL loss weighting. Com- pared to prior work (Lipman et al., 2023; Liu et al., 2023; Tong et al., 2024), model 1 achieves superior FID scores with fewer function evaluations and performs comparably at higher evaluations. Using the adaptive Dopri5 solver further improves scores, highlighting the importance of capturing flow ambiguity. Model 2 increases the KL loss weight, improving performance at higher function evaluations but reducing effecttieness at lower evaluations, likely due to reduced information from latent 2. Model 3, with additional time conditioning, significantly improves FID at low eval- uations and performs best with the adaptive solver. Model 4, incorporating bottleneck sum fusion, delivers robust FID scores across evaluation settings, demonstrating the flexibil- ity of the variational rectified flow objective with different fusion strategies
Similar to the MNIST results in Section 4.3, we observe
7

---

Variational Rectified Flow Matching
NFE/ ssapla
# Params.
5
100
1000
Adaptive
I-CFM (Tong et al., 2024)
36.5M
168.654
35.489
13.788
5.288
4.461
3.643
3.659
V-RFM (adaptive norm, 21, 2e-3)
37.2M
135.275
28.912
13.226
5.382
4,430
3.642
3.545
2 V-RFM (adaptive norm, æ1, 5e-3)
37.2M
159.940
35.293
14.061
5.265
4.349
3.582
3.561
3 V-RFM (adaptive norm, 41 + L, 5e-3)
37.2M
117.666
27.
13.632
5,512
4.484
3.614
3,478
4 V-RFM (bottleneck sum, + t, 20-3)
37.0M
104,634
25.841
13.508
5.618
4.540
3.596
3.520
Table Following Tong al. (2024), we train the same UNet model and reported the FID scores for our method and the baselines using both fixed-step Euler and adaptive-step Dopris ODE solvers. The baselines checkpoint was directly taken from Tong et al. (2024). We present four model variants of V-RFM, which differ in fusion mechanism, posterior model input, and KL loss weight.
Model
Params (M)
Training Steps
FID 4
DIT-XL
675
400K
19.5
SiT-XL
675
400K
17.2
V-SIT-XL
677
400K
14.6
SiT-XL
675
800K
13.1
V-SIT-XI
677
800K
10.6
(a) 20
(b) z
(c) 2
SiT-XT eig=155
675
400K
5.40
V-SiT-XI xfg=15
677
400K
4.91
Figure 7: Varying zo while keeping the latent fixed. Im- ages at the same position across panels share the same while images within a panel share the same latent sampled from the prior distribution.
SiT-XLdg=115
675
800K
3.43
V-SiT-XL<1g-15
677
800K
3.22
Table 4: FID-50K score evaluation of class-conditional gen- eration ImageNet 256 x 256, comparing the baselines (DiT-XI,, SiT-XII) with our proposed model V-SiT-XI.
clear patterns in color and content for the generated samples T1. demonstrating a degree of controllability. Figure 7 visu- alizes three sets of images (c)- Each set is conditioned on different latent z, while the starting noise zo varies across individual images within each set. The same noise Co is applied to images at the same grid location across all subplots. Images conditioned on the same latent exhibit con- sistent color patterns, while images the same grid location display similar content, as highlighted in the last row.
we use the Euler-Maruyama sampler with the SDE solver and 250 integration steps, as described by Ma et al. (2024).
Following the evaluation protocol of Ma et al. (2024), we randomly generate 50K images from the models and report the FID scores in Table 4. V-SiT-XL consistently outper- forms both DiT-XL and SiT-XL, achieving gains under the same training conditions. with and without classifier-f free guidance. These results underscore the importance of mod- eling multi-modality in the velocity vector field, which con- tributes to a substantial improvement in generation quality, particularly in the large-scale high-resolution data domain. Additionally, we analyze the model's performance across different training iterations and varying numbers of function evaluations, presenting the findings in Figure 8. The results reveal a consistent performance boost, further highlighting the effectiveness of our approach.
4.5. ImageNet
To assess efficacy on large-scale data, we use ImageNet 256 256 data and SiT-XL (Ma et al., 2024), a recent transformer-based model that has shown strong results in strictly follow
image generation. For fair comparison, the original training recipe in the open-source SiT repository and replicate the training process from the SiT paper, while introducing our model, V-SIT-XL, by substituting the classic rectified flow loss with the variational rectified flow loss in Equation (5). The posterior model qa also utilizes an SiT transformer architecture but with half the number blocks. In the final layer, the features are average pooled and passed through an MIP layer to predict 24 and og. We sample the latent variable = from the posterior during training and from the prior distribution during inference. This latent variable z is then processed by two MLP layers and fused with the velocity network vy adaptive normalization. By default,
5. Related Work
Generative modeling has advanced significantly in the last decade, thanks in part due to seminal works like genera- tive adversarial nets (Goodfellow et al., 2014), variational auto-encoders (Kingma & Welling, 2014), and normalizing flows (Rezende & Mohamed, 2015).
More recently, score matching (Song & Ermon, 2019) and
8

---

Variational Rectified Flow Matching
inherent in the usually employed ground-truth flow fields.
25.0
V-8ITXL (Ouss)
(eaadre)
30
BT-- Buseird V-BITAL (Ows))
22.5
Structurally similar to this idea is work by Preechakul et al. (2022). In a first stage, an autoencoder is trained to compress images into latent space. The resulting latents then serve as a conditioning signal for diffusion model training in a second stage. Note, this two-stage approach doesn't It directly model ambiguity in the tuta-time space. In spirit similar work by Pandey et al. (2022). A VAE and a diffusion model are trained in two separate stages, with the goal to enable controllability of diffusion models. Related is also work by Eijkelboom et al. (2024) which focuses on flow matching only for categorical data, achieving compelling results on graph generation tasks. Another related work by Vahdat et al. (2021) uses a VAE to map raw data zo into latent space 2o. with the VAE jointly trained with score-based generative modeling (SGM). Unlike our approach, SGM still faces ambiguity issues due to its use of uni-modal Gaussian distribution.
FD40K
20.0
-2
15.0
15
250K
503K
403K SOOK GCOK. 7006 BJDK
50
100
150
250
250
Traniing Steps
Numbor of Function Exalustions
Figure 8: FID-50K score over training iterations and number of function evaluations. Our model, V-SiT-XL, consistently achieves a better FID score compared to SiT-XL trained with classic rectified flow matching.
diffusion models (Ho et al., 2020) were introduced. They can be viewed as augmenting variational auto-encoders hi- erarchically (Luo, 2022) while restricting involved distribu- tions to be Gaussian. Notably, and analogously to classic discrete normalizing flows, the number of hierarchy levels, i.e., the number of time steps, remained discrete, which introduced complications.
Related work by Xiao et al. (2022) replaces the Gaussian model in the denoising step with a multimodal distribution. Unlike our method, a conditional GAN with separate dis- criminator models the distribution. But GANs face mode collapse and stability issues. In contrast, our method uses rectified flow matching, preserving the maximum likelihood benefits. Xu et al. (2023) notes that the posterior distribution is multi-modal. To model this distribution, the paper reduces training target variance using a reference batch. In contrast, our method directly models this multi-modal posterior via a recognition model.
Flow matching (Lipman et al., 2023) was introduced re- cently as a compelling alternative to avoid some of these complications. It formulates an ordinary differential equa- tion (ODE) in continuous time. This ODE connects source distribution to a target distribution. Solving the ODE via forward integration through time permits to obtain samples from the target distribution, essentially by 'moving' samples from the known source distribution to the target time along a learned velocity field.
Concurrently, Zhang et al. (2025) also study method to model multi-modal velocity vector fields. In this paper, we discuss how to use a lower-dimensional latent space to en- able modeling of the velocity distribution via a variational approach. Differently, Zhang et al. (2025) study use of a hierarchy of ordinary differential equations. The variational approach enables to capture semantics while use of hi- erarchy of ordinary differential equations permits to more accurately model the velocity distribution.
To learn the velocity field, various mechanisms to interpolate between the source distribution and the target distribution have been considered (Lipman al., 2023; Liu al., 2023; Tong et al., 2024). Rectified flow matching emerged as a compelling variant, which linearly interpolates between samples from the two distributions. For instance, it was used to attain impressive results on large scale data (Ma et al., 2024; Esser et al., 2024). Different from other techniques, linear interpolation encourages somewhat straight flows, which simplifies numerical solving of the ODE.
6. Conclusion
The importance of straight flows was further studied in ReFlow (Liu et al., 2023), which sequentially formulates multiple ODEs and learns velocity fields by adjusting the interpolations and 're-training.' Consistency models (Song et al., 2023; Kim et al., 2023; Yang et al., 2024) strive for straight flows by modifying the loss to encourage self- consistency across timesteps. More details and comparisons are provided in Appendix B and C.I.
We study Variational Rectified Flow Matching, a framework which enables to model the multi-modal velocity vector fields induced by the ground-truth linear interpolation be- tween source and target distribution samples. Encouraging results can be obtained on low-dimensional synthetic and high-dimensional image data.
Acknowledgements: We thank Zhizhen Zhao for insightful feedback.
While the aforementioned works aim to establish straight flows either via 're-training' or 're- parameterizinng of an already existing flow, differently, in this work we study the results of enabling a rectified flow to capture the ambiguity
9

---

Variational Rectified Flow Matching
Impact Statement
Kapelyukh, I., Vosylius, V., and Johns, E. Dall-e-bot: In- troducing web-scale diffusion models to roboties. IEEE Robotics and Automation Letters, 2023.
This paper presents work whose goal to advance the field of machine learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Kim, D., Sony, A., Lai, C.-H., Liao, W.-H., Murata, N., Takida, Y., He, Y., Mitsufuji, Y., and Ermon, S. Consis- tency trajectory models: Learning probability flow ode trajectory of diffusion. Proc. NeurIPS, 2023.
References
Albergo, M. and Vanden-Eijnder E. Building normalizing flows with stochastic interpolants. In Proc. ICLR, 2023.
Kingma, D. and Welling, M. Auto-Encoding Variational Bayes, In Proc. ICLR, 2014.
Albergo, M., Boffi, N., and Vanden-Eijnden, E. Stochastic Interpolants: A unifying framework for flows and diffu- sions. In arXiv preprint arXiv:2303.08797, 2023.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, Gradient- based learning applied document recognition. Proceed- ings of the IEEE, 1998.
Lipman, Y., Chen, R., Ben-Hamu, H., Nickel, M., and Le, M. Flow Matching for Generative Modeling. In Proc. ICLR, 2023.
Chen, R., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equations. In Proc. NeurlPS, 2018.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. In Proc. ICLR, 2023.
Eijkelboom, F., Bartosh, G., Naesseth, C., Welling, M., and van de Meent, J.-W. Variational Flow Matching for Graph Generation. In arXiv preprint arXiv:2406.04884 04444000044,,, 2024.
Luo, C. Understanding diffusion models: A unified perspec- tive. In arXiv preprint 1XX1::2208.119000 2022.
Esser, P. Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., Podell, D., Dockhorn, T., English, Z., Lacey, K., Good- win, A., Marek, Y., and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. In Proc. ICML, 2024.
Ma, N., Goldstein, M., Albergo, M., Boffi, N., Vanden- Eijnden, E., and Xie, S. SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Inter- polant Transformers. In arXiv preprint arXiv:2401.08540. 2024.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. In Proc. NeurlPS, 2014.
Nguyen, B., Nguyen, B., and Nguyen, A. Bellman opti- mal stepsize straightening of now-matching models. In The Twelfth International Conference on Learning Repre- sentations, 2024.
Grathwohl, W., Chen, R., Bettencourt, J., Sutskever, I., and Duvenaud, D. FFJORD: Free-form continuous dynamics for scalable reversible generative models. In Proc. ICLR, 2018.
Pandey, K., Mukherjee, A., Rai, P., and Kumar, A. Dif- VAE: Efficient, controllable and high-fidelig gener- ation from low-dimensional latents. In arXiv preprint arXiv:2201.00308, 2022.
Guo, Z., Liu, J., Wang, Y., Chen, M., Wang, D., Xu, D., and Cheng, J. Diffusion models in bioinformatics and computational biology. Nature reviews bioengineering, 2024.
Preechakul, K., Chatthee, N., Wizadwongsa, S., and Suwa- janakorn, S. Diffusion autoencoders: Toward meaning- ful and decodable representation. In Proc. CVPR, 2022.
He, K., Zhang, X., Ren, S., and Sun, J. Deep Residual Learning for Image Recognition. happs//arxisoorgaab////12.3333 2015.
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. In Proc. ICML., 2015.
Skilling, The eigenvalues of mega-dimensional matrices. Maximum Entropy and Bayesian Methods, 1989.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. In Proc. NeurIPS, 2020.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In Proc. ICLR, 2021a.
Hutchinson, M. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Com- munications in Statistics-Simulatton and Computation, 1990.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Proc. NewrIPS, 2019.
10

---

Variational Rectified Flow Matching
Song, Y., Sohl-Dickstein, J., Kingma, D., Kumar, A., Ermon, S., and Poole, B. Score-Based Generative Modeling Through Stochastic Differential Equations. In Proc. ICLR, 2021b.
Song, Y., Shen, L.,, Xing. L., and Ermon, S. Solving inverse problems in medical imaging with score-based generative models. In Proc. ICLR, 2022.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, Consis- tency models. In International Conference on Machine Learning, pp. 3221 1-32252. PMLR, 2023.
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang. Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improv- ing and generalizing flow-based generative models with minibatch optimal transport. TMLR, 2024.
Vahdat, A., Kreis, K., and Kautz, J. Score-based Generative Modeling in Latent Space. Proc. NeurIPS, 2021.
Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion GANs. In Proc. ICLR, 2022.
Xu, Y., Tong. S., and Jaakkola, T. Stable Target Field for Reduced Variance Score Estimation in Diffusion Models. In Proc. ICLR, 2023.
Yan, H., Liu, X., Pan, J., Liew, H., Liu, Q., and Feng, J. Perflow: Piecewise rectified flow universal plug- and-play accelerator. arXiv preprint wrx.v:2405.07510, 2024.
Yang. L.. Zhang, Z., Zhang, Z., Liu, X., Xu, M., Zhang, W., Meng, C., Ermon, S., and Cui, B. Consistency flow matching: Defining straight flows with velocity consis- tency. arXiv preprint arXiv:2407.02338. 2024.
Zhang, Y., Yan, Y., Schwing, A., and Zhao, Towards Hierarchical Rectified Flow. In Proc. ICLR, 2025.

---

Variational Rectified Flow Matching
Appendix: Variational Rectified Flow Matching
This appendix is structured as follows: in Appendix A we show that our approach maintains the marginal distribution; in Appendix we discuss additional related work; in Appendix we provide additional experimental analysis; in Appendix we provide more implementation details; in Appendix E list additional qualitative results.
A. On Preserving the Marginal Data Distribution
We obtain samples by numerically solving the ordinary differential equation
du (2 z)dt with ~
= p==)))=N(z)))
This differs slightly from Theorem 3.3 of Liu al. (2023) because the velocity vg depends on latent variable z drawn from standard Gaussian. However, Theorem 3.3 of Liu et al. (2023) can be extended to fit this setting as follows.
First, note that we have s² (x,,42 E|Xx Z] where X, and Z are random variables corresponding to instances x; and 2.
Incorporating the velocity field depending on the latent variable 2 into the transport problem defined in Equation (2) and taking an expectation over the latent variable, we obtain the continuity equation
Pr + div [vo,,, L, =)(pr) = 0.
(6)
Following Liu al. (2023), one show equivalence to the following equality, which uses any compactly supported continuously differentiable test function h:
[[,(),)] = (V,, X VVV) (X, 0) = Ex
Z)]]
dt
Concretely, equivalence can be shown via
0 0= 0=Ez EEz ( h(p: + div(u" (X,,,, 2)p:)
d [[A(X,)] - Ex (V(()) (Xe, Z)]].
= dt
Note, different from Liu et al. (2023), in our case U, driven by a velocity field v( 2) that depends latent variable. Averaging over instantiations of the random latent variable Z leads to the same marginal velocity that appears in the continuity equation (Equation (6)). Therefore, we solve the same equation with the same initial condition (Xo = Ug). Equivalence follows if the solution to Equation (6) is unique.
B. Additional Related Work Discussion
2.0
Here, we discuss related work aimed at improving the sample efficiency of diffu- sion and flow matching models, either via consistency modeling or via distillation.
0,8
0.0
Consistency models. Consistency models, such as those by Song et al. (2023) and Yang et al. (2024), enforce self- celfisonsit across timesteps, ensuring trajectories map back to the same initial point. Moreover, Kim et al. (2023) ensure consistent trajectories for probability flow ODEs. While consistency models focus on im- proving results via trajectory alignment if few function evaluations are used, they don't model the multi-modal ground-truth velocity distribution, which is our goal.
0.5
0,4
0.0
0.2
To illustrate this, we train the recently developed consistency flow matching model proposed by Yang et al. (2024) (which improves upon work by Song et al. (2023) and Kim et al. (2023); both are not flow matching based; it also improves upon distillation work by Nguyen et al. (2024)) on the data for which V-RFM results are presented in Figures 2 and 13. Specifically, we used the publicly available baseline. We obtain the results illustrated Figure 9. As expected, we observe that classic consistency modeling does not capture the multi-modal velocity distribution, unlike the proposed V-RFM.
0.0
255
0.0
2.5
0.0
-3
Figure Velocity distribution con- sistency flow matching (Yang et al., 2024).
https://gihuuh.omm/aaa/////// Angg // flow matching
12

---

Variational Rectified Flow Matching
10
True Log-Likelihood
Wasserstein Distance
Parzen Window
Concistnccy
-1
Log-Ukkelloook
VRFM (ours)
2
10-1
-3
Bectified Comiisinnyy FM
Cono
VRPM (0urs)
VRPH (Ours]
-4
2
8
10
16
20
10
16
20
10-2
10
16
20
Evaluation Steps
Evaluation Steps
Evaluation Steps
Figure 10: Additional quantitative evaluation with the consisteney flow matching baseline on synthetic 1D data. Higher
values are better for True and Parzen Window Log-Likelihood, while lower values are preferred for Wasserstein Distance.
Figure 11: Additional quantita- tive evaluation with the consis- tency flow matching baseline on synthetic 2D data. Metries are av- eraged over three runs with differ- ent random seeds.
True Log Liketihoodd
0.0
Parzen Window Log Likellhood
0.0
-2.5
-2.5
5.0
5.0
Rocuries VRFM
-7.
Consitenty 1ArM
-7.5
-10.0
2
4.
10
20
-10.0
2 4
8 10
16
20
Evaluation Steps
Evaluation Stops
Furthermore, we conduct additional experiments with consistency flow matching across multiple datasets, summarizing the results in Appendix C. We observe that the consistency flow matching method performs well in the low function evaluation regime (i.e., NFE 2 5), but its performance degrades as the NFEs increase. Most notably, its best performance across all NFEs does not surpass that of classic rectified flow matching our proposed variational rectified flow matching. Based on the empirical evidence and the key differences in capturing multi-modal velocity distributions, we believe consistency models are orthogonal to our proposed variational formulation. Therefore, we find it exciting to explore future research on combining variational flow matching with consistency models, which is beyond the scope of this paper.
Distillation. Nguyen et al. (2024) perform distillation by optimizing step sizes in pretrained flow-mathhing models to refine traeectories and improve training dynamics. Moreover, Yan al. (2024) perform distillation by introduceing a piecewise rectified flow mechanism accelerate flow-based generative models. Note, both methods distill useful information from pretrained model, either by using dynamic programming to optimize the step size or by applying reflow to straighten trajectories, i.e., they focus on distilling already learned models. In contrast, our V-RFM focuses on learning via single-stage training, directly from ground-trute truth data, and without use of pre-trained deep net, a flow-mathhing model, which captures a multi-modal velocity distribution. More research on the distillation of V-RFM model is required to assess how multi- modality can be maintained in the second distillation step. We think this is exciting future research, which is beyond the scope of this paper.
Reflow. We have added comparison with 2/3-Rectified Flow via Reflow (Liu al., 2023) in Table 5. We find that while strong FID scores in the low-NFE regime are achieved, it does the cost of limiting peak performance high NFE. We emphasize that Reflow is supplementary technique applied on top of Rectified Flow Matching (RFM))-rimarrly aimed at fast sampling rather than improved sample quality. It also requires N times longer training and a significantly larger fine-tuning dataset, where N denotes the number of Reflow rounds. These differences make a direct comparison with our V-RFM less fair. Hence, RFM without Reflow a more appropriate baseline. Additionally, Reflow can be applied to our method as well, potentially improving results the cost of increased training overhead.
C. Additional Experimental Results and Analysis
C.1. Comparison to Consistency Flow Matching
We conduct additional experiments to compare our approach with consistency models across multiple datasets. For this, we use the recently developed consistency flow matching model from Yang et al. (2024) as representative baseline, as it advances earlier consistency modeling efforts by Song et al. (2023); Kim et al. (2023) and distillation work by Nguyen et al. (2024). Specifically, we used the publicly available implementation."
The results are summarized as follows: Synthetic ID data in Figure 10, Synthetic 2D data in Figure 11, MNIST data in
http:://ggihuu..oo/aaagggg////// 0ng08 8/consistency flow matching
13

---

Variational Rectified Flow Matching
FID score evaluation for the MNIST experiment, including the additional consistency flow matching baseline. Our model with latent dimension of 2 outperforms the base- lines, except at 2 evaluation steps where Consistency FM performs best. Note, the latent dimension of 2 is chosen for a controllability analysis rather than being optimized for FID score improvement.
Figure 12:
100
Rectified FM Consistency
VRRM (Ouns)
25
10
10c
Evaluation Steps
NFE/ / sample
# Params.
2
5
10
50
100
1000
Adaptive
OT-FM
36.5M
166.655
36.188
14.3%
5.557
4.640
3.822
3.655
(Lipman et al., 2023; Tong et al., 2024)
I-CFM
36,5M
168.654
35,489
13.788
5.288
4,461
3,643
3.659
(Liu al., 2023; Tong et al., 2024)
I-CFM Reflow (Liu al., 2023; Tong et al., 2024)
36,5M
7.512
5.906
5,513
5.283
5,276
5,276
5.275
I-CFM w. Reflow
36.5M
7.559
6.925
6.776
6.729
6.733
6.752
6.752
(Liu et al., 2023; Tong et al., 2024)
Consistency-Fih (Yang et al., 2024)
36.5M
15.758
14.588
24.107
36.800
38.675
40.486
40.711
Consistency-FMM -M
61.8M
5.323
11,412
23.948
36.652
38.680
40.402
40.677
(Yang et al., 2024)
1 V-RFM (adaptive norm, .n, 2c-3)
37.2M
135.275
28.912
13.226
5.382
4.430
3.642
3.545
2 V-RFM (adaptive norm, 21, 5e-3)
37.2M
159.940
35.293
14.061
5,265
4.349
3,582
3,561
3 V-RFM (adaptive norm, æ1 + 1, 5e-3)
37.2M
117.666
27.464
13.632
5,512
4.484
3,614
3,478
4 V-RFM (bottleneck sum, æ1 + 2e-3)
37.0M
104.634
25.841
13.508
5.618
4.540
3.596
3.520
Table Additional quantitative evaluation of the consistency flow matching baseline on CIFAR-10. The consistency flow
matching method performs well in the low function evaluation regime (NFE 5), but its performance degrades as NFEs increase. Notably, its best performance across all NFEs does not surpass that of classic rectified flow matching (OT-FM, I-CFM) our proposed variational rectified flow matching (V-RFM).
Figure 12, and CIFAR-10 10 data in Table These results demonstrate that V-RFM outperforms the consistency flow matching baseline across various evaluation steps for synthetic data, with V-RFM showing superior performance when the number of evaluation steps exceeds for MNIST and 5 for CIFAR-10. Importantly, while consistency flow matching achieves strong performance for a low number of evaluation steps, its best performance still does not surpass that of classic rectified flow matching or our proposed variational rectified flow matching with high number of evaluation steps. This highlights its distinct nature as an orthogonal research direction to our method. As discussed in Appendix B, we believe that combining variational formulations with consistency models presents an exciting avenue for future research, though is beyond the scope of this paper.
C.2. Qualitative Results of Synthetic 1D Experiment
We provide qualitative flow visualizations from the synthetic 1D experiment in Figure 13. Our method effectively captures velocity ambiguity and predicts crossing flows, whereas the baselines produce deterministic outputs.
C.3. Quantitative Results of MNIST Experiment
We evaluate the FID scores of our method using this 2-dimensional conditional latent space and report the results in Figure 12. Despite the small latent dimension, still enables the velocity model U6 to achieve better FID scores than the basclines, except at evaluation steps where consistency flow matching (Yang et al., 2024) performs best.
C.4. Inception Score Evaluation of CIFAR-10 Experiment
We evaluate the Inception Score of our model trained on CIFAR- -00 data and present results in Table 6. This score quantifies the distribution of predicted labels for the generated samples. Compared to the vanilla rectified flow baseline, our method
14

---

Variational Rectified Flow Matching
1.00
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.00
-3
2
0
1
3
2
1
0
1
3
2
(a) Ground Truth
(b) Rectified FM
(c) Consistency FM
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.00
0.00
-3
-2
0
3
-2
1
-3 -2
2
X
(d) Ours (z0n)
(e) Ours (x1)
(f) Ours (2:+)
1.00
0.75
0.50
0.25
0.00
x
(g) Ours (zo + x1 + xz.)
Figure 13: 1D flow visualization for uni-modal Gaussian to bi-modal Gaussian.
consistently achieves higher Inception Scores, reflecting improved diversity in the generated samples.
C.5. Ablation on Posterior Model Size
We conducted ablations to study the impact of varying the size of the encoder ga. reducing it to 6.7% and 17.5% of its original size. The results reported in Table 7 demonstrate that our model maintains comparable performance across these variations, highlighting the flexibility and robustness of our approach.
C.6. Reconstruction Loss Visualizations
We present the reconstruction loss curves for our model and the baseline trained on MNIST and CIFAR-10 data in Figure 14. We observe better reconstruction losses of our model compared vanilla rectified flow, indicating that the predicted velocities more accurately approximate the ground-truth velocities.
D. Implementation Details
D.I. Synthetic data
In the rectified flow baseline, the velocity network vy features separate encoders for time and data æ. Each encoder consists of sinusoidal positional encoding layer followed by two MLP layers with GeLU activation. The resulting time and data embeddings are concatenated and passed into four-layer MIP, also utilizing GeL.U activations. Both the positional
15

---

Variational Rectified Flow Matching
NFE/ / sample
10
50
100
1000
Adaptive
I-CFM
2.786
7.143
8.326
8.770
8.872
9.022
9.041
(Liu et al., 2023; Tong et al., 2024)
V-RFM (adaptive norm, æi, 2e-3)
3.943
7.728
8.499
8.973
9.050
9.168
9.171
2 V-RFM (adaptive norm, æ1, 5e-3)
3.083
7.202
8.342
8.868
8.997
9.166
9.183
3 V-RFM (adaptive norm, 21 + f, 5c-3)
4.460
7.930
8,583
9.007
9.104
9.220
9.238
V-RFM (bottleneck sum, 211 + f, 20-3)
4.831
7.996
8.529
9.062
9.150
9.293
9.308
Table 6: Inception Score evaluation of our method compared to the baseline on CIFAR-10, using fixed-step Fuler and
adaptive-step Dopri5 ODE solvers. Higher scores indicate better performance.
NFE/ / sample
2
5
10
50
100
1000
Adaptive
OT-FM
166.655
16.188
14.396
5.557
4.640
3.822
3.655
(Lipman et al., 2023; Tong el al., 2024)
I-CFM
168.654
35.489
13.788
5.288
4.461
3.643
3.659
(Liu et al., 2023;
Tong al., 2024)
1 V-RFM-L (100% Posterior Model)
135.275
28,912
13.226
5.382
4,430
3.642
3,545
2 V-RFM-M (17.5% Posterior Model)
135.983
30.106
13.783
5.486
4.500
3.697
3.607
3 V-REM-S (6.7% Posterior Model)
144.676
31.224
13.406
5.289
4.398
3.699
3.639
Table 7: We use the same flow matching model vg and pair it with different sizes of encoders qs during training while maintaining the exact same Euler and adaptive-step Dopri5 ODE solvers.
We report the FID scores for our method and the baseline using both fixed-step
embedding and hidden dimensions of the encoder and decoder are set to 64. The training batch size is 1000, and we employ the standard rectified flow objective, i.e., we compute the current data via V = (1 - t) + trı, the ground truth velocity via v(zu, 21, = I1 Fo. and we use the L2 loss for supervision.
For consistency flow matching, we adopt the same velocity network vo and modify the loss function incorporate the velocity consistency loss proposed by Yang et al. (2024). We find the hyperparametes settings suggested by the publicly available codebase to work best. Specifically, we use At 1 x 10-3, Nusemennts = 2. and boundary = 0.0 for the first training stage, transitioning to boundary = 0.9 in the second stage. Additionally, the loss weighting factor a is set to 1 x 10-5. For complete implementation details, we kindly direct readers to the open-source repository which we used to obtain the reported results.
In both cases, the AdamW optimizer is used with the default weight decay and learning rate of 10-3, over a total of 20,000 training iterations.
In our variational flow matching approach, the velocity network uo incorporates an additional latent encoding module comprising three MLP layers with a hidden dimension of 128. The conditional latent embedding z is concatenated with the embeddings for time f and data x. The decoder maintains the same structure as the baseline, with the first MLP layer adjusted to accommodate the increased channel input. For the posterior model qa we employ a similar architecture, designing separate encoder for each possible input selected from [xo, x2, t). Each encoder consists of a sinusoidal positional encoder layer followed by two MLP layers with GeLU activation. The output embeddings are concatenated along the channel dimension and processed through three MLP layers to produce the predicted Ps and Øg- The latent dimension of is set to 4 for 1D experiments and 8 for 2D experiments. During training, we utilize the reparameterization trick to sample from the predicted posterior distribution; during inference, the posterior model is omitted, and sampling is performed from a unit variance Gaussian prior distribution. The loss is defined as the sum of the rectified flow reconstruction loss and the KL divergence loss, with the KL loss weighted at 1.0 for the ID experiments and 0.1 for the 2D experiments. We employ AdamW W as the optimizer with a learning rate of x 10-3 and train the two networks do and vo jointly for 20,000 iterations.
D.2. MNIST
In the rectified flow baseline, the velocity network vy uses separate encoders for time t and data x. The time t encoder consists of a sinusoidal positional encoding layer followed by two MLP layers with SiLU activation. The data x encoder includes a convolutional in-projection layer, five consecutive ResNet (He al., 2015) blocks (each consisting of two
https///iihuuh.oo//aaanniigigii 8/consistency flow matching
16

---

Variational Rectified Flow Matching
MNIST
CIFAR-10
12x
reconsiruciin loes.
2=10"
Baseline Ours
reconstructin loss
Whyl vn MmM
mo Wnw
17+10
1,8x
Baseline Ours
1.4x100
20000
40000
60000
8O00D
100000
100000
200000 itrrtions
300000
400000
ilerations
Figure 14: Reconstruction loss for MNIST (left) and CIFAR-10 (right). We observe lower reconstruction losses for the variational formulation, indicating a better fit.
convolutional layers with a kernel size of 3, group normalization, and SiLU activation), followed by a convolutional out out-projection layer. The time and data embeddings are concatenated and passed to a decoder composed of a convolutional in- nn-projection layer, five consecutive ResNet blocks, and a convolutional out out-projection layer with a kernel size of and an output channel of 1. The hidden dimension is set to 64. MNIST data normalized the [-1, 1] range. We adopted the consistency velocity loss from the consistency flow matching baseline used for synthetic data experiments. We train the network for 100,000 iterations using the Adam optimizer with a learning rate of 1 10-3 and batch size of 256.
In our variational flow matching approach, the velocity network 0g includes an additional latent encoding module consisting of a sinusoidal positional encoding layer followed by two MLP layers with SiLU activation. The conditional latent embedding z is concatenated with the embeddings for time and data x. The decoder structure mirrors the baseline, with the first in-projection layer adjusted to handle the increased channel imput. The posterior model qa follows a similar architecture, with separate encoders for each input [1, The resulting embeddings are concatenated and passed through a decoder consisting of convolutional in- un-projection layer, followed by three consecutive interleaving ResNet blocks and average pooling layers. The final hidden activation is flattened and processed by two linear MLP layers to predict the 1D latent 2 with a dimension of 2. The two networks are trained jointly for 100,000 iterations using the AdamW optimizer with a learning rate of x 10-3 and a batch size of 256. The KL. loss weight is set to 1 x 10-3.
D.3. CIFAR-10
For the rectified flow baseline, we directly use the OT-FM and I-CFM models from (Tong et al.. 2024) and evaluate their performance under different NFEs. For the consistency flow matching model, we take the public implementation from (Yang et al., 2024) and integrate the consistency loss into the same I-CFM model, naming it Consistency- Additionally, we evaluate the original model from (Yang al., 2024) with a larger parameter count, referring to it as Consistency-FNN FM-
For V-RFM model variants, we adopt the I-CFM model from (Tong al., 2024) and add modules to incorporate conditional signals from ID latent For both conditioning mechanisms discussed in Section 4.4, the sampled latent is processed through two MLP layers with SiLU activation, with both hidden and output dimensions set to 512.
In the adaptive norm variant, the latent embedding is combined with the time embedding from fi to regress the learnable scale and shift parameters y and 33 for the adaptive group norm layers. For the bottleneck sum variant, the latent is added to the bottleneck feature of ug- Since the lowest spatial resolution of the baseline network is 4 x 4, the 1D latent spatially repeated and fused with the bottleneck feature via a weighted sum. To ensure effective use of the latent, we assign weighting of 0.9 to the latent and 0.1 to the original velocity feature.
The posterior model le shares a similar encoder structure to up but omits the decoder. To achieve greater spatial compression, we increase the number of downsampling blocks, predicting features at a x 1 spatial resolution. The base channel size is set to 16. Both networks are trained jointly for 600,000 iterations using the Adam optimizer with a learning rate of 2 x 10 and a batch size of 128. The KL loss weighting is presented alongside the results in Table 3.
D.4. ImageNet
We build upon the open-source SiT-XL model (Ma et al., 2024) by incorporating additional modules to integrate conditional signals from the sampled ID latent variable 2. The sampled latent is processed through two MLP layers with SiLU activation, with both the hidden and output dimensions set to 1152. The processed latent then directly added to the original
17

---

Variational Rectified Flow Matching
conditional latent which contains timestep and class label information. The resulting conditional feature is used predict the learnable scale and shift parameters, Y and 8, for the adaptive group normalization layers.
The posterior model shares the SiT-XL architecture but uses only half the number of transformer blocks. To achieve greater spatial compression, we apply an average pooling layer to compress the latent representation into ID vector, which is then processed by an MLP layer to predict j&g and dg. The base channel size is set to 1 1152, the patch size to 2, and the number of heads to 16. Both networks are trained jointly for 800,000 iterations using the AdamW optimizer with learning rate of X 10-4 and global batch size of 256. The KL loss weight is set x 10 3 and the posterior model 9s takes #1 as input. To ensure a fair comparison, we strictly adhere to the original training recipe of SiT (Ma et al., 2024), i.e., we don't tune learning rate, decay warm-up schedules, AdamW parameters, or employ additional data augmentation or gradient clipping during training.
E. Qualitative Results
E.1. CIFAR-10
We present qualitative results of our model trained on CIFAR-IO data in Figure 15.
E.2. ImageNet
We present qualitative results of our model trained on ImageNet data in Figure 16.
18

---

Variational Rectified Flow Matching
Figure 15: Randomly selected samples generated from our model trained on CIFAR-10 data.
19

---

Variational Rectified Flow Matching
RVIEAEH
Figure 16: Randomly selected samples generated from our model trained on ImageNet data.
20


Reviews:
[{"review_id": "a4bj3M0xr1", "reviewer": "Reviewer_8Z7u", "summary": "This paper proposes Variational Rectified Flow Matching (V-RFM), a generative model that integrates Variational Autoencoders (VAEs) with Rectified Flow Matching (RFM). Unlike conventional RFM, which struggles to capture the multimodal nature of the ground-truth velocity vector field and learns only a single averaged direction, V-RFM introduces an encoder-based architecture to enable modeling of multimodal velocity fields. This capability allows V-RFM to theoretically achieve straighter sampling trajectories compared to traditional RFM frameworks.\n\n## update after rebuttal\nI thank the authors for their response and I will maintain my score as Weak Accept. I suggest that the authors add discussion and evidence on training stability and convergence in an updated paper. Also, it would be valuable to report more results in an updated version.", "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": "Thanks a lot for your detailed feedback and for recognizing V-RFM's novelty and promise as a generative model. See below for answers to your questions:\n\n\n**1. Training stability of V-RFM. VAE training is not as stable as RFM because it may encounter the posterior collapse problem.**\n\nGreat question. During our studies we observed that stability can be controlled by the architecture. Specifically, as stated in **L353-357**, bottleneck sum, which fuses the latent $z$ with activations of the velocity network, and adaptive normalization, which explicitly scales and offsets the latent $z$ at multiple layers of the velocity network, are effective at ensuring that the latent variable $z$ sampled from the posterior is not ignored. As shown in **App Fig 15**, the reconstruction losses of V-RFM remain consistently lower than the baseline, demonstrating that the latent variable contributes meaningfully to reducing the reconstruction loss. Furthermore, **Fig 6 and 7** confirm that modifying the latent $z$ alters the predicted image, further verifying that our latent representation remains informative and utilized throughout training. Lastly, we conducted an ablation study on the size of the posterior model, summarized in **App Tab 5**. The results show that performance remains consistent across variations, indicating that even with a very small encoder (6.7% of its original size), the latent information remains informative and helps achieve competitive FID scores.\n\n\n**2. More baselines in Tables 1 and 2, such as flow matching and 1/2/3-Rectified Flow.**\n\nNote, the OT-FM/I-CFM baselines in **Tab 1** and the SiT baseline in **Tab 2** employ the flow matching objective with differences in the parameterization of $x_t$ and the conditional vector field $u_t$.  Specifically, OT-FM parameterizes $x_t$ as $\\mathcal{N}(t x_1, 1-(1-\\sigma) t)$, while I-CFM defines it as $\\mathcal{N}(t x_1+(1-t)x_0, \\sigma)$. The conditional vector field $u_t$ is $\\frac{x_1-(1-\\sigma)x_t}{1-(1-\\sigma)t}$ for OT-FM, while it is simply $(x_1-x_0)$ for I-CFM.  We also note that 1-Rectified Flow is equivalent to I-CFM.\n\nWe have added a comparison with 2/3-Rectified Flow via Reflow. We find that while strong FID scores in the low-NFE regime are achieved, it does so at the cost of limiting peak performance at high NFE. We emphasize that Reflow is a supplementary technique applied on top of Rectified Flow Matching (RFM)—primarily aimed at fast sampling rather than improved sample quality. It also requires $N$ times longer training and a significantly larger fine-tuning dataset, where $N$ denotes the number of Reflow rounds. These differences make a direct comparison with our V-RFM less fair. Hence, RFM without Reflow is a more appropriate baseline. Additionally, Reflow can be applied to our method as well, potentially improving results at the cost of increased training overhead. We will clarify these points in the final version and include additional qualitative comparisons for a more precise evaluation.\n\n|Methods    |# Params    |2    |5    |10    |50    |100    |1000    |Adaptive|\n|:---------:|:--------:|:--------:|:-----------:|:-----------:|:-----------:|:-----------:|:-----------:|:-----------:|\n|RFM w.o. Reflow / I-CFM    |36.5M    |168.654    |35.489    |13.788    |5.288    |4.461    |3.643    |3.659|\n|RFM w. 1 Reflow            |36.5M    |7.512    |5.906    |5.513    |5.283    |5.276    |5.276    |5.275|\n|RFM w. 2 Reflow            |36.5M    |7.559    |6.925    |6.776    |6.729    |6.733    |6.752    |6.752|", "claims_and_evidence": "Yes, the major claim of this submission is supported by experiments.", "methods_and_evaluation_criteria": "Yes, the proposed method makes sense for the current problem.", "theoretical_claims": "I checked the correctness of the proof of Claim 1 of the main paper. The proof looks correct to me.", "experimental_designs_or_analyses": "I suggest adding more baselines in the experiments in Tables 1 and 2, such as flow matching and 1/2/3-Rectified Flow, which can increase the soundness/validity of experimental designs. It is also recommended to increase qualitative comparison with those baselines.", "relation_to_broader_scientific_literature": "No.", "essential_references_not_discussed": "The paper provides a full discussion of related work.", "other_strengths_and_weaknesses": "Strengths\n+ The proposed V-RFM is novel for me, and is promising as a generative model.\n+ The introduction of encoder enables V-RFM to have the ability to infer latent codes given the data samples.\n\nWeaknesses\n- I am more concerned about the training stability and effectiveness of the method as a generative model. V-RFM introduces VAE to RFM. Relatively speaking, the training stability and effectiveness of VAE are not as good as RFM because it may encounter the posterior collapse problem. However, this paper does not seem to discuss whether the encoder will encounter this problem and how it affects the generation effect.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Please see Weaknesses and Experimental Designs Or Analyses.", "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": "AU8BfR6HvN", "reviewer": "Reviewer_Sjgu", "summary": "The paper introduces Variational Rectified Flow Matching (VRFM), a novel approach that integrates techniques from Rectified Flow Matching (RFM) and Variational Autoencoders (VAEs). This design aims to address the vector ambiguity issue inherent in the original RFM method. Through extensive experiments, the authors demonstrate that VRFM improves data generation quality, producing samples that more closely align with ground truth compared to standard RFM. Additionally, the proposed method effectively mitigates vector ambiguity to a significant extent. Empirical evaluations on benchmark image datasets, including CIFAR-10 and ImageNet, reveal that VRFM consistently achieves superior Fréchet Inception Distance (FID) scores compared to baseline models, highlighting its effectiveness in generative modeling.", "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": "Thanks a lot for your detailed feedback and for recognizing our well-written paper, extensive experiments, and comprehensive evaluation of performance. We also appreciate the acknowledgment of V-RFM’s effectiveness in addressing velocity ambiguity and its superior performance. Below, we address questions:\n\n**1. The method requires additional parameters to compute the latent representation during training, increasing computational complexity.**\n\nYes, extra computation is used during training. It is essential to extract latent information within the velocity network, ultimately enhancing generation quality and expressiveness. As discussed in **Sec 3.3 and 4**, the latent encoding network is not used during inference—we directly sample $z$ from the prior distribution. Furthermore, the increase in parameters is minimal (e.g., less than 0.3% on ImageNet), and the impact on inference speed is negligible while delivering superior results. The ablation study on the size of the posterior model summarized in **App Tab 5** shows that performance remains consistent across variations, demonstrating the robustness and flexibility of our approach. This allows users to balance training efficiency and runtime quality based on their computational constraints.\n\n**2. While the proposed method outperforms baseline, it lags behind the current SOTA.**\n\nThe primary contribution of our work is not in surpassing the current state-of-the-art, but rather in introducing a methodological innovation, i.e., capturing the multimodal velocity field, which offers new avenues for improvement in the field. While our method lags behind the state-of-the-art, it can be combined with innovations put forth by current SOTA SiT-XL/2 + MG [1] and SiT-XL/2 + REPA [2], both built upon the SiT framework. As detailed in our experiments (**Sec 4.5, L364-367 right column**), we strictly followed the original training recipe from the open-source SiT repository and replicated the process outlined in the SiT paper for a fair comparison. Our experimental results provide empirical evidence of the effectiveness of our approach.\n\n[1] Tang, Z. et al. (2025). *Diffusion Models without Classifier-free Guidance*.\n\n[2] Yu, S. et al. (2024). *Representation alignment for generation: Training diffusion transformers is easier than you think*.\n\n**3. A definitive ground truth mapping may not exist.**\n\nA definitive “ground-truth” does not exist indeed. The (“ground-truth”) velocities form a distribution for every $(x_t, t)$-location. To see this we independently sample from both the source distribution and the target data distribution, calculate the rectified flow interpolants for each pair, and visualize them in **Fig 1** and the velocity distribution in **Fig 3(a)**. Our method models this velocity distribution at every $(x_t, t)$-location, while standard rectified flow matching cannot capture it. In the abstract and the main paper, we use quotation marks around “ground-truth” to emphasize this distinction. We will correct any missing instances in the final version. \n\n**4. Provide a rigorous theoretical proof to substantiate the claim that the proposed method resolves vector ambiguity.**\n\nIn **Sec 3**, we show that the proposed approach leads to a mixture model for the velocity distribution (**L179-180**). A mixture model is theoretically capable of capturing multi-modality. Following classic expectation maximization or variational inference, we introduce the recognition model, derive the lower bound of the marginal likelihood for an individual data point (**L189-192**), and present the variational flow matching objective (**L197-199**). To further substantiate our approach, in **App A**, we show how to prove that the distribution learned by the variational objective preserves the marginal data distribution. Empirically, on 1D data we visualize the learned velocity distribution in **Fig 3**, showing that the method indeed learns the velocity ambiguity. Further, during training on high-dimensional data, we observe that our method achieves better velocity reconstruction losses (**App Fig 15**) compared to standard rectified flow, indicating that the predicted velocities more accurately approximate the “ground-truth” velocities.\n\n**5. How does the choice of latent variable affect the generative path.**\n\nWe studied the role of $z$ for MNIST data (**Fig 6**) and CIFAR-10 data (**Fig 7**). As noted in **Sec 4.3 and 4.4**, we observe clear patterns in the generated samples based on $z$. Specifically, images conditioned on the same latent $z$ exhibit consistent color patterns, while images at the same grid location show similar content. These observations validate the effectiveness of the latent variable $z$ in influencing and controlling the generated samples.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "I reviewed the proofs quickly, and they appear to be correct.", "experimental_designs_or_analyses": "yes, there are not issues.", "relation_to_broader_scientific_literature": "The proposed Variational Rectified Flow Matching (VRFM) builds upon previous work in Rectified Flow Matching (RFM) and Variational Autoencoders (VAEs). A key limitation of RFM is vector ambiguity, which can hinder generative performance. VRFM addresses this issue by integrating a variational framework, a well-established technique in generative modeling for learning more effective latent representations. Experimental results show that VRFM achieves improved Fréchet Inception Distance (FID) scores compared to RFM, indicating enhanced generative quality beyond what was previously achievable with RFM alone.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths:\n\t1.\tThe paper is well-written, with clear and easy-to-follow explanations.\n\t2.\tThe authors conduct extensive experiments to evaluate the proposed method comprehensively.\n\t3.\tThe proposed approach effectively addresses vector ambiguity and demonstrates superior generative performance compared to baseline models.\n\nWeaknesses:\n\t1.\tThe method requires additional parameters to compute the latent representation during training, increasing computational complexity.\n\t2.\tSome claims rely on empirical observations; providing stronger theoretical proofs would further strengthen the paper.\n\t3.\tWhile the proposed method outperforms baseline models, it still lags behind the current state-of-the-art models.", "other_comments_or_suggestions": "No", "questions_for_authors": "1.\tIn Figure A, the visualization depicts the ground truth data, including the source data distribution, target data distribution, and the mapping between them. I assume that the mapping between source and target data points is randomly generated, meaning any point in the source could potentially correspond to any point in the target. If this assumption is correct, then a definitive ground truth mapping may not exist. Could the authors clarify how they obtained the ground truth mapping used in the visualization?\n2.\tThe claim that the proposed method resolves vector ambiguity is primarily supported by visualizations using toy data. Would it be possible to provide a rigorous theoretical proof to substantiate this claim?\n3.\tHow does the choice of latent variable affect the generative path? Specifically, would the generative trajectory be different when using different numbers of latent variables during inference?", "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": "No"}}, {"review_id": "tTXsiefdxo", "reviewer": "Reviewer_7qJ3", "summary": "# Update\n\nIn the rebuttal, the authors have addressed many of my questions and criticism. I feel that the main issue has not been adequately addressed, so I decided not change the score.\n\nTo elaborate, I feel that the method adds complexity to diffusion models. The added complexity has to serve some purposes in order for it to be worthwhile. In light of the paper, the added complexity has two practical benefits.\n\n(1) It improves the scores of over the base models.\n(2) It allows a form of conditional sampling.\n\nFor (1), while improvements are consistent across sampling steps, they are not very pronounced. \n\n* For the CIFAR-10, the scores are very close to baselines from NFE=10 onward. For NFE=2 or 5, improving from 166 to 104 or 36 to 25 cannot really be consider a practical improvement because the images produced by both the V-RFM and baselines are still of low quality. Distillation techniques, only the other hand, reduces 166 to a 1-digit figure at NFE=2 and NFE=5. This is what I meant when I said improvements cannot be compared to distillation techniques.\n\n* For ImageNet datasets, while there is quite a significant gap when not using CFG, the gap becomes gradually smaller as training becomes longer (although the authors show that the percentage improvement still incrasese) and is significantly reduced when CFG is applied. This means that techniques already employed on non-variational models are already quite effectively, and one has to wonder whether adding a variational component to the model would worth the trouble.\n\nBenefit (2), on the other hand, is much more interesting to me because it is a feature that a normal diffusion model has: VAE-style latent codes. This is something that I feel *worths the trouble* of making the model more complex, so I think stressing this benefit should become a bigger part of the paper. However, from reading the rebuttal, while the authors did experiments on interpolating the latents, it seems they have not investigated how to use latent codes to control the outputs with more degrees of certainly. As a result, it is unlikely that the final version of the paper would contain more material in this direction.\n\nBecause of these concerns, I decided not to change my evaluation.\n\n# Old Summary \n\nThe paper proposes \"variational rectified flow matching,\" an extension to (rectified) flow matching. The latter trains a neural network $v_\\theta(x,t)$  that predicts the expected value of velocities induced from velocity fields that continuously transforms one Gaussian distribution (whose mean comes from a \"source\" distribution) to another Gaussian distribution (whose mean comes from a \"target\" distribution). The paper observes that there is a distribution of velocity vectors at $p(v|x,t)$ each $(x,t)$ point and explores modeling them as a part of training the flow matching model instead of just estimating the distribution's means as is done by standard rectified flow matching.\n\nTo model the velocity distribution at each point $(x,t)$, the paper casts the flow matching model as a latent variable model, much like a variational autoencoder. The flow matching model nows accept a latent variable $z$ and becomes $v_\\theta(x,t,z)$. The latent is supposed to come from a prior distribution $p(z)$, taken to be the standard multivariate Gaussian distribution. To train the model, one needs to model the conditional distribution $q(z|x,t)$, which the paper models with an encoder network, muck like what VAE does.\n\nThe distribution, conditioned on the latent $z$, is modeled as a Gaussian distribution around the predicted value: $p(v|x,t,z) = \\mathcal{N}(v; v_\\theta(x,t,z), I)$. This gives $p(v|x,t) = \\int \\mathcal{N}(v; v_\\theta(x,t,z), I) p(z)\\, \\mathrm{d}z$ where $p(z)$ is the prior distribution of $z$, which is taken to be the standard Gaussian distribution. The flow matching matching, together with the encoder, can be trained like a VAE with a loss derived from the ELBO of $\\log p(v|x,t)$. The KL divergence term between $q(z|x,t)$ and $p(z)$ remains the same, but the reconstruction term in the VAE loss is replaced by the conditional flow matching loss instead. To sample with a flow matching model trained this way, one must first sample $z$ from $p(z)$, and then one can use the flow matching model to generate a sample as usual with the exception of feeding $z$ to it at every integration step.\n\nThe paper demonstrates that the method has several benefits.\n\n(1) It yielded better evaluation scores on various datasets and models architectures compared to vanilla flow matching. In particular, the gap is wider when low number of NFEs are used to generate samples or when the training time is shorter.\n\n(2) It yields flow matching models that can model distributions at each  $(x,t)$ point better. In aggretate, such models can model sampling trajectories, and these trajectories seem to be less curved than the non-intersecting trajectories of vanilla flow matching models.\n\n(3) By varying $z$ at test time, one can control the generation output.", "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": "Thanks for feedback and for highlighting our theoretical contributions and strong results across datasets and models. \n\n**1. Quantify the average curvature of 2D trajectories**\n\nWe calculated the curvature for 2D data results (**Sec 4.2**) and find significantly lower curvature for our method:\n\n| | Mean/Max Curvature|\n| - | - |\n|Baseline (rectified flow)|\t21.03/171.35|\n|Ours|\t0.98/4.23|\n\n**2: How to obtain desired digits/brightness via $z$. Latent interpolation would strengthen the claim**\n\nGreat suggestion. We conducted interpolation experiments and summarize findings verbally as images can't be uploaded. For MNIST, interpolating latents leads to smoothly transitioning digits (e.g., 1 → 7 → 8 → 3 → 2 → 0). Note, **Fig 6** illustrates that each digit corresponds to a specific latent $z$. Intermediate digits emerge naturally when interpolating. For CIFAR-10, we observe analogous effects—interpolating between two latents leads to smooth transitions in brightness and color patterns.\n\n**3: Marginal data distribution preservation lacks a proof**\n\nThe proof in **App A** refers to results established by Liu et al. to provide adequate credit. This may require readers to reconstruct parts of the proof using work of Liu et al. To make the paper self-contained, we will expand it to add the missing steps (derivation of $d/dt E[h(X_t)]$ and equivalence of $0 = E_Z \\int(...))$.\n\n**4. Formulations of OT-FM and I-CFM are almost equivalent. Redundant comparison with I-CFM**\n\nOT-FM/I-CFM differ in $x_t$ and the conditional vector field $u_t$.  Specifically, OT-FM defines $x_t$ as $\\mathcal{N}(t x_1, 1-(1-\\sigma) t)$ and $u_t$ is $\\frac{x_1-(1-\\sigma)x_t}{1-(1-\\sigma)t}$. I-CFM defines $x_t$ as $\\mathcal{N}(t x_1+(1-t)x_0, \\sigma)$ and $u_t$ as $(x_1-x_0)$.  SiT uses the rectified flow objective, which is equivalent to I-CFM. We'll follow the suggestion to remove OT-FM for consistency.\n\n**5. Figure 12 of MNIST FID score should be table**\n\nWe'll convert **Fig 12** into a table. \n\n**6. Discussions of related work**\n\nThanks for highlighting these works. We'll add a full discussion. Below is a brief summary:\n\n1. Denoising Diffusion GANs replace the Gaussian model in the denoising step with a multimodal distribution. Unlike our method, a conditional GAN with separate discriminator models the distribution. But GANs face mode collapse and stability issues. In contrast, our method uses rectified flow matching, preserving the maximum likelihood benefits.\n2. Score-based Generative Modeling uses a VAE to map raw data $x_0$ into latent space $z_0$, with the VAE jointly trained with score-based generative modeling (SGM). Unlike our approach, SGM still faces ambiguity issues due to its use of a uni-modal Gaussian distribution.\n3. Stable Target Field notes that the posterior distribution is multi-modal. To model this distribution, the paper reduces training target variance using a reference batch. In contrast, our method directly models this multi-modal posterior via a recognition model.\n\n**7. Improvements between SiT-XL and V-SiT-XL diminish as training continues**\n\nA diminishing gap is expected as absolute values decrease, a trend also presented in Fig 2 of the SiT paper (comparing the gap of DiT v.s. SiT). However, the relative improvement remains strong. We extend training to 1200k steps and report both absolute and percentage improvements. We observe the percentage improvements increase with more training iterations. \n\n||200k\t|400k\t|600k\t|800k\t|1200k\t|800k (cfg=1.5)\t|1200k (cfg=1.5)|\n| - | - | - | - | - | - | - | - |\n|SiT-XL\t|26.09\t|17.84\t|14.77\t|13.15\t|11.26\t|3.43\t|2.97|\n|V-SiT-XL\t|23.34\t|14.60\t|12.00\t|10.62\t|8.97\t|3.22\t|2.76|\n|abs diff\t|2.75\t|3.24\t|2.78\t|2.53\t|2.29\t|0.21\t|0.20|\n|percent diff\t|10.53%\t|18.16%\t|18.79%\t|19.24%\t|20.31%\t|6.12%\t|6.86%|\n\n**8. Improvements are small, not comparable to distillation**\n\nWe respectfully disagree. Our improvement is solid, while strictly following the open-source SiT training. Results are consistent with the SiT-to-DiT improvement, demonstrating a comparable level of progress (19.5 for DiT-XL, 17.2 for SiT-XL, and 14.6 for V-SiT-XL at 400k steps, as shown in  **Tab 2**).\n\nAlso note, our key contribution is to model the velocity distribution. We find this to consistently improve evaluation metrics across all NFEs. While distillation methods may show improvements for low NFEs, our method achieves better results across both low and high NFEs.\n\nAdditionally, as discussed in **App B**, V-RFM focuses on single-stage training to capture a multi-modal velocity distribution from “ground-truth” data without leveraging pre-trained models. Exploring distillation for V-RFM is an exciting avenue for future research, particularly when the interest is to improve results for low NFEs.\n\n**9. Replace \"data-domain-time-domain\"**\n\nGreat suggestion. We'll revise. \n\n\n**10. ImageNet performance below 250 NFEs**\n\n**Fig 8** shows those results, revealing a consistent boost, further highlighting our method's effectiveness.", "claims_and_evidence": "I believe claim (1) is supported by enough evidence as the paper contains experiments on 5 datasets, and three types of architecture. Performance gaps on ImageNet generation without guidance are quite significant.\n\nClaim (2) is supported by showing that, in the 1D dataset, the vanilla flow matching model often collapes the velocity distribution. The experiment on the 2D dataset clearly shows in Figure 4(c) that the proposed method can model intersecting trajectories. However, while Figure 4(b) shows trajectories that seem to be more curved than those in Figure 4(c), it is better to quantify the average curvature of the trajectories and show the numbers along with the pictures.\n\nFor Claim (3), the paper shows that the generated samples changed when $z$ is changed in Figure 6 (MNIST dataset) and Figure 7 (CIFAR-10) datset. While one must accept that the outputs do change, it is not quite clear whether these changes are useful or intuitive. In Figure 6, different areas of the unit square seems to correspond to different digits, but the paper does not explicity show how one can obtain the desired digits through controlling $z$. In Figure 7, different latent codes seem to yield different overall brightness of the outputs, but it is unclear whether one can arbitrarily control the brightness though varying $z either. Several simple experiments where the latent codes interpolated to get the desired outcomes would make this claim stronger.", "methods_and_evaluation_criteria": "The method seems to make sense from the problem at hand.\n\nThe 1D and 2D datasets are used to effectively show the ability of the trained models to better capture velocities of distributions. However, I think it is better to quantify the average curvature of 2D trajectories instead of just showing pictures.\n\nMNIST, CIFAR-10, and ImageNet are widely used datasets to benchmark generative models. The paper also uses appropriate architectures for these datasets.\n\nThe metrics (log-likelihood for 1D/2D datsets and FID for image datsets) also make sense.", "theoretical_claims": "Section 3 is easy to read and seems sound. However, the most important theoretical claim that the paper's training method preserves the marginal data distribution lacks a full formal proof. What is provided in Appendix A is a proof sketch where what is to be proven is supported by statesments such as \"one can show equivalence\" and \"equivalence can be shown via\" without any work being shown. A reader would have to go to Liu's paper and follow all the logic by themself. I suggest the authors write down the proof in Appendix A for completeness.", "experimental_designs_or_analyses": "To my understanding, the paper compares its training method, variational rectified flow matching (VRFM), against two other training methods:\n\n(1) vanilla flow matching (OT-FM), proposed by Lipman et al. (2023), and \n(2) independent coupling flow matching (I-CFM), proposed by Tong et al. (2024).\n\nThe mathematical formulation of these algorithms are slightly different, and they are almost equivalent if their source distributions are the standard Gaussian $\\mathcal{N}(0,I)$. As a result, I find the comparison with I-CFM for the CIFAR-10 dataset redundant. In fact, the numbers of OT-FM and I-CFM in Table 1 are very close. Moreover, the comparison with I-CFM is only available for the CIFAR-10 dataset.\n\nThe paper would feel more consistent if either (a) comparison with I-CFM is removed or (b) it also provides comparison with I-CFM training method for other datasets.", "relation_to_broader_scientific_literature": "The paper propses a new extension to flow mathcing that allows sampling paths to be chosen based on a latent vector. In a sense, it is an interesting way to combine VAE with flow matching model.", "essential_references_not_discussed": "(1) I believe that the idea of modeling distributions of values inside a diffusion sampling process has been explored previously, and this paper is one instance of it. An example that comes to mind is the Denosing Diffusion GANs paper by Xiao et al. [1], which uses conditional GANs to model the denoising distribution at each step of the diffusion process.\n\n(2) There is another way to combine a diffusion model with a VAE, and it involves using the former to model the latent space of the latter [2].\n\n(3) The opposite idea to the one proposed in the paper is to regard the distribution of target values to match the neural network's output against as a kind of noise and seeks to eliminate it. Stable Target Field by Xu et al. implements this idea. [3]\n\n1. Zhisheng Xiao, Karsten Kreis, Arash Vahdat. Tackling the Generative Learning Trilemma with Denoising Diffusion GANs. ICLR 2022.\n2. Arash Vahdat, Karsten Kreis, Jan Kautz. Score-based Generative Modeling in Latent Space. NeurIPS 2021.\n3. Yilun Xu, Shangyuan Tong, Tommi Jaakkola. Stable Target Field for Reduced Variance Score Estimation in Diffusion Models. ICLR 2023.", "other_strengths_and_weaknesses": "I believe this paper presents a new and interesting formulation of flow matching models. However, I do not feel that its benefits are compelling. Being able to model the velocity distribution is clearly a novelty, but rather a conceptual one. The most concrete one improvements in metrics which diminish as the number of function evalutions become larger. Still, for image datasets such as CIFAR-10 and MNIST (and perhaps ImageNet), these improvements are small and not comparable to improvements achieved by distillation methods.\n\nAnother benefit claimed by the paper is the ability to control the output through the latent code $z$. However, to make the paper stronger, I think the paper should do more experiments to highlight this aspect. This can include a simple method to sample a specific number from the MNIST dataset or a way to control the brightness of CIFAR-10 samples.", "other_comments_or_suggestions": "I suggest replace the term \"data-domain-time-domain\" with \"data-time space\" or \"$(x,t)$-space,\" which should be more concise.", "questions_for_authors": "(1) In Table 2, it seems that performance gap between SiT-XL (the baseline) and V-SiT-XL seem to diminish as training becomes longer. Can you show the FID scores at 1200K steps and/or 1600K steps to confirm that the gap still exists there as well?\n\n(2) It would be interesting to see how ImageNet models perform at NFEs lower than 250. Please include those numberse in Table 2 if possible.", "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": "mj8dm5EDnH", "reviewer": "Reviewer_hji9", "summary": "This paper introduces a variational rectified flow matching method. Instead of learning a deterministic mean velocity at time $t$, the paper explicitly models a distribution over the velocity $v_t$, grounding the approach in VAE theory.", "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": "Thanks for detailed feedback and for recognizing the importance of addressing ambiguity in the marginal velocity field, and our strong results.\n\n**1. Paper's unique contributions**\n\nWe study a method for capturing multi-modal velocity vector fields. We show that incorporating an unobserved continuous latent variable z via a variational formulation (akin to a VAE) enables the velocity model to learn a multi-modal vector field. Experiments across diverse data and models show that our method outperforms classic approaches.\n\n**2. How many z used for large-scale experiments. Averaging**\n\nTo remain fair and as shown in Algorithm 2, we sample a single $z$ per data point and keep it fixed throughout integration. No averaging is performed on $z$ or predicted velocity.\n\n**3. Impact of KL weight**\n\nWe summarize the impact of the KL weight $\\lambda$ based on our experimental findings:\n\n1. The model successfully captures velocity ambiguity and predicts crossing flows when $\\lambda$ is in a reasonable range (in [0.1,10.0]).\n2. When $\\lambda$ is large (e.g., 100.0), the latent model is forced to equal a standard Gaussian. Hence, the latent z contains minimal useful information. Hence, the velocity network behaves similar to one obtained using classic rectified flow. This is also apparent in the loss: the KL loss diminishes, and the velocity reconstruction loss is comparable to the baseline loss. The resulting flow cannot capture ambiguity.\n3. When $\\lambda$ is small (e.g., 0.01), the model can exploit excessive information from the latent. This leads to a very low velocity reconstruction loss but a very high KL loss. The resulting flow appears as straight lines, but the endpoint distributions do not match the target data due to the mismatch between the predicted posterior and prior.\n\nIn our experiments, we didn't tune the KL regularization weight much, but instead scaled the KL loss with the dimension of the latent variable. E.g., for ImageNet SiT experiments, we directly used the KL weight that we employed for CIFAR-10.\n\n**4.  Implementation of $q_\\phi $and $v_\\theta$**\n\nAs stated in **Sec 4.4 (L342 - 343)**, $q_\\phi$ and $v_\\theta$ share a similar structure but are separate nets. We will clarify this to avoid confusion. Regarding the increase in parameters, as described in **Sec 3.3 and Sec 4 (L262, 348, 372 right column)**, during inference, $q_\\phi$ is not used. Instead, we sample the latent variable from a prior. The only increase in parameters comes from the two MLP layers to fuse the latent z in $v_\\theta$. This design ensures that our velocity network remains comparable in size to the baseline, with less than a 2% parameter increase for CIFAR-10 and less than a 0.3% increase for ImageNet.\n\n**5. Code Snippet**\n\nWe will release code and models. Note, we provided more implementation details in **App D**. We are happy to address further questions during the rebuttal phase. \n\n**6. More baselines: consistency models, shortcut models, HRF**\n\nConsistency model:\nA detailed comparison to consistency models, particularly distillation models, is included in **App B**. We used the recently developed consistency flow matching [1]. It improves upon consistency models [2] and is more closely related to flow matching. We summarized the results in **App C.1**. Our key findings:\n\n* The consistency flow matching model performs well at low function evaluation regimes (i.e., with NFEs of 2 or 5).\n* Its performance degrades as NFEs increase.\n* Its best performance across all NFEs remains below classic rectified flow matching and our variational rectified flow matching.\n\nWe also highlight an exciting future research direction: combining variational flow matching with consistency models, which could further enhance results.\n\nShortcut model:\nWe evaluate the Shortcut Model (XL), trained for 800k iterations, using the FID score and following the same evaluation protocol used in **Tab 2**. Our results show that our method consistently outperforms it.\n\n| | Params (M) |FID|\n| - | - | - |\n|SiT-XL| 675| 13.1|\n|Shortcut-XL| 676| 19.752 (128 NFE, reproduced)/19.630 (250 NFE, reproduced)|\n|V-SiT-XL| 677| 10.6|\n|SiT-XL (cfg=1.5)| 675| 3.43|\n|Shortcut-XL (cfg=1.5)| 676| 3.8 (128 NFE, from paper)/4.709 (128 NFE, reproduced)/4.707 (250 NFE, reproduced)|\n|V-SiT-XL (cfg=1.5) |677| 3.22|\n\nHierarchical Rectified Flow (HRF): This concurrent work also aims to model multi-modal velocity and acceleration fields but uses a hierarchical rectified flow. Their method requires multiple integrations during inference, making it slower than our approach. Also, HRF does not support semantic disentangling of flows, as demonstrated in our **Fig 6 and 7** for MNIST and CIFAR-10.\n\n[1] Yang, L. et al. (2024) Consistency flow matching.\n\n[2] Song, Y. et al. (2023) Consistency models.\n\n**7. || should be used instead of | in KL**\n\nThanks for spotting this typo, we'll fix.\n\n**8. resources for SiT-XL on ImageNet 256**\n\nWe used 8 H100 GPUs and trained the model for about 3.5 days.", "claims_and_evidence": "Partially, the baselines do not include recent methods. See the detailed comments in the Weaknesses section.", "methods_and_evaluation_criteria": "Yes. They make sense.", "theoretical_claims": "Yes, I checked Claim 1 and the proof.", "experimental_designs_or_analyses": "Yes. I checked the expriments on synthetic data and CIFAR10 and ImageNet.", "relation_to_broader_scientific_literature": "The paper addresses the ambiguous nature of the marginal velocity field, which is critical for the performance of FM models. Previous methods tackling this challenge mainly focus on using different noise-data couplings or distillation-based methods to straighten the trajectories. This paper takes an orthogonal approach.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "### Strengths:\n\n1.  The paper addresses an important problem in flow matching: the ambiguity of the marginal velocity field.\n\n2.  The method demonstrates strong performance compared to the baseline OT-FM and I-CFM methods.\n\n### Weaknesses:\n\n1.  Novelty: VAEs are already well-established, which positions this work as an application of VAEs to flow matching.  It would be beneficial if the authors could further highlight the paper's unique contributions beyond the direct application of VAEs.  (To be clear, this is more of an open-ended question than a significant limitation. If VAEs are shown to be a good tool for the problem the authors address, the paper can still be considered strong.)\n\n2.  Sampling in Algorithm 1: In Algorithm 1, how many samples of $z$ are used for large-scale experiments (e.g., ImageNet)? How does the sample size affect performance? Furthermore, during inference, are the sampled $z^{(i)}$ first averaged and then fed into the velocity predictor, or are they individually fed into the velocity predictor to predict $v^{(i)}$, which are then averaged to obtain the final velocity?  Clarifying this is crucial, as these two approaches have significantly different implications for inference speed.\n\n3.  KL Divergence Regularization: Based on experience, VAEs are often sensitive to the choice of the KL divergence regularization weight (denoted as $\\lambda$).  This sensitivity is also apparent in Table 1. The authors should provide an analysis of the impact of $\\lambda$ on the learned velocity field and discuss the intuition behind its effect.\n\n4.  Implementation Details: The authors did not provide source code, preventing a deep dive into implementation details. Specifically, there are several options for implementing the pipeline:\n    (A) The encoder $q_{\\phi}$ shares parameters with *a part of* the velocity predictor $v_{\\theta}$.  In this case, the initial part (i.e., the first several layers) of the velocity predictor could be used to predict $\\mu_{\\phi}$ and $\\sigma_{\\phi}$, while the latter part predicts the velocity field.\n    (B) The encoder $q_\\phi$ uses a similar structure to $v_\\theta$ but is a separate network, and they are jointly trained.\n\n    If (B) is the case, it could introduce a substantial number of extra parameters, potentially limiting scalability. However, Table 1 shows that the total number of parameters is 37M, compared to 36.5M for the baseline.  This discrepancy needs further clarification.\n\n5.  Related to point 4, it would be very helpful if the authors could provide a code snippet illustrating the implementation of lines 342-357.\n\n6.  Limited Baselines: The baselines only include vanilla OT-FM and I-CFM.  Other relevant methods, such as distillation-based consistency models (Song et al. 2023) and shortcut models (Frans et al. 2024), which may also address velocity field ambiguity by encouraging \"stepping over\" ambiguous regions of $x_t$ with merged steps. If possible, these methods should also be considered. Including a broader range of baselines would further strengthen the paper. E.g. ``Towards Hierarchical Recitfied Flow'' from Zhang et al. (ICLR 2025) addressed a similar multimodal problem.", "other_comments_or_suggestions": "### Suggestions:\n1. In Eq. (4), $||$ should be used as separator for $q$ and $p$ in the KL divergence in stead of $|$.\n\n### Minor questions:\n1. What resources have you used for training SiT-XL on ImageNet-1k with 256x256 resolution? How long is the training schedule?", "questions_for_authors": "See the questions in the Weaknesses.", "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: After the rebuttal and discussion phases, the paper received scores of 4, 3, 3, and 3, which exceed the expected threshold for acceptance. After briefly reviewing the comments and the authors' responses, I believe the paper meets the acceptance criteria for ICML.

However, I strongly encourage the authors to address the unresolved concerns raised by the reviewers—such as the theoretical explanations pointed out by Reviewer Sjgu, the remaining issues raised by Reviewer 7qJ3, and the request for code release made by Reviewer hji9. These revisions would help strengthen the paper and ensure its contributions are more clearly presented.
Decision: Accept (poster)