Original: C ONTINUOUS W ASSERSTEIN -2 B ARYCENTER E STIMATION WITHOUT M INIMAX O PTIMIZATION
Revised: C ONTINUOUS W ASSERSTEIN -2 B ARYCENTER E STIMATION WITHOUT M INIMAX O PTIMIZATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport.
Revised: Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport.

Original: In this paper, we present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures, which are not restricted to being discrete.
Revised: In this paper, we present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures, which are not restricted to being discrete.

Original: While past approaches rely on entropic or quadratic regularization, we employ input convex neural networks and cycle-consistency regularization to avoid introducing bias.
Revised: While past approaches rely on entropic or quadratic regularization, we employ input convex neural networks and cycle-consistency regularization to avoid introducing bias.

Original: As a result our approach does not resort to minimax optimization.
Revised: As a result our approach does not resort to minimax optimization.

Original: We provide theoretical analysis on error bounds as well as empirical evidence of the effectiveness of the proposed approach in low-dimensional qualitative scenarios and high-dimensional quantitative experiments.
Revised: We provide theoretical analysis on error bounds as well as empirical evidence of the effectiveness of the proposed approach in low-dimensional qualitative scenarios and high-dimensional quantitative experiments.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Wasserstein barycenters have become popular due to their ability to represent the average of probability measures in a geometrically meaningful way.
Revised: Wasserstein barycenters have become popular due to their ability to represent the average of probability measures in a geometrically meaningful way.

Original: Techniques for computing Wasserstein barycenters have been successfully applied to many computational problems.
Revised: Techniques for computing Wasserstein barycenters have been successfully applied to many computational problems.

Original: In image processing, Wasserstein barycenters are used for color and style transfer (Rabin et al., 2014; Mroueh, 2019), and texture synthesis (Rabin et al., 2011).
Revised: In image processing, Wasserstein barycenters are used for color and style transfer (Rabin et al., 2014; Mroueh, 2019), and texture synthesis (Rabin et al., 2011).

Original: In geometry processing, shape interpolation can be done by computing barycenters (Solomon et al., 2015).
Revised: In geometry processing, shape interpolation can be done by computing barycenters (Solomon et al., 2015).

Original: In online machine learning, barycenters are used for aggregating probabilistic predictions of experts (Korotin et al., 2019b).
Revised: In online machine learning, barycenters are used for aggregating probabilistic predictions of experts (Korotin et al., 2019b).

Original: Within the context of Bayesian inference, the barycenter of subset posteriors converges to the full data posterior, thus enabling efﬁcient computational methods based on ﬁnding the barycenters (Srivastava et al., 2015; 2018).
Revised: Within the context of Bayesian inference, the barycenter of subset posteriors converges to the full data posterior, thus enabling efﬁcient computational methods based on ﬁnding the barycenters (Srivastava et al., 2015; 2018).

Original: Fast and accurate barycenter algorithms exist for discrete distributions (see Peyré et al.
Revised: Fast and accurate barycenter algorithms exist for discrete distributions (see Peyré et al.

Original: for a survey), while for continuous distributions the situation is more difﬁcult and remains unexplored until recently (Li et al., 2020;
Revised: for a survey), while for continuous distributions the situation is more difﬁcult and remains unexplored until recently (Li et al., 2020;

Original: Fan et al., 2020;
Revised: Fan et al., 2020;

Original: Cohen et al., 2020).
Revised: Cohen et al., 2020).

Original: The discrete methods scale poorly with the number of support points of the barycenter and thus cannot approximate continuous barycenters well, especially in high dimensions.
Revised: The discrete methods scale poorly with the number of support points of the barycenter and thus cannot approximate continuous barycenters well, especially in high dimensions.

Original: In this paper, we present a method to compute Wasserstein-2 barycenters of continuous distributions based on a novel regularized dual formulation where the convex potentials are parameterized by input convex neural networks (Amos et al., 2017).
Revised: In this paper, we present a method to compute Wasserstein-2 barycenters of continuous distributions based on a novel regularized dual formulation where the convex potentials are parameterized by input convex neural networks (Amos et al., 2017).

