Hierarchical Rectified Flow Matching with Mini-Batch Couplings

Yichi Zhang, Yici Yan, Alex Schwing, Zhizhen Zhao

Keywords: Generative Model, Flow Matching, Rectified Flow

Abstract: Flow matching has emerged as a compelling generative modeling approach that is widely used across domains. During training, flow matching learns to model a velocity field. At inference, to generate samples, 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. Each level of the hierarchy of ODEs captures the distribution of the next level, just like vanilla flow matching uses the velocity field to capture a multi-modal data distribution. While this hierarchy enables to model multi-modal distributions at any hierarchy level, the complexity of the modeled distributions 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.

Primary Area: generative models

Submission Number: 12412