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

We apologize for the lack of background information on Dirichlet energy. The Dirichlet energy represents how sensitive the distribution $p(\cdot \mid c)$ is to perturbations in $c$ with respect to the Wasserstein distance $W_2$. In fact, the following holds:

$\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 is equivalent to yielding posteriors that are not unnaturally sensitive in the Wasserstein sense. The experiment in §7.1 verifies this effect. Specifically, when generating the distribution of unobserved conditions using COT-FM, the distribution changes significantly from $c=(0.25,0)$ to $c=(0.5,0)$. In contrast, using MMOT-EFM reduces the extent of this change.

• 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 ci one can simply do the joint training over (xi,j,ci).

We apologize for the misunderstanding regarding the work of Wildberger et al., 2023. The contribution of the authors lies in their focus on the fact that the objective function of the FM for conditional generation only requires joint samples $(c,x)$, and they generated these joint samples in a Bayesian manner, such as $x \sim p(x)$ and $c \sim p(c \mid x)$.

However, the treatment of conditions $c$, aside from the aforementioned joint sampling, is almost identical to (OT-C)FM as described in §7. Therefore, if we apply their results to our setup, we can expect similar experimental outcomes to those of (OT-C)FM.

We have clarified this in the revised manuscript.

We appreciate your thorough review of our manuscript and your understanding of the motivation behind our research. Below, we address your concerns in detail. We would be grateful if you could reconsider our rating based on these clarifications.

For Weakness

• 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=ty+(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.

Indeed, the primary motivation of our paper is to address conditional generation in scenarios where the conditions are sparsely observed, a situation where "the classical conditional algorithms are not applicable." For example, in the case of style transfer between two known conditions $c_1$ and $c_2$, as you mentioned, one can simply use a conditional FM network that receives the two conditions (source and target) as inputs. However, if either $c_1$ or $c_2$ is unavailable, this FM network cannot be used. Furthermore, even if all conditions $c_1, c_2, \dots, c_M$ are known, and we wish to perform style transfer among them, we would need to train the FM network approximately $\binom{M}{2}$ times, which becomes inefficient as the number of conditions $M$ increases. One of the advantages of our EFM approach is that it can handle the aforementioned intractable situations in FM with a single training of the matrix field model.

• The scaling in Nc 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.

As you pointed out, the runtime for MMOT-EFM can become significantly longer compared to the baselines when $N_c$ and the condition dimension are large. However, the computation of MMOT should be independent of matrix field learning, which means that the runtime can be significantly reduced by optimizing the implementation, e.g., by introducing parallel computing.

Regarding the runtime, the 4-hour runtime includes the evaluation phase. A significant amount of time was spent on evaluation, and the implementation needs to be fully optimized, which contributes to the longer runtime. We will provide the net runtime without the evaluation phase at a later date. As for the experiment on MNIST with a continuous condition space, such as inpainting, our method may not be well suited from a computational complexity perspective due to the high dimensionality of the condition vector.

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

We have deleted these parts.

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

We have just submitted the code as supplementary material.

We appreciate your interest in the EFM concept. Below, we address your concerns regarding the molecular generation experiment:

For Weakness

• 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]

The experimental design is motivated by the specific objectives outlined in §1. Our method demonstrates a distinct advantage in scenarios where multiple molecules $x$ share the same label value, and the label value $c$ is sparse. Standard chemical properties, which are continuous and have fewer molecules with identical labels, are not suitable for this experiment. Consequently, we selected a scenario involving the number of bonds, where only two label values are available.

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

The term FM in §7 refers to the OT-CFM method. We apologize for any confusion caused by the lack of clarity. This has been clarified in the revised manuscript.

For Questions

• Could you please justify the ZINC-250k experimental design?

ZINC-250k is a widely utilized molecular database of computationally designed compounds, created by G'omez-Bombarelli et al. (2018). It is a more realistic dataset that includes drug-like, commercially available molecules, compared to QM9, which is a standard dataset for molecule generation. We anticipate that similar results would be obtained if trained on other datasets