Original: Our algorithm is straightforward without introducing bias (e.g. Li et al.
Revised: Our algorithm is straightforward without introducing bias (e.g. Li et al.

Original: or requiring minimax optimization (e.g. Fan et al.
Revised: or requiring minimax optimization (e.g. Fan et al.

Original: This is made possible by combining a new congruence regularizing term combined with cycle-consistency regularization (Korotin et al., 2019a).
Revised: This is made possible by combining a new congruence regularizing term combined with cycle-consistency regularization (Korotin et al., 2019a).

Original: As we will show in the analysis, thanks to the properties of
Revised: As we will show in the analysis, thanks to the properties of

Original: Wasserstein-2 distances, the gradients of the resulting convex potentials “push” the input distributions close to the true barycenter, allowing good approximation of the barycenter.
Revised: Wasserstein-2 distances, the gradients of the resulting convex potentials “push” the input distributions close to the true barycenter, allowing good approximation of the barycenter.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: We denote the set of all Borel probability measures on R D with ﬁnite second moment by P 2 ( R D ) .
Revised: We denote the set of all Borel probability measures on R D with ﬁnite second moment by P 2 ( R D ) .

Original: We use P 2 , ac ( R D ) ⊂ P 2 ( R D ) to denote the subset of all absolutely continuous measures (w.r.t.
Revised: We use P 2 , ac ( R D ) ⊂ P 2 ( R D ) to denote the subset of all absolutely continuous measures (w.r.t.

Original: the
Revised: the

Original: Lebesgue measure).
Revised: Lebesgue measure).

Original: Wasserstein-2 distance.
Revised: Wasserstein-2 distance.

Original: For P , Q ∈ P 2 ( R D ) , the Wasserstein-2 distance is deﬁned by
Revised: For P , Q ∈ P 2 ( R D ) , the Wasserstein-2 distance is deﬁned by

Original: [Equation]
Revised: [Equation]

Original: where Π( P , Q ) is the set of probability measures on R D × R D whose marginals are P , Q , respectively.
Revised: where Π( P , Q ) is the set of probability measures on R D × R D whose marginals are P , Q , respectively.

Original: This deﬁnition is known as Kantorovich’s primal form of transport distance (Kantorovitch, 1958).
Revised: This deﬁnition is known as Kantorovich’s primal form of transport distance (Kantorovitch, 1958).

Original: The Wasserstein-2 distance W 2 is well-studied in the theory of optimal transport (Brenier, 1991;
Revised: The Wasserstein-2 distance W 2 is well-studied in the theory of optimal transport (Brenier, 1991;

Original: McCann et al., 1995).
Revised: McCann et al., 1995).

Original: In particular, it has a dual formulation (Villani, 2003):
Revised: In particular, it has a dual formulation (Villani, 2003):

Original: [Equation]
Revised: [Equation]

Original: where the minimum is taken over all the convex functions (potentials) ψ : R D → R ∪ {∞} , and ψ ( y ) =
Revised: where the minimum is taken over all the convex functions (potentials) ψ : R D → R ∪ {∞} , and ψ ( y ) =

Original: max x ∈ R D (cid:0) (cid:104) x, y (cid:105) − ψ ( x ) (cid:1) :
Revised: max x ∈ R D (cid:0) (cid:104) x, y (cid:105) − ψ ( x ) (cid:1) :

Original: R D → R ∪ {∞} is the convex conjugate of ψ (Fenchel, 1949), which is also a convex function.
Revised: R D → R ∪ {∞} is the convex conjugate of ψ (Fenchel, 1949), which is also a convex function.

Original: The optimal potential ψ ∗ is deﬁned up to an additive constant.
Revised: The optimal potential ψ ∗ is deﬁned up to an additive constant.

Original: Brenier (1991) shows that if P does not give mass to sets of dimensions at most D − 1 , then the optimal plan π is uniquely determined by π ∗ =
Revised: Brenier (1991) shows that if P does not give mass to sets of dimensions at most D − 1 , then the optimal plan π is uniquely determined by π ∗ =

Original: [ id R D , T ∗ ] (cid:93) P , where T ∗ : R D
Revised: [ id R D , T ∗ ] (cid:93) P , where T ∗ : R D

Original: → R D is the unique solution to the Monge’s problem
Revised: → R D is the unique solution to the Monge’s problem

Original: [Equation]
Revised: [Equation]

Original: The connection between T ∗ and the dual formulation (2) is that T ∗ = ∇ ψ ∗ , where ψ ∗ is the optimal solution of (2).
Revised: The connection between T ∗ and the dual formulation (2) is that T ∗ = ∇ ψ ∗ , where ψ ∗ is the optimal solution of (2).

Original: Additionally, if Q does not give mass to sets of dimensions at most D − 1 , then T ∗ is invertible and
Revised: Additionally, if Q does not give mass to sets of dimensions at most D − 1 , then T ∗ is invertible and

Original: [Equation]
Revised: [Equation]

Original: In particular, the above discussion applies to the case where P , Q ∈ P 2 ,ac ( R D ) .
Revised: In particular, the above discussion applies to the case where P , Q ∈ P 2 ,ac ( R D ) .

Original: Wasserstein-2 barycenter.
Revised: Wasserstein-2 barycenter.

Original: Let P 1 , . .
Revised: Let P 1 , . .

Original: . , P N ∈ P 2 ,ac ( R D ) .
Revised: . , P N ∈ P 2 ,ac ( R D ) .

Original: Then, their barycenter w.r.t.
Revised: Then, their barycenter w.r.t.

Original: weights α 1 , . .
Revised: weights α 1 , . .

Original: . , α N ( α n > 0 and (cid:80)
Revised: . , α N ( α n > 0 and (cid:80)

Original: Nn =1 α n = 1 ) is
Revised: Nn =1 α n = 1 ) is

Original: [Equation]
Revised: [Equation]

Original: Throughout this paper, we assume that at least one of P 1 , . .
Revised: Throughout this paper, we assume that at least one of P 1 , . .

Original: . , P N ∈ P 2 ,ac ( R D ) has bounded density.
Revised: . , P N ∈ P 2 ,ac ( R D ) has bounded density.

Original: Under this assumption, P is unique and absolutely continuous, i.e., P ∈ P 2 ,ac ( R D ) , and it has bounded density (Agueh & Carlier, 2011, Deﬁnition 3.6 & Theorem 5.1).
Revised: Under this assumption, P is unique and absolutely continuous, i.e., P ∈ P 2 ,ac ( R D ) , and it has bounded density (Agueh & Carlier, 2011, Deﬁnition 3.6 & Theorem 5.1).

Original: For n ∈ { 1 , 2 , . . . , N } , let ( ψ ∗ n , ψ ∗ n ) be the optimal pair of (mutually) conjugate potentials that transport P n to P , i.e., ∇ ψ ∗ n (cid:93) P n = P and ∇ ψ ∗ n (cid:93) P = P n .
Revised: For n ∈ { 1 , 2 , . . . , N } , let ( ψ ∗ n , ψ ∗ n ) be the optimal pair of (mutually) conjugate potentials that transport P n to P , i.e., ∇ ψ ∗ n (cid:93) P n = P and ∇ ψ ∗ n (cid:93) P = P n .

Original: Then { ψ ∗ n } satisfy
Revised: Then { ψ ∗ n } satisfy

Original: D (cid:88) n = α n ∇ ψ ∗ n ( x ) = x and

Original: D (cid:88) n= α n ψ
Revised: N (cid:88) n

Original: [Equation]
Revised: [Equation]

Original: n
Revised: and

Original: [Equation]

Original: = (cid:107)
Revised: N (cid:88)

Original: [Equation]
Revised: [Equation]

Original: 5) for all x ∈ R D (Agueh & Carlier, 2011; Álvarez-Esteban et al., 2016).
Revised: for all x ∈ R D (Agueh & Carlier, 2011; Álvarez-Esteban et al., 2016).

Original: Since optimal potentials are deﬁned up to a constant, for convenience, we set c = 0 .
Revised: Since optimal potentials are deﬁned up to a constant, for convenience, we set c = 0 .

Original: The condition (5) serves as the basis for our algorithm for computing Wasserstein-2 barycenters.
Revised: The condition (5) serves as the basis for our algorithm for computing Wasserstein-2 barycenters.

Original: We say that potentials ψ 1 , . . .
Revised: We say that potentials ψ 1 , . . .

Original: , ψ N are congruent w.r.t.
Revised: , ψ N are congruent w.r.t.

Original: weights α 1 , . .
Revised: weights α 1 , . .

Original: , α n if their conjugate potentials satisfy (5), i.e., (cid:80) Dn =1 α n ψ n ( x ) = (cid:107) x (cid:107) 2 2 for all x ∈ R D .
Revised: , α n if their conjugate potentials satisfy (5), i.e., (cid:80) Dn =1 α n ψ n ( x ) = (cid:107) x (cid:107) 2 2 for all x ∈ R D .

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Most algorithms in the ﬁeld of computational optimal transport are designed for the discrete setting where the input distributions have ﬁnite support; see the recent survey by Peyré et al.
Revised: Most algorithms in the ﬁeld of computational optimal transport are designed for the discrete setting where the input distributions have ﬁnite support; see the recent survey by Peyré et al.

Original: for discussion.
Revised: for discussion.

Original: A particular popular line of algorithms are based on entropic regularization that gives rise to the famous Sinkhorn iteration (Cuturi, 2013; Cuturi & Doucet, 2014).
Revised: A particular popular line of algorithms are based on entropic regularization that gives rise to the famous Sinkhorn iteration (Cuturi, 2013; Cuturi & Doucet, 2014).

Original: These methods are typically limited to a support of 10 5 − 10 6 points before the problem becomes computationally infeasible.
Revised: These methods are typically limited to a support of 10 5 − 10 6 points before the problem becomes computationally infeasible.

Original: As a consequence, discrete  methods  cannot provide precise approximation of continuous barycenters in high dimensions, since a large number of samples is needed; see experiments in Fan et al.
Revised: Similarly, discrete barycenter methods (Cuturi & Doucet, 2014), particularly the ones that rely on a ﬁxed support for the barycenter (Dvurechenskii et al., 2018; Staib et al., 2017), cannot provide precise approximation of continuous barycenters in high dimensions, since a large number of samples is needed; see experiments in Fan et al.

Original: (2020, §4.3) for an example.
Revised: (2020, §4.3) for an example.

Original: Thus we focus on the existing literature in the continuous setting.
Revised: Thus we focus on the existing literature in the continuous setting.

Original: Computation of Wasserstein-2 distances and maps.
Revised: Computation of Wasserstein-2 distances and maps.

Original: Seguy et al.
Revised: Genevay et al.

Original: (2017) propose one of  the ﬁrst methods for continuous computational transport.
Revised: (2016) demonstrate the possibility of computing Wasserstein distances given only sample access to the distributions by parameterizing the dual potentials as functions in the reproducing kernel Hilbert spaces.

Revised: Based on this realization, Seguy et al.

Original: They tackle (2) by approximating the potentials with neural networks, using entropic or L 2 regularization w.r.t.
Revised: (2017) propose a similar method but use neural networks to parameterize the potentials, using entropic or L 2 regularization w.r.t.

Original: P × Q to keep the potentials approximately conjugate.
Revised: P × Q to keep the potentials approximately conjugate.

Original: The transport map is recovered from optimized potentials via barycentric projection.
Revised: The transport map is recovered from optimized potentials via barycentric projection.

Original: As we note in §2, W 2 enjoys many useful theoretical properties.
Revised: As we note in §2, W 2 enjoys many useful theoretical properties.

Original: For example, the optimal potential ψ ∗ is convex, and the corresponding optimal transport map is given by ∇ ψ ∗ .
Revised: For example, the optimal potential ψ ∗ is convex, and the corresponding optimal transport map is given by ∇ ψ ∗ .

Original: By exploiting these properties, Makkuva et al.
Revised: By exploiting these properties, Makkuva et al.

Original: (2019) propose a minimax optimization algorithm for recovering transport maps, using input convex neural networks (ICNNs)
Revised: (2019) propose a minimax optimization algorithm for recovering transport maps, using input convex neural networks (ICNNs)

Original: (Amos et al., 2017) to approximate the potentials.
Revised: (Amos et al., 2017) to approximate the potentials.

Original: An alternative to entropic regularization is the cycle-consistency regularization proposed by Korotin et al. (2019a).
Revised: An alternative to entropic regularization is the cycle-consistency regularization proposed by Korotin et al. (2019a).

Original: It uses the property that the gradients of optimal dual potentials are inverses of each other.
Revised: It uses the property that the gradients of optimal dual potentials are inverses of each other.

Original: The imposed regularizer requires integration only over the marginal measures P and Q , instead of over P × Q as required by entropy-based alternatives.
Revised: The imposed regularizer requires integration only over the marginal measures P and Q , instead of over P × Q as required by entropy-based alternatives.

Original: Their method converges faster than the minimax method since it does not have an inner optimization cycle. 
Revised: Their method converges faster than the minimax method since it does not have an inner optimization cycle. Xie et al.

Revised: (2019) propose using two generative models with a shared latent space to implicitly compute the optimal transport correspondence between P and Q .

Revised: Based on the obtained correspondence, the authors are able to compute the optimal transport distance between the distributions.

Original: Computation of Wasserstein-2 barycenters.
Revised: Computation of Wasserstein-2 barycenters.

Original: A few recent techniques tackle the barycenter problem (4) using continuous rather than discrete approximations of the barycenter:
Revised: A few recent techniques tackle the barycenter problem (4) using continuous rather than discrete approximations of the barycenter:

Original: • M EASURE - BASED ( GENERATIVE ) OPTIMIZATION :
Revised: • M EASURE - BASED ( GENERATIVE ) OPTIMIZATION :

Original: Problem (4) optimizes over probability measures This can be done
Revised: Problem (4) optimizes over probability measures.

Original:  using the generic algorithm by Cohen et al.
Revised: This can be done using the generic algorithm by Cohen et al.

Original:  who employ generative networks to compute barycenters w.r.t.
Revised: (2020) who employ generative networks to compute barycenters w.r.t.

Original: arbitrary discrepancies
Revised: arbitrary discrepancies.

Original: They test their method with the maximum mean discrepancy (MMD) and Sinkhorn divergence
Revised: They test their method with the maximum mean discrepancy (MMD) and Sinkhorn divergence.

Original: This approach suffers from the usual limitations of generative models such as mode collapse
Revised: This approach suffers from the usual limitations of generative models such as mode collapse.

Original: Applying it to W 2 barycenters requires estimation of W 22 ( P n , P ) .
Revised: Applying it to W 2 barycenters requires estimation of W 22 ( P n , P ) .

Original: Fan et al.
Revised: Fan et al.

Original: (2020) test this approach using the minimax method by Makkuva et al.
Revised: (2020) test this approach using the minimax method by Makkuva et al.

Original: , but they end up with a challenging min-max-min problem.
Revised: , but they end up with a challenging min-max-min problem.

Original: • P OTENTIAL - BASED OPTIMIZATION :
Revised: • P OTENTIAL - BASED OPTIMIZATION :

Original: Li et al.
Revised: Li et al.

Original: (2020) recover the optimal potentials { ψ ∗ n } via a non-minimax regularized dual formulation.
Revised: (2020) recover the optimal potentials { ψ ∗ n } via a non-minimax regularized dual formulation.

Original: No generative model is needed: the barycenter is recovered by pushing forward measures using gradients of potentials or by barycentric projection.
Revised: No generative model is needed: the barycenter is recovered by pushing forward measures using gradients of potentials or by barycentric projection.

Original:  M ETHODS
Revised:  M ETHODS

Original: Inspired by Li et al.
Revised: Inspired by Li et al.

Original: we use a potential-based approach and recover the barycenter by using gradients of the potentials as pushforward maps.
Revised: we use a potential-based approach and recover the barycenter by using gradients of the potentials as pushforward maps.

Original: The main differences are: (1) we restrict the potentials to be convex, (2) we enforce congruence via a regularizing term, and (3) our formulation does not introduce bias, meaning the optimal solution of our formulation gives the true barycenter.
Revised: The main differences are: (1) we restrict the potentials to be convex, (2) we enforce congruence via a regularizing term, and (3) our formulation does not introduce bias, meaning the optimal solution of our formulation gives the true barycenter.

Original: 4. D ERIVING THE D UAL P ROBLEM
Revised: 4. D ERIVING THE D UAL P ROBLEM

Original: Let P be the true barycenter.
Revised: Let P be the true barycenter.

Original: Our goal is to recover the optimal potentials { ψ ∗ n , ψ ∗ n } mapping the input measures P n into P .
Revised: Our goal is to recover the optimal potentials { ψ ∗ n , ψ ∗ n } mapping the input measures P n into P .

Original: To start, we express the barycenter objective (4) after substituting the dual formulation (2): 
Revised: To start, we express the barycenter objective (4) after substituting the dual formulation (2): [Equation]

Original: [Equation]
Revised: [Equation]

Original: The minimum is attained not just among convex potentials { ψ n } , but among congruent potentials (see discussion under (5)); thus, we can add the constraint that { ψ n } are congruent to (6).
Revised: The minimum is attained not just among convex potentials { ψ n } , but among congruent potentials (see discussion under (5)); thus, we can add the constraint that { ψ n } are congruent to (6).

Original: Hence,
Revised: Hence,

Original: N 
Revised: N (cid:88)

Original: [Equation]
Revised: [Equation]

Original: To transition from (6) to (7)
Revised: To transition from (6) to (7)

Original: [Equation]
Revised: [Equation]

Original: we used the fact that for congruent { ψ n } we have
Revised: we used the fact that for congruent { ψ n } we have

Original: [Equation]
Revised: [Equation]

Original: We call the value inside the minimum in (7) the multiple correlation of { P n } with weights { α n } w.r.t.
Revised: We call the value inside the minimum in (7) the multiple correlation of { P n } with weights { α n } w.r.t.

Original: potentials { ψ n } .
Revised: potentials { ψ n } .

Original: Notice that the true barycenter P appears nowhere on the right side of (7).
Revised: Notice that the true barycenter P appears nowhere on the right side of (7).

Original: Thus the optimal potentials { ψ ∗ n } can be recovered by solving the following
Revised: Thus the optimal potentials { ψ ∗ n } can be recovered by solving the following

Original: [Equation]
Revised: [Equation]

Original: 4. I MPOSING THE C ONGRUENCE C ONDITION
Revised: 4. I MPOSING THE C ONGRUENCE C ONDITION

Original: It is challenging to impose the congruence condition on convex potentials.
Revised: It is challenging to impose the congruence condition on convex potentials.

Original: What if we relax the congruence condition?
Revised: What if we relax the congruence condition?

Original: The following theorem bounds how close a set of convex potentials { ψ n } is to { ψ ∗ n } in terms of the difference of multiple correlation.
Revised: The following theorem bounds how close a set of convex potentials { ψ n } is to { ψ ∗ n } in terms of the difference of multiple correlation.

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: Let P ∈ P 2 ,ac ( R D ) be the barycenter of P 1 , . .
Revised: Let P ∈ P 2 ,ac ( R D ) be the barycenter of P 1 , . .

Original: . , P N ∈ P 2 ,ac ( R D ) w.r.t.
Revised: . , P N ∈ P 2 ,ac ( R D ) w.r.t.

Original: weights α 1 , . .
Revised: weights α 1 , . .

Original: . , α N .
Revised: . , α N .

Original: Let { ψ ∗ n } be the optimal congruent potentials of the barycenter problem.
Revised: Let { ψ ∗ n } be the optimal congruent potentials of the barycenter problem.

Original: Suppose we have B -smooth 1 convex potentials { ψ n } for some B ∈ [0 , + ∞ ] , and denote ∆ =
Revised: Suppose we have B -smooth 1 convex potentials { ψ n } for some B ∈ [0 , + ∞ ] , and denote ∆ =

Original: MultiCorr ( { α n , P n } | { ψ n } ) − MultiCorr ( { α n , P n } | { ψ ∗ n } ) .
Revised: MultiCorr ( { α n , P n } | { ψ n } ) − MultiCorr ( { α n , P n } | { ψ ∗ n } ) .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: Here (cid:107) · (cid:107) µ denotes the norm induced by inner product in Hilbert space L 2 ( R D → R D , µ ) .
Revised: Here (cid:107) · (cid:107) µ denotes the norm induced by inner product in Hilbert space L 2 ( R D → R D , µ ) .

Original: We call the second term on the left of (9) the congruence mismatch .
Revised: We call the second term on the left of (9) the congruence mismatch .

Original: We prove this in Appendix B. Note that if the congruence mismatch is non-positive, then
Revised: We prove this in Appendix B. Note that if the congruence mismatch is non-positive, then

Original: [Equation]
Revised: [Equation]

Original: where the last inequality of (10) follows from (Korotin et al., 2019a, Lemma A.2).
Revised: where the last inequality of (10) follows from (Korotin et al., 2019a, Lemma A.2).

Original: From (10), we conclude that for all n ∈ { 1 , . . .
Revised: From (10), we conclude that for all n ∈ { 1 , . . .

Original: , N } , we have W 22 ( ∇ ψ n (cid:93) P n , P ) ≤ B ∆ α n .
Revised: , N } , we have W 22 ( ∇ ψ n (cid:93) P n , P ) ≤ B ∆ α n .

Original: This shows that if the congruence mismatch is non-positive, then ∆ , the difference in multiple correlation, provides an upper bound for the Wasserstein-2 distance between the true barycenter and each pushforward ∇ ψ n (cid:93) P n .
Revised: This shows that if the congruence mismatch is non-positive, then ∆ , the difference in multiple correlation, provides an upper bound for the Wasserstein-2 distance between the true barycenter and each pushforward ∇ ψ n (cid:93) P n .

Original: This justiﬁes the use of ∇ ψ n (cid:93) P n to recover the barycenter.
Revised: This justiﬁes the use of ∇ ψ n (cid:93) P n to recover the barycenter.

Original: Notice for optimal potentials, the congruence mismatch is zero.
Revised: Notice for optimal potentials, the congruence mismatch is zero.

Original: Thus to penalize positive congruence mismatch, we introduce a regularizing term
Revised: Thus to penalize positive congruence mismatch, we introduce a regularizing term

Original: [Equation]
Revised: [Equation]

Original:  Because we take the positive part of the integrand of (9) to get (11) and that the right side of (9) is non-negative, we have 
Revised: [Equation] Because we take the positive part of the integrand of (9) to get (11) and that the right side of (9) is non-negative, we have [Equation]

Original: [Equation]
Revised: [Equation]

Original: for all convex potentials { ψ n } .
Revised: for all convex potentials { ψ n } .

Original: On the other hand, for optimal potentials { ψ n } = { ψ ∗ n } , the inequality turns into equality, implying that adding the regularizing term 1 · R P 1 ( { α n } , { ψ n } ) to (8) will not introduce bias – the optimal solution still yields { ψ ∗ n } .
Revised: On the other hand, for optimal potentials { ψ n } = { ψ ∗ n } , the inequality turns into equality, implying that adding the regularizing term 1 · R P 1 ( { α n } , { ψ n } ) to (8) will not introduce bias – the optimal solution still yields { ψ ∗ n } .

Original: However, evaluating (11) exactly requires knowing the true barycenter P a priori.
Revised: However, evaluating (11) exactly requires knowing the true barycenter P a priori.

Original: To remedy this issue, one may replace P with another absolutely continuous measure τ · (cid:98) P ( τ ≥ 1 and (cid:98) P is a probability measure) whose density bounds that of P from above almost everywhere.
Revised: To remedy this issue, one may replace P with another absolutely continuous measure τ · (cid:98) P ( τ ≥ 1 and (cid:98) P is a probability measure) whose density bounds that of P from above almost everywhere.

Original: In this case,
Revised: In this case,

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Hence we obtain the following regularized version of (8) where { ψ ∗ n } is the optimal solution:
Revised: Hence we obtain the following regularized version of (8) where { ψ ∗ n } is the optimal solution:

Original: [Equation]
Revised: [Equation]

Original: Selecting a measure τ · (cid:98) P is not obvious.
Revised: Selecting a measure τ · (cid:98) P is not obvious.

Original: Consider the case when { P n } are supported on compact sets
Revised: Consider the case when { P n } are supported on compact sets

Original: X 1 , . . .
Revised: X 1 , . . .

Original: , X N ⊂ R D and P 1 has density upper bounded by h < ∞ .
Revised: , X N ⊂ R D and P 1 has density upper bounded by h < ∞ .

Original: In this scenario, the barycenter density is upper bounded by h · α − D(Álvarez-Esteban et al., 2016, Remark 3.2).
Revised: In this scenario, the barycenter density is upper bounded by h · α − D(Álvarez-Esteban et al., 2016, Remark 3.2).

Original: Thus, the measure τ · (cid:98) P supported on ConvexHull ( X 1 , . . .
Revised: Thus, the measure τ · (cid:98) P supported on ConvexHull ( X 1 , . . .

Original: , X N ) with this density is an upper bound for P .
Revised: , X N ) with this density is an upper bound for P .

Original: We will address the question of how to choose τ, (cid:98) P properly in practice in §4.4.
Revised: We will address the question of how to choose τ, (cid:98) P properly in practice in §4.4.

Original: 4. E NFORCING C ONJUGACY OF P OTENTIALS P AIRS
Revised: 4. E NFORCING C ONJUGACY OF P OTENTIALS P AIRS

Original: Throughout this subsection, we assume the upper bound ﬁnite measure τ · (cid:98) P of the P is known.
Revised: Throughout this subsection, we assume the upper bound ﬁnite measure τ · (cid:98) P of the P is known.

Original: The optimization problem (13) involves not only the potentials { ψ n } , but also their conjugates { ψ n } .
Revised: The optimization problem (13) involves not only the potentials { ψ n } , but also their conjugates { ψ n } .

Original: This brings practical difﬁculty since evaluating conjugate potentials is hard (Korotin et al., 2019a).
Revised: This brings practical difﬁculty since evaluating conjugate potentials is hard (Korotin et al., 2019a).

Original: Instead we parameterize potentials ψ n and ψ n separately using input convex neural networks (ICNN) as ψ † n and ψ ‡ n respectively.
Revised: Instead we parameterize potentials ψ n and ψ n separately using input convex neural networks (ICNN) as ψ † n and ψ ‡ n respectively.

Original: We add an additional cycle-consistency regularizer to enfore the conjugacy of ψ † n and ψ ‡ n as in Korotin et al. (2019a).
Revised: We add an additional cycle-consistency regularizer to enfore the conjugacy of ψ † n and ψ ‡ n as in Korotin et al. (2019a).

Original: This regularizer is deﬁned as
Revised: This regularizer is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: Note that R P n 2 ( ψ † n , ψ ‡ n ) = 

Revised: Note that R P n 2 ( ψ † n , ψ ‡ n ) =

Original: this condition is necessary for ψ † n and ψ ‡ n to be conjugate with each other.
Revised: this condition is necessary for ψ † n and ψ ‡ n to be conjugate with each other.

Original: Also, it is a sufﬁcient condition for convex functions to be conjugates up to an additive constant.
Revised: Also, it is a sufﬁcient condition for convex functions to be conjugates up to an additive constant.

Original: We use one-sided regularization.
Revised: We use one-sided regularization.

Original: In our case, computing the regularizer of the other direction (cid:107)∇ ψ † n ◦ ∇ ψ ‡ n − id R D (cid:107) 2 P is infeasible, since P is unknown.
Revised: In our case, computing the regularizer of the other direction (cid:107)∇ ψ † n ◦ ∇ ψ ‡ n − id R D (cid:107) 2 P is infeasible, since P is unknown.

Original: If fact, Korotin et al. (2019a) demonstrates that such one-sided condition is sufﬁcient.
Revised: If fact, Korotin et al. (2019a) demonstrates that such one-sided condition is sufﬁcient.

Original: In this way we use 2 N input convex neural networks for { ψ † n , ψ ‡ n } .
Revised: In this way we use 2 N input convex neural networks for { ψ † n , ψ ‡ n } .

Original: By adding the new cycle consistency regularizer into (13), we obtain our ﬁnal objective: min
Revised: By adding the new cycle consistency regularizer into (13), we obtain our ﬁnal objective: min

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Note that we express the aproximate multiple correlation by using both potentials { ψ † n } and { ψ ‡ n } .
Revised: Note that we express the aproximate multiple correlation by using both potentials { ψ † n } and { ψ ‡ n } .

Original: This is done to eliminate the freedom of an additive constant on { ψ † n } that is not adressed by cycle regularization.
Revised: This is done to eliminate the freedom of an additive constant on { ψ † n } that is not addressed by cycle regularization.

Original: We denote the entire objective as MultiCorr (cid:0) { P n } | { ψ † } , { ψ ‡ } ; τ, (cid:98) P , λ (cid:1) .
Revised: We denote the entire objective as MultiCorr (cid:0) { P n } | { ψ † } , { ψ ‡ } ; τ, (cid:98) P , λ (cid:1) .

Original: Analogous to Theorem 4.1, we have following result showing that this new objective enjoys the same properties as the unregularized version from (8).
Revised: Analogous to Theorem 4.1, we have following result showing that this new objective enjoys the same properties as the unregularized version from (8).

Original: Theorem 4.2.
Revised: Theorem 4.2.

Original: Let P ∈ P 2 ,ac ( R D ) be the barycenter of P 1 , . .
Revised: Let P ∈ P 2 ,ac ( R D ) be the barycenter of P 1 , . .

Original: . , P N ∈ P 2 ,ac ( R D ) w.r.t.
Revised: . , P N ∈ P 2 ,ac ( R D ) w.r.t.

Original: weights α 1 , . .
Revised: weights α 1 , . .

Original: . , α N .
Revised: . , α N .

Original: Let { ψ ∗ n } be the optimal congruent potentials of the barycenter problem.
Revised: Let { ψ ∗ n } be the optimal congruent potentials of the barycenter problem.

Original: Suppose we have τ, ˆ P such that τ ≥ 1 and τ · (cid:98) P ≥ P .
Revised: Suppose we have τ, ˆ P such that τ ≥ 1 and τ · (cid:98) P ≥ P .

Original: Suppose we have convex potentials { ψ † n } and β ‡ -strongly
Revised: Suppose we have convex potentials { ψ † n } and β ‡ -strongly

Original: convex and B ‡ -smooth convex potentials { ψ ‡ n } with 0 < β ‡ ≤ B ‡ < ∞ and λ >
Revised: convex and B ‡ -smooth convex potentials { ψ ‡ n } with 0 < β ‡ ≤ B ‡ < ∞ and λ >

Original: B † 2( β ‡ ) 2 .
Revised: B † 2( β ‡ ) 2 .

Original: Then
Revised: Then

Revised: MultiCorr (cid:0) { α n , P n

Revised: [Equation]

Original: MultiCorr (cid:0) { α n , P n } | { ψ † n } , { ψ ‡ n } ; τ, (cid:98) P , λ (cid:1) ≥ MultiCorr (cid:0) { α n , P n } | { ψ ∗ n } (cid:1) .
Revised: | { ψ † n } , { ψ ‡ n } ; τ, (cid:98) P , [Equation] (cid:1) ≥ MultiCorr (cid:0) { [Equation]

Original: Denote ∆ =
Revised: Denote ∆ =

Original: MultiCorr (cid:0) { α n , P n } | { ψ † n } , { ψ ‡ n } ; τ, (cid:98) P , λ (cid:1) − MultiCorr (cid:0) { α n , P n } | { ψ ∗ n } (cid:1) .
Revised: MultiCorr (cid:0) { α n , P n } | { ψ † n } , { ψ ‡ n } ; τ, (cid:98) P , λ (cid:1) − MultiCorr (cid:0) { α n , P n } | { ψ ∗ n } (cid:1) .

Original: Then for all n ∈ { 1 , . . . , N } , we have
Revised: Then for all n ∈ { 1 , . . . , N } , we have

Original: [Equation]
Revised: [Equation]

Original: Informally, Theorem 4.2 states that the better we solve the regularized dual problem, (14) the closer we expect each ∇ ψ † n (cid:93) P n to be to the true barycenter P in W 2 .
Revised: Informally, Theorem 4.2 states that the better we solve the regularized dual problem, (14) the closer we expect each ∇ ψ † n (cid:93) P n to be to the true barycenter P in W 2 .

Original: It follows from (15) that our ﬁnal objective (14) is unbiased : the optimal solution is obtained by { ψ ∗ n , ψ ∗ n } .
Revised: It follows from (15) that our ﬁnal objective (14) is unbiased : the optimal solution is obtained by { ψ ∗ n , ψ ∗ n } .

Original:  P RACTICAL A SPECTS AND O PTIMIZATION P ROCEDURE
Revised: 4. P RACTICAL A SPECTS AND O PTIMIZATION P ROCEDURE

Original: In practice, even if the choice of τ, (cid:98) P does not satisfy τ · (cid:98) P ≥ P , we observe the pushforward measures ∇ ψ † n (cid:93) P n often converge to P .
Revised: In practice, even if the choice of τ, (cid:98) P does not satisfy τ · (cid:98) P ≥ P , we observe the pushforward measures ∇ ψ † n (cid:93) P n often converge to P .

Original: To partially bridge the gap between theory and practice, we dynamically update the measure (cid:98) P so that after each optimization step we set (for γ ∈ [0 , 1] )
Revised: To partially bridge the gap between theory and practice, we dynamically update the measure (cid:98) P so that after each optimization step we set (for γ ∈ [0 , 1] )

Original: [Equation]
Revised: [Equation]

Original: i.e., the probability measure (cid:98) P (cid:48) is a mixture of the given initial measure (cid:98) P and the current barycenter estimates {∇ ψ † (cid:93) P n } .
Revised: i.e., the probability measure (cid:98) P (cid:48) is a mixture of the given initial measure (cid:98) P and the current barycenter estimates {∇ ψ † (cid:93) P n } .

Original: For the initial (cid:98) P one may use the barycenter of {N ( µ P n , Σ P n ) } .
Revised: For the initial (cid:98) P one may use the barycenter of {N ( µ P n , Σ P n ) } .

Original: It can be efﬁciently computed via an iterative ﬁxed point algorithm (Álvarez-Esteban et al., 2016; Chewi et al., 2020).
Revised: It can be efﬁciently computed via an iterative ﬁxed point algorithm (Álvarez-Esteban et al., 2016; Chewi et al., 2020).

Original: During the optimization, these estimates become closer to the true barycenter and can thus improve the congruence regularizer (12).
Revised: During the optimization, these estimates become closer to the true barycenter and can thus improve the congruence regularizer (12).

Original: We use mini-batch stochastic gradient descent to solve (14) where the integration is done by MonteCarlo sampling from input measures { P n } and regularization measure (cid:98) P , similar to Li et al.
Revised: We use mini-batch stochastic gradient descent to solve (14) where the integration is done by MonteCarlo sampling from input measures { P n } and regularization measure (cid:98) P , similar to Li et al.

Original: The detailed  procedure is given in Algorithm 1 in  the appendix.
Revised: We provide the detailed optimization procedure (Algorithm 1) and discuss its computational complexity in Appendix A. In Appendix C.3, we demonstrate that the impact of the considered regularization on our model: we show that cycle consistency and the congruence condition of the potentials are well satisﬁed.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We compare our method [C W 2 B] with the potential-based method [CR W B] by Li et al.
Revised: We compare our method [C W 2 B] with the potential-based method [CR W B] by Li et al.

Original: (2020) (with Wasserstein-2 distance and L 2 -regularization) and with the measure-based generative method [SC W 2 B] by Fan et al.
Revised: (2020) (with Wasserstein-2 distance and L 2 -regularization) and with the measure-based generative method [SC W 2 B] by Fan et al.

Original: All considered methods recover 2 N potentials { ψ † n , ψ ‡ n } ≈ { ψ ∗ n , ψ ∗ n } and approximate the barycenter as pushforward measures {∇ ψ † n (cid:93) P n } .
Revised: All considered methods recover 2 N potentials { ψ † n , ψ ‡ n } ≈ { ψ ∗ n , ψ ∗ n } and approximate the barycenter as pushforward measures {∇ ψ † n (cid:93) P n } .

Original: Regularization in [CR W B] allows access to the joint density of the transport plan, a feature of their method that we do not consider here.
Revised: Regularization in [CR W B] allows access to the joint density of the transport plan, a feature of their method that we do not consider here.

Original: The method [SC W 2 B] additionally outputs a generated barycenter g(cid:93) S ≈ P where g is the generative network and S is the input noise distribution.
Revised: The method [SC W 2 B] additionally outputs a generated barycenter g(cid:93) S ≈ P where g is the generative network and S is the input noise distribution.

Original: To assess the quality of the computed barycenter, we consider the unexplained variance percentage deﬁned as UVP (˜ P )
Revised: To assess the quality of the computed barycenter, we consider the unexplained variance percentage deﬁned as UVP (˜ P )

Original: = 100 W 22 (˜ P , P ) 1 / 2 Var ( P ) % .
Revised: = 100 W 22 (˜ P , P ) 1 / 2 Var ( P ) % .

Original: When UVP ≈ 0% , ˜ P is a good approximation of P .
Revised: When UVP ≈ 0% , ˜ P is a good approximation of P .

Original: For values ≥ 100% , the distribution ˜ P is undesirable: a trivial baseline P 0 = δ E P [ y ] achieves UVP ( P 0 )
Revised: For values ≥ 100% , the distribution ˜ P is undesirable: a trivial baseline P 0 = δ E P [ y ] achieves UVP ( P 0 )

Original: = 100% .
Revised: = 100% .

Original: Evaluating UVP in high dimensions is infeasible: empirical estimates of W 22 are unreliable due to high sample complexity
Revised: Evaluating UVP in high dimensions is infeasible: empirical estimates of W 22 are unreliable due to high sample complexity

Original: (Weed et al., 2019).
Revised: (Weed et al., 2019).

Original: To overcome this issue, for barycenters given by ∇ ψ † n (cid:93) P n we use L 2 -UVP deﬁned by
Revised: To overcome this issue, for barycenters given by ∇ ψ † n (cid:93) P n we use L 2 -UVP deﬁned by

Original: [Equation]
Revised: [Equation]

Original: (17) where the inequality in brackets follows from (Korotin et al., 2019a, Lemma A.2).
Revised: where the inequality in brackets follows from (Korotin et al., 2019a, Lemma A.2).

Original: We report the weighted average of L 2 -UVP of all pushforward measures w.r.t.
Revised: We report the weighted average of L 2 -UVP of all pushforward measures w.r.t.

Original: the weights α n .
Revised: the weights α n .

Original: For barycenters given in an implicit form g(cid:93) S , we compute the Bures-Wasserstein UVP deﬁned by
Revised: For barycenters given in an implicit form g(cid:93) S , we compute the Bures-Wasserstein UVP deﬁned by

Original: [Equation]
Revised: [Equation]

Original: where B W 22 ( P , Q ) =
Revised: where B W 22 ( P , Q ) =

Original: W 22 (cid:0) N ( µ P , Σ P ) , N ( µ Q , Σ Q ) (cid:1) is the Bures-Wasserstein metric and we use µ P , Σ P to denote the mean and the covariance of a distribution P (Chewi et al., 2020).
Revised: W 22 (cid:0) N ( µ P , Σ P ) , N ( µ Q , Σ Q ) (cid:1) is the Bures-Wasserstein metric and we use µ P , Σ P to denote the mean and the covariance of a distribution P (Chewi et al., 2020).

Original: It is known that
Revised: It is known that

Original: B W 2 2 lower-bounds W 2 2 (Dowson & Landau, 1982), so the inequality in the brackets of (18) follows.
Revised: B W 22 lower-bounds W 22 (Dowson & Landau, 1982), so the inequality in the brackets of (18) follows.

Original: A detailed discussion of the adopted metrics is given in Appendix C.2.
Revised: A detailed discussion of the adopted metrics is given in Appendix C.2.

Original: 5. H IGH -D IMENSIONAL L OCATION -S CATTER E XPERIMENTS
Revised: [Figure]

Original: In this section, we consider N = 4 with ( α 1 , . . .
Revised: In this section, we consider N = 4 with ( α 1 , . . .

Original: , α 4 )
Revised: , α 4 )

Original: = (0 . 1 , 0 . 2 , 0 . 3 , 0 . 4) as weights.
Revised: = (0 . 1 , 0 . 2 , 0 . 3 , 0 . 4) as weights.

Original: We consider the location-scatter family of distributions (Álvarez-Esteban et al., 2016, §4) whose true barycenter[Figure] can be computed.
Revised: We consider the location-scatter family of distributions (Álvarez-Esteban et al., 2016, §4) whose true barycenter can be computed.

Original: Let P
Revised: Let P

Original: ∈
Revised: ∈

Original: P 2 , ac and deﬁne the following location-scatter family of distributions
Revised: P 2 , ac and deﬁne the following location-scatter family of distributions

Original: F ( P 0 )
Revised: F ( P 0 )

Original: = { f S,u (cid:93) P 0 | S ∈ M + D × D , u ∈ R D } , where f S,u : R D
Revised: = { f S,u (cid:93) P 0 | S ∈ M + D × D , u ∈ R D } , where f S,u : R D

Original: → R D is a linear map f S,u ( x ) =
Revised: → R D is a linear map f S,u ( x ) =

Original: Sx + u with positive deﬁnite matrix S ∈ M + D × D .
Revised: Sx + u with positive deﬁnite matrix S ∈ M + D × D .

Original: When { P n } ⊂ F ( P 0 ) , their barycenter P is also an element of F ( P 0 ) and can be computed via ﬁxed-point iterations (Álvarez-Esteban et al., 2016).
Revised: When { P n } ⊂ F ( P 0 ) , their barycenter P is also an element of F ( P 0 ) and can be computed via ﬁxed-point iterations (Álvarez-Esteban et al., 2016).

Original: Figure 1a shows a 2-dimensional location-scatter family generated by using the Swiss roll distribution as P 0 .
Revised: Figure 1a shows a 2-dimensional location-scatter family generated by using the Swiss roll distribution as P 0 .

Original: The true barycenter is shown in Figure 1b.
Revised: The true barycenter is shown in Figure 1b.

Original: The generated barycenter g(cid:93) S of [SC W 2 B] is given in Figure 1c.
Revised: The generated barycenter g(cid:93) S of [SC W 2 B] is given in Figure 1c.

Original: The pushforward measures ∇ ψ † n (cid:93) P n of each method are provided in Figures 1d, 1e, 1f, respectively.
Revised: The pushforward measures ∇ ψ † n (cid:93) P n of each method are provided in Figures 1d, 1e, 1f, respectively.

Original: In this example, the pushforward measures ∇ ψ n (cid:93) P n all reasonably approximate P , whereas the generated barycenter g(cid:93) S of [SC W 2 B] (Figure 1c) visibly underﬁts.
Revised: In this example, the pushforward measures ∇ ψ n (cid:93) P n all reasonably approximate P , whereas the generated barycenter g(cid:93) S of [SC W 2 B] (Figure 1c) visibly underﬁts.

Original: For quantitative comparison, we consider two choices for P 0 : the D -dimensional standard Gaussian distribution and the uniform distribution on [ −√ 3 , + √ 3] D .
Revised: For quantitative comparison, we consider two choices for P 0 : the D -dimensional standard Gaussian distribution and the uniform distribution on [ −√ 3 , + √ 3] D .

Original: Each P n is constructed as f S Tn Λ S n , 0 (cid:93) P 0 ∈
Revised: Each P n is constructed as f S Tn Λ S n , 0 (cid:93) P 0 ∈

Original: F ( P 0 ) , where S n is a random rotation matrix and Λ is diagonal with entries [ 12 b 0 , 12 b 1 , . .
Revised: F ( P 0 ) , where S n is a random rotation matrix and Λ is diagonal with entries [ 1 2 b 0 , 1 2 b 1 , . .

Original: , 2] where
Revised: , 2] where

Original: b = D − 1 √ 4 .
Revised: b = D − 1 √ 4 .

Original: We consider only centered distributions (i.e. zero mean) because the barycenter of non-centered { P n } ∈ P 2 , ac ( R D ) is the barycenter of { P (cid:48) n } shifted by (cid:80) Nn
Revised: We consider only centered distributions (i.e. zero mean) because the barycenter of non-centered { P n } ∈ P 2 , ac ( R D ) is the barycenter of { P (cid:48) n } shifted by (cid:80) Nn

Original: =1 α n µ P n , where { P (cid:48) n } are centered copies of { P n } (Álvarez-Esteban et al., 2016).
Revised: =1 α n µ P n , where { P (cid:48) n } are centered copies of { P n } (Álvarez-Esteban et al., 2016).

Original: Results are shown in Table 1 and 2. 
Revised: Results are shown in Table 1 and 2. [Table]

Original: In these experiments, our method outperforms [CR W B] and [SC W 2 B].
Revised: In these experiments, our method outperforms [CR W B] and [SC W 2 B].

Original: For [CR W B], dimension ∼ 16 is the breakpoint: the method does not scale well to higher dimensions.
Revised: For [CR W B], dimension ∼ 16 is the breakpoint: the method does not scale well to higher dimensions.

Original: [SC W 2 B] scales with the increasing dimension better, but its errors L 2 -UVP and B W 22 -UVP are twice as high as ours.
Revised: [SC W 2 B] scales with the increasing dimension better, but its errors L 2 -UVP and B W 22 -UVP are twice as high as ours.

Original: This is likely due to the generative approximation and the difﬁcult min-max-min optimization in [SC W 2 B].
Revised: This is likely due to the generative approximation and the difﬁcult min-max-min optimization in [SC W 2 B].

Original: [Table]
Revised: For completeness, we also compare our algorithm to the proposed in Cuturi & Doucet (2014) which approximates the barycenter by a discrete distribution on a ﬁxed number of free-support points. In our experiment, similar to Li et al. (2020), we set 5000 as the support size. As expected, the B W 22 -UVP error of the method increases drastically as the dimension grows and the method is outperformed by our approach.

Original: [Table]
Revised: To show the scalability of our method with the number of input distributions N , we conduct an analogous experiment with a high-dimensional location-scatter family for N = 20 . We set α n = 2 n N ( N +1) for n = 1 , 2 , ..., 20 and choose the uniform distribution on [ −√ 3 , + √ 3] D as P 0 and construct distributions P n ∈ F ( P 0 ) as before. The results for dimensions 32, 64 and 128 are provided in Table 3.

Original: [Table]
Revised: Similar to the results from Tables 1 and 2, we see that our method outperforms the alternatives.

Original: 5. S UBSET P OSTERIOR A GGREGATION
Revised: 5. S UBSET P OSTERIOR A GGREGATION

Original: We apply our method to aggregate subset posterior distributions.
Revised: We apply our method to aggregate subset posterior distributions.

Original: The barycenter of subset posteriors converges to the true posterior (Srivastava et al., 2018).
Revised: The barycenter of subset posteriors converges to the true posterior (Srivastava et al., 2018).

Original: Thus, computing the barycenter of subset posteriors is an efﬁcient alternative to obtaining a full posterior in the big data setting (Srivastava et al., 2015; Staib et al., 2017; Li et al., 2020).
Revised: Thus, computing the barycenter of subset posteriors is an efﬁcient alternative to obtaining a full posterior in the big data setting (Srivastava et al., 2015; Staib et al., 2017; Li et al., 2020).

Original: Analogous to (Li et al., 2020), we consider Poisson and negative binomial regressions for predicting the hourly number of bike rentals using features such as the day of the week and weather conditions.
Revised: Analogous to (Li et al., 2020), we consider Poisson and negative binomial regressions for predicting the hourly number of bike rentals using features such as the day of the week and weather conditions.

Original: We consider the posterior on the 8-dimensional regression coefﬁcients for both Poisson and negative binomial regressions.
Revised: We consider the posterior on the 8-dimensional regression coefﬁcients for both Poisson and negative binomial regressions.

Original: We randomly split the data into N = 5 equally-sized subsets and obtain 10samples from each subset posterior using the Stan library (Carpenter et al., 2017).
Revised: We randomly split the data into N = 5 equally-sized subsets and obtain 10samples from each subset posterior using the Stan library (Carpenter et al., 2017).

Original: This gives the discrete uniform distributions { P n } supported on the samples.
Revised: This gives the discrete uniform distributions { P n } supported on the samples.

Original: As the ground truth barycenter P , we consider the full dataset posterior also consisting of 10 5 points.
Revised: As the ground truth barycenter P , we consider the full dataset posterior also consisting of 10 5 points.

Original: We use B W 22 -UVP (˜ P , P ) to compare the estimated barycenter ˜ P (pushforward measure ∇ ψ † n (cid:93) P n or generated measure g(cid:93) S ) with the true barycenter.
Revised: We use B W 22 -UVP (˜ P , P ) to compare the estimated barycenter ˜ P (pushforward measure ∇ ψ † n (cid:93) P n or generated measure g(cid:93) S ) with the true barycenter.

Original: The results are in Table 3.
Revised: The results are in Table 4.

Original: All considered methods perform well (UVP < 2% ), but our method outperforms the alternatives.
Revised: All considered methods perform well (UVP < 2% ), but our method outperforms the alternatives.

Original: 5. C OLOR P ALETTE A VERAGING
Revised: 5. C OLOR P ALETTE A VERAGING

Original: For qualitative study, we apply our method to aggregating color palettes of images.
Revised: For qualitative study, we apply our method to aggregating color palettes of images.

Original: For an RGB image
Revised: For an RGB image

Original: I , its color palette is deﬁned by the discrete uniform distribution P ( I ) of all its pixels ∈ [0 , 1] 3 .
Revised: I , its color palette is deﬁned by the discrete uniform distribution P ( I ) of all its pixels ∈ [0 , 1] 3 .

Original: Forimages {I n } we compute the barycenter P of each color palette P n = P ( I n ) w.r.t.
Revised: Forimages {I n } we compute the barycenter P of each color palette P n = P ( I n ) w.r.t.

Original: uniform weights α n = 13 .
Revised: uniform weights α n = 13 .

Original: We apply each computed potential ∇ ψ † n pixel-wise to I n to obtain the “pushforward” image ∇ ψ † n (cid:93) I n .
Revised: We apply each computed potential ∇ ψ † n pixel-wise to I n to obtain the “pushforward” image ∇ ψ † n (cid:93) I n .

Original: These “pushforwad” images should be close to the barycenter P of { P n } .
Revised: These “pushforward” images should be close to the barycenter P of { P n } .

Original: [Figure]
Revised: The results are provided in Figure 2. Note that the image ∇ ψ † 1 (cid:93)

Original: The results are
Revised: I 1 inherits certain attributes of images I 2 and I 3 : the sky becomes bluer and the trees becomes greener. On the other hand, the sunlight in images ∇ ψ † 2 (cid:93) I 2 , ∇ ψ † 3 (cid:93) I 3 has acquired an orange tint, thanks to the dominance of orange in I 1 .

Original: B P ROOFS
Revised: [Table] [Table] [Table] [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: T

Revised: [Equation]

Revised: The numerical procedure for solving our ﬁnal objective (14) is given below.

Revised: for t = 1 , 2 , . . .

Revised: do

Revised: Sample batches X n ∼ P n for all n = 1 , . . .

Revised: , N ; 2.

Revised: Compute the pushforwards Y n = ∇ ψ θ n (cid:93)X n for all n = 1 , . . .

Revised: , N ; 3.

Revised: Sample batch Y 0 ∼ (cid:98) P ; 4.

Revised: Compute the Monte-Carlo estimate of the congruence regularizer: L Congruence := 1 K ·

Revised: [Equation]

Revised: , where γ 0 = γ and γ n = α n · (1 − γ ) for n = 1 , 2 , . .

Revised: . , N ; 5.

Revised: Compute the Monte-Carlo estimate of the cycle-consistency regularizer: L Cycle :

Revised: = 1 K N (cid:88) n = α n (cid:20) (cid:88) x ∈ X n (cid:107)∇ ψ ω n ◦ ∇ ψ θ n ( x ) − x (cid:107)(cid:21) ; 6.

Revised: Compute the Monte-

Revised: Carlo estimate of multiple correlations:

Revised: [Equation]

Revised: Compute the total loss:

Revised: [Equation]

Revised: Perform a gradient step over { θ n , ω n } by using

Revised: [Equation]

Revised: end

Revised: Parametrization of the potentials.

Revised: To parametrize potentials { ψ θ n , ψ ω n } , we use DenseICNN (dense input convex neural network) with quadratic skip connections; see (Korotin et al., 2019a,

Revised: Appendix B.2).

Revised: As an initialization step, we pre-train the potentials to satisfy

Revised: [Equation]

Revised: and

Revised: [Equation]

Revised: Such pre-training provides a good start for the networks: each ψ θ n is approximately conjugate to the corresponding ψ ω n .

Revised: On the other hand, the initial networks { ψ θ n } are approximate congruent according to (5).

Revised: Computational Complexity.

Revised: For a single training iteration, the time complexity of both forward (evaluation) and backward (computing the gradient with respect to the parameters) passes through the objective function (14) is O ( NT ) .

Revised: Here N is the number of input distributions and T is the time taken by evaluating each individual potential (parameterized as a neural network) on a batch of points sampled from either P n or (cid:98) P .

Revised: This claim follows from the well-known fact that gradient evaluation ∇ θ h θ ( x ) of h θ : R D

Revised: → R , when parameterized as a neural network, requires time proportional to the size of the computational graph.

Revised: Hence, gradient computation requires computational time proportional to the time for evaluating the function h θ ( x ) itself.

Revised: The same holds when computing the derivative with respect to x .

Revised: Then, for instance, computing the term ∇ ψ ‡ n ◦ ∇ ψ † n ( x ) in (14) takes

Revised: O ( T ) time.

Revised: The gradient of this term with respect to θ also takes O ( T ) time: Hessian-vector products that appear can be calculated in O ( T ) time using the famous Hessian trick, see Pearlmutter (1994).

Revised: In practice, we compute all the gradients using automatic differentiation.

Revised: We empirically measured that for our DenseICNN potentials, the computation of their gradient w.r.t.

Revised: input x , i.e., ∇ ψ † ( x ) , requires roughly 3-4x more time than the computation of ψ † ( x ) .

Revised: B P ROOFS

Original: In this section, we prove our main Theorems 4.1 and 4.2.
Revised: In this section, we prove our main Theorems 4.1 and 4.2.

Original: We use L 2 ( R D → R D , µ ) to denote the Hilbert space of functions f : R D → R D with integrable square w.r.t.
Revised: We use L 2 ( R D → R D , µ ) to denote the Hilbert space of functions f : R D → R D with integrable square w.r.t.

Original: a probability measure µ .
Revised: a probability measure µ .

Original: The corresponding inner product for f 1 , f 2 ∈ L 2 ( R D → R D , µ ) is denoted by
Revised: The corresponding inner product for f 1 , f 2 ∈ L 2 ( R D → R D , µ ) is denoted by

Original: [Equation]
Revised: [Equation]

Original: where (cid:104) f 1 ( x ) , f 2 ( x ) (cid:105) is the Euclidean dot product.
Revised: where (cid:104) f 1 ( x ) , f 2 ( x ) (cid:105) is the Euclidean dot product.

Original: We use (cid:107) · (cid:107) µ = (cid:112) (cid:104)· , ·(cid:105) µ to denote the norm induced by the inner product in L 2 ( R D → R D , µ ) .
Revised: We use (cid:107) · (cid:107) µ = (cid:112) (cid:104)· , ·(cid:105) µ to denote the norm induced by the inner product in L 2 ( R D → R D , µ ) .

Original: We also recall a useful property of lower semi-continuous convex function ψ : R D
Revised: We also recall a useful property of lower semi-continuous convex function ψ : R D

Original: → R :
Revised: → R :

Original: [Equation]
Revised: [Equation]

Original: (19) which follows from the fact that
Revised: (19) which follows from the fact that

Original: [Equation]
Revised: [Equation]

Original: We begin with the proof of Theorem 4.1.
Revised: We begin with the proof of Theorem 4.1.

Original: Proof.
Revised: Proof.

Original: We consider the difference between the estimated correlations and true ones:
Revised: We consider the difference between the estimated correlations and true ones:

Original: [Equation]
Revised: [Equation]

Original:  where we twice use (19) for f = ψ n and f = ψ ∗ n .
Revised: [Equation] where we twice use (19) for f = ψ n and f = ψ ∗ n .

Original: We note that
Revised: We note that

Original: N
Revised: N

Original: [Equation]
Revised: [Equation]

Revised: n =

Original: (21) where we use of change-of-variable formula for ∇ ψ ∗ n (cid:93) P n = P and (5) [Equation]
Revised: [Equation] [Equation] where we use of change-of-variable formula for ∇ ψ ∗ n (cid:93) P n = P and (5).

Original: Analogously,
Revised: Analogously,

Original: N
Revised: N

Original: [Equation]
Revised: [Equation]

Revised: Thus

Original: [Equation]
Revised: [Equation]

Original: Since each ψ n is B -smooth, we conclude that ψ n is 1 B -strongly convex, see (Kakade et al., 2009).

Original: Thus, we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: or equivalently
Revised: or equivalently

Original: [Equation]
Revised: [Equation]

Original: We integrate (24) w
Revised: We integrate (24) w.r

Original: [Equation]

Original: r

Original: [Equation]
Revised: [Equation]

Original: t. P n and sum over n = 1 , 2 , . . .
Revised: t. P n and sum over n = 1 , 2 , . . .

Original: , N with weights α n :
Revised: , N with weights α n :

Original: [Equation]
Revised: [Equation]

Original: We note that
Revised: We note that

Original: [Equation]
Revised: [Equation]

Original: Now we substitute (25), (26), (21) and (22) into (20) to obtain (9).
Revised: Now we substitute (25), (26), (21) and (22) into (20) to obtain (9).

Original: Next, we prove Theorem 4.2.
Revised: Next, we prove Theorem 4.2.

Original: Proof.
Revised: Proof.

Original: Since ψ ‡ n is β ‡ strongly convex, its conjugate ψ ‡ n is 1 β ‡ -smooth, i.e. has 1 β ‡ -Lipschitz gradient ∇ ψ ‡ n (Kakade et al., 2009).
Revised: Since ψ ‡ n is β ‡ strongly convex, its conjugate ψ ‡ n is 1 β ‡ -smooth, i.e. has 1 β ‡ -Lipschitz gradient ∇ ψ ‡ n (Kakade et al., 2009).

Original: Thus, for all x, x (cid:48) ∈ R D :
Revised: Thus, for all x, x (cid:48) ∈ R D :

Original: [Equation]
Revised: [Equation]

Original: We substitute x (cid:48)
Revised: We substitute x (cid:48)

Original: = ∇ ψ 
Revised: = ∇ ψ ‡ n ◦ ∇ ψ † n

Original: [Equation]
Revised: [Equation]

Original: n ◦ ∇ ψ † n ( y ) = 
Revised: y ) = [Equation]

Original: (cid:0) ∇ ψ ‡ n (cid:1) − 1 ◦ ∇ ψ † n ( y ) and obtain:
Revised: ∇ ψ ‡ n (cid:1) − 1 ◦ ∇ ψ † n ( y ) and obtain:

Original: [Equation]
Revised: [Equation]

Original: Since the function ψ ‡ n is B ‡ -smooth, we have for all x ∈ R D : 
Revised: Since the function ψ ‡ n is B ‡ -smooth, we have for all x ∈ R D : ψ

Original: [Equation]
Revised: [Equation]

Original: that is equivalent to:
Revised: that is equivalent to:

Original: [Equation]
Revised: [Equation]

Original: We combine (28) with (27) to obtain
Revised: We combine (28) with (27) to obtain

Original: [Equation]
Revised: [Equation]

Original: For every n = 1 , 2 , . . .
Revised: For every n = 1 , 2 , . . .

Original: , N we integrate (29) w.r.t.
Revised: , N we integrate (29) w.r [Equation]

Original:  P n and sum up the corresponding cycle-consistency regularization term:
Revised: t. P n and sum up the corresponding cycle-consistency regularization term:

Original: [Equation]
Revised: [Equation]

Original: We sum (30) for n = 1 , 2 , . . .
Revised: We sum (30) for n = 1 , 2 , . . .

Original: , N w.r.t.
Revised: , N w.r.t.

Original: weights α n to obtain:
Revised: weights α n to obtain:

Original: N
Revised: N

Original: [Equation]
Revised: [Equation]

Original: We add τ · R (cid:98) P 1 ( { ψ ‡ n } ) to both sides of (31) to getMultiCorr (cid:0) { [Equation] n
Revised: We add τ · R (cid:98) P 1 ( { ψ ‡ n } ) to both sides of (31) to get

Original: [Equation]

Original: P

Original: [Equation]

Original: } | { ψ

Original: [Equation]

Original: [Equation]

Original: ψ

Original: [Equation]

Original: MultiCorr
Revised: MultiCorr

Original: [Equation]
Revised: [Equation]

Original: We substract MultiCorr ( { α n , P n } | { ψ ‡ n } ) from both sides and use Theorem 4.1 to obtain 
Revised: We substract MultiCorr ( { α n , P n } | { ψ ‡ n } ) from both sides and use Theorem 4.1 to obtain ∆ ≥ −

Original: [Equation]
Revised: [Equation]

Original: [Equation] In transition from (33) to (34)
Revised: n

Original: [Equation]
Revised: [Equation]

Original:  we explot the fact that the sum of the ﬁrst term of (32) with the regularizer τ · R (cid:98) P 1 ( { ψ ‡ n } ) .
Revised: In transition from (33) to (34), we explot the fact that the sum of the ﬁrst term of (32) with the regularizer τ · R (cid:98) P 1 ( { ψ ‡ n } ) .

Original: Since λ >
Revised: Since λ >

Original: B † 2( β ‡ ) 2 , from (34) we immediately conclude ∆ ≥ 0 ; i [Equation] e., the multiple correlations upper bound (15) holds true.
Revised: B † 2( β ‡ ) 2 , from (34) we immediately conclude ∆ ≥ 0 ; i.e., the multiple correlations upper bound (15) holds true.

Original: On the other hand, for every n = 1 , 2 , . . .
Revised: On the other hand, for every n = 1 , 2 , . . .

Original: , N we have
Revised: , N we have

Original: [Equation]
Revised: [Equation]

Original:  We combine the second part of (35) with (27) integrated w.r.t.
Revised: and [Equation] We combine the second part of (35) with (27) integrated w.r.t.

Original: P n :
Revised: P n :

Original: [Equation]
Revised: [Equation]

Original: Finally [Equation] we use the triangle inequality for (cid:107) · (cid:107) P n and conclude[Equation]
Revised: Finally, we use the triangle inequality for (cid:107) · (cid:107) P n and conclude

Original: i

Original: [Equation]
Revised: [Equation]

Original: e.,
Revised: i.e.,

Original: [Equation]
Revised: [Equation]

Original: where the ﬁrst inequality follows from (Korotin 2019a
Revised: where the ﬁrst inequality follows from (Korotin et al., 2019a, Lemma A.2).

Original: C E XPERIMENTAL DETAILS AND EXTRA RESULTS
Revised: C E XPERIMENTAL DETAILS AND EXTRA RESULTS

Original: In this section, we provide experimental details and additional results.
Revised: In this section, we provide experimental details and additional results.

Original: In Subsection C.1, we demonstrate qualitative results of computed barycenters in the 2-dimensional space.
Revised: In Subsection C.1, we demonstrate qualitative results of computed barycenters in the 2-dimensional space.

Original: In Subsection
Revised: In Subsection

Original: C.2, we discuss used metrics in more detail.
Revised: C.2, we discuss used metrics in more detail.

Original: In Subsection C.3, we list the used hyperparameters of our method (C W 2 B) and methods [SC W 2 B], [CR W B].
Revised: In Subsection C.4, we list the used hyperparameters of our method (C W 2 B) and methods [SC W 2 B], [CR W B].

Original: C. A DDITIONAL T OY E XPERIMENTS IN 2D
Revised: C. A DDITIONAL T OY E XPERIMENTS IN 2D

Original: We provide additional qualitative examples of computed barycenters of probability measures on R 2 .
Revised: We provide additional qualitative examples of computed barycenters of probability measures on R 2 .

Original: In Figure 3, we consider the location-scatter family F ( P 0 ) with P 0 =
Revised: In Figure 3, we consider the location-scatter family F ( P 0 ) with P 0 =

Original: Uniform [ −√ 3 , √ 3] D .
Revised: Uniform [ −√ 3 , √ 3] D .

Original: In principle, all the methods capture the true barycenter.
Revised: In principle, all the methods capture the true barycenter.

Original: However, the generated distribution g(cid:93) S of [SC W 2 B] (Figure 3c) provides samples that lies outside of the actual barycenter’s support (Figure 3b).
Revised: However, the generated distribution g(cid:93) S of [SC W 2 B] (Figure 3c) provides samples that lies outside of the actual barycenter’s support (Figure 3b).

Original: Also, in [CR W B] method, one of the potentials’ pushforward measure (top-right in Figure 3e) has visual artifacts.
Revised: Also, in [CR W B] method, one of the potentials’ pushforward measure (top-right in Figure 3e) has visual artifacts.

Original: [Figure]
Revised: [Figure]

Original: In Figure 4, we consider the Gaussian Mixture example by (Fan et al., 2020).
Revised: In Figure 4, we consider the Gaussian Mixture example by (Fan et al., 2020).

Original: The barycenter computed by [SC W 2 B] method (Figure 4b) suffers from the behavior similar to mode collapse. 
Revised: The barycenter computed by [SC W 2 B] method (Figure 4b) suffers from the behavior similar to mode collapse. [Figure] [Figure]

Original: C. METRICS
Revised: M

Revised: [Figure]

Original: The unexplained variance percentage (UVP) (introduced in Section 5) is a natural and straightforward metric to assess the quality of the computed barycenter.
Revised: The unexplained variance percentage (UVP) (introduced in Section 5) is a natural and straightforward metric to assess the quality of the computed barycenter.

Original: However, it is difﬁcult to compute in high dimensions: it requires computation of the Wasserstein-2 distance.
Revised: However, it is difﬁcult to compute in high dimensions: it requires computation of the Wasserstein-2 distance.

Original: Thus, we use different but highly related metrics L 2 -UVP and B W 22 -UVP.
Revised: Thus, we use different but highly related metrics L 2 -UVP and B W 22 -UVP.

Original: To access the quality of the recovered potentials { ψ † n } we use L 2 -UVP deﬁned in (17).
Revised: To access the quality of the recovered potentials { ψ † n } we use L 2 -UVP deﬁned in (17).

Original: L 2 -UVP compares not just pushforward distribution ∇ ψ † n (cid:93) P n with the barycenter P , but also the resulting transport map with the optimal transport map ∇ ψ ∗ n .
Revised: L 2 -UVP compares not just pushforward distribution ∇ ψ † n (cid:93) P n with the barycenter P , but also the resulting transport map with the optimal transport map ∇ ψ ∗ n .

Original: It bounds UVP ( ∇ ψ † n (cid:93) P n ) from above, thanks[Figure] to (Korotin et al., 2019a, Lemma A.2).
Revised: It bounds UVP ( ∇ ψ † n (cid:93) P n ) from above, thanks to (Korotin et al., 2019a, Lemma A.2).

Original: Besides, L 2 -UVP naturally admits unbiased Monte Carlo estimates using random samples from P n .
Revised: Besides, L 2 -UVP naturally admits unbiased Monte Carlo estimates using random samples from P n .

Original: For measure-based optimization method, we also evaluate the quality of the generated measure g(cid:93) S using Bures-Wasserstein UVP deﬁned in (18).
Revised: For measure-based optimization method, we also evaluate the quality of the generated measure g(cid:93) S using Bures-Wasserstein UVP deﬁned in (18).

Original: For measures P , Q whose covariance matrices are not degenerate, B W 22 is given by
Revised: For measures P , Q whose covariance matrices are not degenerate, B W 22 is given by

Original: [Equation]
Revised: [Equation]

Original: Bures-Wasserstein metric compares P , Q by considering only their ﬁrst and second moments.
Revised: Bures-Wasserstein metric compares P , Q by considering only their ﬁrst and second moments.

Original: It is known that B W 22 ( P , Q ) is a lower bound for W 22 ( P , Q ) , see (Dowson & Landau, 1982).
Revised: It is known that B W 22 ( P , Q ) is a lower bound for W 22 ( P , Q ) , see (Dowson & Landau, 1982).

Original: Thus, we have B W 22 -UVP ( g(cid:93) S ) ≤ UVP ( g(cid:93) S ) .
Revised: Thus, we have B W 2 2 -UVP ( g(cid:93) S ) ≤ UVP ( g(cid:93) S ) .

Original: In practice, to compute B W 22 -UVP ( g(cid:93) S ) , we estimate means and covariance matrices of distributions by using 10 5 random samples. 
Revised: In practice, to compute B W 2 2 -UVP ( g(cid:93) S ) , we estimate means and covariance matrices of distributions by using 10 5 random samples. C. C YCLE C ONSISTENCY AND C ONGRUENCE IN P RACTICE

Revised: To assess the effect of the regularization of cycle consistency and the congruence condition in practice, we run the following sanity checks.

Revised: For cycle consistency, for each input distribution P n we estimate (by drawing samples from P n ) the value (cid:107)∇ ψ ‡ n ◦ ∇ ψ † n ( x ) − x (cid:107) 2 P n / Var ( P n ) .

Original: C. T RAINING H YPERPARAMETERS
Revised: This metric can be viewed as an analog of the L 2 -UVP that we used for assessing the resulting transport maps. In all the experiments, this value does not exceed 2%, which means that cycle consistency and hence conjugacy are satisﬁed well. For the congruence condition, we need to check that (cid:80) Nn

Revised: =1 α n ψ † n ( x ) = (cid:107) x (cid:107) 2 / 2 .

Revised: However, we do not know any straightforward metric to check this exact condition that is scaled properly by the variance of the distributions.

Revised: Thus, we propose to use an alternative metric to check a slightly weaker condition on gradients, e.g., that (cid:80) Nn

Original: {+=
+}
Revised: {+=
+}

Revised: α n ∇ ψ † n ( x ) = x .

Revised: This is weaker due to the ambiguity of the additive constants.

Revised: For this we can compute (cid:107) (cid:80) Nn

Original: {+=
+}
Revised: {+=
+}

Revised: α n ∇ ψ † n ( x ) − x (cid:107) 2 P / Var ( P ) , where the denominator is the variance of the true barycenter.

Revised: We computed this metric and found that it is also less than 2% in all the cases, which means that congruence condition is mostly satisﬁed.

Revised: C. T RAINING H YPERPARAMETERS

Original: The code is written using the PyTorch framework.
Revised: The code is written using the PyTorch framework.

Original: The networks are trained on a single GTX 1080Ti.
Revised: The networks are trained on a single GTX 1080Ti.

Original: C.3. W ASSERSTEIN -2 C ONTINUOUS B ARYCENTERS (C W 2 B, OUR METHOD )
Revised: C.4. W ASSERSTEIN -2 C ONTINUOUS B ARYCENTERS (C W 2 B, OUR METHOD )

Original: Regularization .
Revised: Regularization .

Original: We use τ = 5 and ˆ P = N (0 , I D ) in our congruence regularizer τ · R ˆ P 1 .
Revised: We use τ = 5 and ˆ P = N (0 , I D ) in our congruence regularizer τ · R ˆ P 1 .

Original: We use λ = 10 for the cycle regularization λ
Revised: We use λ = 10 for the cycle regularization λ

Original: · R P n 2 for all n = 1 , 2 , . . .
Revised: · R P n 2 for all n = 1 , 2 , . . .

Original: , N .
Revised: , N .

Original: Neural Networks (Potentials) .
Revised: Neural Networks (Potentials) .

Original: To approximate potentials { ψ † n , ψ ‡ n } in dimension D , we use
Revised: To approximate potentials { ψ † n , ψ ‡ n } in dimension D , we use

Original: DenseICNN [2; max(64 , 2 D ) , max(64 , 2 D ) , max(32 , D )] with CELU activation function.
Revised: DenseICNN [2; max(64 , 2 D ) , max(64 , 2 D ) , max(32 , D )] with CELU activation function.

Original: DenseICNN is an input-convex dense architecture with additional convex quadratic skip connections.
Revised: DenseICNN is an input-convex dense architecture with additional convex quadratic skip connections.

Original: Here 2 is the rank of each input-quadratic skip-connection’s
Revised: Here 2 is the rank of each input-quadratic skip-connection’s

Original: Hessian matrix.
Revised: Hessian matrix.

Original: Each following number max( · , · ) represents the size of a hidden dense layer in the sequantial part of the network.
Revised: Each following number max( · , · ) represents the size of a hidden dense layer in the sequantial part of the network.

Original: For detailed discussion of the architecture see (Korotin et al., 2019a,
Revised: For detailed discussion of the architecture see (Korotin et al., 2019a,

Original: Section B.2).
Revised: Section B.2).

Original: Training process .
Revised: Training process .

Original: We perform training according to Algorithm 1 of Appendix A. We set batch size
Revised: We perform training according to Algorithm 1 of Appendix A. We set batch size

Original: K = 1024 and balancing coefﬁcient γ = 0 .
Revised: K = 1024 and balancing coefﬁcient γ = 0 .

Original: We use Adam optimizer by (Kingma & Ba, 2014) with a ﬁxed learning rate 10 − 3 .
Revised: We use Adam optimizer by (Kingma & Ba, 2014) with a ﬁxed learning rate 10 − 3 .

Original: The total number of iterations is set to 50000.
Revised: The total number of iterations is set to 50000.

Original: C.3. S CALABLE COMPUTATION OF W ASSERSTEIN B ARYCENTERS (SC W 2 B)
Revised: C.4. S CALABLE COMPUTATION OF W ASSERSTEIN B ARYCENTERS (SC W 2 B)

Original: Generator Neural Network .
Revised: Generator Neural Network .

Original: For the input noise distribution of the generative model we use S = N (0 , I D ) .
Revised: For the input noise distribution of the generative model we use S = N (0 , I D ) .

Original: For the generative network g : R D
Revised: For the generative network g : R D

Original: → R D we use a fully-connected sequential ReLU  [max(100 , 2 D [Equation] , max
Revised: → R D we use a fully-connected sequential ReLU network with hidden layer sizes [max(100 , 2 D ) , max(100 , 2 D ) , max(100 , 2 D )] .

Original: [Equation]

Original: 100 , 2 D

Original: [Equation]

Original: , max

Original: [Equation]

Original: 100 , 2 D )] .

Original: Before the main optimization, we pre-train the network to satisfy g ( z )
Revised: Before the main optimization, we pre-train the network to satisfy g ( z )

Original: ≈ z for all z ∈ R D .
Revised: ≈ z for all z ∈ R D .

Original: This has been empirically veriﬁed as a better option than random initialization of network’s weights.
Revised: This has been empirically veriﬁed as a better option than random initialization of network’s weights.

Original: Neural Networks (Potentials) .
Revised: Neural Networks (Potentials) .

Original: We used exactly the same networks as in Subsection C.3.1.
Revised: We used exactly the same networks as in Subsection C.4.1.

Original: Training process .
Revised: Training process .

Original: We perform training according to the min-max-min procedure described by (Fan et al., 2020, Algorithm 1).
Revised: We perform training according to the min-max-min procedure described by (Fan et al., 2020, Algorithm 1).

Original: The batch size is set to 1024 .
Revised: The batch size is set to 1024 .

Original: We use Adam optimizer by (Kingma & Ba, 2014) with ﬁxed learning rate 10 − 3 for potentials and 10 − 4 for generative network g .
Revised: We use Adam optimizer by (Kingma & Ba, 2014) with ﬁxed learning rate 10 − 3 for potentials and 10 − 4 for generative network g .

Original: The number of iterations of the outer cycle ( min -max-min) number of iterations is set to 15000.
Revised: The number of iterations of the outer cycle ( min -max-min) number of iterations is set to 15000.

Original: Following (Fan et al., 2020), we use 10 iterations per the middle cycle (min- max -min) and 6 iterations per the inner cycle (min-max- min ).
Revised: Following (Fan et al., 2020), we use 10 iterations per the middle cycle (min- max -min) and 6 iterations per the inner cycle (min-max- min ).

Original: C.3. C ONTINUOUS R EGULARIZED W ASSERSTEIN B ARYCENTERS (CR W B)
Revised: C.4. C ONTINUOUS R EGULARIZED W ASSERSTEIN B ARYCENTERS (CR W B)

Original: Regularization .
Revised: Regularization .

Original: [CR W B] method uses regularization to keep the potentials conjugate.
Revised: [CR W B] method uses regularization to keep the potentials conjugate.

Original: The authors impose entropy or L 2 regularization w.r.t.
Revised: The authors impose entropy or L 2 regularization w.r.t.

Original: some proposal measure ˆ P ; see (Li et al., 2020, Section 3) for more details.
Revised: some proposal measure ˆ P ; see (Li et al., 2020, Section 3) for more details.

Original: Following the source code provided by the authors, we use L 2 regularization (empirically shown as a more stable option than entropic regularization).
Revised: Following the source code provided by the authors, we use L 2 regularization (empirically shown as a more stable option than entropic regularization).

Original: The regularization measure ˆ P is set to be the uniform measure on a box containing the support of all the source distributions, estimated by sampling.
Revised: The regularization measure ˆ P is set to be the uniform measure on a box containing the support of all the source distributions, estimated by sampling.

Original: The regularization parameter (cid:15) is set to 10 − 4 .
Revised: The regularization parameter (cid:15) is set to 10 − 4 .

Original: Neural Networks (Potentials) .
Revised: Neural Networks (Potentials) .

Original: To approximate potentials { ψ † n , ψ ‡ n } in dimension D
Revised: To approximate potentials { ψ † n , ψ ‡ n } in dimension D

Original: , we use fullyconnected sequential ReLU neural networks with layer sizes given by [Equation]
Revised: , we use fullyconnected sequential ReLU neural networks with layer sizes given by [max(128 , 4 D ) , max(128 , 4 D ) , max(128 , 4 D )] .

Original: We have also tried using DenseICNN architecture, but did not experience any performance gain.
Revised: We have also tried using DenseICNN architecture, but did not experience any performance gain.

Original: Training process .
Revised: Training process .

Original: We perform training according to (Li et al., 2020, Algorithm 1).
Revised: We perform training according to (Li et al., 2020, Algorithm 1).

Original: We set batch size to 1024 .
Revised: We set batch size to 1024 .

Original: We use Adam optimizer by (Kingma & Ba, 2014) with ﬁxed learning rate 10 − 3 .
Revised: We use Adam optimizer by (Kingma & Ba, 2014) with ﬁxed learning rate 10 − 3 .

Original: The total number of iterations is set to 50000.
Revised: The total number of iterations is set to 50000.