Multimarginal Generative Modeling with Stochastic Interpolants
Authors: Michael Samuel Albergo
Nicholas Matthew Boffi
Michael Lindsey
Eric Vanden-Eijnden
Venue: ICLR 2024 poster (2024)
Keywords: multi-marginal
unsupervised learning
generative modeling
measure transport
optimal transport
Primary area: generative models
TL;DR: We introduce a method to generalize flow-based and diffusion based generative models to map between K distributions instead of two, revealing multiway-correspondences between densities.

Abstract: Given a set of $K$ probability densities, we consider the multimarginal generative modeling problem of learning a joint distribution that recovers these densities as marginals. The structure of this joint distribution should identify multi-way correspondences among the prescribed marginals. We formalize an approach to this task within a generalization of the stochastic interpolant framework, leading to efficient learning algorithms built upon dynamical transport of measure. Our generative models are defined by velocity and score fields that can be characterized as the minimizers of simple quadratic objectives, and they are defined on a simplex that generalizes the time variable in the usual dynamical transport framework. The resulting transport on the simplex is influenced by all marginals, and we show that multi-way correspondences can be extracted. The identification of such correspondences has applications to style transfer, algorithmic fairness, and data decorruption.  In addition, the multimarginal perspective enables an efficient algorithm for optimizing the dynamical transport cost in the ordinary two-marginal setting. We demonstrate these capacities with several numerical examples.

Published as a conference paper at ICLR 2024
MULTIMARGINAL GENERATIVE MODELING WITH STOCHASTIC INTERPOLANTS
Michael S. Albergo
Nicholas M. Boffi
Center for Cosmology and Particle Physics
Courant Institute of Mathematical Sciences
New York University
New York University
New York, NY 10003, USA
New York, NY 10012, USA
elbergo@nyu.edi
boffiecims nyu..du
Michael Lindsey
Eric Vanden-Eijnden
Department of Mathematics University of California, Berkeley Berkeley, CA 94720, USA
Courant Institute of Mathematical Sciences
New York University
New York, NY 10012, USA
seddsey@math.berreee. edu
eve2@cims edu
ABSTRACT
Given set of probability densities, we consider the multimarginal generative modeling problem of learning a joint distribution that recovers these densities as marginals. The structure of this joint distribution should identify multi-way correspondences among the prescribed marginals. We formalize an approach this task within generalization of the stochastic interpolant framework, leading to efficient learning algorithms built upon dynamical transport of measure. Our generative models are defined by velocity and score fields that can be characterized as the minimizers simple quadratic objectives, and they are defined on simplex that generalizes the time variable in the usual dynamical transport framework. The resulting transport on the simplex is influenced by all marginals, and we show that multi-way correspondences can be extracted. The identification of such correspondences has applications to style transfer, algorithmic fairness, and data decorruption. In addition, the multimarginal perspective enables an efficient algorithm for reducing the dynamical transport cost in the ordinary two-marginal setting. We demonstrate these capacities with several numerical examples.
INTRODUCTION
Generative models built upon dynamical transport of measure, in which probability densities connected by Icarnable transformation, underlic many recent advances in unsupervised learning (Rombach et al., 2022; Dhariwal & Nichol, 2021). Contemporary methods such as normalizing flows (Rezende & Mohamed, 2015) and diffusions (Song et al., 2021b) transform samples from one density Pu into samples from another density P1 through an ordinary or stochastic differential equation (ODE/SDE). In such frameworks one must learn the velocity field defining the ODE/SDE. One effective algorithm for learning the velocity is based on the construction of stochastic interpolant, a stochastic process that interpolates between the two probability densities at the level of the individual samples. The velocity can be characterized conveniently as the solution of a tractable square loss regression problem. In conventional generative modeling, one chooses po to be analytically tractable reference density, such as the standard normal density, while P1 some target density of interest, accessible only through dataset of samples. In this setting, general stochastic interpolant I, can be written as
Hr === ao()) + ai (1)=1,
(1)
where Ey ~ Po and Æ1 ~ P1 and we allow for the possibility of dependence between 2g and ø1-
Meanwhile au and a1 (t) are differentiable functions of t E [0. 1] such that (0) (1) = and an (1) a (0) = 0. These constraints guarantee that zı=o = In and 2;=1 = 21 by construction. It was shown in Albergo & Vanden-FFijnden (2023); Albergo et al. (2023) that if X, is the solution of

---