[G´omez-Bombarelli et al., ACS Central Science, 2018]

We appreciate your interest in our EFM theory. We hope that the responses to your questions below will address your concerns comprehensively.

For Weakness

• It would be beneficial to elucidate how this method extends to unseen conditions more clearly. Lines 402-405 touch on this, but further emphasis on this point would be valuable. Designing a conditional generative model that technically extends to unseen conditions is straightforward, but ensuring that the model extends in a reasonable manner is more challenging. EFM has the potential to guide this extension, and further exploration of this point would be appreciated.

Our method, MMOT-EFM, extends to unseen conditions by ensuring that the target conditional distribution $p(x \mid c)$ is "smooth." Specifically, we aim to minimize the Dirichlet energy as defined in Equation (3.2). Intuitively, the Dirichlet energy quantifies the sensitivity of the distribution $p(x \mid c)$ with respect to the conditioning vector $c$. Minimizing the Dirichlet energy implies that the sensitivity to the condition $c$ is not excessively large.

While there are numerous methods of extrapolation, it is reasonable to assume an inductive bias that the sensitivity of natural data (e.g., molecules) to conditions (e.g., chemical properties) is not unnaturally large. Therefore, our method addresses extrapolation by training a model such that the data to be extrapolated adheres to this inductive bias of low sensitivity.

We would like to highlight that this type of inductive bias has been historically employed in generative models to prevent overfitting and stabilize generative processes; see, for example, Miyato et al. in ICLR, 2018.

We have revised the manuscript to further emphasize this point in §1, illustrating how MMOT-EFM guides the extension to unseen conditions in a reasonable manner.

• The algorithm is not yet applicable to real-world scenarios. While the authors acknowledge this, it remains a significant limitation of the work's impact. The molecule experiment is limited in terms of comparisons to existing work and the overall training setup.

Our method has yet to be evaluated on real-world datasets. However, the ZINC-250k dataset used for conditional molecule generation includes drug-like commercially available molecules, making it more realistic compared to well-known datasets such as QM9.

• Much of the theoretical statements are direct extensions from prior work.

While it is true that Theorem 3.4 is a direct extension of the fundamental theorem in FM, the theory developed in §3 as a whole should not be seen as a straightforward derivation from existing FM theory. For instance, Proposition 3.1, which introduces a method for utilizing the generalized continuity equation (3.1) for generation and style transfer, is a novel technique within the context of generative models.

For Questions

• When is MMOT-EFM and EFM in general expected to outperform COT-FM / Bayesian-FM? Although there is a brief explanation of the differences in assumptions, it is challenging to understand the benefits of assuming piecewise continuity of $p(x|c)$ versus measurability. How does this compare to prior works in general?

Our EFM approach is expected to perform better when there is prior knowledge that the distribution $p(x \mid c)$ we aim to generate changes continuously with respect to $c$. Specifically, it is effective when the sensitivity of $p(x \mid c)$ with respect to the condition $c$, quantified by the Dirichlet energy, is known to be small.

In contrast, COT-FM and Bayesian-FM do not incorporate such prior knowledge regarding the continuity or sensitivity of $p(x \mid c)$ with respect to $c$.

This distinction is evident in the generated results, as shown in Figure 4(b). When generating the distribution of unobserved conditions using COT-FM, the distribution changes significantly from $c=(0.25,0)$ to $c=(0.5,0)$. Conversely, with MMOT-EFM, the extent of this change is mitigated.

We have included this explanation in the revised manuscript.

For Overall

• This work presents an intriguing idea with potential to enhance understanding of how these models generalize to unseen conditions. However, this aspect is not theoretically explored. Additionally, the current method does not scale to practical settings. Further investigation into when the assumptions behind this method are valid relative to other methods would significantly strengthen this work. A deeper understanding of how this relates to prior literature and when this method is preferable would likely change my opinion of this work.

