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

Authors: Noboru Isobe, Masanori Koyama, Jinzhe Zhang, Kenji Fukumizu, Kohei Hayashi
Primary area: generative models
Keywords: Flow Matching, Generative Model

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.

Decision: reject
Reviewer scores: 3, 6, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1729370314185

thread_type:
official_review

<thread_object>
creator: reviewer_a59p

<content>
summary:
This paper introduces extended flow matching a new flow matching based method, that is designed for conditional generation. For this, the authors make use of the generalized continuity equation by Lavenant. The authors show that their proposed loss indeed has the correct gradients, i.e., regresses onto the true velocity field of the generalized continuity equation. The algorithm consists "learning" an interpolation via kernel regression (which is needed since "straight paths" are not the only viable solution anymore), and then regressing onto a flow matching loss where the is now matrix-valued. This is a generalization of the usual inverse problems framework of flow matching. Further, the authors showcase the effiacy of their algorithm via a toy example and conditional molecular generation.

strengths:
I find the motivation very clear: Sometimes we already know the posteriors for several conditions (for instance in molecular dynamics, where we obtain some posterior samples via MCMC), and want to "smartly" interpolate between the conditions, i.e., learn a generative model which walks along "generalized "geodesics in the space of probability measures. I also like that the authors were very rigorous in their theorems and motivation for the developed algorithm.

weaknesses:
However, the glaring weakness is that there is not clear cut numerical use case shown. I would like to see a not toyish example where we actually need several conditions and the transport between them. Usually, in the classical inverse problems works there is an implicit geodesic path taken where $y_t = t y + (1-t)y$, since one does not need to alter the condition if posterior sampling is the ultimate goal. If one wants to do style transfer (which seems to be the second motivation of this paper), then one can simply use a conditional FM network which receives the two conditions (source and target) as inputs. Therefore, while theoretically neat I am not convinced of why the generalized continuity equation and a network which moves efficiently also in the condition space, is advantageous. The authors can convince me by providing a clear example where either i) the classical conditional algorithms are not applicable or ii) this approach significantly outperforms the other flow matching models. 

I also have some smaller concerns. 

1) The scaling in $N_c$ and condition dimension seems to be bad. can you provide the run times for the molecular example also for the baselines? it only says in the appendix that they were completed within 4 hours, but I expect the baselines to train much quicker. Also latent space of a VAE is pretty low dimensional. Please provide training your conditional flow matching model on MNIST (no VAEs..), where the condition space is not discrete (i.e., for instance inpainting). Even if this does not fit your motivation, I would like to see the results in such a more standard example and this would improve my confidence in the scalability. 

2) Appendix D5 and F are empty (or almost empty). 

3) you do not seem to provide any code. I find the algorithm description to be not perfectly clear, there I would very strongly suggest that you at least publish code for the toy example. 

4) I believe that the example 7.1 is meaningless. You construct a random example with sparse conditions. Then you show, that your algorithm performs better on the OOD. But basically you can construct an inverse problem which aligns with your in distribution posteriors and does anything else on the OOD data. Of course I am aware that your point is that your algorithm is minimizing the Dirichlet energy and you measure the distribution induced by this. However, it is not clear to me if this is the theoretically optimal thing to do (wrt to Wasserstein). I am guessing that your algorithm computes something like Wasserstein barycenters weighted by some distance to the known conditions? Please clarify why the minimization of the generalized Dirichlet energy should yield theoretically sound posteriors. 

5) The manuscript is sloppy at times when discussing related work. "The authors in (Wildberger et al., 2023; Atanackovic et al., 2024) developed FM-based models to estimate the posterior distribution when the prior distribution p(c) of conditions is known. In contrast, our approach tackles situations where the conditions can only be sparsely observed, and the prior distribution is unknown."

The prior distribution p(c) is not known in (Wildberger et al, 2023). They are only able to sample from the joint distributions (c,x), but this does not mean that you can evaluate it. Further, their algorithm can very easily be adapted to the setting you described. If one has posterior samples for sparse conditions $c_i$ one can simply do the joint training over $(x_{i,j}, c_i)$.

6) when style transfer is one of the main modes of motivation, I would also like to see an example of it. 

Overall, I appreciate the idea and think that it has merits, but the execution prevents me from accepting it in the current form. I would love to see a practical example, where the main motivation of your algorithm becomes clear. Furthermore, providing a more standard inverse problem on MNIST (with no encoder/decoder) and a continuous condition space would show me that your algorithm at least somewhat scales.  If these problems are discussed/solved, then I am willing to raise my score.