Published as a conference paper at ICLR 2024
the ordinary differential equation (ODE)
X, ((t,Xi).
with velocity b defined by
b(t, x) = Ba = x] (o( B[roolr 2] dr (1) E 21 by = x3
and initial condition Xa drawn from Po. then X, matches z; in law for all times t E [0, 1]. Hence we can sample from P1 by generating samples from Pu and propagating them via the ODE (2) to time t 1. Equivalent diffusion processes depending on the score can also be obtained by the introduction of noise into the interpolant (Albergo et al., 2023).
significant body of research both flows and diffusions has studied how to (t and a1 (t) that reduce the computational difficulty of integrating the resulting velocity field b (Karras et al., 2022).
In this work, first observe that the decomposition of the velocity field of (1) into conditional expectations of samples from the marginals and #1 suggests more general definition of process (a)) defined not with respect the scalar / [0, 1] but with respect to an interpolation coordinate (oo, a1),
z (a) ag + 01
By specifying curve (t), we recover the interpolant in (1) with the identification 2; z (t))))). We use the generalized perspective in (4) to identify the minimal conditions upon a so that the density p(a, of v((a) is well-defined for all a, reduces to po(x)) for (1,0), and to pi (ar) for (0, 1). These considerations lead to several new advances, which we summarize as our main contributions:
1. We show that the introduction of an interpolation coordinate decouples the learning problem of estimating given bt, from the design problem for a path att) governing the transport. We use this to devise an optimization problem over curves at)) with the Benamou-Brenier transport cost, which gives a geometric algorithm for selecting a performant a.
2. By lifting to a higher-dimensionn a generic paradigm for multimarginal generative modeling. To this end, we derive a generalized probability flow valid among K + 1 marginal densities po, PI..... PK.. PK. whose corresponding ve- locity field b(1, æ)) is defined via the conditional expectations E E[ [=<(a) x] for = 0, 1,..., K and a curve (t). We characterize these conditional expectations as the minimizers of square loss regression problems, and show that one gives access to the score of the multimarginal density.
space,
the stochastic interpolant framework to devise
3. We show that the multimarginal framework allows us solve (K) marginal transport problems using only K marginal vector fields. Moreover, this framework naturally learns multi-wal correspondences between the individual marginals, detailed in Section 3. The method makes possible concepts like all-to-all image-to-in tomi-agg translation, where we observe an emergent style transfer amongst the images.
4. Moreover, in contrast to existing work, we consider multimarginal transport from novel generative perspective, demonstrating how to generate joint samples (Zo....,Ik) matching prescribed marginal densities and generalizing beyond the training data.
The structure of the paper is organized follows. In Section 1.1, review some related works in multimarginal optimal transport and two-marginal generative modeling. In Section 2, we introduce the interpolation coordinate framework and detail its ramifications for multimarginal generative modeling. In Section 1. we formulate the path-length minimization problem and illustrate its application. In Appendix B, we consider a limiting case where the framework directly give one-step maps between any two marginals. In Section 3, we detail experimental results in all-to-all image translation and style transfer. We conclude and discuss future directions in Section
1.1 RELATED WORKS
Generative models and dynamical transport Recent years have seen an explosion of progress in generative models built upon dynamical transport of measure. These models have roots as early (Tabak & series of steps (Rezende & Mohamed, 2015; Dinh et al., 2017; Huang et al., 2016; Durkan et al., 2019), while modern models typically formulated via continuous-time transformation. A
2010; Tabak & Turner, 2013), and originally took form discrete
2

---

Published as a conference paper at ICLR 2024
particularly notable example of this type is score-based diffusion (Song et al., 2021c;a; Song & Ermon, 2019), along with related methods such as denoising diffusion probabilistic models (Ho et al., 2020: Kingma et al., 2021), which generate samples by reversing stochastic process that maps the data into samples from a Gaussian base density. Methods such as flow matching (Lipman et al., 2022; Tong et al., 2023), rectified flow (Liu, 2022; Liu al., 2022), and stochastic interpolants (Albergo & Vanden-Eijnden, 2023; Albergo et al., 2023) refine this idea by connecting the target distribution to an arbitrary base density, rather than requiring Gaussian base by construction of the path, and allow for paths that are more general than the one used in score- based diffusion. Importantly, methods built transformations that posit connection between densities often lead more efficient learning problems than alternatives. Here, we continue with this formulation, but focus on generalizing the standard framework to multimarginal setting.
upon continuous-
Multimarginal modeling and optimal transport Multimarginal problems are typically studied from an optimal transport perspective (Pass, 2014b), where practitioners are often interested in computation of the Wasserstein barycenter (Cuturi & Doucet, 2014; Agueh Carlier, 201 1: Altschuler & Boix-Adsera) The barycenter is thought to lead to useful generation of combined features from set of datasets (Rabin al., 2012), but its computation is expensive, leading the need for approximate algorithms. Some recent work has begun to fuse the barycenter problem with dynamical transport, and has developd algorithms for computation based on the diffusion Schrödinger bridge (Noble et al., 2023).
significant body of work multimarginal optimal transport has concerned the existence of an optimal transport plan Monge type (Gangho & Šwiech, 1998; Pass, 2014a), as well as generalized notions (Friesecke & Vögler, 2018). A Monge transport plan is a joint distribution of the restricted form (31) considered below and in particular yields set compatible deterministic correspondences between marginal spaces. We refer to such a set of compatible correspondences as multi-way correspondence. Although the existence of Monge solutions in multimarginal optimal transport is open in general (Pass, 2014a), in the setting of the quadratic cost that arises in the study of Wasserstein barycenters, in fact Monge solution does exist under mild regularity conditions (Gangbo & Swiech, 1998).
We do not compute multimarginal optimal transport solutions in this work, but nonetheless we will show how to extract an intuitive multi-way correspondence from our learned transport. We will also show that if the multimarginal stochastic interpolant is defined using coupling of Monge type, then this multi-way correspondence is uniquely determined.
2 MULTIMARGINAL STOCHASTIC INTERPOLANTS
In this section, we study stochastic interpolants built upon an interpolation coordinate, as illustrated in (4). We consider the multimarginal generative modeling problem, whereby we have access to dataset of n samples {at 1= from each of K + densities with ri. ~ Pk. We wish to construct generative model that enables us to push samples from any Pj onto samples from any other Pk- By setting po to be tractable base density such Gaussian, we may use this model to draw samples from any marginal; for this reason, we hereafter assume Po to be a standard normal. We denote by
K the K-simplex in I K+1 i.e., A = {a- = (ap.... agk E R Ek=o Og = 1 and og 2 O for ls = O.... K}. We begin with a useful definition of a stochastic interpolant that places a E A K; Definition (Barycentric stochastic interpolant). Given K + 1 probability density functions {px K with full support on RRd,, the barycentric stochastic interpolant æ(a) with a (oo...., E A is the stochastic process defined as
xa))
TROK::
(5)
where (I1,..., II) are jointly drawn from probability density p(x1,'.. ,KK) such that
Vk: 1,..,, and we set 20 (O, Ida) independent of (I1 xK))
K JE plx,,.... IK )dxı 1 dx pa (xx):
(6)
3

---

Published as a conference paper at ICLR 2024
The barycentric stochastic interpolant emerges natural extension of the wwo-marginal inter- polant (1) with the choice (u(t) = 1 - and 01 = t. In this work, we primarily study the barycentric interpolant for convenience, but we note that our only requirement in the following discussion is that Li=o a O. This condition ensures that æ(a) always contains a contribution from some 2k, and hence its density p(a. -) does not collapse to a Dirac measure at zero for any a. In the following, we classify p(a..) -) the solution to a set of continuity equations.
Theorem (Continuity equations). For all E A the probability distribution of the barycentric stochastic interpolant x(a)) has a density p(ar, I) which satisfies the K + equations
Pa,p(a, + V (9x((,,))((a x)) =0, k o...., K.
(7)
Above, each (s( x)) is defined as the conditional expectation
gk I x) B[ra(=(a) m), 0,.0..... K,
(8)
where [[zn(=(a)- = z] denotes an expectation over Po(zz))(( 1(x1,,,......... .x) conditioned on the event a(())= = The score along each two-marginal path connected to pa is given by
Yao # 10: Ve log pla, x) = a0000=-aaaaa 90(a,
Moreover, each gk is the unique minimizer of the objective
(ūne) S&a Ellin(a, x(a))) - 2n Iu (ax, x(a)))
k =
(10)
where the expectation IE is taken over (zo...., ,kk ~ pa ( TT).
The proof of Theorem is given in Appendix A: it relies on the definition of the density p(a, x) that says that, for every suitable test function Φ(x), we have
J ((x))(a, = - M (=(0))) xK )dro drk.
(11)
where x(a)= = Ek= Ig. Taking the derivative of both sides with respect to og using the chain rule well 2 Tk. gives us
Jee ((x)0n x )dx = [
J IK (ola(()))o(()))t)).............. oo drk.
(12)
Integrating the right hand side by parts and using the definition of the conditional expectation implies (7).
As we will see, the loss functions in (10) are amenable to empirical estimation over a dataset of samples, which enables efficient learning of the gs- The resulting approximations can be combined according to (14) to construct multimarginal generative model. In practice, we have the option to parameterize a single, weight- shared g(a, = (gn(a, ......... ........e gK (a, x²) as a function from numerical experiments, we proceed with the former for efficiency.
A (k-1) x 2 + R x'dd or to parameterize the g A 1) x Rd + Bd individually. our
Theorem provides relations between derivatives of the density p(a, .) with respect to a and with respect to that involve the conditional expectations Iu (a,.x)) By specifying curve a(t) that traverses the simplex, we can use this result to design generative models that transport directly from any one marginal to another, or that are influenced by multiple marginals throughout the generation process, as we now show.
Corollary 2 (Transport equations). Let {er} represent the standard basis vectors of IR K+1 and let a [0, 1] + A denote a differentiable curve satisfying a0)) = e; and a(1) = e; for any i = O...., K. Then the barycentric stochastic interpolant z(a(t)) ) has probability density p(t, z)
x) that satisfies the transport equation
a,plt, + v (((t,)))(tx))) x)) = 0,
vit=0,z) pc(x),
p(t= (d=1,2x) = p,(2), (13)
where we have defined the velocity field
K Σ
b(t, x)
áx ((()))
(14)
k=
4

---

Published as a conference paper at ICLR 2024
In addition, the probability flow associated with (13) given by
X₂ b(t,
(15)
satisfies Xt=1 Pj for any Xt=o ~ ps. and vice-versa.
Corollary 2 is proved in Appendix A, where we also show how to the information about the score in (9) to derive generative models based on stochastic dynamics. The transport equation (13) is simple consequence of the identity (, x) = Ex=o (x,t))a p(a(t), æ). in tandem with the equations in (7). It states that the barycentric interpolant framework leads to a generative model defined for any path on the simplex, which can be used to transport between any pair of marginal
densities.
Note that the two-marginal transports along fixed edges in (13) reduce to set of K independent processes, because each Is takes values independent from the rest of the simplex when restricted to an edge. In fact, for any edge between Pi and Pj with i.j o k. the marginal vector field go ga(a, x) Elar (r(a)) = x] = E[ra], because the conditioning event provides no additional information about Fą- In practice, however, expect imperfect learning and the implicit bias of neural networks to lead to models with -vr-maaginal transports that are influenced by all data used during training. We summarize few possibilities with the multimarginal generative modeling setup in Section 2.1.
2.1 OPTIMIZING THE PATHS TO LOWER THE TRANSPORT COST
Importantly, the results above highlight that the learning problem for the velocity b is decoupled from the choice of path at) on the simplex. We can consider the extent to which this flexibility can lower the cost of transporting a given Pi to another Pj when this transport is accomplished by solving (15) subject to (13). In particular, we are free to parameterize the 0% in the expression for the velocity b(t, x) given in (14) so long as the simplex constraints on a are satisfied. We use this to state the following corollary.
Corollary 3. The solution
C(à) min à
E
Σ
(x(0))(((0)
Pat
(16)
gives the transport with least path length in Wasserstein-2 metric over the class of velocities b(t. x) Ek âu (t)gx (á(t), z). Here gk is given by (8), z(a) = _k= oyTK,, the expectation is taken over (Ko, Z1 .......... xx) )x== ~ Po p, and the minimization over all paths E C. ([0, 1)) such that à(t) E A K for all t [0,
Note the objective in (16) can be estimated efficiently without having to simulate the probability flow given by (15) by sampling the interpolant. Note also that this gives a way to reduce the transport directly without having to solve the max-min problem presented (Albergo & Vanden- Eijnden, 2023). We stress however that this class of velocities is not sufficient to achieve optimal transport, such guarantee requires the use of nonlinear interpolants (cf. Appendix D of Albergo & Vanden-Eijnden (2023)). An example learnable parameterization is given in Appendix C.I. In addition optimization over the path, the transport equation (13) enables generative pro- cedures that are unique to the multimarginal set- ting, such as transporting any given density to the barycenter of the K + 1 densities, or trans- port through the interior of the simplex.
Algorithm 1: Learning each 9%
Input: model 0k. coupling p(z TRR) gradient steps Ng, loss function La as defined in (10)
for j j Ng do
draw (Nu. IK) po(zo))((1, ,KK)
3 NUMERICAL EXPERIMENTS
draw a ~
make s0(a) Σt CAATR
Take gradient step with respect to L end
In this section, numerically explore the char- acteristics, behavior, and applications of the mul- timarginal generative model introduced above.
Return: Ôk-
5

---

Published as a conference paper at ICLR 2024
Example
Example Application
Simplex representation
2-marginal
standard generative modeling
P1 - P2
3-marginal
two-marginal with smoothing
P1 P2
K-marginal all-to-all image translation/style transfer
P: Pj
Table I: Characterizing various generative modeling aims with respect to how their transport appears on the simplex. We highlight a sampling of various other tasks that can emerge from generic multimarginal generative modeling. Here, two-margnn owo-marginal with smoothing corresponds to the addition of Gaussian marginal, as done (Albergo et al., 2023), with a path through the simplex.
Lourned patb
Pall Insete
Seandard Ineer path
820
602
L00
Tritiing step
Figure Direct optimization of (t)) over parametric class to reduce transport cost in the 2-marginal learning problem of a Gaussian to the checkerboard density. Left: The initial and final âg, å₃ learned over 300 optimizations steps on (16). Center: The reduction in the path length over training. Right: Time slices of the probability density p(t,, a)) corresponding to learned interpolant with learned à as compared to the linear interpolant a = t.
We start with the numerical demonstration of the path minimization in the restricted class defined in Corollary 3. Following that, we explore various pathways of transport the simplex illustrate various characteristics and applications of the method. We explore what happens empirically with the respect transport on the edges of the simplex and from barycenter. We also consider empirically if probability flows following (15) with different paths specified by a(t) but the same endpoint ((1) on give meaningfully similar but varying samples. Algorithms for learning the vector fields gk and learning least-cost parametric path (if desired) on the simplex are given in Algorithm 1 and Algorithm respectively.
3.1 OPTIMIZING THE SIMPLEX PATH a(t)
Algorithm 2: Learning a path (())
As shown in Section 2, the multimarginal for- mulation reveals that transport between distri- butions can be defined in path-independent manner. Moreover, Corollary highlights that the path can be determined a-posteriori via an optimization procedure. Here we study to what extent this optimization can lower the transport cost, and in addition, we consider what effect has on the resulting path at the qualitative level of the samples. We note that while we focus here on the stochastic interpolant framework consid- ered in (5), score-based diffusion (Song et al.,
Input: model ğk. model path a(t), coupling
p( gradient steps Naups. loss function C(ā) as defined in (16) .g do po(2o))(Z1, ... ... ,²x)
for j ,,
draw (......-rp draw t ~ Unif compute 6(t)
make x<a(t)) Σn an (t)ara
Take gradient step with respect to C(â)
end Return: &(t).
2021c) fits within the original interpolant formalism presented in (1). This means that the frame-

---

Published as a conference paper at ICLR 2024
work that consider can also be used optimize the diffusion coefficient g(t) for diffusion models, providing an algorithmic procedure for the experiments proposed by Karras et al. (2022).
A numerical realization of the effect of path optimization is provided in Fig. 1. We train interpolant on the two-marginal problem of mapping a standard Gaussian to the two-dimensional checkerboard density seen in the right half of the figure. We parameterize the og and 01 in this case with Fourier expansion with m = 20 components normalized to the simplex; the complete formula is given in Appendix C.I. A visualization of the improvements from the learned a for fixed number of function evaluations is also given in the appendix. We note that this procedure yields oo and an that are qualitatively very similar those observed by Shaul al. (2023). who study how minimize kinetic cost of Gaussian paths to the checkerboard.
3.2 ALL-TO-ALL IMAGE TRANSLATION AND STYLE TRANSFER
natural application of the multimarginal framework is image-to-image translation. Given K image distributions, every image dataset is mapped every other in way that can be exploited, for example, for style transfer. the distributions have some distinct characteristics, the aim is to find a map that connects samples z; and zj from P. and pj in a way that visually transforms aspects of one into another. We explore this in two cases. The first is a class to class example on the MNIST dataset, where every digit class is mapped every other digit class. The second is a case using multiple common datasets marginals: the AFHQ-2 animal faces dataset (Choi et al., 2020), the Oxford flowers dataset (Nilsback Zisserman, 2006), and the CelebA dataset (Zhang al., 2019). 1). All three are set to resolution 64 64. For each case, we provide pictorial representations that indicate which marginals map where using either simplex diagrams or Petrie polygons representing the edges the higher-order simplices. For all image experiments, use the U-Net architecture made popular in (Ho et al., 2020).
5
5 5
5
4
3 3
4
4
3
3
1
2
2
3
L
2 2
2
Figure 2: Left: Generated MNIST digits from the same Gaussian sample xp ~ Po. with marginals (Pu and 6 digit classes). To is visualized in the center of the image collection time t = 0, and the perimeter corresponds to transport to the edge of the simplex at time vertices color-coded. A Petrie polygon representing the 6-simplex, with arrows denoting transport from the Gaussians along edges to the color-coded marginals, clarifies the marginal endpoints. Right: Demonstrating the impact of learning with the larger simplex. Top row: learning just the simplex edge from 0 to 3. Middle: Learning on all the simplex edges from 0 through 5. Bottom: Learning on the entire simplex constructed from 0 through and not just the edges.
with
Qualitative ramifications of multimarginal transport We use mappings between MNIST digits as simple testbed to begin a study of multimarginal mappings. For visual simplicity, we consider the case of set of marginals containing the digits 0 through and a Gaussian. The inclusion of the Gaussian marginal facilitates resampling generate from any given marginal, in addition to the ability to map from any P to any other Pj.
In the left of Fig. 2, we show the generation space across the simplex for samples pushed forward from the same Gaussian sample towards all other vertices of the simplex. Directions from the center to the colored circles indicate transport towards of the vertices corresponding to the different digits. We note that this mapping smooth amongst all the generated endpoint digits, allowing for a natural translation toward every marginal.

