Multimarginal Generative Modeling with Stochastic Interpolants

We thank reviewer yfUs for their feedback. We are glad that the reviewer found our presentation clear and our experiments comprehensive. The reviewer, however, had concerns about the intrerpretability, utility, and novelty that we now address.

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.

Naïve extension: 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 betwen two distributions is itself defined via a time-dependent process (for example, a forward SDE that needs 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 muti-way using a vector of $\alpha=(\alpha_1,\ldots,\alpha_K)$.

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

We thank reviewer 2Wkn for their feedback. We are glad that the reviewer found our multi-marginal stochastic interpolant framework unique and intriguing, and noted its connection with multi-marginal OT. The reviewer, however, had concerns about the interpretability, utility, and novelty, of the approach, which we now address.

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.

Utility: The approach offers possibilities that cannot be achieved by current generative models that proceed between only two marginals. Because the method gives access to the vector fields $g_k(\alpha,x)$ for all values of $\alpha$, we can a-posteriori change the interpolation path by varying the functional form of $\alpha(t)$ and thereby produce many different generative models (all unbiased) without having to perform any retraining. In turn, this allows us to use these models for multiple tasks, or to optimize them according to several suitable criteria (OT being just one possible choice), as already discused in the first point about interpretability above.

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)$.

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.

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

Thanks again for your valuable insights.

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

Clarifying how we solve equation (16), which is the path minimization loss. This optimization is done separately from learning the vector fields $g_k(\alpha,x) = \mathbb{E} [x_k | x(\alpha)=x]$ with $x(\alpha) = \sum_{k=1}^K \alpha_k x_k$. Once we have learned the $g_k$, we can a-posteriori change the interpolation path by varying the functional form of $\alpha(t)$. For example, if we want to transport from marginal $\rho_i$ to marginal $\rho_j$, we can take any $\alpha(t)$ with initial condition $\alpha_k(0) = \delta_{k,i}$ and a final condition $\alpha_k(1) = \delta_{k,j}$. Any parameterization of $\alpha(t)$ that meets these boundary conditions is sufficient. In the Appendix we present one such parameterization via an expansion of Fourier coefficients, which has the advantage that it has very few parameters; alternatively we could also use a neural network to represent $\hat \alpha(t)$. With a parameterization $\hat \alpha(t)$ and the learned $g_k$ in hand, we can evaluate equation (16) in a simulation free manner just by sampling the interpolant $x(\alpha) \sim \rho(\alpha, x)$ and performing gradient descent on the parameters of $\hat \alpha(t)$. This optimization is extremely cheap and typically costs a minute fraction of the cost of learning the vector fields. We will make this clear with the addition of a pseudocode algorithm in the revision.