Improving Flow Matching by Aligning Flow Divergence

TL;DR

This work identifies a fundamental bottleneck in flow-based generative modeling: minimizing the FM loss alone does not guarantee accurate probability-path learning due to divergence gaps between the learned and exact flows. By deriving a PDE-based error characterization for the probability paths and establishing a TV-based bound that ties this error to a divergence-gap term, the authors motivate a joint training objective called flow and divergence matching (FDM). FDM combines conditional flow matching (CFM) with a conditional divergence matching (CDM) term, including an efficient squared version, to align both the flow and its divergence without sacrificing sampling efficiency. Empirically, FDM improves density estimation, DNA sequence design, dynamical-system trajectory sampling, and video generation across multiple datasets and diffusion-path choices, demonstrating meaningful gains in likelihoods and sample quality. The approach offers a practical, scalable path to more accurate flow-based generative models with broad applicability in scientific modeling and sequential data generation.

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}.

Figures (7)

• Figure 1: Experiments of training an FM model using CFM for sampling the 1D Gaussian mixture distribution in equation (\ref{['eq:mixture-Gaussian']}). The left panel shows that the conditional divergence loss ${\mathcal{L}}_{\rm CDM}$ in equation (\ref{['eq:cfm_div']}) is much larger than the CFM loss ${\mathcal{L}}_{\rm CFM}$, and the right panel shows the significant gap between the exact distribution ($p_{\rm Data}$) and the distribution learned through FM ($\hat{p}_{\rm FM}$).
• Figure 2: Snapshots for probability paths at $t=0.6,0.85$, and $1$ (left to right). First/Second row: FM/FDM vs. data distribution.
• Figure 3: Comparison of probability paths over time learned by FM (left) vs. FDM (right).
• Figure 4: Generated samples from FM and FDM using the optimal transport (OT) path trained on the checkerboard dataset.
• Figure 5: Histograms of the constraint value $C({\bm{x}}_1)$ where ${\bm{x}}_1$ is an event trajectory computed by the Dormand-Prince ODE solver or sampled from the model with event guidance. The unguided sampling histograms are shown in Appendix \ref{['appendix:additional-results:dynamics-forecasting']}.

Theorems & Definitions (17)

• Proposition 3.0
• Corollary 3.0
• Theorem 3.1
• Theorem 4.1
• Remark 4.2
• Proposition 1.0
• proof : Proof of Proposition \ref{['prop:error']}
• Corollary 1.0
• proof : Proof of Corollary \ref{['cor:error-formula']}
• Theorem 1.1