questions:
see weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730475028270

thread_type:
official_review

<thread_object>
creator: reviewer_cdwg

<content>
summary:
To achieve extrapolation beyond observed conditions, the authors proposed Extended Flow Matching (EFM) framework that is developed upon the Conditional generative modeling. Specifically, the authors introduced a novel algorithm called MMOT-EFM derived from the Multi-Marginal Optimal Transport (MMOT). In the experiments, the authors showed improved MAE over compared FM-based methods.

strengths:
1. The Extended Flow Matching sounds novel and the authors show the newly introduced conditional components in Fig.1, which is quite intuitive.

2. I like the well-structured theoretical discussion from FM to EFM, this can help domain experts grasp the main contribution and difference between the existing OT-CFM and the proposed MMOT-EFM

weaknesses:
1. I feel concerned about the experimental design. For instance, the authors introduce a rather usual setting (Appendix 1300-1306). Though it aligns well with the synthetic point cloud experiments, it is quite different from the common practice [1]. 

[1] Ketata M A, Gao N, Sommer J, et al. Lift Your Molecules: Molecular Graph Generation in Latent Euclidean Space[J]. arXiv preprint arXiv:2406.10513, 2024.

2. I think critical experiments against highly related OT-CFM methods are missing in this version.  

Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian
Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models
with minibatch optimal transport. arXiv preprint 2302.00482, 2023b.

questions:
1. Could you please justify the ZINC-250k experimental design?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730647524817

thread_type:
official_review

<thread_object>
creator: reviewer_1cpw

<content>
summary:
This paper proposes an extension to flow matching to conditional generation on unknown (but related) conditions using a flow on both the data space and the condition space. A variant of this based on multi-marginal optimal transport is proposed as an extension to optimal transport conditional flow matching. 2D and conditional molecular generation experiments are performed showing conditional generation.

strengths:
* Understanding how to extend current generative models to more general conditionals (especially unobserved conditionals) is an important problem particularly in the sciences. 
* I enjoyed the symmetry of the presentation of first standard flow matching and OT-CFM settings followed by EFM and MMOT-EFM settings. Table 1 is great to understand the difference to OT-CFM. 
* To the best of my knowledge the theory is correct and answers some of my questions on how one might generalize flow matching to condition-augmented spaces.

weaknesses:
* It would be great to make clearer to the reader how this method extends to unseen conditions. I think lines 402-405 kind of get at this, but I would have loved to see more emphasis on this point. It is very easy to design a conditional generative model that technically extends to unseen conditions, but it is much more difficult to enforce that that model extends in a reasonable way. EFM has the potential to guide that extension and I would love to see that point explored further.
* The algorithm is not yet useful in real applications. While the authors also acknowledge this, it’s still a large limitation of the impact of this work. The molecule experiment is extremely limited in terms of comparisons to existing work and overall training setup.
* Much of the theoretical statements are direct extensions from prior work.

questions:
When is MMOT-EFM and EFM in general expected to work better than COT-FM / Bayesian-FM? I know there is a short explanation on the differences in assumptions but it is difficult for me to translate what is gained when making a piecewise continuous assumption on p(x|c) vs. a measurability assumption. It’s not clear to me how this compares to these prior works in general.

Small comments that don’t affect the score: 
There appears to be an unfinished section D.5 in the appendix. 
GG-EFM isn’t defined in the main text. 
I didn’t understand the distinction between p_c and p_{0,c} line 170. 
Typo on line 311 “ton he”
Shr\”odinger to Schr\”dinger line 425
The source points in Figure 4 b and c (and corresponding appendix figs) are essentially invisible (grey against a grey background). It would be **really nice** to fix this. 


### Overall
I think this work presents an interesting idea with promise to understand how these models generalize to unseen conditions. However, this is not explored theoretically. In addition the current method does not scale to practical settings at the moment. I think further investigation as to when the assumptions behind this method make sense relative to other methods would greatly strengthen this work. A better understanding of how this relates to prior literature and when this method is preferable would likely change my opinion of this work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730658677204

thread_type:
official_review

<thread_object>
creator: reviewer_mdmu

<content>
summary:
Flow matching can generate different data distribution given different desired property conditions. The authors proposed the extended flow matching (EFM) which introduces a continuous mapping from a joint continuous space of time and property conditions to corresponding data distribution, which enables smooth shifts between different conditions. The authors also extended optimal transport to Multi-Marginal Optimal Transport (MMOT) for multiple property conditions. They validated their method on a 2D toy model and conditional molecular generation.

