Variational Rectified Flow Matching

Thanks for your time and feedback.

Reviewer Comment 1: unclear why reducing ambiguity would inherently lead to straighter flows.

Response: We demonstrate that resolving the ambiguity problem allows our model to better match the ground-truth (GT) flow, which is linear by definition in rectified flow. Specifically, a lower loss in Equation (5) indicates that the predicted flow more closely aligns with the GT flow, inherently leading to straighter trajectories.

During training, we observed that our method achieves better velocity reconstruction losses (Appendix F, Figure 10) compared to vanilla rectified flow, indicating that the predicted velocities more accurately approximate the GT velocities, which are linear. Furthermore, in the experimental section, we provide strong empirical evidence that our method outperforms baseline approaches, especially when the number of function evaluations (NFEs) is small on multiple benchmarks (Synthetic, MNIST, CIFAR-10, ImageNet) in the experiments section and the newly added Appendix G-J. This highlights the practical benefits of our approach in efficiently modeling flows while addressing the velocity ambiguity issue.

Reviewer Comment 2: proofs demonstrating that the learned distribution from VRFM preserves the marginal data distribution.

Response: We include the derivation in the newly added Appendix E.

Reviewer Comment 3: empirical evaluation is restricted to MNIST and CIFAR-10, which limits the generalizability of the findings.

Response: To answer, we conduct additional experiments on the ImageNet 64x64 dataset. The training setup and architecture is exactly identical to our CIFAR-10 training, i.e., no additional hyperparameter tuning or cherry-picking. The only changes: increasing the number of iterations to 800k and adjusting the batch size to 128 to accommodate the larger training set. The resulting FID scores are summarized below and added in Appendix I, Table 4. We observe that our method improves upon the baseline, even in this large-scale real-world dataset. These results demonstrate the scalability and effectiveness of our approach in handling more complex data while maintaining its advantages over baseline methods.

Thank you for providing the detailed explanations, proofs, and additional experiments. I appreciate the effort to address my initial concerns, and I find the variational perspective introduced for rectified flow both novel and interesting.

That said, I believe there are still limitations in the practical performance of VRFM, particularly at small NFEs. The additional ImageNet experiments, while helpful, reveal an FID score that remains noticeably behind state-of-the-art models (e.g., FID below 3). Even at larger NFEs (e.g., 1000), the performance does not reach the performance of strong baselines. Given that VRFM’s primary contribution is in reducing ambiguities at intersections and producing straighter flows, I anticipated a more substantial improvement that could bridge the gap with state-of-the-art methods.

In conclusion, while I find the proposed direction intriguing and promising, the current approach seems to require further advancements to meet the standard of a top-tier conference like ICLR. Based on the improvements provided in the rebuttal and the potential of this direction, I am increasing my rating to 6.

I acknowledge that VRFM improves vanilla rectified flow under comparable model sizes, and I appreciate the efforts to highlight this advancement. However, I maintain my current score for two primary reasons:

1. The results remain unconvincing, particularly at low NFEs. Given that the proposed method argues that it reduces ambiguity and learns straighter flows, I expect a more substantial performance gain in low NFE regimes. The significant performance gap between low and high NFEs falls short of these expectations.

2. The main claim of this paper is improved performance in accelerated settings, as mentioned in the abstract and introduction section. To substantiate this claim, I believe it is necessary to include comparisons against other baselines that also target improvements in accelerated settings. Additionally, I believe it is within scope to apply the reflow procedure to VRFM, as it aligns with the goal of extending the capabilities of rectified flow.

In conclusion, while I appreciate the contributions of this work, these considerations prevent me from adjusting my score further at this stage.

We thank the reviewer for detailed feedback and for acknowledging the improvements VRFM achieves over vanilla rectified flow under comparable model sizes. Below, we answer the remaining comments:

Reviewer comment 1: significant performance gap between low and high NFEs

As stated in the abstract (L21), our goal is to capture multimodal velocity vector fields in flow matching. We show that the proposed formulation is able to achieve this goal. We further demonstrate that capturing a multimodal velocity vector field leads to consistent improvements over classic flow matching across all NFE regimes (as shown in Tables and Figures in the main paper). Note, we don’t intend to reduce the performance gap between low and high NFEs. Instead, we aim for improvements across all NFE regimes, which the proposed method achieves.

Reviewer comment 2: main claim of this paper is improved performance in accelerated settings

As stated in the abstract (L21) and the introduction (L49), our (main) claim is to model the multimodality/ambiguity of the velocity field in flow matching. The empirical results (Figures 1, 3, 10) demonstrate that we successfully achieve this goal. We also observe our method to attain compelling results compared to classic rectified flow across all NFE regimes (including low NFEs) and across datasets including synthetic, MNIST, CIFAR-10, and ImageNet.

Reviewer comment 3: comparisons against accelerated baselines and apply reflow to VRFM

