Variational Rectified Flow Matching

Authors: Pengsheng Guo, Alex Schwing
Primary area: generative models
Keywords: Flow Matching, Diffusion Model, Generative Model

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. Further, averaging leads to integration paths that are more curved while making it harder to fit the target distribution. In contrast, the studied variational rectified flow matching is able to capture the ambiguity in flow directions. We show on synthetic data, MNIST, and CIFAR-10 that the proposed variational rectified flow matching leads to compelling results with fewer integration steps.

Decision: reject
Reviewer scores: 5, 6, 3, 6

Reviews and discussion:
<thread>
creation date:
1730118451303

thread_type:
official_review

<thread_object>
creator: reviewer_y61o

<content>
summary:
This paper presents Variational Rectified Flow Matching (VRFM), a framework that improves classic rectified flow matching by incorporating multi-modal velocity vector fields based on a variational perspective. Previous flow matching approaches average out directions, leading to curved integration paths and hindering accurately fitting the target distribution. VRFM, by contrast, captures the multi-modality of flow directions, thus preserving directional diversity. Experimental results on synthetic data, MNIST, and CIFAR-10 demonstrate that VRFM achieves promising results with fewer integration steps.

strengths:
- Significance: This paper addresses a notable limitation in flow matching models, specifically the ambiguity at path intersections that results in curved sampling trajectories. By tackling this ambiguity, the proposed approach demonstrates clear improvements over existing rectified flow models, particularly in low NFE settings.

- Originality and clarity: The paper is well-written and easy to follow, clearly presenting concepts. Interpreting the flow-matching objective through variational inference to reduce directional ambiguity is conceptually sound, adding a meaningful perspective to the flow-matching framework.

weaknesses:
- In line 53, the authors claim, “This results in trajectories that are more straight”.  It is unclear why reducing ambiguity would inherently lead to straighter flows. Including a theoretical proof or a detailed explanation to clarify this result would strengthen the argument.

- For completeness, the paper should include proofs demonstrating that the learned distribution from VRFM preserves the marginal data distribution, as established in Theorem 3.3 in [1].

- The most concerning part of this paper is limited evaluation and performance compared to the recent papers. The empirical evaluation is restricted to MNIST and CIFAR-10, which limits the generalizability of the findings. Extending the evaluation to additional datasets, such as ImageNet 64x64, would improve the generalizability of the findings. Furthermore, the reported results of VRFM in low NFE regimes (e.g., 104 FID for 2 NFE on CIFAR-10) are less compelling, given the recent advances [1,2,3] in reducing sampling costs in diffusion (or rectified flow) models. For instance, reflow on rectified flow (e.g., 2-rectified flow) achieves a 4.85 FID with a single step [1]. Results showing VRFM’s performance with the reflow technique would provide a more competitive comparison. 

- It would be valuable if the authors could provide results on conditional generation setting.

questions:
- Is there any reason for designing the input of the posterior encoder as [x0, x1, xt, t]?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730541088550

thread_type:
official_review

<thread_object>
creator: reviewer_obzd

<content>
summary:
The paper introduces Variational Rectified Flow Matching as a method to model multi-modal velocity and ambiguity in the data-space-time-space domain. The properties of Variational Rectified Flow Matching are studied and validated through experiments with visualizations on low-dimensional synthetic data. Compelling results are demonstrated on the synthetic data, MNIST, and CIFAR-10 datasets.

strengths:
1. The paper presents a valuable observation: because the vector field is parameterized via a Gaussian at each data-domain-time-domain location, ambiguity cannot be captured.

2. The analytical experiments with visualizations are well-executed and contribute significantly to validating the theoretical analysis.

weaknesses:
1. Current evaluations are too weak and this paper lacks enough experiments on more real-world and complex datasets, such as the AFHQ, CelebA and ImageNet datasets. The evaluations on these benchmarks are necessary for demonstrating the effectiveness of the proposed method.

2. The authors should sufficiently discuss these missing related works [1][2][3][4][5], and compare with them in the experiments.

[1] Nguyen B, Nguyen B, Nguyen V A. Bellman Optimal Stepsize Straightening of Flow-Matching Models[C]//The Twelfth International Conference on Learning Representations. 2024.

[2] Song, Yang, et al. Consistency models. arXiv preprint arXiv:2303.01469 (2023).

[3] Yang, Ling, et al. Consistency flow matching: Defining straight flows with velocity consistency. arXiv preprint arXiv:2407.02398 (2024).