---

Published as a conference paper at ICLR 2024
5
4
3
Figure 3: Left: An illustration of how different transport paths the simplex reach final samples that have similar content. Top row: cat sample is transformed into celebrity. Middle row: The same cat is pushed the flower marginal. Bottom row: The new marginal flower sample is then pushed to a celebrity that maintains meaningful semantic structure from the celebrity generated from along the other path the simplex.
On the right of this figure, we illustrate how training on the simplex influences the learned gener- ative model. The top row corresponds to learning just in the two-marginal case. The middle row corresponds learning only on the edges of the simplex, and is equivalent learning multiple two-marginal models in one shared representation. The third corresponds to learning on the entire simplex. We observe here (and in additional examples provided Appendix C.2) that learning on the whole simplex empirically results in image translation that better preserves the characteristics of the original image. This can best be seen in the bottom row, where the three naturally emerges from the shape of the zero.
In addition to translating along the marginals, one can try to forge correspondence between the marginals by pushing samples through the barycenter of the simplex defined by a(t) ( (R4T..... R ). This a unique feature of the multimarginal generative model, and we demonstrate its use the right of Fig. to construct other digits with correspondence to a given sample
We next turn to the higher-dimensional image generation tasks depicted in Fig. 4, where train 4-marginal model (K 3) to generate and translate among the celebrity, flowers, and animal face datasets mentioned above. In the left of this figure present results generated marginally from the Gaussian vertex for model trained on the entire simplex. Depictions the transport along edges of the simplex are shown the right. We show in the right half of the figure that image-to-ln translation between the marginals has the capability to pick up coherent form, texture, and color from sample from sample from marginal and bring it to another (e.g. the wolf maps to flower with features of the wolf). Additional examples are available in Appendix C.2.
As final exploration, consider the effect of taking different paths the simplex the same vertex. As mentioned the related works and Theorem 4, a coupling of the form (31) would
Figure Left: Marginally sampling a K multimarginal model comprised of the AFHQ, flowers, and CelebA datasets at resolution 64 x 64. Shown to the right of the images is the corresponding path taken on the simplex, with (0) eo starting at the Gaussian Po and ending one of P1: P2 P3. Right: Demonstration of style transfer that emerges naturally when learning multimarginal interpolant. With a single shared interpolant, we flow from the AFHQ vertex P2 to the flowers vertex Pi to the CelebA vertex P3- The learned flow connects images with stylistic similarities.
8

---

Published as a conference paper at ICLR 2024
lead to sample generation independent of the interior of the path a(t) and dependent only its endpoints. Here. we do not enforce such a constraint, and instead examine the aath-dependence of an unconstrained multimarginal model. In the left of Fig. 3, we take a sample of a cat from the animals dataset and push to generated celebrity. In the middle row, we take the same cat image and push it toward the flower marginal. We then take the generated flower and push it towards the celebrity marginal. We note that there are similarities in the semantic structure of the two generated images, despite the different paths taken along the simplex. Nevertheless, the loop does not close, and there are differences in the generated samples, such as in their hair. This suggests that different paths through the space of measures defined on the simplex have the potential provide meaningful variations in output samples using the same learned marginal vector fields yk.
4 OUTLOOK AND CONCLUSION
In this work, we introduced a generic framework for multimarginal generative modeling. To do so, used the formalism of stochastic interpolants, which enabled to define a stochastic process on the simplex that interpolates between any pair of K densities at the level of the samples. We showed that this formulation decouples the problem of learning a velocity field that accomplishes given dynamical transport from the problem of desigming path between two densities, which leads to a simple minimization problem for the path with lowest transport cost over specific class of velocity fields. We considered this minimization problem in the two-marginal case numerically, and showed that can lower the transport cost in practice. It would be interesting to apply the method to score-based diffusion, where significant effort has gone into heuristic tuning of the signal to noise ratio (Karras et al., 2022), and see if the optimization can recover or improve upon recommended schedules. In addition, explored how training multimarginal generative models impacts the generated samples in comparison to two-marginal couplings, and found that the incorporation of data from multiple densities leads to novel flows that demonstrate emergent properties such as style transfer. Future work will consider the application of this method to problems measurement decorruption from multiple sources.
REFERENCES
Martial Agueh and Guillaume Carlier. Barycenters in the Wasserstein Space. SIAM Journal Mathematical Analysis, 43(2):904-92) -244, January 2011. ISSN 0036-1410. doi: 10.1137/100805741. URL https: : /ppubs.....//////// sia. org/ /doi/////////// /ooi/1100/0 1...1111110 1111111111000 111.11101110005441111 Industrial and Applied Mathematics.
Publisher: Society for
Michael S. Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants, 2023.
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnnenn Stochastic interpolants: A unifying framework for flows and diffusions, 2023.
Jason M Altschuler and Enric Boix-Adsera. Wasserstein Barycenters can be Computed in Polynomial Time in Fixed Dimension.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. Stargan v2: Diverse image synthesis for multiple domains, 2020.
Marco Cuturi and Arnaud Doucet. Fast Computation of Wasserstein Barycenters, June 2014. URL http:/ / /aaxiiiiiiv///aii//aaaa//aaiavvv orgg/bss/////ooooooo/...ooa abs/ 11100. 4375. arXiv:1310.4377 [stat].
Prafulla Dhariwal and Alex Nichol. Diffussoo models beat gans on image synthesis, 2021.
Laurent Dinh, Jascha Sohk-Dick,,ein, and Samy Bengio. Density Estimation Using Real NVP. In International Conference Learning Representations, pp. 32, 2017.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buc. E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL https: http://///ceeeeeeeeeeeeee/o/ee/e nwwrips.unccmmmm /ttpuppppcccppppe///./////////// 7ac71d433 72ce7728352222330033777222-88882

