Paper: Multimarginal Generative Modeling with Stochastic Interpolants (ICLR-2024)

Reviewer comment (Theory):
Optimizing the vector field on simplex alpha may require delicate parameterizions. It would be better if it is more connected optimal transport.

Author response:
***Optimization on the simplex:*** The approach has precisely the advantage that it allows us to *solve the OT problem in a simulation free way on the simplex.* This is because we precompute (via quadratic regression) the conditional expectations $g_k(\alpha,x) = \mathbb{E}[x_k| x(\alpha)]$ where $x(\alpha) = \sum_{k=1}^K \alpha_k x_k$ and $x_k\sim\rho_k$. Having access to these conditional expectations allows us to calculate the transport cost along any transportation path parameterized by $\alpha(t)$ without having to perform any retraining and without having to solve an ODE: it is completely simulation-free in that respect. In turn, this allows us to optimize the path-parameterization $\alpha(t)$ to reduce the transportation cost. Of course, this only achieves OT within a specific parametric class  of velocities (i.e. those that transport samples on the simplex toward the PDF of any $x(\alpha) =  \sum_{k=1}^K \alpha_k x_k$), but it is still a nontrivial advance.

Reviewer comment (Presentation):
In Table.1, what does the smoothing mean in "two-marginal with smoothing"?

Author response:
***Questions and typos***: Thanks for pointing out a few unclear phrases. In the text, we will clarify that "2-marginal with smoothing" refers to the case of having an interpolant $x(\alpha) = \alpha_0 x_0 + \alpha_1 x_1 + \alpha_2 z$ where $z$ is a Gaussian random variable that is independent of $(x_0, x_1)$. That is, there are three marginals, a density $\rho_0$, a density $\rho_1$, and a Gaussian density $\rho_2 = \rho_z$. For a parametric path in which $\alpha_2(t) > 0$ when $t \neq 0,1$, this corresponds to the Gaussian "smoothed" interpolant from [1]. On the simplex, it resembles the simplicial coordinate moving toward the top vertex of a triangle (which corresponds to $\rho_z$ in the illustration given in the table).

Reviewer comment (Experiments):
There's a lack of clarity on the potential use-cases for this generative model. For instance, if it's geared towards style transfer, it's essential to determine which style gets preserved and how it compares with existing baselines in that domain.

Author response:
***Interpretability:*** The approach rests on the calculation of the conditional expectations $g_k(\alpha,x) = \mathbb{E}[x_k| x(\alpha)=x]$ where $x(\alpha) = \sum_{k=1}^K \alpha_k x_k$ and $x_k\sim\rho_k$. These functions have a direct intepretation: they give the Bayes-optimal estimator of the sample $x_k$ given the information contained in the observation $x(\alpha)$. The multimarginal stochastic interpolant framework turns this static picture into a dynamical one by showing that we can generate samples from $\rho(\alpha)$, the PDF of $x(\alpha)$, at any value of $\alpha$, and transport them along any time-parameterized path $\alpha(t)$ using either an ODE or an SDE with tunable diffusion coefficient. In terms of style transfer, it shows, for example, how a sample from a given distribution (say of cats) can be continuously transformed into a sample from another distribution (say of flowers) in a way that is influenced by additional distributions (via a path through the interior of the simplex), thereby incorporating their "style" along the way. We give some examples of such style-transfering paths in the paper, though we agree that more research is needed to reveal the full potential of the approach. We stress that this notion of style transfer is inherently different than any two-way correspondence, as the learned transport is affected by all the marginals.

Reviewer comment (Novelty):
The novelty of doing stochastic interpolant for multiple marginal distributions is limited given that it is done for the two marginal case. The theoretical result seems to be an extension from the two marginal case to the multi-marginal case as well.

Author response:
***Novelty:*** While the results follow quite naturally from the stochastic interpolant framework, they also constitute a nontrivial generalization of it. In particular, by showing that interpolation can be done in a multi-way fashion, we break from the standard perspective (used e.g. in score based diffusion models) in which the bridging between two distributions is itself defined via a time-dependent process (for example a forward SDE that need then to be reversed in time for generation). The original stochastic interpolant framework already showed that we can decouple the bridging of the densities from the generation process used afterwards; the multimarginal framework that we propose goes a step further in this direction by showing that the bridging itself does not need to use a scalar time variable but rather can be made multi-way using a vector of $\alpha=(\alpha_1,\ldots,\alpha_K)$.

Reviewer comment (Experiments):
Can the authors provide some more justification on why learning such an optimal transport path is useful? If learning the path in between distributions is to facilitate objectives like style transfer, then it should be compared with baselines on that task and show the utility of having this model. Right now, I am struggling to understand what is the use case especially for images when there is not a well-defined joint distribution between the multiple marginals.

Author response:
**Question:**  

- While we agree that the meaning of a joint distribution between images of different type is somewhat unclear, this is a concept that is being explored in the two-marginal case and is at the core of any method that tries to compute the optimal transport path between these two distributions. This is useful, for instance, in terms of function evaluations along the transporting path since it makes this path straighter. Our approach allows for an algorithm that finds a better path which does not require simulating the transport, via equation (16), which is not available in the original interpolant presentation. Note also that this can be achieved *without any retraining* which is a desirable feature, and also applies in the standard two-marginal case.

Reviewer comment (Theory):
It is unclear to me what is the implication of the Monge type coupling theoretical result. Is this shown to be relevant to the experimental results? If not I would recommend it to be put in appendix. It is also worth explaining more about what this result implies and when we would like to have this property.

Author response:
- We agree that the discussion about Monge coupling is somewhat tangential to our main point and will move it to an Appendix, as suggested by the reviewer. As a replacement, to improve the clarity, we are adding pseudocode for the main algorithms presented.

Reviewer comment (Reproducibility):
The overall algorithm of the proposed method might not be very clear or easy to follow, especially for a mathematically less mature audience. I suggest the authors add the pseudocode of the main algorithm to improve clarity.

Author response:
We thank reviewer PiDH for their feedback. We are glad that the reviewer found the work mathematically sound and more feasible than existing methods, and that they were happy with the experimentation.

The reviewer asked us to clarify the algorithm, which we can happily do by:

***Providing pseudocode*** and making a clear distinction between the steps that are specific to: *(i) learning the vector fields* that make possible the multimarginal framework, and *(ii) using them as a generative model* to sample from some initial condition (such as one of the marginals) on the simplex and transport to some final condition (such as one of the other marginals, or, say, to the barycenter).