Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows
Authors: Gabriele Visentin
Patrick Cheridito
Venue: ICML 2025 poster (2025)
Keywords: Optimal transport
Wasserstein barycenter
Normalizing Flows
Primary area: general_machine_learning->scalable_algorithms

Abstract: We present a novel method for efficiently computing optimal transport maps and Wasserstein barycenters in high-dimensional spaces. Our approach uses conditional normalizing flows to approximate the input distributions as invertible pushforward transformations from a common latent space. This makes it possible to directly solve the primal problem using gradient-based minimization of the transport cost, unlike previous methods that rely on dual formulations and complex adversarial optimization. We show how this approach can be extended to compute Wasserstein barycenters by solving a conditional variance minimization problem. A key advantage of our conditional architecture is that it enables the computation of barycenters for hundreds of input distributions, which was computationally infeasible with previous methods. Our numerical experiments illustrate that our approach yields accurate results across various high-dimensional tasks and compares favorably with previous state-of-the-art methods.

Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows
Gabriele Visentin Patrick Cheridito
Abstract
applications in problems of shape interpolation (Solomon et al., 2015), image interpolation (Lacombe et al., 2023; Simon & Aberdam, 2020), color translation (Rabin et al., 2014), style translation (Mroueh, 2019), Bayesian subset posterior estimation (Srivastava et al., 2015), clustering in Wasserstein space (Del Barrio et al., 2019; Ho et al., 2017), and fairness (Chzhen et al., 2020; Gouic et al., 2020).
We present a novel method for efficiently com- puting optimal transport maps and Wasserstein barycenters in high-dimensional spaces. Our ap- proach uses conditional normalizing flows to ap- proximate the input distributions as invertible pushforward transformations from a common la- tent space. This makes possible to directly solve the primal problem using gradient-based minimization of the transport cost, unlike previ- ous methods that rely on dual formulations and complex adversarial optimization. We show how this approach can be extended to compute Wasser- stein barycenters by solving conditional vari- ance minimization problem. A key advantage of our conditional architecture is that it enables the computation of barycenters for hundreds of input distributions, which was computationally infeasi- ble with previous methods. Our numerical experi- ments illustrate that our approach yields accurate results across various high- dine-siinnii tasks and compares favorably with previous state-of-the-art methods.
Traditional numerical methods work with discrete or dis- cretized probability measures (Peyré & Cuturi, 2019), via linear programming (Anderes et al., 2016) and entropy reg- ularization (Cuturi & Doucet, 2014; Solomon et al., 2015). They yield accurate results for discrete measures on low- dimensional spaces but scale poorly in the number of sup- port points, so that their performance degrades for contin- uous distributions, especially in nigh-dimensional settings (Fan al., 2020; Korotin et al., 2021b).
These limitations have generated substantial interest in nu- merical methods for OT and Wasserstein barycenters prob- lems for continuous distributions. While other methods are reviewed Section 4, we note here that almost all of them solve the Kantorovich dual of the OT problem and employ neural networks to parametrize either the dual potentials (Li et al., 2020) or the transport maps (Kolesov et al., 2024), which often results in complex bi-level (Kolesov et al., 2024; Korotin et al., 2019; Xie et al., 2019; Lu et al., 2020) or tri- level (Fan et al., 2020) adversarial learning.
1. Introduction
Optimal transport (OT) and Wasserstein barycenters are active research areas with applications in economics, opera- tions research, statistics, physics, PDE theory, and machine learning (Peyré & Cuturi, 2019; Villani. 2021). OT has been successfully applied to generative modeling (Arjovsky et al., 2017; Petzka al., 2017: Wu et al., 2018: Liu et al., 2019; Cao et al., 2019 Leygonie al., 2019) and domain adaptation (Luo al., 2018; Shen et al., 2018; Xie al., 2019), while Wasserstein barycenters, by providing nat- ural notion of average for probability distributions, found
Contributions. We propose a different approach. Our main contributions are the following:
1.
We propose a method that bypasses complex adversar- ial training by directly solving the primal formulation of the OT problem. This approach directly yields both OT maps and values. In addition, it enables an intuitive reformulation of the Wasserstein barycenter problem expected conditional variance minimization prob- lem.
'epartment of Mathematics, ETH Zurich, Switzerland. Correspondence to: visentin <eabricle.visentin@e math.ethz.ch>.
In our setup, OT maps are bijective. Therefore, we approximate them with conditional normalizing flows, which flexible bijective models and therefore natural choice in this context.
Gabriele Visentin
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
For barycenter problems, our method directly yields

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
OT maps and their inverses between the input dis- tributions and the barycenter without requiring addi- tional computations. Moreover, it provides generative model of the barycenter which allows for direct sam- pling without querying the input distributions.
the case p 2, which is of most practical interest. Let (Hs)sEs S (2) be family of distributions on Rd indexed by finite set 8 and (w,, E S) non-negative
weights such that E.es W, = 1. Then the w-weighted Wasserstein-2 barycenter is given by:
4. Our model's conditional architecture scales well in the number of input distributions and can be used to compute Wasserstein barycenters of hundreds of input distributions, which is computationally infeasible with existing methods.
J = arg
arginl inl MEPa(R") SES
Σ w. (,,
(3)
Under these assumptions, the Wasserstein-2 barycenter jú exists, is unique and is absolutely continuous with respect to the Lebesgue measure (Aguch & Carlier, 2011; Brizzi et al., 2025).
2. Background and Preliminaries
In this section we recall the OT problem (Section 2.1) and the Wasserstein barycenter problem (Section 2.2).
3. Proposed Method
Notation. We denote by || || the Euclidean norm on TRd and by (P) the space LP B((R)), equipped with P(RR)) denotes the space of all Borel probability measures on Rd, (P(R))) is the space of all μ E P(Rd) such that Sox [|=]))/(dd)) < oo, while (RRd) is the space of all (PRR)) that are absolutely continuous with respect to the Lebesgue measure. Finally, we define P (R.) ==
(Section 3.1). Then we derive our main theoretical results (Section 3.2) and present their algorithmic implementation In this section, we first introduce some additional notation
the norm [[f]l := =S{. [[(()))p((dz))
(Section 3.3). The proofs of all lemmas and theorems can be found in Appendix B.
3.1. Notation
Given two Borel probability measures μ4, 1/ E P(R") and Borel function f Rd + Rd, , we say that a Borel function f: + Rd is a (n., v-)inverse of f if f f(x) x for a-almost all x E Rd and f(x)) = x for v-almost all x E Rd
P, n (IR).
2.1. Optimal Transport
The Monge-formulation (Monge, 1781) of the OT problem between two distributions fu, V E P(RR) for p-cost
We denote by B(,,v) the set of all Borel functions f Rd - Rd such that f# and admits (1,, vi-inverse. In Lemmas A.3 and A.4 below we show some basic facts about this set. In particular, we show that if μ,, v E Pp (IR)) and f E B(p, v), then [[ belongs to (*),), the (H, v)- inverse f v-almost surely unique and it belongs to B(v, jl). A special case when fs and have non-vanin non-vanishing densities with respect to the Lebesgue measure on Rd, then B v) is just the set of all almost everywhere bijective pushforward maps that push μ Furthermore, if (p,, v1), then its (J,, -niverse is almost everywhere unique and coincides almost everywhere with the inverse map f-1, whenever the latter exists.
Wp(u,v) inf TAKEE
|| - (()))))))(ddz))
Rd
where T i μu denotes the pushforward measure of a under T. This problem is difficult to handle and does not always have a solution. Therefore, it was relaxed by Kantorovich (1942) as follows:
WH(,))= inf / |/ | 11--yl17(60 (lx, dy), TEI((e.)
where I'(H,)) is the set of all couplings between μ and v, Le. the set of all probability measures in P(Rd Rd) with marginals y and v.
3.2. Thcoretical Results
Bijective pushforward maps arise naturally in OT theory, as shown in the following version of Brenier's 's theorem, which follows from Gangbo McCann (1996).
However, under some mild regularity assumptions, prob- lems (1) and (2) are equivalent (see Lemma 3.1 below), and ", (p, v) defines metric on the space (P((RR)) known as the Wasserstein-p metric.
Lemma 3.1. Let p > 1 and v E P (R..c(Rd)) Then, Kan- torovich's 's problem (2) has a unique solution 7y E I v). Moreover, Y is of the form (id, T) Ty yt where T B v) a p-almost surely unique solution to Monge's problem (1).
2.2. Wasserstein Baryeenters
Wasserstein barycenters are Fréchet means on the Wasser- stein spaces (P,(R.)), Wp). In the following we focus on
The next theorem allows us to reformulate the OT problem
2

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
with p-cost as a constrained LP (A))mmiininiattonn problem for given latent distribution À P Pp (Re))
has a solution, and
2.
For every solution J : I x S + Id of problem (4), the w-weighted Wasserstein-2 barycenter of (Ha))eS is equal to the distribution of h(()), where h RR - Rd is given by the weighted sum
Theorem 3.2. Let X, M, E Pp.ac(R.)) for some p 1. Then:
W WE(H) v) min 1 (l))
)EB().H) g)B(A.))
Furthermore, if and are solutions of the optimization problem above and f is a (A. iiiiveee of f. then g o f is an optimal transport map between y and v.
h(2)
Σ W. (z) SES
3.3. Implementation
Theorem 3.2 can be understood as a generalization of the closed-form solution of OT in one dimension using quantile functions. Recall, that for μr, 1 E (P()), one has:
Our approach uses conditional normalizing flows, After introducing them in Section 3.3.1, we present our algorithms for OT maps (Section 3.3.2) and Wasserstein barycenters (Section 3.3.3). In Appendix C we provide all necessary background information on normalizing flows.
2) v) = 11qu - qu qu112=(0011)
where ([0, 1J) is the L space over [0, 1] equipped with the uniform distribution and qu is the quantile function of ja (see, for instance, Remark 2.30 in Peyré & Cuturi (2019)). Theorem 3.2 shows that for μ4, V E (B we can replace the uniform distribution with any absolutely continuous latent distribution À and the comonotonicity constraint on qu and qu (which is implicit in the definition of quantile function) with closeness in (A)).
3.3.1. CONDITIONAL NORMALIZING FLOWS.
Given a finite set of input distributions (H3)sES E Prac and a reference distribution À E Pp.co we use conditional normalizing flows to parametrize a map ʃ R X - Rad such that ((; E B(,, v.)) for all E S, and we use this parametrization to solve the optimization problems in Theorem 3.2 and Theorem 3.3 with gradient descent in parameter space. Teshima et al. (2020) have shown that coupling- normalizing flows are universal approximators for bijections with the caveat that their con- struction relies on ill-conditioned networks (Koehler et al., 2021: Draxler al., 2024). On the other hand, Draxler et al. (2024) have introduced well-conditioned flow which can approximate arbitrary distributions but not necessarily every bijection.
Our next result shows how to compute Wasserstein-2 barycenters as weighted averages of certain bijective push- forward maps that minimize an expected conditional vari- ance criterion. Consider non-empty finite set S (1....,1)) n and a collection of probability distributions μts E (R2)) E S, together with weights w, 2 0, s E S, such that Exes W. = 1. Then the w- w-iehttted Wasserstein- 2 barycenter of (ua)ses is the unique j E P2..c (TR") sat- isfying Equation (3). In the following theorem, we realize ū as the distribution of a d-dimensional random vector of the form I(Z) for Borel function h : Rdd -> Rd and a d-dimensional random vector Z defined on an underlying probability space (22, F, IP) that also carries an independent S-valued random variable S such that P(S s) wa for all s E S. The expectation and variance corresponding to IP are denoted by E and Var, respectively.
We choose the standard d-dimensional Gaussian distribu- tion as latent distribution À and we enforce the pushforward constraint f(
+) A = f via conditional likelihood maxi- mization. The conditioning variable E S may be R-valued, taking finitely many values (as in Section 5.2.4, where S is finite grid in [0. 1)) one-hot encoded (as in all other nu- merical experiments), depending on how regular we expect the dependence of μ., on s to be.
Theorem 3.3. Let S be a non-empty finite set and consider a collection of probability distributions μa E P2. s E S, and weights ws 2 0, 8 E S, satisfying Exes Ws = Let Z be a d-dimensional random vector with an arbitrary distribution X E a2.ac and S an independent S-valued random variable taking the values with probabilities Was.
We design and implement our own conditional versions of Real NVP (Dinh et al., 2016) and Glow (Kingma & Dhariwal, 2018) with multi-scale architecture. In the case of Real NVP, we use a conditional affine coupling layer with the following coupling function:
E S. Then, the following hold:
Se(=9,0) (=A) = A @ exp(a(< 8)) + b(= s),
1. The minimization problem
where the parameters ((:,)) = (a(.; .), ((,))))) are feedfor- ward neural networks, with an additional residual connec- tion for the conditioning input s. A similar conditional ar- chitecture, but without the residual connection, was already proposed by Atanov et al. (2019), but in our experiments,
d
min x
E (Var (S((2,S) 2)]
(4)
f(.,) pz)
1=1
for all E
3

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
found that adding this residual connection improves the performance. In the case of Glow, we obtain a conditional architecture simply by concatenating the conditioning vari- able pixel-wise with the latent input image.
values of the optimal likelihood loss and 12-cost. In all our numerical experiments we chose exponentially decreas- ing schedule; see Appendix D for full description of the hyperparameters used in each numerical experiment.
At the end of training, the OT map from f4r to μız can be efficiently retrieved by performing an inverse pass from fi to À followed by a forward pass from to μ2, i... is given by x + felfo (z, 1).
3.3.2. COMPUTING OT MAPS
Our method for computing OT maps using Theorem 3.2 is presented in Algorithm The algorithm is given for the case of quadratic cost (i.e. p- = 2). but it can also be used for OT costs of the form c(zz, y) h(x y) for strictly convex function h: _ + [0, oo) satisfying the conditions (H1)-(H3) of Gangbo & McCann (1996).
3.3.3. COMPUTING WASSERSTEIN- 2 BARYCENTERS.
The algorithm for computing Wasserstein-2 barycenters builds upon Theorem 3.3 and is presented in Algorithm 2.
As shown in Algorithm 1, we learn the source distribution μ = μi and the target distribution V = #2 using a joint conditional normalizing flow with conditioning variable (1, In principle, it would be possible use two separate normalizing flows for μ4 and U2. But in our exper- iments, conditional model resulted in better performance, probably due to the fact that this architecture introduces smooth dependence on the conditioning variable & E {0, 1}. which facilitates learning pushforward maps that are close in L (A)..
Algorithm Wasserstein-2 Barycenter via Conditional Nor- malizing Flows.
Input: input distributions (Ha)sES accessible by sam- pling; weights (wa)ses and weighting measure (ss) == ws; conditional normalizing flow fa(::-) R2d x X - Rd with initialized parameters @ and latent distribution A; number of iterations T; learning rate yi decreasing weights (Gi)
Output: learned Wasserstein-2 barycenter aad x p. where h(z) = Eees (z,
Algorithm 1 OT Map via Conditional Normalizing Flows. Input: input distributions μ1 and μ22 accessible by sam- conditional normalizing flow fo(:-) : Rd x (1, + Ed with initialized parameters 0 and latent distribution X; number of iterations T; learning rate m decreasing weights (G))=1-
for / = 1 to T do
sample S ~ a,, X ~ us. Z ~ X; compute model likelihood on (X, S):
pling:
pe (X, S) - ( (X, S)) |det (det(D. " (X, compute -cosstt:
L3( (Z, - Exes wafo (2 8) })?;
Output: OT map from P1 to P2 given by x H So (ar 2).
update model parameters @ by gradient descent: 0 + 0 - v/ (- log(po (X) + GL S)) end for
for t 1 to T do
sample S ~ Uniform((1, 2)), X ~ us. Z ~ X;
compute model likelihood on S):
Also in this case we employ conditional normalizing flow architecture and train using a sequence of decreas- ing weights. At the end of training, we obtain a genera- tive model of the barycenter given by Ay for the func- tion h(2) = Eves W, s). Additionally, the OT map from any input distribution He, to the barycenter is given by x h(So (x. s)), while the OT map from a sampled barycenter datapoint h(2)) to an input distribution Me is given by x + fo (xx), s) = fo(2.s).
5)(X, = Pz( (X)))) /det D D S)):: compute LL-cost:
ZZ = Ilfo (Z, 1)- - fz.
update model parameters 0 by gradient descent: 0 x- 0 - Vo (- log(po + (Z))
end for
Decreasing weights. Both Algorithm 1 and Algorithm 2 require as input sequence of decreasing weights ( for the cosst This is motivated by the fact that we are trying to solve a multi-objective optimization problem with two competing objectives: the pushforward maps need to be as close as possible in (A) while still being far enough to satisfy their pushforward constraints. Naively minimizing both the -.o-tost and the negative log-likelihood in each gra- dient step leads to poor performance. Instead, we gradually decrease the importance of the L c,2-cost, thus guaranteeing that at convergence the pushforward constraints are satis- fied. The optimal choice of weights depends on the relative
4. Related Work
Traditional numerical methods tackle OT between discrete or discretized distributions with linear programming (An- deres et al., 2016; Peyré & Cuturi, 2019). The Sinkhorn- Knopp algorithm (Sinkhorn & Knopp, 1967; Cuturi, 2013) yields good results for entropy-regularizee discrete OT prob- lems in low dimensions. But discrete methods loose accu- racy for continuous distributions and become infeasible in high-dimensional settings. This has lead to continuous
4

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
methods, most of which tackle dual formulation of the OT problem with neural networks kernel expansions; see e.g. Korotin et al. (2021a).
ODE model, for the optimal coupling, but they enforce the pushfoward constraint using the Kantorovich-Rubinste duality for the Wasserstein-1 distance, which results in minimax optimization problem. Both these models do not address the problem of computing Wasserstein barycenters.
Methods for Wasserstein barycenters can be divided into non-generative and generative. While non-generative mod- els are limited to transporting existing samples from the input distributions, generative models output a learned barycenter distribution which can be used for sampling.
A Advantages of generative models. Depending on the application of interest, generative models may offer sub- stantial benefits over non-generative ones. For instance, in shape interpolation, the density of the generative model is precisely the interpolating shape and does not require any surface reconstruction from point clouds. In style transfer the generative model can generate previously unseen sam- ples (e.g. new MNIST digits in a given style), as opposed to just transporting already existing ones. In fair regression the generative model effectively performs distributional re- gression for the fair premium, which allows, for instance, the construction of confidence intervals, the estimation of any statistic of the barycenter distribution, as well as further transformations of the distribution; e.g. the barycenter can be shifted to have mean zero, which is useful, for instance, in acturrill applications.
Non-generative models. An early example of such model has been given by Li et al. (2020), who exploit the dual formulation of the entropic regularized Wasserstein barycenter problem and paramatrize the potentials using neural networks. Their method produces only approximate barycenters, due to the bias introduced by the entropic regu- larization, and outputs only the learned potentials, instead of the OT maps. Korotin et al. (2021b) parametrize potentials with input-convex neural networks (ICNNs) (Amos et al., 2017) and tackle the dual OT problem by imposing a con- gruence and a cycle-consistency condition. The congruence condition requires fixing a dominating measure that may be difficult to choose a priori, while there is evidence that the sycle-consistency condition (Amos et al., 2017) and the use of ICNNs (Korotin et al., 2021a) may be suboptimal. Kolesov et al. (2024) present a model that can be applied to generic cost functionals (both regularized and exact), but relies on bi-level adversarial learning and does not provide a generative model of the barycenter.
Since our model is generative, in Section 5.2 we compare with other state-of-the oorthe--rt generative models, specifically SCW (Fan et al., 2020) and WIN (Korotin et al., 2019).
5. Numerical Experiments
We showcase the performance of our method in a series of numerical experiments. Section 5.1 presents results for the computation of OT maps in high-dimensional settings using Algorithm 1, while Section 5.2 deals with the computation of Wasserstein-2 barycenters on the following data sets: the Swiss roll data set (Section 5.2. two high-dimensional location-scatter data sets (Section 5.2.2), the MNIST data set (Section 5.2.3), a high-dimensional Gaussian data set with a large number of input distributions (Section 5.2.4). and a real-life dataset for multivariate fair regression (Sec- tion 5.2.5). Hyperparameter choices and all other implemen- tation details are given in Appendix D. Our source code is available online at ht tttpps:/ //git/bbbitt///t/ com/ hin/gvise gvisen/ No
Generative models. More closely related to our approach are generative models of Wasserstein barycenters. Also for this class of models the dominant approach is to start from the dual formulation and to solve a tri-level or bi-level optimization problem. Fan et al. (2020) solve the barycen- ter problem by computing Wasserstein distances through a min-max-min tri-level optimization problem using ICNNs, which may be unstable and may result in under-training of the generative model of the barycenter (Korotin et al., 2021b). Korotin et al. (2022) use the fixed point algorithm of Álvarez-Esteban et al. (2016) and the reversed maximin neural OT solver by Dam et al. (2019) to obtain VAE-based generative model of the barycenter through an iterative pro- cedure. Unfortunately, their fixed point algorithm is not guaranteed to converge. Additionally, the model does not output the OT maps between the input distributions and the barycenter, which therefore, need to be computed in additional step. Finally, we mention two early generative methods that share some similarities with our approach. Lu et al. (2020) also tackle directly the primal formulation of the OT problem, but enforce the pushforward constraint by adversarial learning (due to their choice of GAN as generative model) and the bijectivity constraint through a cycle-consistency condition. The model of Xie et al. (2019) also uses bijective generative model, specifically neural
Location-scatter families. Location-scatter families have become standard benchmarks for Wasserstein barycenter models (Korotin et al., 2019; 2021b; Kolesov et al., 2024). They are families of distributions obtained by transforming base distribution μl0 E P(R through invertible affine trans- formations fM. : Rdd + Rd of the form fM... = M - u. for a positive definite matrix M E Rdxd and a vector u E Rd. The Wasserstein-2 barycenter of finitely many distributions belonging to a given location-scater family can be computed numerically using a fixed-point algorithm (Alvarez-Esteba et al., 2016). In the following experiments,
5

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
whenever we mention a location-scatter data set, we will implicitly reproduce the experimental set-up of Korotin et al. (2021b). In particular, we fix S = {1, 3, 4), vector of weights w = (w,).ES = (0.4, 0.3, 0.2, 0.1), and choose as base distribution 140 either the Swiss roll distribution (d = 2), the standard Gaussian distribution . .V I) on Id or the uniform distribution Uniform([-v] /3, v. We then generate four distributions (pe)&ES from the base distribu- tion by setting μs gM where fM, is an affine transformation with positive definite matrix M, == R2 A R&s
10.0%
7.5%
5.0%
distance Relative
2.5%
0.0%
-2.5%
and zero shift, R, is a random rotation matrix sampled uni- formly from the Haar measure on (O))) and A E Raxd is diagonal matrix with diagonal (2,0, 1 ..... 1....b0--11 1) for b = 1)
-5.0%
-7.5%
-10.0%
Metrics. As proposed in Korotin al. (2021b), we use the unexplained variance percentage measure the relative difference between a target distribution y and its estimate j::
d-16
d=32 d4
Figure I. W2 relative error on high-dimensional Gaussian data.
UVP (H, ()) == 100 WH(H, P) (%).
Var (11
solutions are known (see Remark 2.11 in Peyré & Cuturi (2019)). We report the -VVUP and BW?-UVD metrics in Table and Table 2, for Gaussian and uniform distributions respectively, together with their standard deviation from ten trials. In terms of Wasserstein-2 -2 distance, in Figure we re- port box plot of the (signed) relative error of the estimated distance computed on high-dimensional Gaussian data for ten random initializations of our model. All results show very good performance in high-dimensions
In a high-dimensional setting it's not possible to compute UVP (u. (i) exactly, due to the intractability of (2 (.). Nevertheless, as shown in Lemma A.2 of Korotin et al. (2019), if μ T4 no and jú- T4 po. then the following upper bound holds:
|T (1)) (U,
UVP (V, A) 100
Var Po))
5.2. Computing Wasserstein-2 Barycenters
Additionally, as shown by Dowson & Landau (1982), the following lower bound holds:
5.2.1. SWISS ROLL DATASET
BW3-UVP (u,A) == 100 BW3 ())
UVP (n, A), Var(p)
We begin by evaluating Algorithm 2 qualitatively on two-dimensional location-scatter data set with Swiss roll base distribution, S = 2, 3, 4} and weights w = (0.4, 0.3. 0.2, 0.1). Figure shows samples from the true input distributions (Ma)ees (first row) and samples gener- ated from our trained model (second row). At convergence, all input distributions have been learned well, which im- plies that the pushforward constraints are fully satisfied. In Figure 3 we compare samples from the true barycenter and the learned barycenter, i.e. hg Unlike other generative models (see Figure 1 in Korotin et al. (2021b) for a com- parison with SCW,B/ ), our learned barycenter distribution is sharp and reproduces the highly non-linear structure of the Swiss roll distribution. We also verify in Figure 4 that the learned OT maps correctly transport samples from the input distributions to the barycenter via the transformations x + h(fo (Ir s)). for all E S.
where (W A)) denotes the wassersnnnn metric on Gaussian distributions, computed using the means and covariance matrices of μ and j..
When solving an OT problem between Yi and #2 with known OT map T, we evaluate the performance of our estimated optimal transport map T by computing C2-UVP (T,Tipp 11) and BW2-UVP (H2, fiz). For barycen- ter problems, we compute the metries -V- (Ta, P.,) p.) for each E S (where T, and T, are the true and esti- mated OT maps from Ss to the barycenter ī7) and then report the mean C2-UVP metric, averaged over all s E S using the weights (wa)ses. Additionally, we report the BW?-UVP (p, fi) metric, where ji is the true barycenter and j is the model estimate.
5.1. Computing OT Maps
5.2.2. HIGH- NIGH-DIMENSIONAL EXPERIMENTS
We first validate Algorithm by computing OT maps be- tween high-dimensional distributions, for which close-form
We now compare our method's performance to other state-of-the thett-- models on two location-scatte benchmark
6

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
Table 1. -VP and B UVP metries for our method when computing OT maps between high- Gaussian distributions.
METRIC
2=
d=
64
-VV
0.004 = 0.001
0.086 = 0.001
0.204 = 0.002
0.623 = 0.004
2,389 = 0.018
1.658 = 0.005
3.102 + 0.009
BWG-UVP
0,002 = 0.001
0.015 0.005
0.008 = 0.001
0.014 = 0.001
0.034 = 0.002
0.054 0.002
0.088 +: 0.002
Table 2. 2 and BW3-UVP metries for our method when computing OT maps between high-dimensional uniform distributions.
METRIC
d=2
d= 64
2V-UVP
0.54 + 0.006
1.506 = 0,019
2.046 = 0.031
1.404 = 0.006
4.255 = 0.016
3.833 = 0.014
7.137 + 0.012
BW3-UVP
0.007 0.001
0.029
0.003
0.047 = 0.004
0.036 = 0.003
0.704 = 0.009
0.738 + 0.01
1.223 + 0.005
Figure Samples from the true input distributions (first row) and from the input distributions as learned by our model (second row).
Figure 4. Samples from the input distributions μa, (first row) and their image after being transported to the barycenter (second row).
1 + M1
p + U2
Figure 5. Sample from learned barycenter transported to the two input distributions.
Figure 3. True barycenter (left) and learned barycenter (right).
datasets across a range of input dimensions (from d = 2 to d = 128): a Gaussian data set (see Table 3) and a uni- form one (see Table 4). Our model has better performance than the competitors, SCW,B (Fan et al., 2020) and WIN (Korotin et al., 2022), with better or equal training time (for more details, see Appendix D).
each barycenter sample is the average of images from the imput measures (Korotin et al., 2019). In Figure 5 we show samples from the learned barycenter and their images when transported to the two input distributions. While Figure 6 and Figure 7 show the inverse operation on samples from the true input distributions.
We also train model using as input distributions the em- pirical distributions of all ten digits. Thanks to our model's conditional architecture, the training time scales well in the number of input distributions, therefore solving this Wasser- stein barycenter task did not take more training time than the previous one, Figure 8 we show samples from the
5.2.3. MNIST DIGITS
Next, we test our model on image data using the MNIST data set, using our conditional Glow architecture. We first compute the barycenter of the empirical distributions of the digits zero and one. We recover the eell-nnown result that
7

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
Table 3. -VVV and BV3 metrics on location -catter Gaussian data (from Table in Korotin et al. (2019) for SCW,B and WIN).
METRIC
METHOD
d=2
d=
d=
d' 16
d
128
<² UVP
OURS
0.025 0.0
0.079 1: 0.002
0.076 1: 0.001
0.1 1: 0.001
0.382 0.001
0.544 1: 0.001
1.55 1: 0.002
SCW,B
0.070
0.090
0.160
0.280
0.430
0.590
1.280
BWQ UVP
WIN
0.010
0.020
0.010
0.080
0.110
0.230
0.380
OURS
0.005 0.001
0.004
0.001
0.004 + 0.001
0.007 0.001
0.024 0.001
0.046 + 0.001
0.094 +: 0.001
Table -V-UVP and BW2-UVP metrics on location-scatee uniform (from Table Korotin et al. (2019) SCW,B WIN).
METRIC
METROD
d=
d'
8
d 16
d= 32
d= 64
d 128
2V-
OuRS
0.998
0.399 0.005
+ 0.003
0.697 0.002
1.301 0.002
2.428 0.002
5.506 + 0.004
SCW,B
0.290
0,460
0.600
1.380
-WV-
WIN
0,040
0.060
0.060
0.080
0.110
0.270
0.460
OURS
0.01 + 0.004
0.053 # 0.007
0.033 + 0.005
0.058 + 0.003
0.081 : 0.002
0.159 + 0.002
0.362 + 0.003
H1
n 11
O
O
O
f - H22
/
/
1
Figure Sample from the first input distribution transported to the barycenter.
2
2
2
2
2
2
2
2
/ー μ4
3
3
8 4
3
3
3
8 4 6 6
3 4 5
ps
4
4
4
4
4 5
1
H2
Ā Ôー μ6
6
6
5
5
U2
6
6
66666
6
6
6
6
Åー Ps
7 777777
7
7
7
7 8 9
7 8
Figure 7. Sample from the second input distribution transported to the barycenter.
8
8
8
8
8
8
A - H10
9 9
9
9 9
9
9
learned barycenter and their images when transported to ten input distributions.
We illustrate a style translation application of our model in Figure 9, where we transport samples from source distribution (e.g. the distribution of zeros), through the latent space, to a target distribution (c.g. the distribution of ones). Under this transformation, stylistic features of the zero digit such its slant and thickness are preserved, showing that the latent space in model encodes structural features of the input data and can be used to transfer them.
Figure 8. Sample from learned barycenter transported to all input distributions.
our model on high-dimensional Gaussian data sets (d = 64) with various numbers of input distributions (n 8, 64, 128). Different input distributions are indexed by S, the uniform grid on [0, π with n points, and are zero-mean Gaussian distributions with covariance Σ,
E,, where En is diagonal matrix with diagonal (2, 1/2....... 1/2) and R(s) is a rotation matrix describing a rotation in the first two coordinates by an angle s E C [0, x|. In this way we can generate large number of input distributions that vary smoothly in the univariate parameter s. We then com-
5.2.4. LARGE NUMBER OF INPUT DISTRIBUTIONS
In this section, we leverage the conditional architecture of our model to compute Wasserstein-2 barycenters for a large number of input distributions. Specifically, we evaluate
8

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
features X) and the percentage of the population that is African-American (sensitive feature S with range [0, 1].
0
11
As shown in Table 6, standard regression leads to strong correlation between the first predicted variable and the sen- sitive feature for several correlation measures, while our fair regression achieves almost perfect uncorrelatedness. In all cases, we fail to reject the null hypothesis of no association.
H1 U2
1
1
f
/
1
2
4
2 3
2
H1 - 11
3 3
3
7
H1 μ14
4s
4 4
4
4
4 5 6
Table 6. Correlation between regressor and sensitive feature.
46
5
8
8
5
5
5
CORRELATION COFFFICIENT
STANDARD
FAIR REGRESSION
REGRESSION
6
4
6
6
6
6
6
μ1 17
PEARSON
0.43
0.0003 (P-VALUE: 0.99)
SPEARMAN
0.47
0.0058 (P-VALUE: 0.80)
μ1 μg
7
7 7
7
7
7
7
7 8
KENDALL 0.32
0.0039 (P-VALUE: 0.80)
μ1 + μ9
8
8
8
8
8
8
8
M1 - N10
9
8
9
9
9
9 9
9
We point out that this experiment requires computing the barycenter of 100 input distributions (the cardinality of the range of S, rounded up to the nearest percentage point), which would be computationally challenging for any other numerical Wasserstein barycenter method.
Figure 9. Sample from first input distribution transported to all other input distributions.
6. Conclusion
pute the Wasserstein-2 barycenters of these distributions using Algorithm 2. The results presented in Table 5 and show degradation in performance as the number of input distributions increases.
We have introduced new method for computing OT maps in high-dimensional settings using conditional normaliz- ing flows. Unlike previous methods, we avoid complex adversarial optimization and directly solve the primal OT problem using gradient-base -hased minimization the transport cost. We have shown how the approach can be extended compute Wasserstein barycenters by solving conditional variance minimization problem, and we have compared to state-of- sitee-of-hheee the-hhrrt competitor models on several benchmarks. Our approach yields accurate results and can be used to com- pute Wasserstein barycenters for hundreds of input distribu- tions, which was computationally infeasible with previous methods.
Table Performance as number of input distributions increases.
METRIC
4
16
64
128
LVVVV 0.057
0.119
0.081
0.073
BW2-UVP
0.053
0.031
0.016
0.051
5.2.5. MULTIVARIATE FAIR REGRESSION
Fair regression, in the sense of demographic parity, looks for a regression function f(X,S) S) that minimizes the cost E||lY f(x, {) such that f(X, is independent of S. In this context S denotes sensitive features, such as race or gender. with respect to which discrimination is not allowed. The solution to this constrained optimization is precisely Wasserstein-2 barycenter of the conditional distributions of Y given (Chzhen al., 2020; Gouic al., 2020). To demonstrate the applicability of our model multivariate fair regression, work the dataset "Communities and Crime" (Redmond, 2002), which is benchmark dataset in the fairness literature, and we regress the target vari- ables "percentage of officers assigned to drug units" and "total number of violent crimes per 100k population" (target variable Y) on 127 socio-economic features inon-sensitive
7. Limitations
A limitation of our approach is that Algorithm 2 applies only to Wasserstein-2 barycenters. Its extension to Wasserstein-p barycenters for * 2 is currently still an open research question. Furthermore, training might be computationally intensive on image datasets much larger than the ones we tested on, i.e. for > 1000.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
specifically highlighted here.
Dam, N., Hoang, Q. Le, T., Nguyen, T. D., Bui, H., and Phung. D. Three-player Wasserstein GAN via amortised duality. In International Joint Conference Artificial Intelligence 2019, pp. 2202 2008. Association for the Advancement of Artificial Intelligence (AAAI), 2019.
References
Agueh, M. and Carlier, G. Barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 43(2): 904-924, 2011.
Del Barrio, E., Cuesta-Albertos, J. A., Matrán, C., and Mayo-iscar, A. Robust clustering tools based on optimal transportation. Statistics and Computing. 29:139-160 2019.
Alvarez-Esteban, P. C., Del Barrio, E., Cuesta--lbertoss J., and Matrán, C. A fixed-point approach to barycenters in Wasserstein space. Journal of Mathematical Analysis and Applications, (41(2)4444-72 2016.
Dinh, L., Krueger, D., and Bengio, Y. NICE: Non-linear Independent Components Estimation. arXiv preprint arXiv:1410.8511 2014.
Amos, B., Xu, L., and Kolter, J. Z. Input convex neural net- works. In International Conference on Machine Learning, pp. 155-155. PMLR, 2017.
Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density estima- tion using Real NVP. arXiv preprint 1rXiv:1605.08800 2016.
Anderes, E., Borgwardt, S., and Miller, J. Discrete Wasser- stein barycenters: Optimal transport for discrete data. Mathematical Methods of Operations Research, 84:389- 409, 2016.
Dowson, D. and Landau,
B. The Fréchet distance between multivariate normal distributions. Journal of Multivariate Analysis, 12(3):450
1982.
Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein gen- erative adversarial networks. In International Conference on Machine Learning. pp. 214-223, PMLR, 2017.
Draxler, F., Wahl, S., Schnörr, C., and Köthe, U. On the universality of normalizing flows. arXiv preprint arXiv:2402.06578, 2024.
and coupling-base based
Atanov, A., Volokhova, A., Ashukha, A., Sosnovik, I., and Vetrov, D. Semi-conditional normalizing flows for semi- supervised learning. arXiv preprint 1XXiv:9005.00500 920, 2019.
Durrett, R. Probubility: Theory and Examples, volume 49, Cambridge University Press, 2019.
Bogachev, V. I. and Ruas, M. A. S. Measure Theory, vol- ume 1. Springer, 2007.
Fan, J., Taghvaei, A., and Chen, Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.0444: 2020.
Brizzi, C., Friesecke, G., and Ried, T. p- D-Wasserstein barycenters. Nonlinear Analysis, 251:113698 2025.
Folland, G. B. Real Analysis: Modern Techniques and their Applications, volume 40. John Wiley & Sons, 1999.
Brown, B. C., Caterini, A. L., Ross, B. L., Cresswell, C., and Louiza-Ganem, G. Verifying the union of manifolds hypothesis for image data. arXiv preprint arXiv:2207.02862 2022.
Gangbo, W. and McCann, R. J. The geometry of opti- mal transportation, Acta Mathematica, 177(2):1331161, 1996.
Cao, J., Mo, L., Zhang, Y., Jia, K., Shen, C., and Tan, M. Multi-marginal Wasserstein GAN. Advances in Neural Information Processing Systems, 32, 2019.
Gong, S., Boddeti, V. N., and Jain, A. K. On the intrinsic dimensionality of image representations. Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 3987-3996, 2019.
Chzhen, E.,, Denis, C., Hebiri, M., Oneto, I.., and Pontil, M. Fair regression with Wasserstein barycenters. Ad- vances in Neural Informaionn Processing Systems, 33: 7321-7331, 2020.
Gouic, T. L., Loubes, J.-M., and Rigollet, P. Projec- tion to fairness in statistical learning. arXiv preprint arXiv:2005. 10000.11720 17200, 2020.
Cuturi, M. Sinkhorn Distances:
Lightspeed Computation of Optimal Transport. Advances in Neural Information Processing Systems, 26, 2013.
Ho, N., Nguyen, X., Yurochkin, M.. Bui, H. H., Huynh, V. and Phung, D. Multilevel clustering via Wasserstein means. In International Conference Machine Learn- ing. pp. 1501-1509, PMLR, 2017.
Cuturi, M. and Doucet, A. Fast computation of Wasserstein barycenters. In International Conference on Machine Learning, pp. 685-693, PMLR, 2014.
Kantorovich, V. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS), volume 37, pp. 199-201, 1942.
10

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
Kingma, D. and Dhariwal, P. Glow: Generative flow with invertible 1x1 convolutions. Advances in Neural Information Processing Systems, 31, 2018.
Loaiza-Ganem, G., Ross, B. L.., Hosseinzadeh, R., Caterini, A. I.., and Cresswell, J. C. Deep generative models through the lens of the manifold hypothesis: A survey and new connections. Transactions on Machine Learning Research, 09, 2024.
Kobyzev, I., Prince, S. J., and Brubaker, M. A. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine In- telligence, 43(1 1):3944-3979, 2020.
Lu, G., Zhou, Z., Shen, J., Chen, C., Zhang, W., and Yu, Y. Large-scale optimal transport via adversarial training with cycle-consistency, arXiv preprint 2XXi:200330066335 2020.
Koehler, F., Mehta, V., and Risteski, A. Representational aspects of depth and conditioning in normalizing flows. In International Conference on Machine Learning, pp- 5628-5636. PMLR, 2021.
Luo, Y., Zhang, S.-Y., Zheng, W.-L., and Lu, B.-L. WGAN domain adaptation for eeg-based emotion recognition. In International Conference on Neural Information Process- ing. pp. 275-286, 2018.
Kolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., and Korotin, A. Estimating barycenters distributions with neural optimal transport. In International Conference Machine Learning, pp- 25016-25041, PMLR, 2024.
Monge, G. Mémoire sur la théorie des déblais et des rem- blais. Mem. Math. Phys. Acad. Royale Sci., pp. 666-704, 1781.
Mrouch, Y. Wasserstein style transfer. arXiv preprint arXiv: 1900.9922.. 2019.
Korotin, A., Egiazarian, V., Asadulaev, A., Safin, A., and Burnaev, E. Wasserstein-2 generative networks. arXiv preprint arXiv:/909 190.::999.11082, 2019.
Papamakarios, G., Nalisnick, E., Rezende, J., Mohamed, S., and Lakshminarayanan, B. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57): 2021.
Korotin, A., Li, L.,, Genevay, A., Solomon, J. M., Filippov, A., and Burnaev, E. Do neural optimal transport solvers work? A continuous W benchmark. Advances in Neural Information Processing Systems, 34:14593- 14605, 2021a.
Petzka, H., Fischer, A., and Lukovnicov, D. On the regularization of Wasserstein GANs. arXiv preprint arXiv: 170p:0899..844 2017.
Korotin, A., Li, I.., Solomon, J., and Burnaev, F. Continuous Wasserstein-2 barycenter estimation without minimax optimization. arXiv preprint 1rXiv:2102.01752, 2021b.
Peyré, G. and Cuturi, M. Computational optimal trans- port: With applications data science. Foundations and Trends@ in Machine Learning, 11(5-6):355-607, 2019.
Korotin, A., Egiazarian, V.,
Li, L., and Burnaev, E. Wasser- stein iterative networks for barycenter estimation. Ad- vances in Neural Information Processing Systems, 35: 15672-15686, 2022.
Pope, P., Zhu, C., Abdelkader, A., Goldblum, M., and Gold- stein, T. The intrinsic dimension of images and its impact learning. arXiv preprint urXiv:2104.08894 2021.
Rabin, J., Ferradans, S., and Papadakis, N. Adaptive color transfer with relaxed optimal transport. In 2014 IEEE International Conference on Image Processing (ICIP), pp. 4855-48855 IEFE, 2014.
Lacombe, J., Digne, J., Courty, N., and Bonneel, N. Learn ing to generate Wasserstein barycenters. Journal of Math- ematical Imaging and Vision, 65(2):354-3700 2023.
Redmond, M.
Communities and Crime. UCI
Leygonie, J., She, J., Almahairi, A., Rajeswar, S., and Courville, A. Adversarial computation of optimal trans- port maps. arXiv preprint arXiv: 190,01,,900099),2200
Machine Learning Repository, 2002. DOI:
http://bo.oo//1022333/33333
Shen, J., Qu, Y., Zhang, W., and Yu, Y. Wasserstein distance guided representation learning for domain adaptation. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 4058 0058-4065.. 2018
Li, L., Genevay, A., Yurochkin, M., and Solomon, J. M. Continuous regularized Wasserstein barycenters. Ad- vances in Neural Information Processing Systems, 33: 17755-17765, 2020.
Simon, D. and Aberdam, A. Barycenters of natural images constrained Wasserstein barycenters for image morphing, Proceedings of the IEEE/CVF Conference on Com- puter Vision and Pattern Recognition, pp. 7910-7919, 2020.
Liu, H., Gu, X., and Samaras, D. Wasserstein GAN with quadratic transport cost. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4832- 4841, 2019.

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
Sinkhorn, R. and Knopp, P. Concerning nonnegative ma- trices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343-348, 1967.
Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional
Wasserstein distances: Efficient optimal transportation geometric domains. ACM Transactions Graphics (ToG), 14
Srivastava, S., Cevher, V., Dinh, Q. and Dunson, D. WASP: Scalable Bayes via barycenters of subset posteriors. In Artificial Intelligence and Statistics, pp. 912-920. PMLR, 2015.
Stimper, V., Liu, D., Campbell, A., Berenz, V., Ryll, L., Schölkopf, B., and Hernández-Lobato, J. M. normflows: A PyTorch Package for Normalizing Flows. Journal of Open Source Softwree 8(86):5361, 2023. doi: 10.21105/ 1oss.05361. URL joss. 05361.
Teshima, T., Ishikawa, I., Tojo, K., Oono, K., Ikeda, M., and Sugiyama, M. Coupling- invertible neural networks are universal diffcomorphism approximators. Advances in Neural Information Processing Systems, 33:3362-3373, 2020.
Villani, C. Topics in Optimal Transportation, volume 58.
American Mathematical Soc., 2021.
Wu, J., Huang. Z., Thoma, Acharya, D., and Van Gool, L. Wasserstein divergence for GANs. In Proceedings of the European Conference Computer Vision (ECCV), pp. 653 03-668,222 2018.
Xie, Y., Chen, M., Jiang, H., Zhao, T., and Zha, H. On scalable and efficient computation of large scale opti- mal transport. In International Conference Machine Learning, pp.
PMLR, 2019.
12

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
A. Auxiliary Results
The following are some standard results from measure and integration theory that are used in this paper.
Lemma A.1 (Change of variables). Let f:
X + Y measurable function between two measurable spaces (X, A) and (Y, B). p probability measure on (X, A), then measurable finction g Y satisfies LL' (f##1) if and only if g f E (p)). Moreover, in this case has
J ((W)(FWW)dd)) J s(f(z)))(dzx)
Proof. See Theorem 3.6.1 in Bogachev & Ruas (2007).
Lemma A.2. Let f Rd Rd be Borel function and μ. Borel probability measure RR Then f 0 j-almost surely if and only if { [1(()]lu(dx>= = 0.
Proof. Sa (l( = if and only fx) = p-almost surely (see, for instance, Proposition 2.16 in Folland (1999), which the case and only if ((x) is the zero vector p-almost surely.
Lemma A.3. Let u, E (RR)) for some 1. If E B(u, v) and is a (M, vi-iverrse of J, then
① E E
( is v-a.s. unique (μ, iiivere of f.
(ii) f B),, [lf[l (V(v) and is p.a.s. unique (v, -iiverse of
Proof. (i f E B(,,v) Then by Lemma A.1.
J [lf()))Pp(dr)) = f [Ivl(Px(dy) < co.
(ii) Assume / is another (1, iinnvese of f, then
|1f f)) x11 p(dx)
(fo o f = id, p-a.s.)
3d
I(( (x)) (J (()))l
(oo = = p-a.s.)
[ (If(z) - f(y)ll v(dy).
(change of variable f(z), see Lemma A.1)
which, by Lemma A.2, implies that f-f v-a.s.
(iii) We first show that fœv= yr. Let B(RR)), then
x()+(A)) = -P(F-((--(A))
(fpp=v)
-I((fo()((A))
=A(A).
( id, pra.s.)
Since f is the iinveree of f, clearly satisfies the definition of a (v, iinverre of therefore f belongs to B(P, Finally, [| E LP(v) simply follows from f B(v, p) and (i).
Lemma A.4. Let X, A. E PA(R"). If B(A,)) and 9 E B(H,v), then g of E B(),)). Furthermore, if f is (), n) -ivvere of If and ğ is a (n, innvrre of g, then ff ğ is a (), in-iveesee of g o
13

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
Proof. Clearly (g + 9# (f -4
Furthermore.
[(fo - (f g)(il v(dy)
|[f [(f (f g))( x(dy)
(g = id v-a.s.)
[[fogog)(z) - f(x)) p(dz))
(change of variables = g(y))
l (x) f(x))l u(dx)
(--.s.)
[fo.goo (2) z|| X(dz),
(change of variables f(x))
which shows that fogogof id A-almost surely by Lemma A.2. An analogous computation shows that go fofoo id p-almost surely.
B. Proofs
Proof of Lemma 3.1. Existence of unique solution 7 E ((,, v) to Kantorovich's problem (2) and the fact that y is of
the form (id, () for a measurable map T Rd + Rd pushing μ to v, follows from Theorem 3.7 of Gangbo McCann (1996). That p-almost surely unique and belongs to (2u, v) is consequence of Theorem 4.5 of Gangbo McCann (1996).
The proofs of Theorem 3.2 and Theorem 3.3 rely on the following lemma.
Lemma B.1. Consider probability measures X, H4, Prpac ( for some p 1. Then B(A, p) is non-empty, and for every f E p),, there exists a E v) such that
WH(K = { [1((z) - g(=))I = - |12 g (()))P((())))
(5)
where j is any (2,z) -invverse of f. In particular, o an optimal transport map between μ and v.
Proof. It follows from Theorems 3.7 and 4.5 Gangbo & McCann (1996) that B(A, 11) is non-empty and there exists a T E B(,,v) such that
f || - T(I) 11==)(())IP u(dr) = inf /
=ell(p.)) 3u || 3/1" π (dx, dy) = (H,
So, for a given f E B(A, u), the mapping T o in B(A, v) by Lemma A.4 and, by the change-o--variall--foolll
J/L [1(()) ( = / [1f(z) To f()) 1 (G-T(())PPp(a)) (P,
Since T - go y-almost surely, this shows (5), and g f is an optimal transport map between y and which completes the proof of the lemma.
Proof of Theorem 3.2. We know from Lemma 3.1 that there exists optimal transport map T E B(p, v). So, for any E B(A, 11) and B(A,)) we have
|l g (1 = [ ((o)P(a))))( / [(=)(()))Ppda)))
J
Wy(,,))
(6)
where f is any (2, iniieees of and we used the fact that o pushes f to by Lemma A.4.
14

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
On the other hand, it follows from Lemma B.1 that there exist f ((), and y E B(A,)) such that the inequality in (6) becomes an equality, which completes the proof.
Proof of Theorem 3.3. It follows from Brizzi et al. (2025) that the w-weiggttd Wasserstein-2 barycenter jı of (Ha)sES
belongs to a2 (. Moreover, we obtain from the definitions of the barycenter ji and the Wasserstein-2 distance that
Σ BES
w. (p, S Σ W. W2 S Σ Ws / [1b(z) #ES
(2,,s))
BES
for all Borel maps h : Rd + Rd and f Rd x hand, we know from Proposition B.1 that there exist E B(X, (J) and f(;,s) E B(X, u.), E S, such that
S - Rd such that f(, 8) transports À to μa for every E S. On the other
(s) - ( [1{(z) f(z, 8) A(d=) for 8 E
and therefore,
- ((E,)) Σ W. (2, = Σ w. w, Σ Ws | -1(3))
SSS
#ES
SES
In particular, minimizes
Σ
Rd wg [1(()) f(2, E - [ f(2,
#ES
condicitoonl expectation over all Borel maps from dR Rd, which implics that is given by the component-wisc conditional
h(Z) - =((
Σ
w. (Z,
SES
see, e.g. Durrett (2019). Moreover, minimizes
E [E 12 f(2,)))
Σ E Var 2)]
over all mappings f: Rd S Rd satisfying f(, E B(A, p.) for all s S. This proves (i).
Now, assume f Rd S Rdd is another function such that f(:,s) E B,), for all S and
Var(F(Z,,)) 2) Σ E (Vr((f. 2)].
Then, for i(2) = E[Í(2,5) 2 = Exes w. (2,( 8), one has
Σ
m. 22 (h Σ
[lā(2) - {(=,5))12 Σ B Var(ff 2)]
ws
Judd
SES
#ES
d
E [Var(f{(
2) = Σ W. [1(2)) s)l3>)))>>(>>)) = Σ w. (,,
SES
#ES
But since j is the unique w-weighted Wasserstein-2 barycenter, we conclude that d A = th, which shows (ii).
15

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
C. Background on Normalizing Flows
Our method relies on conditional normalizing flows to compute OT maps and Wasserstein barycenters. For the reader's convenience, we collect here all the necessary background on normalizing flows. For more information, we point the reader Kobyzev al. (2020) and Papamakarios et al. (2021), which are two recent, well-written survey papers.
Normalizing flows are generative models used to learn high-dimensional data distributions by transforming simple latent distribution (e.g. standard normal distribution) through sequence of diffeomorphisms. A key advantage of normalizing flows is that, unlike other generative models such as VAEs and GANs, density evaluations of the model distribution are efficient, thanks to the change of variables formula, which allows efficient learning by likelihood maximization. More precisely, if Z is an Rd-valued random variable with density pz and ʃ R2d + Rd is a composition of diffeomorphisms == fi ..... ... o fN, then the random variable = f(Z) has density:
p(x) = p2 (x)) det(DD.f
(7)
where f-1 = -N .... o fi and
[del(D2s (x))
II del del (D2f (x)))-
The challenge in designing good normalizing flows consists in finding flows that have a tractable Jacobian and are easy to invert, but are expressive enough to be able to approximate interesting data distributions.
Real NVP. Real NVP was the first normalizing flow model with competitive performance on benchmark image data sets (Dinh et al., 2016). It leverages coupling layers, first introduced by Dinh et al. (2014), provide an economical way of representing highly expressive transformations. In coupling layer the input is first split into two components (a*, 2 then 2 is transformed using simple parametric bijective transformation, whose parameters depend (possibly in highly non-linear way) on NB, as follows:
x BB = B
fo(z0)(21)
where fe(-) is parametric bijection with parameters 0, called a coupling function. In particular, Real NVP uses affine coupling layers with log-scale and shift parameters, 0(= = (a(z b(2( b)) resulting in a coupling function of the form:
fu(ze) = A o exp(a(= b(zB),
where a(-) and b(-) are feedforward or convolutional neural networks.
Stacking many affine coupling layers and permutation layers (i.e. bijective transformations 2 + T(2), for random permutation æ fixed at initialization) results in highly complex, non-linear transformations.
Glow. Kingma& & Dhariwal (2018) proposed a normalizing flow model with improved performance image data. They replaced the random permutation layers of Real NVP with invertible Ixl convolutions, thereby allowing the model to learn the permutations, instead of randomly fixing them at initialization. They also added activation normalization layers, which adapt the batch normalization technique to the small mini-batch regime typical of training on high-resolution image data.
Multi-scale architecture. A common problem of normalizing flows is that the dimension of the latent space must match the dimension of the target distribution be learned, due to the bijectivity constraint. In high-dimensional settings, this results in very deep and large models, that may be difficult train. Multi-scale architectures solve this problem by exploiting the fact that empirical data distributions, especially for image data, tend to concentrate around low dimensional manifolds (Gong et al., 2019; Pope et al., 2021; Brown et al., 2022; Loaiza-Gant Ganem et al., 2024). This is done by dividing the model's flows into L separate scales or levels, as shown in Figure 10, consisting of K layers each, with latent dimensions being gradually added at each subsequent scale. In this way, only a dow-dimensional latent vector & is pushed through the full depth of the model, while the remaining dimensions undergo simpler, easier-to-train transformations, capturing finer details.
16

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
K 1d
id
K
Figure 10. scati-scale architecture with L scales of K flows each.
D. Experimental details
Table 7 presents all hyperparameter choices for our model in all numerical experiments. The architecture type (either conditional Real NVP conditional Glow) reports the number of hidden neurons in each hidden layer of the conditioner network. L and K denote the number of scales and flows per scale in the multi-scale architecture.
Table 7. Hyperparameter choices for model for all numerical experiments.
DATASET
OPT
ITERS
BS
TYPE
K
WEIGHTS
SWISS ROLL
ADAM
10-4
5
10
10
REAL
NVP
32
LOGSPACE [0,-4]
664
32
FOR 4,8, FOR d 32
HIGH-DIM GAUSSIAN (d = 128)
ADAM
10-3
10
10
REAL
NVP
logz(d)
16
LOGSPACE [0, -2-2]
(64, 64]
8
32
FOR = 2
HIGH-DIM UNIFORM
ADAM
10
10a
10a
REAL (64, 64]
NVP log
16
FOR 4, 8,
LOGSPACE [0, -0]
(d = 128)
8
32
MNIST
ADAM
10-4
5
10
32
GLOW,
256
16
LOGSPACE
CHANNELS
[1, [1,-2]
LARGE NUMBER OF IN-
ADAM
10-3
10
10a
REAL 664, 64]
NVP logg(d)
32
LOGSPACE
PUT DISTRIBUTIONS
[0, -2]
All numerical experiments were run on an NVIDIA GeForce RTX 4090 GPU with 24 GB of memory, except for the MNIST data set experiment, which Python, is both GPU and CPU-compatible, and builds on PyTorch, the normfl
NVIDIA RTX 6000 Ada with 48 GB of memory. Our implementation is written in floloww Lows package by Stimper et al. (2023) and
the code repository by Korotin et al. (2021b).
The training times for the location-scatter experiments Section 5.2.2 are shown in Table and Table 9. We emphasize, though, that in practice all methods converge in approximately 10-20 minutes, with longer times needed only for top-notch performance. On the MNIST data set our method converges in approximately 30 minutes.
17

---

Computing Optimal Transport Maps and Wasserstein Baryeenters Using Conditional Normalizing Flows
Table 8. Training times high-dimensional location-scatter Gaussian data.
METHOD
d = 2
d= 4
d=
16
d = 32
d 64
= 128
SCW, B
114M 9s
110M51S
108M21s
108M 4s
108M52s
109M36s
114M56S
WIN
60M 9s
57M56S
58M32s
75M53s
57M26s
92M54S
93M28S
OURS
23M32S
23M58S
35M28S
47M
30M 1s
35M59s
42m
Table 9. Training times on high-dimensional location-scatter Uniform data.
METHOD
d
d
8
=
:22
d 128
S WIN
114M
110M51S
108M215
108M
108M52s
109M36s
114M56S
90M 1 7s
58M55s
59M45s
60M45s
60M 8s
93M16S
103M 13s
OURS
110M IS
114M36S
110M46S
111M14S
110M11S
108M55s
11 MM49
18


Reviews:
[{"review_id": "5WNdGEMci3", "reviewer": "Reviewer_U9wD", "summary": "The paper proposes a new algorithm to compute the Wasserstein barycenter of a set of distributions and, by extension, the optimal transport between any pair of the given distributions. The method rests on a new representation of the barycenter objective in terms of conditional normalizing flows. Since this objective can be readily approximated by empirical averages and variances, the resulting algorithm is very simple and elegant. Experiments illustrate the algorithm's outcomes on the Swiss roll and the MNIST datasets and demonstrate its superiority over the competition in systematic comparisons.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, \n\nwe have carefully addressed all your comments below. According to the ICML review guidelines, we are not allowed to submit a revised manuscript at this point of the review process, but below we discuss in detail the edits we will make based on your suggestions in the camera-ready version of the paper, if accepted.\n\n**1. Can you demonstrate the usefulness of Wasserstein barycenters for downstream tasks?**\n\nWasserstein barycenters have found application in many downstream tasks, such as shape interpolation, image interpolation, color translation, style translation, Bayesian subset posterior estimation, clustering in Wasserstein space, and fairness (for references for each application, see the introduction of the paper).\n\nOur interest in real-life applications to fairness in insurance prompted us to develop a method with good scalability properties in the number of input distributions, which is absent in the existing literature and one of the main contributions of our paper. Since the method is of independent interest for the broader Wasserstein barycenter community, in this paper we focus on introducing it and showcasing its competitive performance on well-established benchmarks, leaving its application to real-life actuarial datasets for upcoming work.\n\n**2. Samples from the MNIST barycenter look like a superposition of 10 digits with the same style. Is the Wasserstein barycenter a useful notion?**\n\nThe notion of Wasserstein barycenter is in general non-trivial. For instance, the Wasserstein-2 barycenter of two Dirac measures $\\delta_x$ and $\\delta_y$ (for $x,y \\in \\mathbb{R}^d$) is not the measure $\\frac{1}{2} \\delta_x + \\frac{1}{2} \\delta_y$, but rather the Dirac measure centered at their midpoint, i.e. $\\delta_{\\frac{x+y}{2}}$. Thanks to this property, Wasserstein barycenters can provide very natural-looking averages of 2d/3d images, *when each image is modelled as an input distribution in 2d/3d space*, as can be seen in Figures 1 and 7 of (Solomon et al, 2015) or Figures 1 and 2 of (Cuturi et al, 2014). \n\nRecent works, instead, focus on image datasets in which *input distributions are not themselves images, but rather empirical distributions of images* (as in the MNIST barycenter task in our Section 5.2.3). While these benchmarks may be useful to showcase model scalability, they are of very limited practical interest (the barycenter images are just pixel-wise averages of input images) and do not correspond to any meaningful image interpolation task. We agree with the reviewer and acknowledge that there is a pressing need for the Wasserstein barycenter community to develop more meaningful benchmark tasks in high dimensions, possibly moving away from image data. We hope our future work on actuarial datasets will contribute to this.\n\n**3. It is never specified over which distributions expectations and variances are taken. This makes theorem 4.3 hard to understand.**\n\nWe follow the standard convention of assuming that there is an underlying probability space on which all random variables are defined (this can be done without loss of generality, see Lemma 8.16 of (Kallenberg, 2002)). All expectations are then well-defined Lebesgue integrals with respect to this probability measure. We will add this assumption explicitly in the paper.\n\n**4. Line 205: \"s may be univariate discrete\". Do you mean \"univariate continuous\"?**\n\nBy \"univariate discrete'', we actually meant an $\\mathbb{R}$-valued random variable taking finitely many values (such as the discrete uniform distribution). To avoid misunderstandings, we will edit the text in l.205 to read \"[...] the conditioning variable $s \\in \\mathcal{S}$ may be $\\mathbb{R}$-valued and taking finitely many values (as in Section 5.2.4, where $\\mathcal{S}$ is the uniform grid on $[0, \\pi]$ with $n$ points) or one-hot encoded (as in all other numerical experiments).''\n\n**5. The authors might want to consider putting the proofs of theorems 4.2 and 4.3 in the main paper. They are not very long, the paper seems to have substantial unused white space, and it would make the presentation more self-contained.**\n\nThis is a good suggestion. We will do this, provided we can respect the ICML page limit for the revised paper.\n\nThanks also for reporting the typos, missing definitions and missing parameter values. We will correct them all in the revised paper.\n\n**References**\n\n- Solomon, Justin, et al. \"Convolutional wasserstein distances: Efficient optimal transportation on geometric domains.\" ACM Transactions on Graphics (ToG) 34.4 (2015): 1-11.\n- Cuturi, Marco, and Arnaud Doucet. \"Fast computation of Wasserstein barycenters.\" International conference on machine learning. PMLR, 2014.\n- Kallenberg, Olav. Foundations of Modern Probability. Springer Science \\& Business Media, 2002.", "claims_and_evidence": "The paper derives and proves the new variant of the objective (theorem 4.3). Since the proof short-cuts some steps by citing results from the literature, its correctness it not readily apparent. The experiments are well designed and support the superiority claim.", "methods_and_evaluation_criteria": "The experiments adapt an evaluation protocol with analytic ground truth from the literature. This makes results directly comparable and supports the superiority claim.", "theoretical_claims": "See above.", "experimental_designs_or_analyses": "See above.", "relation_to_broader_scientific_literature": "The experiments only include two alternative methods. I do not know if there are others that should have been included.", "essential_references_not_discussed": "None found.", "other_strengths_and_weaknesses": "The authors note -- in accordance to earlier findings -- that samples from the MNIST barycenter look like a superposition of 10 digits with the same style. It should be discussed what this means: Is the barycenter a useful notion of the \"average\" of a set of given distributions? Does it have appealing properties beyond minimizing the Wasserstein distance? If yes, are these properties useful to simplify or enable downstream tasks? If no, what's the point of computing barycenters?", "other_comments_or_suggestions": "The authors might want to consider putting the proofs of theorems 4.2 and 4.3 in the main paper. They are not very long, the paper seems to have substantial unused white space, and it would make the presentation more self-contained. \n\nMinor points and typos:\n* The formal definition of the L_p(lambda)-norm is missing. This makes theorem 4.2 hard to understand.\n\n* It is never specified over which distributions expectations and variances are taken. This makes theorem 4.3 hard to understand.\n\n* A closing bracket is missing in the change-of-variables formulas in algorithms 1 and 2.\n\n* Line 205: \"s may be univariate discrete\". Do you mean \"univariate continuous\"?\n\n* Section 5.2.1: The weights w_s are not specified.", "questions_for_authors": "Can you demonstrate the usefulness of Wasserstein barycenters for downstream tasks?", "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": "3hvsz2DZNK", "reviewer": "Reviewer_91Cy", "summary": "The authors propose a new method for finding Wasserstein-2 barycenter via Conditional Normalizing Flows (CNF) as well as computation of Optimal Transport (OT) maps from input distributions to the barycenter. The key advantage of the method is minimization of the primal OT problem by invertible pushforward bijections, avoiding adversarial bi or tri-level optimization problems present  in the previous methods. They demonstrate performance of the method for computing Wasserstein-2 barycenter for Swiss Roll dataset, Gaussians and high-dimensional case for MNIST dataset.\n\n## **Update after rebuttal.**\n\nI have carefully read the authors' responses and would like to raise a few concerns:\n- (minor) Upon re-reading the manuscript, I noticed that the paper (Korotin et al., 2019) is incorrectly cited as representing the WIN approach in several places (lines 142, 385, 394). I believe the authors intended to reference (Korotin et al., 2022). This appears to be a typo and should be corrected, as it may confuse readers.\n-  I am *confused* by the fact that you do not use the time of rebuttal to perform the comparison with some of the non-generative approaches which I pointed to (Kolesov 2024a,b). There are not so many papers on continuous OT barycenter estimation and even less of them are generative (in your classification). Thus, it seems for me that the comparison with the best non-generative approaches in the appropriate experimental setups is (1) meaningful and (2) not so time-consuming. It is even more confusing since your section 4.3.2 is dedicated to the similar idea of computing the barycenter using the conditional normalizing flows, which allow for computation of OT maps between input distributions and barycenter (i.e., finding barycenter points from samples of input distributions), and the experimental section 5.2 tests this conditional model.  At least in experiment with *location-scatter Gaussian* data the comparison is possible and **meaningful**. See my next point for details.\n- I previously referred to the L$_2$-UVP values in Table 1 for the Gaussian experiment, which was incorrect. I actually meant the experiment with location-scatter Gaussian data. As another reviewer noted, the L$_2$-UVP comparison is missing for this case. I believe that, in this experiment, comparing against both generative and non-generative methods would be valuable for situating your approach within the current landscape of continuous barycenter solvers.\n\n- Besides, I am confused by the overall practical validity of the proposed approach, because (1) it works only with the quadratic cost, (2) was not tested on any practical experimental setups considered in SC$\\mathbb{W}_2$B (Fan et al, 2021) or WIN (Korotin et al., 2022) papers or any others. For example, the authors do not test their approach on \"Ave, celeba!\" benchmark dataset considered in WIN paper. It is strange because both competitors SC$\\mathbb{W}_2$B and WIN were tested here (see section 6.1 of WIN's paper). Instead, the authors argue that such experiments are not meaningful. However, given that both SC$\\mathbb{W}_2$B and WIN were evaluated on this benchmark, omitting it weakens the comparison. (Actually, here it is not necessary to use your conditional model - the comparison of the learned barycenters is enough to measure the performance of your approach.)\n\nOverall, I remain unconvinced that the current contribution is substantial enough to warrant publication at this conference.\n\n**References.** \n\nFan, J., Taghvaei, A., and Chen, Y. Scalable computations of wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462, 2020.\n\nKorotin, A., Egiazarian, V., Asadulaev, A., Safin, A., and Burnaev, E. Wasserstein-2 generative networks. arXiv preprint arXiv:1909.13082, 2019.\n\nKorotin, A., Egiazarian, V., Li, L., and Burnaev, E. Wasserstein iterative networks for barycenter estimation. Advances in Neural Information Processing Systems, 35: 15672–15686, 2022.\n\nKolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., and Korotin, A. Estimating barycenters of distributions with neural optimal transport. arXiv preprint arXiv:2402.03828, 2024a.  \n\nKolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Kratsios, A., ... & Korotin, A. (2024b). Energy-guided continuous entropic barycenter estimation for general costs. Advances in Neural Information Processing Systems, 37, 107513-107546.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, \n\nplease find our replies below.\n\n**1. The authors do not discuss limitations of the proposed method.**\n\nWe agree that discussing the limitations is very important. We propose to add discussions of the following limitations: \n1. Alg. 2 is suitable only for Wasserstein-2 barycenters. \n2. Training might be computationally intensive on image datasets larger than the ones we test ($\\gg$1000 dimensions). \n\n**2. Normalizing flows (NFs) have known drawbacks (e.g., constraints, high computational cost, instability).**\n\nIn our experience NFs are stable and reliable generative models, as our results show. They require large latent spaces, due to the **bijectivity constraint**, but the multi-scale architecture makes training in high dimensions unproblematic (see App. D). Training is **stable** with minimal fine-tuning (e.g. only lower learning rate on non-Gaussian data, see App. D). Our method did not require more **compute** than competitors and can run on a single GPU with 12GB. For training times, see question 7.\n\n**3. Is the method appropriate for “Ave, CelebA”?**\n\nYes. Indeed, the Glow model (which we build on in Sec 5.2.3) was trained on the CelebA-HQ dataset with great success (Kingma e al., 2018), but it required one week of training on 40 GPUs (https://github.com/openai/glow/issues/37). We lack such computational resources. We will mention this in the method's limitations (see question 1).\n\nWe claim that our experimental set-up is sufficient to prove the **scability of our method**, because the MNIST dataset ($d=784$) is high-dimensional enough **for all real-life applications** (see paper references): style transfer ($d=500$), color translation ($d=3$), clustering in Wasserstein space ($d=10-200$), fairness: ($d=1-200$).\n\nNotice that \"Ave, CelebA'' is not related to any real-life task (see question 2 by reviewer U9wD).\n\n**4. The authors do not compare their approach with the recent approaches for barycenter estimation, e.g., (Kolesov 2024a,b), which have much better performance than that of WIN (Korotin et al., 2022).**\n\nGenerative methods solve a strictly harder problem than OT map estimation. To keep the comparison fair, we compared **generative models only** (see question 5).\n\nFurthermore, the StyleGAN methods in (Kolesov 2024,a,b) solve an OT problem in the StyleGAN latent space with non-quadratic cost. Since none of the methods in our paper supports generic costs, a **comparison is impossible**.\n\nIn terms of model evaluation, nothing would be gained by adding (Kolesov 2024a,b), because WIN has a very similar performance (see Table 3 in (Kolesov 2024a) and our Tables 2-3 for location-scatter data and Fig 5 (a, b) in (Kolesov 2024b) for MNIST data).\n\n**5. StyleGAN is a manifold-constrained generator in (Kolesov et al, 2024a)  and it directly samples from barycenter.**\n\nWe explain why that's not the case. Barycenter generative models \"output a learned barycenter distribution, which can be directly sampled, while non-generative models are limited to transporting existing samples from the input distributions\" (see Sec 3). The pre-trained StyleGAN $G$ in (Kolesov et al, 2024a,b) (which is an *input* of the method, not an output) is **not** the barycenter (see their Fig 3). A barycenter sample can be obtained only by pushing an input distribution sample $x_k$ through the OT map $T_{k, \\phi}(x_k, z)$ (see their Fig 1 (a,b)). In our method, instead, barycenter samples are obtained by pushing latent samples through the map $h$ (no input distribution query needed).\n\n**6. Could you compare with (Kolesov et al, 2024a) on the MNIST dataset?**\n\nSee questions 4 and 5, above.\n\n**7. What about convergence time of the proposed method?**\n\nTraining time on Gaussian location-scatter data ($d=128$) are: 114m56s (SCWB), 93m28s (WIN), 42m 3s (ours). All methods converge in approximately 10-20 minutes, longer times needed only for top-notch performance. On MNIST our method converges in 30 minutes. We will provide full information on training times.\n\n**8. While the dimensionality of tasks grows, the conditional variance grows too [...] and it leads to [...] instability.**\n\nThe conditional variance scales linearly in the number of dimensions $d$. Let $f(Z,s)$ have mean $m_s$ and covariance $\\Sigma_s$ and denote $\\bar{m} := \\sum_{s \\in \\mathcal{S}} w_s m_s$, then by the law of total variance:\n$$\\sum_{i=1}^d \\mathbb{E}\\left[ \\text{Var}(f_i(Z,S) | Z) \\right] \\le \\sum_{i=1}^d \\text{Var}(f_i(Z,S)) = \\sum_{s \\in \\mathcal{S}} w_s \\rm{Tr}(\\Sigma_s) + \\sum_{s \\in \\mathcal{S}} w_s \\| m_s - \\bar{m}\\|^2 = O(d)$$\nThis growth rate is the same as multivariate regression and is unproblematic.\n\n**References**\n- Kingma, Prafulla. \"Glow: Generative flow with invertible 1x1 convolutions.\" NeurIPS 31 (2018).\n- Kolesov et al. \"Estimating barycenters of distributions with neural OT.\" arXiv:2402.03828 (2024a).\n- Kolesov et al. \"Energy-guided continuous entropic barycenter estimation for general costs.\" NeurIPS 37 (2024b): 107513-107546.", "claims_and_evidence": "Most of the claims are supported by clear evidence. However, some points of the submission remain unclear from me:\n\n- **Limitations of the approach.** The authors do not discuss the limitations of the proposed method. It is an important point since normalizing flows are known to suffer from different drawbacks (e.g., constraints on the architecture, high computational expense) and the proposed approach should have inherited the same issues which is not directly stated in the text;\n\n- **Scalability.** It is stated that the approach is applicable for high-dimensional datasets, however, the only high-dimensional experiment considered un the paper corresponds to the computation of barycenter for MNIST images which have the dimension 28x28. It is not clear how the method behaves of experiments with higher dimensions, see, e.g., “Ave, CelebA” experiment from (Kolesov et al, 2024a).\n\n- **Limited comparison.** In the considered experimental setups, the authors do not compare their approach with the recent approaches for barycenter estimation, e.g., (Kolesov 2024a,b). As it is evident from the results reported in this papers, their performance is much better than that of WIN (Korotin et al., 2022), thus, the comparison with them is crucial for understanding the performance of the proposed approach. Since it is possible to retrieve the OT maps using your approach, such a comparison seems to be valid.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria make sense for the considered problem.", "theoretical_claims": "Yes, I skimmed through the theoretical results and their proofs and do not find any issues.", "experimental_designs_or_analyses": "Experimental results are overall valid. However, I have several concerns which I have written in 'claims and evidence' section.", "relation_to_broader_scientific_literature": "The paper proposes a new approach for finding a barycenter of distributions which avoids min-max optimization objective by using the normalizing flows. As far as I know, the usage of normalizing flows in this context is quite novel.", "essential_references_not_discussed": "The authors have cited most of the related literature.", "other_strengths_and_weaknesses": "**Strengths**:\n\n- This method proposes non-adversarial optimization problem for finding Wasserstein-2 barycenter;\n- The authors  offer new intuitive reformulation of the Wasserstein barycenter as expected conditional variance of pushforward bijective transformations;\n- The authors develop new generative model that is able to sample directly from barycenter.\n\n**Weaknesses**:\n\n- The authors do not mention limitations of their approach in the main text. For example, since normalizing flows have drawbacks as constraints of considered transformations , high computational expense and instability, these  shortcomings are being moved to the proposed method. \n- While the dimensionality of tasks grows, the conditional variance grows too. As a consequence of this, the variance of the loss function grows too and it leads to another point of the method’s instability.\n- The comparison with other approaches has limitations, see previous sections.", "other_comments_or_suggestions": "- The authors claim that (Kolesov et al, 2024) does not provide a generative model of barycenter. However, it is not indeed true. StyleGAN is a manifold-constrained generator in the aforementioned method and it directly samples from barycenter.", "questions_for_authors": "- The authors claim that their approach performs well on high-dimensional data taking MNIST dataset as  a high-dimensional experiment. Is the developed method appropriate for the “Ave, CelebA” experiment (See p.9 from (Kolesov et al, 2024) )?\n- The authors provide comparisons with WIN and SCWB. However, (Kolesov et al, 2024) is also suitable for the finding barycenter in MNIST dataset experiment, could you compare with it?\n- What about convergence time of the proposed method? One would like to see comparisons of convergence times for the proposed approach with SCWB, WIN and  (Kolesov et al, 2024).\n\n**References.**\n\nKolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., and Korotin, A. Estimating barycenters of distributions with neural optimal transport. arXiv preprint arXiv:2402.03828, 2024.", "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": "QQzu2kK1YZ", "reviewer": "Reviewer_u7iE", "summary": "This paper introduces a new way to compute optimal transport maps and Wasserstein barycenters using conditional normalizing flows.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, \n\nthanks for the very constructive feedback! Apologies for our terse replies (due to character limit).\n\n**1. Can you make a plot between true and estimated $W_2$ distance?**\n\nSee here for Gaussian OT ($d=64$): https://pasteboard.co/jk8v0aZgFTom.png. We will add such pictures.\n\n**2. Why is there no L-UVP in Tables 2-3?**\n\nThe **L-UVP values for our model are available** (for $d=126$, L-UVP: 1.5% (gaussian), 5.5% (uniform)). We will add them to Tables 2-3. We left them out because Table 5 in (Korotin et al, 2022) does not report them. If appropriate, we can also report the L-UVP for SCWB and WIN (retrained by us).\n\n**3. The rotation matrix in 5.2.4 is only 2d.**\n\nAny high-dim. rotation is 2d in the right coordinates. Since the algo only sees realizations and has no \"rotational\" prior, the experiment does capture scaling behavior.\n\n**4. Theorem 4.3 lacks proof.**\n\nWe wrote a clearer proof, which we sketch below.\n\nWe know that\n\n$$\\sum_s w_s W^2_2(\\bar{\\mu}, \\mu_s) \\le \\sum_s w_s W^2_2(h_\\\\\\#\\lambda, \\mu_s)  \\le \\sum_s w_s \\\\\\| h - f_s \\\\\\|^2_{L^2(\\lambda)}$$\n\nfor all $h:\\mathbb{R}^d \\to \\mathbb{R}^d$ and all $f_s \\in B(\\lambda, \\mu_s)$. Since $\\bar{\\mu} \\in \\mathcal{P}_{ac}(\\mathbb{R}^d)$, Lem 4.1 and arguments in the proof of Thm 4.2 imply that there exist $h \\in B(\\lambda, \\bar{\\mu})$ and $f_s \\in B(\\lambda, \\mu_s)$, such that \n\n$$ W^2_2(\\bar{\\mu}, \\mu_s) = \\\\\\| h - f_s \\\\\\|^2_{L^2(\\lambda)}$$\n\nand, therefore, achieve equality in the first chain of inequalities. This implies (see paper) that $h(Z) = \\mathbb{E}[f(Z,S)|Z]$ and $f$ minimizes $\\sum_{i=1}^d \\mathbb{E}[\\text{Var}(f_i(Z,S))]$.\n\nAssume that there is another function $\\tilde{f}$ with $\\tilde{f}(\\cdot, s) \\in B(\\lambda, \\mu_s)$ achieving minimal conditional variance. Then for $\\tilde{h}(Z) = \\mathbb{E}[\\tilde{f}(Z,S)|Z]$ one has \n\n$$ \\sum_s w_s W^2_2({\\tilde{h}\\\\\\#\\lambda}, \\mu_s) \\le \\sum_s w_s \\\\\\| \\tilde{h} - \\tilde{f_s} \\\\\\|^2_{L^2(\\lambda)} = \\sum_{i=1}^d \\mathbb{E}[\\text{Var}(\\tilde{f}(Z,S))] = \\sum_{i=1}^d \\mathbb{E}[\\text{Var}(f_i(Z,S))] = \\sum_s w_s W^2_2(\\bar{\\mu}, \\mu_s)$$\n\nThen by def of $\\bar{\\mu}$ we must have that $\\tilde{h}_\\\\\\# \\lambda = \\bar{\\mu}$.\n\n**5. Eq (l. 179) is not used in Algorithm 2.**\n\nThe L2 cost in Alg. 2 is exactly Eq (l. 179), where the conditional expectation is given by Eq (l. 173). \n\n**6. Loss Tradeoff in Alg. 1 always introduces bias.**\n\nLet us discuss the proposed example. If $\\mu_s = \\mathcal{N}(m_s, \\sigma_s^2)$, for $s \\in \\{1, 2\\}$, $\\lambda = \\mathcal{N}(0, 1)$ and $f(z, s|\\theta) = \\hat{\\sigma}_s z + \\hat{m}_s$, for $\\theta = (\\hat{m}_1, \\hat{m}_2, \\hat{\\sigma}_1, \\hat{\\sigma}_2)$, then the model admits four MLEs: $ \\theta^* \\in \\{(m_1, m_2, \\sigma_1, \\sigma_2), (m_1, m_2, -\\sigma_1, \\sigma_2), (m, m_2, \\sigma_1, -\\sigma_2), (m_1, m_2, -\\sigma_1, -\\sigma_2)\\}$.\n\nThe model OT map, $\\hat{T}(x) = f(f^{-1}(x, 1), 2)$, is the right one only for two MLEs. In agreement with Thm 4.2, the right MLEs minimize $ L^2(z |\\theta) = ((\\hat{\\sigma}_1 - \\hat{\\sigma}_2) z + (\\hat{m}_1 - \\hat{m}_2))^2$.\n\nIn the reviewer's example ($\\sigma_1 = \\sigma_2 = 1$ and $m_1 \\neq m_2$), at each $t$ in Alg. 1 we take a GD step towards the minimum of $ -\\log(p(x,s|\\theta)) + \\zeta_t L^2(z |\\theta)$:\n\n$$ \\hat{m^*}_1(t) = \\frac{m_1 + 2 \\zeta_t \\hat{m}_2}{1 + 2 \\zeta_t}, \\quad \\hat{m^*}_2(t) = \\frac{m_2 + 2 \\zeta_t \\hat{m}_1}{1 + 2 \\zeta_t}.$$\n\nBut $(\\hat{m^*}_1(t), \\hat{m^*}_2(t)) \\to (m_1, m_2)$ as $\\zeta_t \\downarrow 0$, therefore at convergence there will be **no bias**. Empirically we do not observe any bias.\n\n**7. The argumentation regarding the annealing structure is incorrect.**\n\nThanks for the insightful remarks. We do agree and we will remove any reference to annealing. We also saw empirically that minimizing the L2 cost with a likelihood regularizer works well, which shows that MLEs may lie on a connected manifold.\n\n**8. UATs for NFs use arbitrarily ill-conditioned networks.**\n\nThis is a valid concern, but it applies to all UATs, which do not provide numerical stability guarantees. It is surely possible to construct distributions for which our networks are ill-conditioned. Empirically, the approach works well.\n\n**9. Other remarks**\n\n- We do have Alg. 1 results for uniform data (e.g. L-UVP=3.8%, BW-UVP=0.7%, for $d=126$). We will add them.\n- Rectified flows also use invertible nets, but do not compute *Wasserstein* OT maps nor barycenters.\n- Alg. 1 can be adjusted for generic $p>1$ in l.202. It's unclear if the same adjustment gives the $W_p$ barycenter in Alg. 2.\n- Any application in the paper intro is feasible for our method. See question 1 by reviewer U9wD.\n- Our method is not suitable for entropic OT, since we model Monge maps, not generic couplings.\n\n**References**\n\n- Korotin et al. \"Wasserstein iterative networks for barycenter estimation.\" arXiv:2201.12245 2023\n- Liu, Chengyue, and Qiang \"Flow straight and fast\" arXiv:2209.03003 (2022).\n- Pooladian et al. \"Multisample flow matching\" arXiv:2304.14772 (2023).", "claims_and_evidence": "By order of appearance:\n\n1. Correctness of OT and barycenter problems: I have doubts on the barycenter derivation, as well as doubts on the construction of the loss, see theoretical claims section below.\n2. Normalizing flows: There is some debate in the literature regarding universality, what papers are the authors referring to?\n3. Yes, the method yields explicit maps, which is a great strength of this method.\n4. Scaling: The presented evidence is convincing.\n\n\n### Update after rebuttal: Main concern fixed\n\nThe barycenter proof was fixed, minor points see below.", "methods_and_evaluation_criteria": "## OT Maps\n\nAlgorithm 1 is only evaluated on Gaussian data, which should be relatively easy to learn for a flow (e.g. Draxler et al. 2022 show that RealNVP flows converge exponentially fast to Gaussian data in terms of number of layers). I think estimating the Wasserstein distance for more complicated tasks would be a great add-on, e.g. pick any of the OT applications in the introduction. I think this is somewhat compensated by the Barycenter experiments, since OT estimation is a prerequisite for the barycenter computation.\n\nCan you make a plot between the true Wasserstein distance and the distance estimated by the model?\n\n## Barycenters\n\nWhy is there no L-UVP in Tables 2 and 3? The displayed loss is only a lower bound on the true metric.\n\nGiven that the rotation matrix in the 64-dimensional experiment in 5.2.4 is 2-dimensional, does the experiment really measure the scaling behavior?\n\n\n## References\n\nDraxler, Felix, Christoph Schnörr, and Ullrich Köthe. \"Whitening convergence rate of coupling-based normalizing flows.\" Advances in Neural Information Processing Systems 35 (2022): 37241-37253.\n\n### Update after rebuttal\n\nQuestions answered, experimental evidence still limited, but enough in my opinion.", "theoretical_claims": "I did not check any statements in terms of where suitable maps exist, i.e. whether they are measurable. Instead I assumed that all densities are continuous wrt. Lebesgue measure.\n\n## Theorem 4.2\n\nTheorem 4.2 is convincing to me.\n\n## Theorem 4.3 lacks proof, in my understanding\n\nLet's first establish common ground, please confirm that I understand the proof of Theorem 4.3 correctly:\n\n1. l. 707-715 show that h exists (since the Barycenter exists).\n2. l. 715-716 argue that h establishes equality in Eq. 4 (since h + f(., s) are the solution of Theorem 4.2 to the OT between the Barycenter and each mu_s).\n3. l. 717-722 show that h minimizes Eq. (l. 720) (if there was a better h, then it would not be the solution to the OT).\n4. l. 722-724: Since h is the solution of that minimization, it is the average over the OT solution for each s (first statement of Thm 4.3).\n5. Because p=2, we can separate the optimization tasks over dimensions.\n\nHowever, I think the proof is missing the most important step: The map that satisfies Eq (l. 173) and Eq (l. 179) points to the barycenter has $h_\\sharp \\lambda = \\bar \\mu$. In addition, it seems that the second equation is not even used in Algorithm 2.\n\nIn other words: if one wants to train a model with these two as their loss, then one also has to show that the solution of this optimization is the barycenter.\n\nThis is my most important criticism, addressing it will increase my rating. \n\n\n## Loss Tradeoff in Alg. 1 always introduces bias\n\nI think that the argumentation regarding the annealing structure is incorrect and of different nature than simulated annealing. This is easy to see for the following counterexample: Let mu_1 and mu_2 be two Gaussian distributions with means $\\pm 1$ and standard deviations $1$, and let the latent lambda also be a Gaussian distribution. Then the loss offers a closed-form solution which reveals that the learned means are off proportional to $\\sqrt{\\chi}$ (I might be wrong in the scaling).\n\nIn other words: The **optimal solution of the loss in Alg. 1 over all bijections will not be the optimal transport solution**. I would conjecture that it comes closer as $\\chi \\to 0$, but it is unclear how fast and what the tradeoff will be if presented arbitrary distributions.\n\nWhy different from annealing? In terms of suboptimal minima, I think that one can always move between different optima by sending infinitesimal packages of mass from one place to another -- so the loss barriers really are vanishing (and this effect seems to transfer to neural network parameterizations, see \"mode connectivity\" as per Garipov et al 2018, and Draxler et al 2018).\n\nI think this will be an easy fix.\n\n## Universality\n\nI think there is a logical gap that the existing literature has not fully closed in terms of universality of normalizing flows. With the statement \"normalizing flows, which are universal approximators for bijections and are thus the most natural generative models to employ in this context\", the authors probably refer to the work by Teshima et al. 2020. However, as Koehler et al. 2021 and Draxler et al. 2024 point out, the underlying proofs use arbitrarily ill-conditioned networks (become arbitrarily ill-conditioned as the error decreases, e.g. Section 5.1 in Draxler et al. 2024).\n\n\n\n## References\n\nGaripov, Timur, et al. \"Loss surfaces, mode connectivity, and fast ensembling of dnns.\" Advances in neural information processing systems 31 (2018).\n\nDraxler, Felix, et al. \"Essentially no barriers in neural network energy landscape.\" International conference on machine learning. PMLR, 2018.\n\nTeshima, Takeshi, et al. \"Coupling-based invertible neural networks are universal diffeomorphism approximators.\" Advances in Neural Information Processing Systems 33 (2020): 3362-3373.\n\nKoehler, Frederic, Viraj Mehta, and Andrej Risteski. \"Representational aspects of depth and conditioning in normalizing flows.\" International Conference on Machine Learning. PMLR, 2021.\n\nDraxler, Felix, et al. \"On the universality of volume-preserving and coupling-based normalizing flows.\" ICML (2024).\n\n\n### Update after rebuttal\n\n- Theorem 4.3 fixed.\n- Bias: The authors argue that weight schedule on loss leads to correct solution, for which they provide additional evidence, which should be included in the paper and is good enough for now.\n- Universality: Reached common ground on understanding of the literature", "experimental_designs_or_analyses": "I am unsure about the usefulness of the UVP metrics and hope for input from another reviewer.", "relation_to_broader_scientific_literature": "Optimal Transport is a principled mathematical framework for finding transport maps between distributions that adhere to external cost. It has a broad set of applications, but finding optimal transport maps in high dimensions remains an open challenge.", "essential_references_not_discussed": "Rectified Flows (and mini-batch flow matching) also try approximating the optimal transport between two distributions. This is not mentioned or compared to in the paper. The nature of the competing loss terms makes this challenging, of course.\n\nReferences:\n\nPooladian, Aram-Alexandre, et al. \"Multisample flow matching: Straightening flows with minibatch couplings.\" arXiv preprint arXiv:2304.14772 (2023).\n\nLiu, Xingchao, Chengyue Gong, and Qiang Liu. \"Flow straight and fast: Learning to generate and transfer data with rectified flow.\" arXiv preprint arXiv:2209.03003 (2022).\n\n\n### Update after rebuttal: Mostly clarified\n\n- The authors clarified the relation to Rectified Flows (Liu et al. 2022), and pointed me to the explanation in that paper.\n- Regarding (Pooladian et al. 2023), the authors did not provide an explicit explanation, but I think the argument concerning rectified flows persists that there is no known cost that is minimized, and so the approaches are different", "other_strengths_and_weaknesses": "The flow construction is innovative and easy to follow, nice work!", "other_comments_or_suggestions": "l. 122 right: Add citation to McCann and Gangbo.\n\nPlease add links from the proofs of theorems in the main text to the appendix.", "questions_for_authors": "Please take the criticism regarding the theory seriously. I will adjust my recommendation when the points are addressed.\n\n## Extension ideas\n\nAre there any of the real-world applications mentioned in the introduction feasible with this new method?\n\nWasserstein barycenters outside p=2? Relatedly, the \"specialization\" of Algorithm 1 to $p=2$ mentioned in the first paragraph of 4.3.2 affects only l. 202, right?\n\nIs it possible to extend the algorithms to entropic OT? I guess this does not make sense since this produces non-unique couplings, right?\n\n\n### Updates after rebuttal: Answers provided to all points", "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": "LMMDH8fnM2", "reviewer": "Reviewer_aK7B", "summary": "An alternative formulation of the p-Wasserstein distance is introduced and it consists in a constrained L^p minimisation problem for a given latent distribution. This alternative formulation allows the authors to both i) directly solve the Wasserstein primal minimization problem via stochastic gradient decent and ii) exploit conditional normalizing flows in order to compute the optimal transport map (already exiting here, since the source and destination measures are absolutely continuous w.r.t. to the Lebesgue measure) as well as the Wasserstein barycenters. Two pseudo-algorithms are well detailed and experiments on two datasets (Swiss roll and Mnist) are used to support the two main claims: the proposed approach allows one to compute barycenters from hundreds of input distributions and it is efficient in high-dimensional spaces.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, \n\nthanks for taking the time to review our submission and for the kind appreciation of our contribution. We have carefully addressed all your comments below. According to the ICML review guidelines, we are not allowed to submit a revised manuscript at this point of the review process, but below we discuss in detail the edits we will make based on your suggestions in the camera-ready version of the paper, if accepted.\n\n**1. Some comparisons with standard OT solver would be beneficial: barycenters similar to those in Figures 3-4 could be computed via standard discrete OT (e.g. in Python Optimal Transport) and it would be interesting to perform qualitative and quantitative comparisons with your method. E.g. what about the running times?**\n\nThe quantitative comparison between discrete OT and neural continuous OT methods has already been investigated extensively in the literature. The consensus is that discrete OT methods perform poorly in high dimensions, due to the required discretization of the input distributions. Specifically, performance degrades above $d=16$ (see, for instance, Figure 5 of (Fan et al, 2020) and Tables 1, 2, and 3 of (Korotin et al, 2021)), meaning that any discrete OT technique is assured to perform poorly on the MNIST 0/1 barycenter task shown in Figure 4 ($d=784$).\n\nOn the other hand, discrete OT methods will certainly achieve good performance on the Swiss roll dataset (Figure 3) and decent computational times. Nevertheless, the point of presenting Figure 3 is not to stake a performance claim, but rather to verify visually that our conditional normalizing flow model performs well on highly non-linear distributions with complex support. The real performance evaluation is done on more high-dimensional problems in the subsequent sections.\n\nIn conclusion, we think that a direct comparison with discrete OT methods in Figures 3 and 4 would not be meaningful.\n\n**2. Consider Table 1, for instance. Since you simulate data from (high-dimensional) Gaussian distributions and you know the actual Wasserstein distance between them, you might calculate the absolute difference between the Wasserstein loss you come up with with your method, from each data sample, and the true one! Why to use the upper/lower bound of another quantity in this case? It is misleading.**\n\nWe completely agree with the reviewer. We re-ran Algorithm 1 on high-dimensional Gaussian data in the last few days and found that indeed our estimated Wasserstein distance converges to the true one very quickly. The plot for $d=64$ can be found here: https://pasteboard.co/jk8v0aZgFTom.png.\n\nThis (and similar plots) will be added to the revised paper.\n\n**3. Moreover, and this is the weak point of the paper in my view, in Tables 2-3, you only compare your methods with Win and $SCW_2B$ in terms of the lower bound $BW_2^2-UVP$. Now, a lower bound is a lower bound, what about the upper bound? Their UVP might be better than yours despite having a worse lower bound.**\n\nThis is a very important point, also raised by another reviewer. We entirely agree. Please see our answer to question 2 by reviewer u7iE.\n\n**4. I fell that Theorem 4.2 might open the way to other approaches, adopting or not Normalizing Flows.**\n\nThis is definitely worth pursuing. We only mention that, since OT maps in Threorem 4.2 are ultimately computed by inverting the model, only invertible architectures are valid candidates for this approach.\n\n**5. Just a couple of parentheses missing in the pseudo-codes such as the one at the very end of line 257.**\n\nThanks. We will add them in the revised paper.\n\n**References**\n\nFan, Amirhossein, and Yongxin. \"Scalable computations of wasserstein barycenter via input convex neural networks.\" arXiv:2007.04462 (2020).\n\nKorotin et al. \"Continuous wasserstein-2 barycenter estimation without minimax optimization.\" arXiv:2102.01752 (2021).", "claims_and_evidence": "From a theoretical point of view, the two main claims of the paper (i.e. Theorems 4.2 and 4.3) are clear and I did not find any error in the proofs. About the experimental results: everything is convincing except for some comparisons, essentially in Tables 2-3  as I will detail below.", "methods_and_evaluation_criteria": "Yes, most of the time the methods/evaluation criteria do make sense, although some comparisons with standard OT solver would be beneficial: barycenters similar to those in Figures 3-4 could be computed via standard discrete OT (e.g. in Python Optimal Transport) and it would be interesting to perform qualitative and quantitative comparisons with your method. E.g. what about the running times?", "theoretical_claims": "I checked the correctness of proofs of part of Lemma A.3 and entirely revised proofs of Theorems 4.2 and 4.3. As I said everything looks correct at my eyes.", "experimental_designs_or_analyses": "The experimental design is quite clear but it coud be improved. Consider Table 1, for instance. Since you simulate data from (high-dimensional) Gaussian distributions and you know the actual Wasserstein distance between them, you might calculate the absolute difference between the Wasserstein loss you come up with with your method, from each data sample, and the true one! Why to use the upper/lower bound of another quantity in this case? It is misleading. Moreover, and this is the weak point of the paper in my view, in Tables 2-3, you only compare your methods with Win and SCW_2B in terms of the lower bound BW_2^2-UVP. Now, a lower bound is a lower bound, what about the upper bound? Their UVP might be better than yours despite having a worse lower bound. Am-I missing something?", "relation_to_broader_scientific_literature": "The relation with previous/related literature is correct, in my view.", "essential_references_not_discussed": "I don't think there are essential references not discusses, but I might be wrong.", "other_strengths_and_weaknesses": "I liked very much this paper. It is well written, mathematics are elegant and the theoretical arguments are solid. I fell that Theorem 4.2 might  open the way to other approaches, adopting or not Normalizing Flows and, by the way, I consider that being able to solve OT problems (distance and barycenters) via standard SGD, in the primal, is an important advance. \n\nAs I sad, however I think that the experimental part could be improved in order to asses the interest of adopting this approach in place of alternatives.", "other_comments_or_suggestions": "Just a couple of parentheses missing in the pseudo-codes such as the one at the very end of line 257.", "questions_for_authors": "I would like to see this paper published at ICML. I will tentatively recommend is for \"weak accept\" but if you can do the one or two edits (Tables 1-2-3) that I suggested for the experimental part or alternatively show me that I am wrong about my concerns (it might be the case) I am ready to raise my note to \"accept\".", "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:

Meta-review: In this study, the authors propose a conditional normalizing flow method to compute optimal transport maps and Wasserstein barycenters. All four reviewers, including AC, agree that the proposed method is interesting and solid in theory. However, AC and Reviewer 91Cy require the authors to polish the experimental part and introduce the implementation details clearly. In summary, AC tends to accept this work.
Decision: Accept (poster)