Improving Flow Matching by Aligning Flow Divergence

2026-01-31

International Conference on Machine Learning

AI Summary

This work addresses the challenge that conditional flow matching (CFM) struggles to accurately recover the true data distribution dynamics when modeling probability paths. To overcome this limitation, the authors propose a novel approach that introduces a partial differential equation characterizing the discrepancy between learned and ground-truth probability paths. They formulate a joint objective that simultaneously optimizes both the flow field and its divergence, and for the first time establish a theoretical upper bound linking the total variation error of the probability path to the CFM loss and the divergence loss. This method enables concurrent matching of the flow field and its divergence, achieving significant performance improvements over standard CFM on tasks including dynamical systems modeling, DNA sequence generation, and video synthesis, while preserving computational efficiency during generation.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 gap between the two 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 sacrificing generation efficiency. 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. Code is available at \href{https://github.com/Utah-Math-Data-Science/Flow_Div_Matching}{Utah-Math-Data-Science}.

Problem

• conditional flow matching
• probability path accuracy
• flow-based generative models
• flow divergence

Innovation

• flow matching
• divergence alignment
• generative modeling
• probability path
• partial differential equation