Improving Flow Matching by Aligning Flow Divergence

Yuhao Huang, Taos Transue, Shih-Hsin Wang, William M Feldman, Hong Zhang, Bao Wang

TL;DR: We establish a new user-friendly total variation bound between the learned and exact probability paths in flow matching, informing efficient algorithms to improve flow matching.

Abstract: Conditional flow matching (CFM) stands out as an efficient, simulation-free approach for training flow-based generative models, achieving remarkable performance for data generation. However, CFM is insufficient to ensure accuracy in learning probability paths. In this paper, we introduce a new partial differential equation characterization for the error between the learned and exact probability paths, along with its solution. We show that the total variation between probability paths is bounded above by a combination of the CFM loss and an associated divergence loss. This theoretical insight leads to the design of a new objective function that simultaneously matches the flow and its divergence. Our new approach improves the performance of the flow-based generative model by a noticeable margin without significantly raising the computational cost. We showcase the advantages of this enhanced training approach over CFM on several important benchmark tasks, including generative modeling for dynamical systems, DNA sequences, and videos.

Lay Summary:

Flow-based generative models are powerful tools for producing complex data such as images, DNA sequences, and videos. A recent method called Conditional Flow Matching (CFM) has improved training efficiency by directly matching learned and exact vector fields, avoiding costly simulations. However, CFM alone can struggle to accurately capture how probabilities evolve over time, which can reduce the reliability of the generated results.

In our research, we developed a new mathematical framework that precisely characterizes the error between the learned and true probability flows. We showed that this error depends not only on the original CFM loss but also on a previously overlooked factor: a mismatch in the divergence of the flow fields.

Motivated by this insight, we designed a new training objective that matches both the flow and its divergence. This simple yet impactful modification leads to significantly better performance across several challenging tasks, including modeling dynamical systems, DNA sequences, and videos.

Our approach retains the speed and scalability of CFM while improving accuracy and reliability, making it a practical and effective upgrade for real-world generative modeling applications.

Link To Code: https://github.com/Utah-Math-Data-Science/Flow_Div_Matching

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Flow Matching, Divergence Matching, Total Variation Bound, Generative Modeling

Submission Number: 9071