strengths:
The theory of integrating property conditions and time in flow matching is highly innovative, and the authors developed MMOT to perform optimal transport within this space.

weaknesses:
The experimental evidence is insufficient.

questions:
Major:

1.	Could the authors explain or give an intuition about the regression in MMOT (Eq. 3.4)?

2.	Could the authors show the extrapolation ability of their methods in a more realistic application of EFM, e.g. style transfer of images?

Minor:

1.	At the end of Line 311, “focus on the” is misspelled as “focus ton he”.

2.	“ConvHull” should be explained.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699526899

thread_type:
official_review

<thread_object>
creator: reviewer_qpjx

<content>
summary:
The authors propose extended flow matching (EFM) for conditional sampling and style transfer using flow matching. EFM consists of 

1. learning a field which also uses the conditioning vector $c$ as input, which the authors call a matrix field. 
2. The authors then integrate the learned field $u(x, t, c)$ along different paths $\gamma: [0, t] \rightarrow [0, 1] \times C$, where $C$ is the set of conditioning vectors. 
    1. For instance, for conditional generation the authors propose integrating along the path $\gamma(t) = (t, c)$, which reduces to conditional flow matching. 
    2. For style transfer, the authors integrate along the path $\gamma(t) = (1, (1-t) c_1 + t c_2)$. Since integrating along $\gamma(t)$ can be out of domain for models learned trained just on pairs $x, c \sim p(x, c)$, the authors propose a learning algorithm such that the field $u_\theta$ also observes such paths during training. 

The authors propose learning such a field $u$ using optimal transport:

1. the authors propose learning an optimal plan similar to [Lipman et al 2023]
2. instead of using linear interpolation between different points on a path, the authors extend the set of paths to include functions belonging to an RKHS.

strengths:
The authors identify an interesting problem: observing the conditioning vector in a number of domains can be hard or expensive. The proposal of integrating along paths between different marginals is also interesting, a similar proposal is studied in [Albergo et al 2023].

weaknesses:
1. While the motivation of EFM was to provide ensure that the learned network $u(x, t, c)$ is smooth with respect to the conditioning vector $c$, the authors do not address how imposing smoothness can allow extrapolation to conditioning vectors not seen during training. 
2. Could the authors explain why the multi-marginal optimal transport approach allows for extrapolating to conditioning vectors not seen during training?
3. The authors should also consider including other works that learn multi-marginal flow models? For instance, [Albergo et al 2023] propose learning multi-marginal flows and present a learning algorithm for optimizing the paths such that the transport cost in $W_2$ metric is minimized. 
4. [Albergo et al 2023] also propose a much more general algorithm for including paths between samples from an arbitrary number of marginal distributions, available during training. 
5. The experiments section can be improved by adding extra text explaining the results and the figures, particularly in figure 4.


[Albergo et al 2023]  Albergo, M.S., Boffi, N.M., Lindsey, M. and Vanden-Eijnden, E., 2023. Multimarginal generative modeling with stochastic interpolants. arXiv preprint arXiv:2310.03695.

questions:
1. Can the authors consider providing definitions before introducing a new notation in the text?
2. What is the effect of defining $\pi$ using plans built using batched samples? Would the vector/matrix field learned change as a function of the batch size? 
3. What kernels do the author use for the RKHS used to construct paths?
4. In lines 212-214 and lines 220-222, can the authors clarify the output of $u$?
5. the discussion about the weak assumption of measurability and continuity of $p(x|c)$ with respect to $c$ requires clarification, particularly since piece-wise continuous functions are measurable as well.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732857503840

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global comments by authors

comment:
We appreciate the detailed feedback on our paper. Many reviewers pointed out the unclear relationship between smoothness and extrapolation, as well as the scalability issues of MMOT-EFM. In response to these comments, we made the following revisions, which have been reflected in the revised manuscript.

---

## How to Impose Smoothness to Allow Extrapolation to Conditions

First, we would like to clarify the meaning of "smoothness," as mentioned in our motivation. **Our goal is to ensure that the conditional distribution $ p(x \mid c) $ which we will generate is "smooth," meaning it minimizes the Dirichlet energy as defined in Equation (3.2).** Intuitively, the Dirichlet energy represents the sensitivity of the distribution $ p(x \mid c) $ with respect to the conditioning vector $ c $. Specifically, it holds that

