Hierarchical Rectified Flow Matching with Mini-Batch Couplings

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 different levels of the hierarchy via mini-batch couplings. We show the benefits of mini-batch couplings in hierarchical rectified flow matching via compelling results on synthetic and imaging data.

The above figure shows an example with 1D data. In this example, the source distribution is a standard Gaussian, while the target distribution is a mixture of two Gaussians with means located at -1 and 1. As shown in the first row, after applying data coupling (depth 1), the velocity distribution (depth 2) collapses into a single-mode Gaussian as the coupling batch size (bs) increases, effectively simplifying the velocity layer's distribution. The number given in the legend refers to the number of used velocity ODE integration steps. Below are some additional experimental results on 2D data, MNIST, CIFAR-10, and CelebA-HQ. We observe the proposed HRF with data coupling (D) and velocity coupling (V) to achieve compelling results.