Extended Flow Matching: a Method of Conditional Generation with Generalized Continuity Equation

TL;DR

The paper tackles conditional generation with continuous conditions by extending flow matching to Extended Flow Matching (EFM), which learns a matrix field to explicitly model how distributions evolve with respect to conditions. By formulating a generalized continuity equation on the condition-augmented domain and introducing a Dirichlet-energy objective, EFM (and its MMOT-EFM variant) enforces a bias toward smoother, more controlled conditioning, and supports both conditional sampling and style transfer. The authors provide theory to justify the matrix-field construction, outline a training algorithm, and demonstrate competitiveness on synthetic 2D data and conditional molecular design. While achieving promising results, they acknowledge the high computational cost of multi-marginal OT and point to future work on efficiency and scaling to more complex condition spaces and unseen conditions.

Abstract

The task of conditional generation is one of the most important applications of generative models, and numerous methods have been developed to date based on the celebrated flow-based models. However, many flow-based models in use today are not built to allow one to introduce an explicit inductive bias to how the conditional distribution to be generated changes with respect to conditions. This can result in unexpected behavior in the task of style transfer, for example. In this research, we introduce extended flow matching (EFM), a direct extension of flow matching that learns a "matrix field" corresponding to the continuous map from the space of conditions to the space of distributions. We show that we can introduce inductive bias to the conditional generation through the matrix field and demonstrate this fact with MMOT-EFM, a version of EFM that aims to minimize the Dirichlet energy or the sensitivity of the distribution with respect to conditions. We will present our theory along with experimental results that support the competitiveness of EFM in conditional generation.

Figures (11)

• Figure 1: Across-condition transfer by the conditional generative model trained with three conditional distributions with two clusters each ($\mu_{c_0}$, $\mu_{c_1}$, $\mu_{c_2}$). When the average sensitivity of the distribution with respect to $c$ (Dirichlet energy) is not optimized, the inner cluster may mix with the outer cluster (left). Meanwhile, if the energy is optimized (right), the transfer would respect the separation of inner vs outer clusters.
• Figure 2: Inferences of FM and EFM.
• Figure 3: Visualization of the flow for (a) conditional generation along $\gamma^{c_1}$ and $\gamma^{c_2}$ (\ref{['alg:gen']}), and (b) style transfer along $\gamma^{c_1\to c_2}$ (\ref{['alg:transfer']}).
• Figure 4: Training distributions for the 2D synthetic experiments (4 conditions, two clusters each)
• Figure 5: Wasserstein distance between GT vs predicted distributions. COT-FM was evaluated with $\beta=5$.

Theorems & Definitions (14)

• proposition 1
• Example 1: Conditional generation
• Example 2: Style transfer
• theorem 1
• Definition 3: A distributional solution of the generalized continuity equation
• proposition 2
• proof
• Remark 4
• Lemma 5: Principled mass alignment lemma
• Lemma 6