$$
\operatorname{Dir}(p)= \text{``}\lim_{\varepsilon\to0}\text{''} C_k \iint_{\Omega \times \Omega}  \frac{W_2^2(p(\cdot \mid c_1), p(\cdot \mid c_2))}{2 \varepsilon^{k+2}} \boldsymbol{1}_{|c_1-c_2| \leqslant \varepsilon} \mathrm{~d} c_1 \mathrm{d} c_2 \quad \text{for} \quad p\colon\Omega\ni c\longmapsto p(\cdot \mid c)\in\mathcal{P}(D),
$$

where $ k $ is the dimension of the condition space $ \Omega $. For the precise meaning of the limit $ \text{``}\lim_{\varepsilon\to0}\text{''} $ and the value of the constant $ C_k $, please refer to [§1.3, Lavenant, 2019].

Thus, **minimizing the Dirichlet energy implies that the sensitivity with respect to the condition $ c $ is not too large.**

Although there are infinite ways of extrapolation, it is reasonable to assume an inductive bias that the sensitivity of data in nature (e.g., molecules) to conditions (e.g., chemical properties) is not unnaturally large. Therefore, our method addresses extrapolation by learning a model such that the data to be extrapolated follows this inductive bias of low sensitivity. We would like to note that this kind of inductive bias has been used throughout the history of generative models as a method to prevent overfitting and stabilize generative models; see, for example, Miyato et al. in ICLR, 2018.

Our experiments in §7 demonstrate that EFM, which minimizes the Dirichlet energy, outperforms methods that do not minimize this energy (such as FM and COT-FM) in terms of generation performance.

In addition, the cost (objective) function used in our multi-marginal optimal transport (MMOT) approach provides an upper bound on the Dirichlet energy; please refer to lines 233-236 and Table 1. Therefore, optimizing the transport plan $ \pi $ through the MMOT approach also minimizes the Dirichlet energy, which in turn reduces the sensitivity of the generated distribution $ p(x \mid c) $ with respect to the conditioning vector $ c $.

---

## Scalability of EFMs

Because we mentioned the complexity of MMOT in the manuscript, many reviewers were concerned about the scalability of the proposed method. **The new experiment we conducted in §7.2 to generate images with continuous conditions is expected to dispel this concern.**

In §7.2 of the revised manuscript, we included the conditional generation of the MNIST dataset with a single ImageNet background, where the two conditions are "color" and "rotation," which form a 4-dimensional condition space $ \Omega=[0, 1]^4 $ (the first dimension is rotation, and the rest are RGB). For training, we used a conditional dataset corresponding to 12 uniform random samples of $ \Omega $. We evaluated the Wasserstein distance from the ground truth (GT) conditional distributions in a manner similar to Figure 4. Note that these conditional distributions are nonlinear in the presence of background. We compared our method with the classifier-free guidance method (Zheng et al. [1]) with different guidance strengths. On this dataset, we see that the Wasserstein error increases monotonically with the distance of the target conditions from the training set of conditions (for example, $ c=[0, 0, 0, 0] $ is very far from the training conditions, and it is much more difficult to realize).

Please note that we had to make color and rotation the choice of conditions in this experiment because the default conditional labels of the image benchmarks, such as MNIST/CIFAR10, are (1) discrete, and (2) there is no impartial metric on the space of label conditions. For example, MNIST is labeled with digits from 0 to 9, but in terms of image generation, 0 is no closer to 1 than it is to 9. The one-hot vector embedding is also not very reflective of actual image generation because 9 is actually much closer to 4 in shape than it is to 3. This is the reason why we limited our experiments to datasets like ZINC-250k of the form $ (x, c(x)) $, where $ c(x) $ is a feature of $ x $.

---

The above explanations have been added in purple highlights in the revised manuscript. We hope that these sections will make the novelty and potential of our research clearer.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a new flow matching algorithm based on a generalized (matrix-valued) continuity equation, with the motivation being unobserved data/conditions. The main claim of the work is that it outperforms existing methods in molecular generation tasks with sparsely observed conditions. 

The main strength is the rigorous development of the proposed algorithm, and clear motivation of the problem. However, the main claim of the paper that the method outperforms existing approaches is not fully validated in the experiments, and it remains unclear whether there is practical setting in which the generalized continuity equation is useful. 

I recommend to reject the paper at this stage, and encourage the authors to take the reviewers details feedback into account for a resubmission.