We appreciate the reviewer’s suggestion to include comparisons against other accelerated baselines. In our experiments, we included consistency flow matching as a baseline (which itself improves upon recent methods) and found that VRFM outperforms it on synthetic datasets (Figure 2, 5), MNIST when the NFE exceeds 2 (Figure 7), and CIFAR-10 when NFEs exceed 5 (General Response).

It’s a great suggestion to apply reflow to VRFM. As noted in our response to another reviewer (https://openreview.net/forum?id=1cM0yQe3pO&noteId=ZvIn52hRIG), reflow is indeed compatible with our framework and we expect further improvements. However, while reflow is a promising extension, numerous other enhancements—such as advanced noise schedules, alternative regularization strategies, or hybrid generative modeling approaches—could similarly improve VRFM’s capabilities. We think including these extensions dilutes the focus of our core contribution and extends the scope beyond what is practical.

The primary objective of this paper is to introduce VRFM and establish its benefits and shortcomings over classic rectified flow across a variety of settings. While we recognize the potential of other advances to further improve performance, we believe that it is more appropriate to focus on the foundational contribution of VRFM and leave additional integrations to future exploration by the community.

Thanks for your time and feedback.

Reviewer Comment 1: lacks enough experiments on more real-world and complex datasets like ImageNet.

Response: To answer, we conduct additional experiments on the ImageNet 64x64 dataset. The training setup and architecture is exactly identical to our CIFAR-10 training, i.e., no additional hyperparameter tuning or cherry-picking. The only changes: increasing the number of iterations to 800k and adjusting the batch size to 128 to accommodate the larger training set. The resulting FID scores are summarized below and in the newly added Appendix I, Table 4. We observe that the same trends hold for our method compared to the baseline models, even in this large-scale real-world dataset. These results demonstrate the scalability and effectiveness of our approach in handling more complex data while maintaining its advantage over baseline methods.

Reviewer Comment 2: whether Variational Flow Matching can enhance performance through “reflow”.

Response: Great question. Technically, it is feasible to integrate reflow into the VFM framework. Specifically, after completing one round of training, we can resample $(x_0,x_1)$ pairs. Unlike classic flow matching, VFM enables us to sample multiple $x_1$ values for a given $x_0$, leveraging the model’s ability to capture multimodal distributions. By consolidating these resampled pairs into a new dataset, we can initiate a subsequent round of VFM training.

1. If ambiguity is modeled nearly perfectly by our variational model in the first round (e.g., for our 1D/2D data), reflow won’t help too much, which is expected as we already attain a very high evaluation score even for very low neural function evaluations (NFE) in those settings.
2. If ambiguity is not modeled perfectly by our variational model in the first round (e.g., for cifar10/imagenet data), reflow will help the model to find more dependent data pairs. As we shown in Figure 8 in the main paper, varying different $z$ will only change the color patterns, while an image’s content is primarily determined by $x_0$. Thus with the new reflow algorithm, we expect to construct an “easier” set for the model to capture, indicating it could be helpful to further improve performance.

This combination is an exciting avenue for future work, and we believe that integrating the two methods could lead to advancements in the performance and scalability of flow-based generative models.

Thanks for your time and feedback.

Reviewer Comment 1: The paper does not clarify how allowing flow trajectories to intersect resolves the ambiguity problem.

Response: As discussed in the Introduction (Lines 41–48) and as illustrated in Figure 1, we note that crossing flows point in different directions at the same location in the data-space-time-space domain. Intuitively, being at the same location and moving in different directions means that flows are crossing. This phenomenon is also observed for the ground-truth velocities. Hence, at those locations, the (ground-truth) flow trajectories point in multiple directions, i.e., they are ambiguous/multimodal.

Classic flow matching methods are unable to capture this type of ambiguity, as they employ a deterministic squared-norm loss that restricts the model to a single velocity vector at any given point.

In contrast, variational rectified flow matching models the multimodal velocity distribution at any point along the flow in the data-space-time-space domain. By allowing different flow directions at the same data-space-time-space point, our approach enables trajectories to intersect and effectively captures the underlying ambiguity.

We revised Line 72 in the paper to clarify this motivation by stating: “Importantly, variational rectified flow matching differs in that it enables to also model ambiguity in the data-space-time-space domain. This enables different flow directions at the same data-space-time-space point, allowing the resulting flows to intersect at that location. ”

Reviewer Comment 2: VAEs are known to sometimes experience mode collapse.

Response: While mode collapse is a well-known issue in GANs, where multiple modes of the data distribution may be ignored, VAEs generally do not exhibit the same degree of collapse due to their explicit probabilistic modeling of the latent space. This probabilistic framework inherently encourages coverage of the data distribution's diversity. Furthermore, as a standard quantitative metric in generative modeling [1,2,3], the Fréchet Inception Distance (FID) score evaluates both the mean and covariance statistics of the generated data compared to the ground-truth. Sample diversity directly impacts the covariance component of the FID score calculation, making it sensitive to issues like mode collapse or a lack of diversity.

To further address the reviewer’s question, we added the inception score to our evaluation. The score explicitly measures the distribution of predicted labels of the generated samples. Compared to the vanilla rectified flow baseline, our method consistently achieves higher Inception Scores, reflecting improved diversity in the generated samples. By reporting both FID and inception scores, we provide a comprehensive assessment of the model's ability to generate diverse and realistic samples. The full table is available in the new Appendix G, Table 2.

[1] Tong, A., FATRAS, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G. and Bengio, Y., Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research.

[2] Peebles, W. and Xie, S., 2023. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 4195-4205).