We appreciate your evaluation of our concept. Indeed, there are scalability issues, and our method may need further development to be applied to general tasks such as image generation. However, in the context of our motivation, such as molecular generation, we have demonstrated its advantages over existing methods in §7.2. Regarding the relation to prior literature, particularly the differences from COT-FM and Bayesian FM, we have clarified these points in our response to your previous question. Please refer to that response for a detailed explanation.

We appreciate your thorough review of our manuscript. Below, we provide detailed responses to your inquiries. We hope that these clarifications will enhance your confidence in your evaluation.

For Questions

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

Equation (3.3) serves as a soft constraint for the boundary conditions in Equation (3.4). The meaning of Equation (3.4) is that the conditional distributions generated by the model should match the known data distributions $\mu_\xi$, $\xi \in A$. This is analogous to the requirement in standard FM that the endpoints of the probability path generated by the model correspond to the source and target distributions. In the case of FM, where only two distributions are considered, this can be achieved by simply connecting them with a straight line, eliminating the need for regression. However, in EFM, we need to consider three or more distributions simultaneously, making it hard to satisfy the constraints strictly. Therefore, we relax the constraints to be soft in this context.

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

Our method is designed to work in situations where the dimensionality of the conditioning vector is relatively low, and the conditions are sparsely observed. Such scenarios are common in the context of molecular generation, as demonstrated in §1 and §7.2. Therefore, EFM is expected to show extrapolation ability in a more practical molecular generation.

We have incorporated the aforementioned explanation into the introduction of the revised manuscript.