---

Published as a conference paper at ICLR 2024
Gero Friesecke and Daniela Vögler. Breaking the curse of dimension in multi marginal kantorovich optimal transport on finite state spaces, SIAM Journal on Mathematical Analysis, 50(4):3996 4019,, 2018.
Wilfrid Gangbo and Andrzej Swiech. Optimal maps for the multidimennsional monge- monge-kantorovich problem. Communications Pure and Applied Mathematics, 51(1):23 45, 2023/09/29 1998.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and Н. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp- 6840-6851. Curran Asso- ciates, Inc., 2020. URL ptatttss:/prroceeeee ngs neerriiipp (GBDcteess00F0000111111101111111 http://dwww...r.......e.ee/////p httpsscccppppr//////11.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Weinberger. Deep Networks with Stochastic Depth. 1/Xiv:1603.0938) [cs].JJly 2016.
Tero Karras, Miika Aittlla, Timo Aila, and Samuli Laine. Elucidating the Design Space of Diffusion- Based Generative Models, October 2022. arXiv:2206.0036/ [cs, stat].
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. On density estimation with diffusion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: http:////ppeeeeeeeeee/eee/eoeeee net / forum? id=2 id=2LdBqxc1 Yv.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, 2022. URL https: /arxivv org/ abs/ 2210 02747.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport, 2022. URL https: : a/aaaaxvvv..g...... org/
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow, 2022. URL https: /ttppv..//axxvv.oorro org/ abs 22009 03003.
Maria-Flena Nilsback and Andrew Zisserman. visual vocabulary for flower classification. IEEE Conference on Computer Vision and Pattern Recognition, volume 2, pp. 1447-1454, 2006.
Maxence Noble, Valentin De Bortoli, Arnaud Doucet, and Alain Durmus. Tree-Based Diffusion "odinger Bridge with Applications to Wasserstein Barycenters, May 2023. URL http: //arxiv. wttp:///ggabaaaaaaeeeaa 11.300..65555 16577. arXiv:2305.16557 [cs, math, stat].
Brendan Pass. Multi-marginal optimal transport and multi-agent matching problems: Uniqueness and structure of solutions. Discrete and Continuous Dynamical Systems, 34(4): 1623-1699 EP-, 2014a.
Brendan Pass. Multi-marginal optimal transport: theory and applications. Technical Report 0026.
arXiv:1406.0003 arXiv, September 2014b. URL arXiv:1406.0026 [math] type: article.
Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernot, Wasserstein Barycenter and Its Application to Texture Mixing. Alfred M. Bruckstein, Bart M. ter Haar Romeny, Alexander M. Bronstein, and Michael M. Bronstein (eds.), Scale Space and Variational Methods in Computer Vision, Lecture Notes in Computer Science, pp. 435-446, Berlin, Heidelberg, 2012. Springer. ISBN 978-3-642-24788-9, doi: 10. 00.1007/978-3--42-277785
37.
Danilo Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows. In International Conference Machine Learning, pp. 1530-1538. PMLR, June 2015.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models, 2022.
Neta Shaul, Ricky T. Q. Chen, Maximilian Nickel, Matt Le, and Yaron Lipman. On kinetic optimal probability paths for generative models, 2023.
10

---

Published as a conference paper at ICLR 2024
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distri- bution. H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL https: h www.eeddiinnnnnnnnneeeeeeeeeeeee 0tw.077755511119996444777777
neurips. cttppppsppppppp/111111 9//19//// 9/file/ pdf.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan (eds.), Advances Neural Information Processing Systems, volume 34, pp- 1415-1428. Curran Associates, Inc., 2021a. URL httttt:://ppooeeeeeeeeee neurnnn...euiiin neurips.............. cc/ Pap.er..-appppp.p..eeee.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben
Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b. URL htthst:s: /op/reppeeviveeeeeeeeeeeeeee nww. formmf
id-PxTIG12RRHS.
Yang Song, Jascha Sohl-Dicke Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021c.
E. G. Tabak and Cristina V. Turner. family of nonparametric density estimation algorithms. Communications Pure and Applied Mathematics, 66(2):145-164, 2013. doi: https://doi.org/ 10.1002/4pa.21425, URL https:////// /nn inelibra ry. wi ley. com/ www.om/oo..........ooooooooo.ooo doi 20i/abb/10..007 cpa. 21423.
Esteban G. Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the log likeehooodd Communications in Mathematical Sciences, 8(1):217-233, 2010. ISSN 15396746, 19450796. doi: 1.44100/CMS.2000.8...o
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. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023.
Honglun Zhang, Wenqing Chen, Jidong Tian, Yongkun Wang, and Yaohui Jin. Show, attend and translate: Unpaired multi-domain image-to-image translation with visual attention, 2019.

---

Published as a conference paper at ICLR 2024
A OMITTED PROOFS
Theorem (Continuity equations). For all a E A the probability distribution of the barycentric stochastic interpolant x(a)) has density p(a, which satisfies the K equations
Pa ((a)() + V.. (1)(a,,))paaa = 0, Above, each gl (o, x) is defined as the conditional expectation
k ....., K.
9k. g x) B[xx/<(a) x],,
k = 0,...... K,
(8)
where E[xx|=(()) = z] denotes an expectation over Pu(zz) p(x1,.............. .x) conditioned on the event x(a)) = The score along each two-marginal path connected to pu is given by
Yao # 0 v log p(o,, x) = - g0
(9)
Moreover, each gk is the unique minimizer of the objective
(āv) = S&a F[(á((a, 2n Ou (ax, x(a)) (a) == K, where the expectation I taken over (Io, (ro...., zk) (0 ((1,,....))
(10)
k= ak its characteristic function is Proof. By definition of the barycentric interpolant x(aa) = Σk given by
Elecholaa)] = /
JRAX 2h 2k p(01.................._.._.. 2K dxx dæke d
(17)
where we used To 1 (II...., and Tp ~ N(O, Ida). The smoothness in k of this expression guarantees that the distribution of sr(a)) has density p(a, x) >> O. Using the definition of s(a)) again, p(a,,) x) satisfies, for any suitable test function Ø R2d - IR:
[ ((x)p(t.z(()dr f (x(a)) p pla xx) e/2 klxol dxo dx K.
(18)
Taking the derivative with respect to ag of both sides we get
1 plz)aaz(a, x)dz
X XRd Tb Vo( ((I,,,,) (2T) -2/2 d/26 6/26- -lmol dxo dow
(19)
B[Tk- Vo(z(a)) = 2] p(a, (la.x)dd x)dx
[[cs/<(a)= = x| (()
where used the chain rule to get the first equality, the definition of the conditional expectation to get the second, and the fact that ((x(a:) = o(x) since we condition on x(a) = x to get the third. Since
E[x (=(a)) = x] = gk (ax x))
(20)
by the definition of gk in (8), we can write (19) as
[ o(z)0a x)dx = / gx(a,z) Vo(z)p(a, x)dx
(21)
This equation is (7) written in weak form.
To establish (9), note that, if oxp > we have
E Toe took-oo -0- ao = =-a 10glN12 iagke 1
(22)
As result, using To 1 (I1,.... IK), have
E Jo 2())]
iig (xxx((a)]
(23)
12

---

Published as a conference paper at ICLR 2024
Using the properties of the conditional expectation, the left-hand side of this equation be written as
B f0
E ( ( [x(()) /2((x) a1) pla,xx)d
Rd
E|zo x(a)= = zeeth pl d)
(24)
go(a, a) r) p(a, x)
where we used the definition of 90 in (8) to get the last equality. Since the right- hand side of (23) is the Fourier transform of -arop(a, 2.), we deduce that
go(a, x)p(ax)))
do Vp(a, oo log, 2(a,
(25)
Since p(a, x² > O, this implies (9) when 0. Finally, to derive (10), notice that we can write
(#x) / E [1 x(a))) - 2σk (a, x(a)) da,
SS
JRd [lecaa x(a)))² 2n ((())(( x] p(a, x dddada
(26)
/A ((a((a,((a))) 2B/r//<(a)) (aaeaaa)llaaa.... [())) p((a,x))drd x))xda
{ x(a)) =(a)) (t. x. E) (a x(a)) p(a, x z)dzda
where we used the definition of gb (8). The unique minimizer of this objective function is (n(a, x) 2))
Corollary 2 (Transport equations). Let {ex} represent the standard basis vectors of TR K+1 and let [0, + A denote a differentiable curve satisfying a(0) e; and a(1) ej for any i = O...., K. Then the barycentric stochastic interpolant z(a(t)) has probability density p(t, x)
x)) that satisfies the transport equation
O,p(t, + V x)p(t,)) = 0,
= 0(z= pc(x),
p(t = 1, = pz(x), (13)
where have defined the velocity field
K
b(t, x)
Σ àu (a()))
(14)
Amm
In addition, the probability flow associated with (13) given by
X, = b(t, X,)
(15)
satisfies X ~ Pj for any X ~ Pi. and vice-versa.
Proof. By definition p(t = 0, x) = 0 p(e,,))) (z(z) and p(t = 1,x) = x()(1) x(c,. = P5 (z)), the boundary conditions in (13) satisfied. To derive the transport equation in (13) use the chain rule as well as (7) to deduce
0,p(t,x) Σ (t))...p(aaz))
Σ an(t)V (o=
(27)
This gives (13) by definition b(t, x) in (14). (15) is the characteristic equation associated with (13) which implies the statement about the solution of this ODE.
[7
13

---

Published as a conference paper at ICLR 2024
Note that, using the expression in (13) for V log p(a(t), x) = log p(t, x)) as well the identity Ap(t, z) = V. (Vp(t, p(t, x)) we can, for any <(t) 2 0, write (13) as the Pekker-Plancls equation (FPE)
0,p(t, + V ((b(t,z) (t)ga)(())o p(a, = c(t)Ap(t,x))
(28)
The SDE associated with this FPE reads
dxf (b(t, (t) (to)) Xf)) dt + 2e(t)dW,
(29)
As result of the property of the solution p(t, x) we deduce that the solutions to the SDE (29) are such that X ~ Pj if X ~ Pt, which results in generative model using diffusion.
Corollary 3. The solution to
K
(1(0ga(((0) F Σ k=0
C(a) = min
(16)
gives the transport with least path length in Wasserstein-2 metric over the class velocities b(t, æ) Ek âu (t) (á(t), z².. Here gk is given by (8), x(a) Σ 04.2R, the expectation is taken over
(Fo, T1... ..... ..... ) Po and the minimization is over all paths C* 1]) such that ā(t) E A for all t E 1].
Proof. The statement follows from the fact that the integral the right-hand side of (16) can be written as
K
= -f" =(((e))(((()))
(()))(())))))))))))
km
(30)
=fo / 16(t, (6(t,z))))))))))()))))) )² plt, x)dzet
where b(t, = Σxn ((())) x) and p(t, 2) solves the transport equation (13) with replaced by b. This expression give the length in Wasserstein-2 metric of the path followed by this density, implying the statement of the corollary.
B DETERMINISTIC COUPLINGS AND MAP DISTILLATION
We now describe an illustrative coupling po(xo))(z1,
xx) for the multimarginal setting in which the probability flow between marginals is given in one step. Let Tp IRd + Rd for = 1...., K be invertible maps such that 2ą T(ro) pk: if zo Po (i.e., each Pk is the pushforward of po by Tx). Assume that the coupling P is of Monge type, i.c.
pl
HK)
II 8( Ta (x(o) k=1
(31)
Clearly, this distribution satisfies (6). For simplicity of notation let us denote by To(x) the identity map. Let us also assume that, for all a E A the map Ek TxTTr(r) invertible, with inverse (()
K Ma. k=0
k = Σ (R(a K=U
Ya E Y IER
R
(32)
In this case, the factors Ok (a, æ) evaluated at eg simply recover the maps Tg. Evaluated e,, they transport samples from to samples from Pk. as shown in the following results:
Theorem 4. Assume that p(æ1...., FK) ) is given by (31) and that (32) holds. Then, under the same conditions as in Corollary 2, the factors gı(a,, x) defined in (8) are given by
gı(a, x)) = (R(xx x))
(33)
In particular, if a Css we have gu (a x)) (x(T, (xx) that if x; Pi then go (ei, vT) Pk-
14

---

Published as a conference paper at ICLR 2024
Proof. If p(z1,..., 2K) is given by (31), have
g x) = Eo (xo)
(xo) = = =o (T(zo) 2o = R(a, x0] = (x((,,,))) (34)
where Egg denotes expectation over zo ~ Po- The relation 3)) (x)) follows from (33)
since R(e,, x) = (x) and the final statement from (31).
In addition, the solution the probability flow ODE (15) with initial data X,=0 -7((2)) is given by
K
X,
Σ 0 (t)) k=
(35)
so that X1=1 (x0) since a(1) Cj by assumption. To check (35), notice that b(t,
Ekzo åu (t)at()(( x)). As a result, using (32),
K
bl1, Σ (1)x (R(a((), (n((0T((R(a(0))) Σ
ow (ITT Σ (1)T)(( =x, (36)
k==0
so that (35) satisfies (15).
C EXPERIMENTAL DETAILS
C.I TRANSPORT REDUCTION EXPERIMENTS
Corollary tells us that, because the problem of learning to transport between any two densities pt Pj on the simplex can separated from a specific path on the simplex, some of the transport costs associated to the probability flow (15) can be reduced by choosing the optimal a(t). a(t) can be parameterized as a learnable function conditional on this parameterization maintaining the simplex constraint >, 0% 1.
In the 2D Gaussian to checkboard example provided the main text, we fulfilled this by taking an o; 0r of the form:
N
ã; = 1+
(
ai sin (n
(37)
x3
N
å = t +
J,n sin (t))
(38)
where dg,n, aj.n are learnable parameters. We also found the square on the sum, though necessary for enforcing a.jj 2 O, could be removed N is not too large for faster optimization. For the case when the simplex contains more than two marginals (k & i.j), the additional og. which do not specify an endpoint i, j of the transport can be parameterized drop the first terms which enforce the boundary conditions so that
N
ann sin(ntt)
(39)
T
(40)
The 0g Vk are normalized such that
6
āx
(41)
Σ
Σ
As such, their time derivatives, necessary for using the ODE (15) as generative model, is given as, for any component 4%
dg áy Emeek am) am) -ā- - 5m
(42)
(Ex āx)²
15

---

Published as a conference paper at ICLR 2024
Learned path
Stundard lincar path
Figure 5: The output of the probability flow (15) realized by for the learned interpolant for the checkerboard problem discussed in Section 2.1 and Appendix C.I. For fixed number of function evaluations, 5 steps of the midpoint integrator, the learned a is quicker to give a more accurate solution.
C.2 ARCHITECTURE FOR IMAGE EXPERIMENTS
We adapt the standard U-Net architecture (Ho al., 2020) to work with vector a describing the "time coordinate" rather than scalar t, which is the conventional coordinate label. The number of output channels from the network is given as image channels K, where each kth slice of # image channels corresponds to the kth marginal vector field necessary for doing the probability flow.
MNIST
JFHO-Flowees--elese 64 x 64 x 3
Dimension
32 x 32
Training point
50,000
15000, 8196, 200000
Batch Size
500
180
Training Steps
3x 105
x 105
Hidden dim
128
256
Attention Resolution
64
64
Learning Rate (LR)
0.0002
0.0002
LR decay (1k epochs)
0.995
0.995
U-Net dim mult
[1,2,2]
[1,2,3,4]
Learned sinusoidal embedding
Yes
Yes
lo2 l when sampling with ODE
[0.0, 1.0]
[0.0, 1.0]
EMA decay rate
0.9999
0.9999
EMA start iteration
10000
10000
GPUs
8
Table 2: Hyperparameters and architecture for image datasets.
C.3 ADDITIONAL IMAGE RESULTS
Here we provide some additional image results, applied to the image-to-image translation.
16

---

Published as a conference paper at ICLR 2024
Figure 6: Random assortment of additional image translation examples built from the simplex of CelebA, the Oxford flowers dataset, and the animal faces dataset all 64 64 resolution.
17

---

Published as a conference paper at ICLR 2024
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
8
8
3
3
0
0
3
3
3
7
3
3
3
3
0
6
3
3
3
3
3
8
3
3
3
3
3
3
3
3
3
3
3
3
26
2
3
3
3
3
Figure Random assortment of examples comparing the class-to-class translation MNIST digits
and 3, where each row in triplet corresponds to: top) model trained only on the edge between ( and middle) model trained on all edges between 0 and 5. bottom) model trained on the entire simplex.
18


Reviews:
[{"review_id": "J22ELomul6", "reviewer": "Reviewer_PidH", "summary": "In this work, the authors propose an extension of stochastic interpolants for multimarginal generative modeling, which involves learning a joint probability distribution that captures multiple marginal probability densities. The proposed framework is able to capture multi-way correspondences.", "questions": "Referring to the weaknesses part, how do you solve (16)? \n\nOther remarks:\n- Incomplete caption for Figure 2 (Right). \n- Page 9 first paragraph: “from a sample from a sample”\n- References: Song et al. (2021b, c) are repeated.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The work proposes a theoretically sound and practical approach based on stochastic interpolants for the multimarginal setting. The overall proposed scheme is mathematically sound and computationally more feasible than existing schemes. Extensive experiments are performed to illustrate the idea and different scenarios of the proposed method.", "weaknesses": "The overall algorithm of the proposed method might not be very clear or easy to follow, especially for mathematically less mature audience. I suggest the authors add the pseudocode of the main algorithm to improve clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YcEvppEiop", "reviewer": "Reviewer_2Wkn", "summary": "The paper introduces a vector-field-based generative model (incorporating ODE/SDE mechanisms) capable of generating samples from multiple marginal distributions. The methodology pivots around a generalized form of stochastic interpolant, specifically a barycentric stochastic interpolant. The stochastic process, in this context, is formulated as a weighted mix of samples from distinct marginal distributions. These weights, restricted to the simplex, allow the model to generate a weighted combination, defined as $x(\\alpha)=\\sum_{k=0}^K x_k \\alpha_k$, where $\\alpha=\\left(\\alpha_0, \\ldots, \\alpha_K\\right) \\in \\Delta^K$. The corresponding probability flows are derived for this multi-weight interpolant path. To identify the optimal path with the lowest transport cost, the model is trained to minimize the Wasserstein-2 metric, considering the path's velocities. Consequently, this learned model allows transitions between various marginal distributions following the optimized path. However, the semantic interpretation of these paths, especially between distributions like images from multiple modes, remains ambiguous. In the experiments,", "questions": "- 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.\n- 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 the what this result implies and when we would like to have this property.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- *Innovative Generative Model*: The paper's proposition of a vector-field-based generative model that utilizes a generalized form of stochastic interpolant to generate from multiple marginals is unique and intriguing.\n- *Optimal Path Identification*: By minimizing the Wasserstein-2 metric, the methodology seeks the most efficient path for transitions, with theoretical groundings.\n- *Barycentric Stochastic Interpolant*: The use of a barycentric stochastic interpolant offers a nuanced approach, ensuring the weights adhere to a simplex and providing a structured method to learn to interpolate samples from different distributions.", "weaknesses": "- *Ambiguous Interpretation*: The paper does not provide a clear semantic understanding of the paths between distributions, making it difficult to ascertain the practical implications or the broader relevance of the methodology. Also, what does multi-marginal optimal transport path mean for image generation? Does it lead to better quality? In this paper, it is not answered.\n- *Unknown Utility*: 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.\n- *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.\n- *Minor*: missing caption for Figure 2 (right).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xDGvP805e9", "reviewer": "Reviewer_yfUs", "summary": "Multimarginal problems are typically studied in the context of optimal transport, by contrast, the authors studied the multimarginal problem using stochastic interpolant with a high-dimensional α in a simplex. The multimarginal framework allows us to (K, 2) marginal transport problems using only K marginal vector fields. Interesting experiments on all-to-all image translation are studied.", "questions": "Table.1, what does the smoothing mean in \"two-marginal with smoothing\" ?\n\n\nMinor:\n\ntypo? \"color from a sample from a sample from one marginal\"", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. by lifting alpha to high-dimensional vector, the work enables to solve (K,2) marginal transport problems using only K marginal vector fields.\n\n2. the overall presentation is clear and the empirical experiments are comprehensive.", "weaknesses": "Optimizing the vector field on simplex alpha may require delicate parameterizions. It would be better if it is more connected optimal transport.\n\nThe theories look like a na\\\"ive extension from the work of \"Stochastic Interpolants: A Unifying Framework for Flows and Diffusions\". The authors can kindly point out some innovations if I may miss some key points.\n\nSection 2.2 may need some proper rewriting to be more clear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)