[4] Yan, Hanshu, et al. Perflow: Piecewise rectified flow as universal plug-and-play accelerator. arXiv preprint arXiv:2405.07510 (2024).

[5] Kim, Dongjun, et al. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279 (2023).

questions:
I am curious whether Variational Rectified Flow Matching can enhance performance through "reflow", similar to classic rectified flow matching as discussed in [1].


[1] Liu, Xingchao, and Chengyue Gong. "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow." The Eleventh International Conference on Learning Representations.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626146997

thread_type:
official_review

<thread_object>
creator: reviewer_qka9

<content>
summary:
The paper introduces a novel framework, Variational Rectified Flow Matching, which addresses the limitations of conventional Rectified Flow (RF) methods in capturing ambiguity in velocity distributions. By incorporating an additional latent variable z drawn from a Gaussian prior, the framework models multiple modes of ambiguity. An encoder is used to derive the posterior distribution p(v∣xt,t,z) at a specific sample xt and time t. This approach is claimed to better capture ambiguity and improve the empirical performance of diffusion model generation.

strengths:
The strengths of the paper are as follows:

- Clear, easy-to-follow presentation with strong empirical performance compared to baseline methods.

weaknesses:
The weaknesses of the paper are listed below:
- The motivation of the paper is unclear, particularly in the Introduction. The statement at line 73, “Importantly, variational rectified flow matching differs in that it enables modeling ambiguity in the data-space-time-space domain, i.e., the goal is a model where flow trajectories can intersect,” does not clarify how allowing flow trajectories to intersect resolves the ambiguity problem.
- As I understand, using an additional latent variable z when modeling the velocity at (xt,t) can partially address the ambiguity problem, as different values of z capture different modes or sources of variation in the velocity distribution. However, VAEs are known to sometimes experience mode collapse, where the same high-density velocities may be generated from multiple modes of z. How does the proposed method handle this issue? To further address this concern, I suggest including a “diversity metric” in the experimental protocol to measure the variety of generated samples. 
- Furthermore, I believe RF can address the ambiguity problem by performing multiple rectifications. When training stabilizes, the optimal velocity at each sample xt at time t becomes unique, eliminating ambiguity. Even without rectification, existing methods such as OT-FM can mitigate ambiguity by improving the coupling between x0 and x1, resulting in less ambiguous directions at (xt, t). Are there any theoretical or methodological benefits of the proposed approach compared to these methods? Without such justification, it’s difficult to attribute the improved performance to the additional variable z.
- Can different values of z affect the visual quality of the generated samples? If x0 is kept constant, does varying z introduce significant variance in the generated samples? Or is there a specific value of z that results in low-quality samples? 
- There is no theoretical guarantee that the proposed approach will achieve a better straight flow when addressing the ambiguity problem.

questions:
Refer to the Weaknesses section for my concerns and questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731258650387

thread_type:
official_review

<thread_object>
creator: reviewer_clob

<content>
summary:
This paper proposes a novel method using variational inference to address the ambiguity of velocity vector fields that classic rectified flow matching fails to capture. By introducing latent variables and modeling ambiguity through a mixture model of velocity vector fields, the method enables more accurate data distribution capture and efficient integration. Experimental results demonstrate that the proposed approach achieves comparable performance to existing methods with fewer steps on synthetic data and image datasets.

strengths:
- The paper is well-structured and easy to read.
- The proposed method integrates VAE and flow matching in a straightforward manner, offering novelty in its ability to learn vector fields with uncertainty. Furthermore, the high performance on MNIST and CIFAR-10 datasets suggests that the hypotheses and approaches of this study are reasonably valid.