• 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 W2 metric is minimized.
• [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.

In fact, our proposed EFM framework indeed incorporates the approach of [Albergo et al., 2023].

More specifically, one of the implementations of EFM, Geo-EFM (supplemented in Appendix F), and [Albergo et al., 2023] both use the same transport plan $\pi$ to learn the matrix (or vector) field. In addition, our interpolator $\bar\psi(c \mid (x_i)_i)$ in Table 1 corresponds to the barycentric stochastic interpolant $x(\alpha) = \sum_i x_i \alpha_i$ in [Equation (5), Albergo et al., 2023]. Here, the interpolation coordinates $\alpha = (\alpha_i)_i \in \Delta^K$ valued in the simplex $\Delta^K$, can be roughly regarded as the condition vector $c \in \Omega$ in our case, i.e., $\Delta^K \approx \Omega$. More precisely, the condition $c$ is given by expanding it on a basis. That is, $\alpha$ is the coefficient when expanding $c$ as $c = \sum_i c_i \alpha_i$ in a basis $(c_i)_i$. Consequently, Geo-EFM, similar to [Albergo et al., 2023], can utilize an arbitrary number of marginal distributions during the training process.

The only difference between Geo-EFM and the method in [Albergo et al., 2023] is that, after learning the vector field, Albergo et al. optimize the path on the space of "conditions" $\alpha \colon [0,1] \to \Delta^K \approx \Omega$ such that the transport cost in the W2 metric is minimized. Thus, we can take the same procedure as above to optimize $\gamma \colon [0,1] \to \Omega$ with the Geo-EFM setting.

Our MMOT-EFM is novel in that it minimizes the transport cost in a complementary way to the optimization of $\gamma$. MMOT-EFM trains a matrix field, which is an extension of a vector field, to also minimize a generalization of the transport cost called Dirichlet energy. This makes it possible to learn a model that transports optimally with only one training of the model without optimizing $\gamma$. We note that there is a computational limitation on the number of marginal distributions (as you pointed out) due to the use of MMOT during training.

Moreover, in our setting, the set of conditions $\Omega$ is continuous, whereas in [Albergo et al., 2023], it is discrete.

In summary, MMOT-EFM and [Albergo et al., 2023] can be seen as complementary approaches. We have mentioned the above in the revision.

• The experiments section can be improved by adding extra text explaining the results and the figures, particularly in figure 4.

We apologize for the inconvenience. Due to page limitations, we could not include sufficient explanation in the submitted manuscript. We have added more explanations to the caption in the revision.

For Questions

• Can the authors consider providing definitions before introducing a new notation in the text?

We apologize for the technical nature of the mathematical equations. We would like to consider your comments on the manuscript. Could you tell us exactly where we introduce the new notation before the definition?

• 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?

Larger batch sizes tend to stabilize the learning process because the changes in the matching $\pi$ per iteration become smaller.

• What kernels do the authors use for the RKHS used to construct paths?

In general, one can employ any nonlinear kernel. In the implementation of this paper, we use the linear kernel.

• In lines 212-214 and lines 220-222, can the authors clarify the output of $u$?

Here, $u$ returns a matrix of size $d \times (1+k)$, where $d$ is the dimension of the data $x \in D$ and $k$ is the dimension of the condition $c \in \Omega$. In general, the output of $u$ is of size $d \times \operatorname{dim} \Xi$ (see Proposition 3.1). In §3.1, $\Xi = I \times \Omega$, so $\operatorname{dim} \Xi = 1+k$.

• The discussion about the weak assumption of measurability and continuity of $p(x \mid c)$ with respect to $c$ requires clarification, particularly since piece-wise continuous functions are measurable as well.

We apologize for the confusion. Our intention was not clearly conveyed. In COT/Bayesian-FM, the conditional distribution $p(x \mid c)$ is assumed to be measurable with respect to $c$, which allows $p(x \mid c)$ to change discontinuously. In contrast, we assume that $p(x \mid c)$ changes continuously, or more precisely, that the sensitivity of $p(x \mid c)$ with respect to the condition $c$ is small. This difference is demonstrated in §7.1, Figure 4(b).

We have included this clarification in the revision.

We appreciate your valuable feedback and apologize for any lack of clarity in our initial explanation. Below, we address your concerns and questions:

For Weaknesses

• 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.
• Could the authors explain why the multi-marginal optimal transport approach allows for extrapolating to conditioning vectors not seen during training?

We apologize for the unclear wording in line 86. 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," which means minimizing 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$. 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 a method to 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$.

Thank you very much for the follow-up, we would like to address your concerns below.

• Can the authors comment on the convergence of the finite-sample approximation of the coupling $\pi$ to the coupling $\pi^\ast$ which minimizes $\mathbb{E}_\pi |x_0−x_1|_2^2$? Any empirical or theoretical analysis would make the case stronger. More specifically, as the authors claim on line 276, can the optimal coupling $\pi^\ast$ be approximated by finite-samples? a proof or asymptotic analysis would be appreciated.

It is known that the finite sample approximation converges asymptotically for any $\pi$ that minimizes $\mathbb{E}_\pi |x_0−x_1|_2^2$, for example [Theorem 5.10, Villani, 2009]. We also conjecture that the multi-marginal case in line 276 converges in a similar way.

• From appendix D.3 it seems the cost of MMOT is prohibitive, therefore the authors propose an approximation. Can the authors discuss the effect of this approximation on the claim that they minimize an upper bound on the Dirichlet energy?

In this approximation, MMOT is applied only at the cluster level, and the coupling between each cluster is conducted using a method of the user’s choice. Indeed, this has an effect of achieving the energy value that is greater than the upper bound shown in the objective function in 3.2 because, in the implementation, the infimum in $\inf_\pi \int_{D^{|A|} \times \Xi } \| \nabla_\xi \phi(\xi | x_A) \|^2 \pi (d x_A) dc$ will be taken not over the space of all joint distributions $\mathcal{P} ( D^{|A|})$, but over its subset of form $\pi_{\mathrm{user}} ( x_A | \lbrace m_{i} \rbrace ) \pi_{\mathrm{cluster}}( \lbrace m_{i} \rbrace )$ where $\pi_{\mathrm{cluster}}$ is computed with discrete MMOT over $\lbrace m_{ik} \rbrace$ and $\pi_{user}$ is chosen by the user.

Indeed, $\pi_{user}$ can be chosen to be couplings that respect the consideration to kinetic energy as well, such as the generalized Geodesics coupling based on optimal transport. Also, by definition, this approximation shall converge to the actual upper bound if we choose $|U_{ik}| = 1$ take the limit of $|B_i| \to \\infty.$

I believe the authors have a typo in their rebuttal, there is no discussion of any Geo-EFM algorithm in appendix F, rather it contains a description of metrics, datasets and baselines.

We are sorry for the typo. Our discussion of Geo-EFM in the currently uploaded version (modified: 21 Nov 2024) is provided in Section E, with the title A REMARK ON GENERALIZED GEODESIC COUPLING(GGC) AND THE SAMPLING OF $\bar{\psi}$.

(Undefined notations at) the paragraph from lines 148-161, including equation 2.3.

Thank you for specifying the spot. We stated the definitions before making the statement in 148-161, just as below:

Let $\psi$ be a random path such that $\psi\colon I \rightarrow D$ is differentiable. Let $Q$ be a distribution over a space $H(I; D) \coloneqq \Set{\psi\colon I \rightarrow D | \psi \text{ is differentiable}}$ of paths that map time $t\in I$ to data $x \in D$, and use $\mu^\psi_t$ to denote $\delta_{\psi(t)}$. With these definitions, we can present $\mu = \mu^Q$ from a random path $\psi$ as

We have revised the phrases with similar expressions when we can.

Can the authors show any examples of such an inductive bias helping in solving any high-dimensional inverse problems? For instance, with the MNIST/CIFAR10, etc datasets? Or class-conditional generation?

The difficulty in providing the efficacy of our method on image benchmarks like MNIST/CIFAR10 is that the conditions in these datasets are (1) discrete and (2) there is no impartial metric on the space of conditions. For example, MNISTS are labeled with digits from 0 to 9, but in terms of image generation, 0 is not closer to 1 than it is to 9. One hot vector embedding is also not too reflective of actual image generation because 9 is indeed much closer to 4 in terms of an image than it is to 3. This is the reason why we restricted our experiments to the dataset like ZINC-250k of the form $(x, c(x))$, where $c(x)$ is a feature of $x$. While there are many datasets of this form in application, for example, in applied fields of science such as economics, biochemistry, and physics, benchmark datasets, models, and publicized architecture are difficult to obtain. (continued to next part)

Can the authors comment on why they have not done even a MNIST generation experiment.

The reason why we did not present the MNIST generation experiment with digit conditions is that there is no continuous dependence on the digits in the coupling between the conditional distributions. The essential purpose of EFM is to generate a good coupling sample of $\lbrace x_c \mid c \in \Omega\rbrace$ under the constraint that $x_c$ changes “smoothly” when $c$ is varied over $\Omega$. Recall that the purpose of this was to generate and interpolate/extrapolate the values of $\mu_c$ when $c$ is not in $C_{\mathrm{train}}$. Therefore, we have chosen an experimental setup that has dependencies on $c$. If the problem is to generate $x_c \sim \mu_c$ for each $c$ independently for each $c \in C_{\mathrm{train}}$ (e.g., MNIST with digit conditions), we do not need to consider the joint distribution, and it makes no difference whether we use MMOT-EFM or Geo-EFM, or even the EFM with random $\pi$ couplings. In particular, when it comes to generating only the 'marginal distribution' for $c \in C_{\mathrm{train}}$, EFM and [Albergo et al.] are theoretically the same as the original OT-CFM with conditions.

it would be useful if the authors can comment on what limitations the method faces when scaling to dimensions > 100.

The effect of dimension on the complexity of EFM is indirect. In general, the computational cost of MMOT/OT for $d$ dimensional particles scales linearly with d. In contrast, it scales with $n^m$, where $n$ is the number of particles and $m$ is the number of marginal distributions. At the same time, the larger the dimension $d$, the slower the batch sample converges to the true distribution; in particular, the expected Wasserstein distance from the empirical distribution to the true distribution scales with $O(n^{-1/d})$ when measured with the $W_1$ distance [3]. Thus, the larger the $d$, the larger the number $n$ of particles it would take to approximate the MMOT/OT with a given precision in terms of the transport cost.

[3] Fournier, and Guillin. "On the rate of convergence in Wasserstein distance of the empirical measure." Probability theory and related fields, 2015.

Even assuming that they have enough model capacity to learn a vector (or matrix) field, what are the implications of using finite-sample approximations of the transport plan?

The implications of using finite-sample approximations are twofold. (1) The quality of the approximated marginal distributions is affected by the batch size. (2) The quality of the approximated coupling (joint distribution) is affected by the batch size.

The first part applies to any batch-based flow-matching method in general because batches are used in the flow matching as the approximation of marginal distributions, which converges to the ground-truth distribution. This also applies to the optimal transport cost and, most likely, to Dirichlet energy. As stated above, the larger the $d$, the greater the size of the batches it would require to obtain a matrix field with the lowest transport energy.

We believe that our response above clarifies the limitations of the method and the impact of mini-batch OT, which are your concerns.

Thank you for your response. We would like to clarify our position further.

So in effect, the authors are not able to upper bound the Dirichlet energy with their objective?

Unbiased estimators of OT plans using finite samples or mini-batches are discussed in the following paper:

• Fatras et al., Learning with minibatch Wasserstein: asymptotic and gradient properties. AISTATS 2020.

It would also be possible to achieve the upper bound of our Dirichlet energy by constructing an estimator in the same manner as they did. In training generative models, we believe in constructing a similar estimator approximately using stochastic gradients.

Their method is applicable to marginals distributions with continuous-valued support, see definition 1.

Firstly, let us clarify the terminology. In [Definition 1, Albergo et al., 2024], the term "support" refers to the set of data points $x$. In both our paper and Albergo et al.'s paper, $x$ takes continuous values. You might be referring to $\Delta^K$ in [Definition 1, Albergo et al., 2024] as the "support." While it is true that $\alpha \in \Delta^K$ takes continuous values, $\alpha$ merely represents a probability vector over a discrete set, making it essentially discrete. In contrast, in our setting, the condition vector $c \in \Omega$ not only takes continuous values but can also represent quantities that vary continuously, such as color or angle. This allows for a more flexible representation of continuously varying conditions.

This is incorrect. In inverse problems, one can define a matrix $A$ and observations $y = Ax + \varepsilon$, where $x$ is the image and $\epsilon$ is mean zero noise, More typically, one can always consider generating labels $y = g(x)$ for a deterministic function $g$.

When considering the relationship $y = Ax + \varepsilon$ for an image $x$ and a categorical label (digit) $y$, if the image $x$ changes continuously, the label $y$ would also change continuously. However, since $y$ is a categorical variable, it is unnatural for it to change continuously. In the context of categorical labels, such as digits, the assumption of a continuous relationship does not hold. This is why we believe there is no impartial metric on the conditions of digits when considering a conditional distribution of images, especially because there are no ground-truth datasets other than those used in training (0, ..., 9).

Did the authors use a pre-trained autoencoder? What model class, were any pre-trained models used, what batch size, the cost of running mmot, how many iterations.

For the image experiment, we used Encoder_ResNet_AE_CIFAR of the Pythae library, fine-tuned on the training distributions over 40 epochs with batch size 128. The general cost of Sinkhorn MMOT is $O(n^m)$, where $n$ is the batch size and $m$ is the number of marginals. Indeed, with this scaling, setting $m= |C_{\mathrm{train}}|$ is prohibitive, where $C_{\mathrm{train}}$ is the set of all conditions we use in training. We also note that this affects memory availability because storing complexity of $O(n^m)$ is infeasible for large $n$.