Hierarchical Rectified Flow Matching with Mini-Batch Couplings
Yichi Zhang, Yici Yan, Alex Schwing, Zhizhen Zhao
TL;DR
This paper introduces a hierarchical rectified flow matching method that uses mini-batch couplings to adapt the complexity of velocity distributions across hierarchy levels, improving generative modeling of multi-modal data.
Contribution
It proposes a novel hierarchical rectified flow matching approach with mini-batch couplings to better model multi-modal velocity distributions in generative models.
Findings
Enhanced modeling of multi-modal velocity distributions.
Improved generative performance on synthetic and imaging data.
Demonstrated benefits of mini-batch couplings in hierarchical flow matching.
Abstract
Flow matching has emerged as a compelling generative modeling approach that is widely used across domains. To generate data via a flow matching model, an ordinary differential equation (ODE) is numerically solved via forward integration of the modeled velocity field. To better capture the multi-modality that is inherent in typical velocity fields, hierarchical flow matching was recently introduced. It uses a hierarchy of ODEs that are numerically integrated when generating data. This hierarchy of ODEs captures the multi-modal velocity distribution just like vanilla flow matching is capable of modeling a multi-modal data distribution. While this hierarchy enables to model multi-modal velocity distributions, the complexity of the modeled distribution remains identical across levels of the hierarchy. In this paper, we study how to gradually adjust the complexity of the distributions across…
Peer Reviews
Decision·Submitted to ICLR 2026
I like the presentation of the approach. The use of mini-batch OT in the two levels of hierarchical rectified flow is intuitive and well-motivated from experiments on bi-modal Gaussian data. The theoretical results are interesting and easy to follow. The experiments do show improvement in lower FID compared to baselines with low NFEs.
Major: Novelty - The paper's main contribution is to use mini-batch OT for training the HRF2 model. The mini-batch OT has been previously used for CFM models (OT-CFM). I understand that this paper extends this to HRF models, but in my view, this is a very limited novelty. Training cost - Flow matching training with the mini-batch OT is expensive. This method adds computing another OT map for the velocity distribution. In addition, HRF2 with velocity coupling requires a simulation from ODE flow
• The paper is well-written and clearly-organized. • The empirical discovery that using mini-batch OT coupling in the current hierarchy level simplifies the distribution at the next level is interesting and could be considered in other scenarios of generative modeling like diffusion models.
• The novelty of the paper is limited and incremental, the usage of minibatch OT coupling in flow matching is already proposed in [1]. The main point of the paper that using mini-batch OT inherently simplifies the velocity distribution is only intuitively explained without further theoretical mathematical evaluation. • The paper lacks further theoretical discussion and comparison between “data coupling” and “joint data and velocity coupling”. • The method is not simulation-free, which increase
The paper extends the exising HRF framework by introducing coupling, showing that the velocity distribution learnt from coupling data can recover the data distrbution as well. Through some 1D experiments, the authors show that coupling can reduce the multi-modality of velocity distribution, which is convincing. Results of large scale experiments also verify the effectiveness of proposed methods.
1. The proposed method is a naive combination of two existing methods: HRF and mini-batch coupling, and provides little new insights. It is a well known result that mini-batch coupling can help straighten the velocity and thus address multi-modality in some sense. The authors didn't provide further useful insights for why mini-batch coupling could benefit HRF other than some 1D numerical experiments. Can the authors develop some theories (even toy example is fine) to demonstrate why mini-batchin
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsFluid Dynamics and Mixing · Data Stream Mining Techniques · Innovative Microfluidic and Catalytic Techniques Innovation