weaknesses:
- The proposed method in this paper introduces latent variables and their inference model, enabling the capture of overlapping vector fields when they occur. However, it is necessary to clarify the setting that assumes such overlapping vector fields. The overlap and ambiguity in question here initially seemed to imply that while the flow from a specific $x_0$​ to $x_1$​ is uniquely determined, there exist different vector fields at a particular time and spatial location that may intersect. However, during inference with the proposed method, a data point $x_0$​ is sampled from the source distribution, followed by the sampling of a latent variable, which is then used as the initial value to integrate an ODE based on a vector field determined by it. This implies that there is uncertainty in the direction from the initial value $x_0$​, meaning there is an assumption that it could lead to a different $x_1$​. Is this setting reasonable? The uncertainty involving deterministic flows crossing and the possibility of tracing different flows from the initial value (i.e., the $x_0$ and $x_1$​ pairings are not unique) appear to be mixed. The authors should clearly distinguish between these and clarify which aspect they are aiming to address.
- Considering the above points, while the proposed method may indeed enable faster transitions to the target due to the learned flow being linear even when vector fields overlap, during inference, $z$ is sampled from the prior and thus is not determined solely by $x_0$. As a result, the model could reach a different $x_1$​, which may not be desirable from the perspective of two-sided conditioning flow matching. 
- The proposed method requires an inference model and employs separate encoders for each of $x_0​,x_1$​, and $x_t$​ with the same structure as the encoder in $v_\theta​$. This implies a significantly larger number of learnable parameters compared to existing models, and although the encoders are not used during inference, it is not entirely fair in terms of parameter count relative to previous research (even if the speed remains similar). Therefore, it would be necessary to evaluate the impact of the size of the inference model’s encoders by modifying their size and investigating how it affects performance.

questions:
- I would like the authors to respond to the points I raised as concerns regarding the above weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732256960983

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers, ACs, and SACs

comment:
We sincerely thank the reviewers for their thoughtful comments and constructive feedback. We apologize for the delay in this rebuttal, as completing the ImageNet experiments required additional computation time. We appreciate the recognition of the strengths of our Variational Rectified Flow Matching (VRFM) work and are particularly grateful for reviewers highlighting:

* The **clarity and structure** of the paper, noted as well-written, easy to read, and easy to follow (cLoB, qkA9, Y61o).
* The **novelty**, **conceptual soundness**, and **valuable observations** of our approach, offering a meaningful perspective to flow matching models and addressing a notable limitation in existing rectified flow models in capturing ambiguity within the velocity vector field (cLoB, obZd, Y61o).
* The **strong empirical performance** demonstrated on multiple datasets and in comparisons with baseline methods (cLoB, qkA9, Y61o).
* The **well-executed analytical experiments** with visualizations that effectively validate our theoretical analysis (obZd).

To address the questions of the reviewers, we conducted additional experiments and expanded our discussions, updating both the main paper and the appendix. The key changes are summarized as follows:

1. Appendix E: Added a theoretical proof demonstrating that VRFM preserves the marginal data distribution.
2. Appendix F: Included visualizations of the reconstruction loss to illustrate that our model better approximates the ground truth velocities compared to baseline methods.
3. Appendix G: Reported the Inception Score on CIFAR-10, showing that our model does not suffer from mode collapse but instead improves sample diversity.
4. Appendix H: Conducted an ablation study on posterior model size, demonstrating that VRFM remains robust across different posterior model sizes.
5. Appendix I: Provided results from ImageNet experiments, demonstrating that our method generalizes well to large-scale datasets.
6. Appendix J: Showcased strong performance in class-conditional generation, highlighting the robustness and flexibility of our method in handling conditional image generation.
7. Appendix K: Expanded the related work discussion, including insights on consistency models and reflow, emphasizing that VRFM addresses the novel velocity ambiguity challenge. We highlight that these related works are orthogonal and complementary to ours, and exploring potential integrations offers an exciting avenue for future research.

We address each reviewer’s specific questions in individual replies.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732525202330

thread_type:
official_comment

<thread_object>
creator: area_chair_vfmy

<content>
title:
Author Reviewer Discussion

comment:
Dear Reviewers,

Thank you for your efforts in reviewing this paper. We highly encourage you to participate in interactive discussions with the authors before November 26, fostering a more dynamic exchange of ideas rather than a one-sided rebuttal.

Please feel free to share your thoughts and engage with the authors at your earliest convenience.

Thank you for your service for ICLR 2025.

Best regards,

AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1733290454125

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers, ACs, and SACs

comment:
We thank all the reviewers for their thoughtful comments and constructive feedback. Throughout the review process, we have carefully addressed all questions by providing detailed explanations, additional experiments, and theoretical justifications. We believe that all questions and issues have been thoroughly resolved, resulting in a more robust and comprehensive presentation of our contributions:

In this paper, we show that the “ground-truth” velocity vector field of flow matching is multimodal, i.e., it can point in different directions at the same location. To model a multimodal velocity vector field, we present Variational Rectified Flow Matching (VRFM), which is based on a variational training objective. Through extensive experiments, we show that VRFM is able to capture a multimodal velocity vector field, while outperforming classic flow matching across a range of benchmarks, including synthetic data, MNIST, CIFAR-10, and ImageNet. We think the proposed method naturally extends classic rectified flow matching and is a meaningful contribution to our community.

