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

Noboru Isobe, Masanori Koyama, Jinzhe Zhang, Kenji Fukumizu, Kohei Hayashi

Keywords: Flow Matching, Generative Model

TL;DR: We propose a novel framework for continuous conditional generative models that extends Flow Matching with a generalized continuity equation.

Abstract:

Conditional generative modeling (CGM), which approximates the conditional probability distribution of data given a condition, holds significant promise for generating new data across diverse representations.
While CGM is crucial for generating images, video, and text, its application to scientific computing, such as molecular generation and physical simulations, is also highly anticipated.
A key challenge in applying CGM to scientific fields is the sparseness of available data conditions, which requires extrapolation beyond observed conditions.
This paper proposes the Extended Flow Matching (EFM) framework to address this challenge.
EFM achieves smooth transitions in distributions when departing from observed conditions, avoiding the unfavorable changes seen in existing flow matching (FM) methods.
By introducing a flow with respect to the conditional axis, EFM ensures that the conditional distribution changes gradually with the condition.
Specifically, we apply an extended Monge--Kantorovich theory to conditional generative models, creating a framework for learning matrix fields in a generalized continuity equation instead of vector fields.
Furthermore, by combining the concept of Dirichlet energy on Wasserstein spaces with Multi-Marginal Optimal Transport (MMOT), we derive an algorithm called MMOT-EFM.
This algorithm controls the rate of change of the generated conditional distribution.
Our proposed method outperforms existing methods in molecular generation tasks where conditions are sparsely observed.

Primary Area: generative models

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6719