[3] Ma, N., Goldstein, M., Albergo, M.S., Boffi, N.M., Vanden-Eijnden, E. and Xie, S., 2024. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv preprint arXiv:2401.08740.

Thank you for your response, which has addressed many of my concerns. While the proposed problem - ambiguity of the flows - is well-motivated, and the proposed method demonstrates strong performance in addressing it, a significant limitation remains in my view. The paper primarily relies on intuition without providing a theoretical guarantee or solid justification for the underlying idea.

Specifically, it is difficult to justify that introducing z fully resolves the ambiguity issue. For instance, what if there exist two values, $z_1$ and $z_2$, such that $v(x_t, t, z_1)$ and $v(x_t, t, z_2)$ exhibit minimal differences? Wouldn't this imply that the ambiguity still persists? I would appreciate further justification on this point.

That said, I recognize the model’s strong performance across datasets, and I am inclined to increase my score to 6, albeit with lower confidence. The lack of a theoretical foundation remains a substantial limitation for an ICLR paper, so I leave it to the ACs/SACs to make a final decision.

Thanks for your time and feedback.

Reviewer Comment 1: The reasonableness of the setting where uncertainty in the direction from $x_0$ could lead to multiple possible outcomes. Might not be desirable from the perspective of two-sided conditioning flow matching.

Response: We appreciate the reviewer highlighting this point. The mapping from $x_0$ to multiple possible $x_1$ is indeed a fundamental aspect of what flow matching aims to model. As discussed in Section 2.2 of our paper, during training, classic flow matching independently samples $x_0$ from the source distribution and $x_1$ from the target distribution. The objective is to learn a mapping between the two distributions. This setup inherently implies that any given $x_0$ can be mapped to multiple $x_1$ values in the target distribution, giving rise to non-unique mappings. The training objective is to match the "ground-truth" velocity vector field for these randomly sampled $(x_0,x_1)$ pairs. Different from prior art, the proposed variational rectified flow formulation acknowledges and models this intrinsic ambiguity. This ambiguity exists across all $t \in \[0, 1)$, including at $t=0$ (at $x_0$), due to the foundational formulation of flow matching.

Quantitatively, we observe better reconstruction losses in our model compared to vanilla rectified flow (Appendix F, Figure 10). Further, we also observe better FID scores at various evaluation steps on multiple datasets (Synthetic, MNIST, CIFAR-10, ImageNet) in the experiments section and in the newly added Appendix G-J. Results indicate that modeling this ambiguity is beneficial for better modeling the ground-truth velocity distribution and target data distribution.

Reviewer Comment 2: Mixing uncertainty from flow crossings and the possibility of tracing different flows from the initial value.

Response: Great point about distinguishing between uncertainty involving flow crossings at $(x_t,t)$ and the possibility of tracing different flows from the initial value $x_0$. To clarify, our variational flow model is designed to handle both aspects through its flexibility in predicting a distribution of velocity vectors rather than a deterministic one, as done in classic flow matching. Specifically, we independently sample $x_0$ from the source data distribution and $z$ from the prior latent distribution. At inference time, for simplicity, we choose to use a single latent variable $z$ sampled from a prior distribution and keep it fixed for the entire trajectory. Doing so enables us to model the ambiguity of the velocity as illustrated in Figure 3 (c)–(f), including potential flow crossings at a given $(x_t, t)$. However, our formulation does not inherently restrict $z$ to being fixed. In fact, to further decouple the distribution of $x_t$ and $z$ at a random time $t$, it is entirely feasible to sample a new $z$ at each time step during inference, allowing us to model uncertainty if required. This flexibility allows for a more dynamic modeling of uncertainty and flow variability if desired.

Reviewer Comment 3: Our model implies a significantly larger number of learnable parameters compared to existing models, it is not entirely fair in terms of parameter count relative to previous research (even if the speed remains similar).

Response: Thanks a lot for highlighting the size of the latent encoder and its potential impact on the parameter count. We note that these encoders are only used during training and do not contribute to the inference-time model complexity. Compared to the baseline, the only additional module introduced is a two-layer MLP to adapt the latent variable $z$ in the flow matching model, which accounts for only 1.3% of the total parameters of the flow model. This ensures that the computational efficiency during inference remains comparable to that of baseline methods. That said, we conduct additional experiments to investigate the impact of varying the encoder size, reducing it to 17.5% (VRFM-M) and 6.7% (VRFM-S) of its original size. The results below demonstrate that our model maintains comparable performance across these variations, highlighting the flexibility and robustness of our approach. The full table and discussion can be found in the newly added Appendix H, Table 3.