We also want to summarize the specific actions we have taken in response to each reviewer’s feedback:

* **Reviewer cLoB** inquired about the rationale behind allowing multiple directions from an initial random sample and the implications of using a larger number of learnable parameters. In response, we provided detailed explanations in the rebuttal, alongside with the reconstruction loss plot (Figure 10, Appendix F), and additional experiments such as a posterior model size ablation study, an inception score analysis, ImageNet studies, and class-conditional generation experiments (Appendices G–J). Since the reviewer never responded, we assume that all questions have been answered satisfactorily. 
* **Reviewer qkA9** asked about the connection between flow trajectory intersections and ambiguity, the potential for mode collapse in VAEs, the uniqueness of the optimal velocity in data-space-time-space, the impact of a latent on generated samples, and theoretical guarantees linking ambiguity resolution to straighter flows. To address these points, we revised the main text (Line 72), provided discussions and justifications, and included analyses such as the Inception score (Table 2, Appendix G), reconstruction loss plot (Figure 10, Appendix F), and theoretical explanations on how a latent captures multimodality. We hope our last response was able to answer the remaining questions.
* **Reviewer obZd** noted the absence of experiments on real-world data like ImageNet and comparison to consistency models. In response, we added experiments on ImageNet (Table 4, Appendix I), analyzed the velocity distribution of consistency models (Figure 4), compared our method with a recent Consistency Flow Matching baseline across multiple settings (Figures 2, 5, and 7), and introduced new CIFAR-10 experiments below. These highlight that VRFM outperforms all baselines when the number of evaluation steps exceeds 2 for MNIST and 5 for CIFAR-10. Since the reviewer didn’t react upon our last messages, we assume all concerns have been addressed.

| NFE / sample| Params. | 2 | 5 | 10| 50 | 100| 1000 | Adaptive |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| OT-FM | 36.5M | 166.655 | 36.188 | 14.396 | 5.557 | 4.640 | 3.822 | 3.655|
| I-CFM | 36.5M | 168.654 | 35.489 | 13.788 | *5.288* | 4.461 | 3.643 | 3.659 |
| Consistency-FM | 61.8M | **5.323** | **11.412** | 23.948 | 36.652 | 38.680 | 40.402 | 40.677
| VRFM (adaptive norm, $x_1$, 2e-3)| 37.2M | 135.275 | 28.912 | **13.226** | 5.382 | *4.430* | 3.642 | 3.545|
| VRFM (adaptive norm, $x_1$, 5e-3)| 37.2M | 159.940 | 35.293 | 14.061 | **5.265** | **4.349** | **3.582** | 3.561|
| VRFM (adaptive norm, $x_1 + t$, 5e-3)| 37.2M  | 117.666 | 27.464 | 13.632 | 5.512 | 4.484 | 3.614 | **3.478**|
| VRFM (bottleneck sum, $x_1 + t$, 2e-3)| 37.0M  | *104.634* | *25.841* | *13.508* | 5.618 | 4.540 | *3.596* | *3.520*|

* **Reviewer Y61o** sought clarity on whether reducing ambiguity leads to straighter flows, proofs demonstrating that VRFM preserves the marginal data distribution, additional experiments on ImageNet, and conditional generation. To address these concerns, we provided a reconstruction loss plot (Figure 10, Appendix F), derivations on preserving the marginal distribution (Appendix E), ImageNet experiments (Tables 4 and 6, Appendix I), and conditional generation results (Table 5, Appendix J). We hope our last response was able to answer the remaining questions.

In summary, we think we have thoroughly addressed all the reviewers' comments. We think our work makes a meaningful contribution to the field. We sincerely thank all reviewers once again for their time and insightful feedback, which have improved the quality of our paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission proposes Variational Rectified Flow Matching, an approach to improve over prior rectified flow methods by better capturing ambiguity and improving the empirical performance of diffusion model generation. Initially, several theoretic points of discussion were raised and experimental evaluation was not sufficient. The rebuttal could address these issues to some extent. After the rebuttal, the reviewers still find that the proposed approach is intuitively motivated while a stronger mathematical foundation would be better. More importantly, the practical performance of the proposed Variational Rectified Flow Matching lags behind which needs to be further analysed.