Towards Marginal Fairness Sliced Wasserstein Barycenter

## Abstract

Abstract The sliced Wasserstein barycenter (SWB) is a widely acknowledged method for efficiently generalizing the averaging operation within probability measure spaces. However, achieving marginal fairness SWB, ensuring approximately equal distances from the barycenter to marginals, remains unexplored. The uniform weighted SWB is not necessarily the optimal choice to obtain the desired marginal fairness barycenter due to the heterogeneous structure of marginals and the non-optimality of the optimization. As the first attempt to tackle the problem, we define the marginal fairness sliced Wasserstein barycenter (MFSWB) as a constrained SWB problem. Due to the computational disadvantages of the formal definition, we propose two hyperparameter-free and computationally tractable surrogate MFSWB problems that implicitly minimize the distances to marginals and encourage marginal fairness at the same time. To further improve the efficiency, we perform slicing distribution selection and obtain the third surrogate definition by introducing a new slicing distribution that focuses more on marginally unfair projecting directions. We discuss the relationship of the three proposed problems and their relationship to sliced multi-marginal Wasserstein distance. Finally, we conduct experiments on finding 3D point-clouds averaging, color harmonization, and training of sliced Wasserstein autoencoder with class-fairness representation to show the favorable performance of the proposed surrogate MFSWB problems.

## 1 Introduction

Wasserstein barycenter [] generalizes "averaging" to the space of probability measures. In particular, a Wasserstein barycenter is a probability measure that minimizes a weighted sum of Wasserstein distances between it and some given marginal probability measures. Due to the rich geometry of the Wasserstein distance [], the Wasserstein barycenter can be seen as the Fréchet mean [] on the space of probability measures. As a result, Wasserstein barycenter has been applied widely to various applications in machine learning such as Bayesian inference [, ], domain adaptation [], clustering [], sensor fusion [], text classification [], and so on. Moreover, Wasserstein barycenter is also a powerful tool for computer graphics since it can be used for texture mixing [], style transfer [], shape interpolation [], and many other tasks on many other domains.

Despite being useful, it is very computationally expensive to compute Wasserstein barycenter. In more detail, the computational complexity of Wasserstein barycenter is $\mathcal{O}(n^{3}\log n)$ when using linear programming [] where $n$ is the largest number of supports of marginal probability measures. When using entropic regularization for optimal transport [], the computational complexity is reduced to $\mathcal{O}(n^{2})$ []. Nevertheless, quadratic scaling is not enough when the number of supports approaches a hundred thousand or a million. To address the issue, sliced Wassserstein barycenter (SWB) is introduced in [] by replacing Wasserstein distance with its sliced variant i.e., sliced Wasseretein (SW) distance. Thank to the closed-form of Wasserstein distance in one-dimension, SWB has a low time complexity i.e., $\mathcal{O}(n\log n)$ which enables fast computation. Combining with the fact that sliced Wasserstein is equivalent to Wasserstein distance in bounded domains [] and sliced Wasserstein does not suffer from the curse of dimensionality [, , , ], SWB becomes a great and scalable alternative choice of Wasserstein barycenter.

In some applications, we might want to find a barycenter that minimizes the distances to marginals and has equal distances to marginals at the same time e.g., constructing shape template for a group of shapes [, ] that can be further used in downstream tasks, exact balance style mixing between images [], fair generative modeling [], and so on. We refer to such a barycenter as a marginal fairness barycenter. Both the Wasserstein barycenter and SWB are defined based on a given set of marginal weights (marginal coefficients), and these weights represent the important levels of marginals toward the barycenter. Nevertheless, a uniform (weights) barycenter does not necessarily lead to the desired marginal fairness barycenter as shown in Figure 1. Moreover, obtaining the marginal fairness barycenter is challenging since such a barycenter might not exist and might not be identifiable given non-global-optimal optimization (Karcher mean problem). To the best of our knowledge, there is no prior work that investigates finding a marginal fairness barycenter.

Figure: Figure 1: The uniform SWB and the MFSWB of 4 Gaussian distributions.
Refer to caption: x1.png

In this work, we make the first attempt to tackle the marginal fairness barycenter problem i.e., we focus on finding marginal fairness SW barycenter (MFSWB) to utilize the scalability of SW distance.

Contribution: In summary, our main contributions are four-fold:

1. We define the marginal fairness SW barycenter (MFSWB) problem which is a constraint barycenter problem, where the constraint tries to limit the average pair-wise absolute difference between distances from the barycenter to marginals. We derive the dual form of MFSWB, discuss its computation, and address its computational challenges.

2. To address the issue, we propose surrogate definitions of MFSWB which are hyperparameter-free and computationally tractable. Motivated by Fair PCA [], we propose the first surrogate MFSWB that minimizes the largest SW distance from the barycenter to the marginals. To address the problem of biased gradient estimation of the first surrogate MFSWB, we propose the second surrogate MFSWB which is the expectation of the largest one-dimensional Wasserstein distance from the projected barycenter to the projected marginals. We show that the second surrogate is an upper bound of the first surrogate and can yield an unbiased gradient estimator. We extend further the second surrogates to the third surrogate by applying slicing distribution selection and showing that the third surrogate is an upper bound of the previous two.

3. We discuss the connection between the proposed surrogate MFSWB problems with the sliced multi-marginal Wasserstein (SMW) distance with the maximal ground metric. In particular, solving the proposed MFSWB problems is equivalent to minimizing a lower bound of the SMW. By showing that the SMW with the maximal ground metric is a generalized metric, we convey that it is safe to use the proposed surrogate MFSWB problems.

4. We conduct simulations with Gaussian and experiments on various applications including 3D point-cloud averaging, color harmonization, and sliced Wasserstein autoencoder with class-fairness representation to demonstrate the favorable performance of the proposed surrogate definitions.

Organization. We first discuss some preliminaries on SW distance, SWB and its computation, and sliced multi-marginal Wasserstein distance in Section 2. We then introduce the formal definition and surrogate definitions of marginal fairness SWB in Section 3. After that, we conduct experiments to demonstrate the favorable performance in fairness of the proposed definitions in Section 4. We conclude the paper and provide some future directions in Section 5. Finally, we defer
the proofs of key results, the discussion on related works, and additional materials in the Appendices.

## 2 Preliminaries

Sliced Wasserstein distance. The definition of sliced Wasserstein (SW) distance [] between two probability measures $\mu_{1}\in\mathcal{P}_{p}(\mathbb{R}^{d})$ and $\mu_{2}\in\mathcal{P}_{p}(\mathbb{R}^{d})$ is:

$$ $\displaystyle\text{SW}_{p}^{p}(\mu_{1},\mu_{2})=\mathbb{E}_{\theta\sim\mathcal {U}(\mathbb{S}^{d-1})}[\text{W}_{p}^{p}(\theta\sharp\mu_{1},\theta\sharp\mu_{2 })],$ (1) $$

where the Wasserstein distance has a closed form in one-dimension which is $\text{W}_{p}^{p}(\theta\sharp\mu_{1},\theta\sharp\mu_{2})=\int_{0}^{1}|F_{
\theta\sharp\mu_{1}}^{-1}(z)-F_{\theta\sharp\mu_{2}}^{-1}(z)|^{p}dz$
where $F_{\theta\sharp\mu_{1}}$ and $F_{\theta\sharp\mu_{2}}$ are the cumulative
distribution function (CDF) of $\theta\sharp\mu_{1}$ and $\theta\sharp\mu_{2}$ respectively.

Sliced Wasserstein Barycenter. The definition of the sliced Wasserstein barycenter (SWB) problem [] of $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$ with marginal weights $\omega_{1},\ldots,\omega_{K}>0$ ($\sum_{i=k}^{K}\omega_{k}=1$) is defined as:

$$ $\displaystyle\min_{\mu}\mathcal{F}(\mu;\mu_{1:K},\omega_{1:K});\quad\mathcal{F }(\mu;\mu_{1:K},\omega_{1:K})=\sum_{k=1}^{K}\omega_{k}\text{SW}_{p}^{p}(\mu, \mu_{k}).$ (2) $$

Computation of parametric SWB. Let $\mu_{\phi}$ be parameterized by $\phi\in\Phi$, SWB can be solved by gradient-based optimization. In that case, the interested quantity is the gradient $\nabla_{\phi}\mathcal{F}(\mu_{\phi};\mu_{1:K},\omega_{1:K})=\sum_{k=1}^{K}
\omega_{k}\nabla_{\phi}\text{SW}_{p}^{p}(\mu_{\phi},\mu_{k})$. However, the gradient $\nabla_{\phi}\text{SW}_{p}^{p}(\mu_{\phi},\mu_{k})=\nabla_{\phi}\mathbb{E}_{
\theta\sim\mathcal{U}(\mathbb{S}^{d-1})}[\text{W}_{p}^{p}(\theta\sharp\mu_{
\phi},\theta\sharp\mu_{k})]=\mathbb{E}_{\theta\sim\mathcal{U}(\mathbb{S}^{d-1}
)}[\nabla_{\phi}\text{W}_{p}^{p}(\theta\sharp\mu_{\phi},\theta\sharp\mu_{k})]$ for any $k=1,\ldots,K$ is intractable due to the intractability of SW with the expectation with respect to the uniform distribution over the unit-hypersphere. Therefore, Monte Carlo estimation is used. In particular, projecting directions $\theta_{1},\ldots,\theta_{L}$ are sampled i.i.d from $\mathcal{U}(\mathbb{S}^{d-1})$, and the stochastic gradient estimator is formed:

$$ $\displaystyle\nabla_{\phi}\text{SW}_{p}^{p}(\mu_{\phi},\mu_{k})\approx\frac{1} {L}\sum_{l=1}^{L}\nabla_{\phi}\text{W}_{p}^{p}(\theta_{l}\sharp\mu_{\phi}, \theta_{l}\sharp\mu_{k}).$ (3) $$

With the stochastic gradient, the SWB can be solved by using a stochastic gradient descent algorithm. We refer the reader to Algorithm 1 in Appendix B for more detail. Specifically, we now discuss the discrete SWB i.e., marginals and the barycenter are discrete measures.

Free supports barycenter. In this setting, we have $\mu_{\phi}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$, $\mu_{k}=\frac{1}{n}\sum_{i=1}^{n}\delta_{y_{i}}$, and $\phi=(x_{1},\ldots,x_{n})$, we can compute the (sub-)gradient with the time complexity $\mathcal{O}(n\log n)$:

$$ $\displaystyle\nabla_{x_{i}}\text{W}_{p}^{p}(\theta\sharp\mu_{\phi},\theta \sharp\mu_{k})=p|\theta^{\top}x_{i}-\theta^{\top}y_{\sigma(i)}|^{p-1}\text{ sign}(\theta^{\top}x_{i}-\theta^{\top}y_{\sigma(i)})\theta,$ (4) $$

where $\sigma=\sigma_{1}\circ\sigma_{2}^{-1}$ with $\sigma_{1}$ and $\sigma_{2}$ are any sorted permutation of $\{x_{1},\ldots,x_{n}\}$ and $\{y_{1},\ldots,y_{n}\}$.

Fixed supports barycenter. In this setting, we have $\mu_{\phi}=\sum_{i=1}^{n}\phi_{i}\delta_{x_{i}}$, $\mu_{k}=\sum_{i=1}^{n}\beta_{i}\delta_{x_{i}}$, $\sum_{i=1}^{n}\phi_{i}=\sum_{i=1}^{n}\beta_{i}$ and $\phi=(\phi_{1},\ldots,\phi_{n})$. We can compute the gradient as follows:

$$ $\displaystyle\nabla_{\phi}\text{W}_{p}^{p}(\theta\sharp\mu_{\phi},\theta\sharp \mu_{k})=\bm{f}^{\star},$ (5) $$

where $\bm{f}^{\star}$ is the first optimal Kantorovich dual potential of $\text{W}_{p}^{p}(\theta\sharp\mu_{\phi},\theta\sharp\mu_{k})$ which can be obtained with the time complexity of $\mathcal{O}(n\log n)$. We refer the reader to Proposition 1 in [] for the detail and Algorithm 1 in [] for the computational algorithm.

When the supports or weights of the barycenter are the output of a parametric function, we can use the chain rule to estimate the gradient of the parameters of the function. For the continuous case, we can approximate the barycenter and marginals by their empirical versions, and then perform the estimation in the discrete case. Since the sample complexity of SW is $\mathcal{O}(n^{-1/2})$ [, , , ], the approximation error will reduce fast with the number of support $n$ increases. Another option is to use continuous Wasserstein solvers [, , ], however, this option is not as simple as the first one.

Sliced Multi-marginal Wasserstein Distance. Given $K\geq 1$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, sliced Multi-marginal Wasserstein Distance [] (SMW) is defined as:

$$ $\displaystyle SMW_{p}^{p}(\mu_{1:K};c)=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1} ,\ldots,\mu_{K})}\int c(\theta^{\top}x_{1},\ldots,\theta^{\top}x_{K})^{p}d\pi( x_{1},\ldots,x_{K})\right],$ (6) $$

where the expectation is under $\theta\sim\mathcal{U}(\mathbb{S}^{d-1})$. When using the barycentric cost i.e., $c(\theta^{\top}x_{1},\ldots,\theta^{\top}x_{K})^{p}=\sum_{k=1}^{K}\beta_{k}
\left|\theta^{\top}x_{k}-\sum_{k^{\prime}=1}^{K}\beta_{k^{\prime}}\theta^{\top
}x_{k^{\prime}}\right|^{p}$ for $\beta_{k}>0\forall k$ and $\sum_{k}\beta_{k}=1$. Minimizing $SMW_{p}^{p}(\mu_{1:K},\mu;c)$ with respect to $\mu$ is equivalent to a barycenter problem. We refer the reader to Proposition 7 in [] for more detail.

## 3 Marginal Fairness Sliced Wasserstein Barycenter

We first formally define the marginal fairness Sliced Wasserstein barycenter in Section 3.1. We then propose surrogate problems in Section 3.2. Finally, we discuss the connection of the proposed surrogate problems to sliced multi-marginal Wasserstein in Section 3.3.

### 3.1 Formal Definition

Now, we define the Marginal Fairness Sliced Wasserstein barycenter (MFSWB) problem by adding marginal fairness constraints to the SWB problem.

###### Definition 1 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, admissible $\epsilon\geq 0$ for $i=1,\ldots,K$ and $j=i+1,\ldots,K$, the Marginal Fairness Sliced Wasserstein barycenter (MFSWB) is defined as:

$$ $\displaystyle\min_{\mu}\frac{1}{K}\sum_{k=1}^{K}\text{SW}_{p}^{p}(\mu,\mu_{k}) \text{ s.t. }\frac{2}{(K-1)K}\sum_{i=1}^{K}\sum_{j=i+1}^{K}|\text{SW}_{p}^{p}( \mu,\mu_{i})-\text{SW}_{p}^{p}(\mu,\mu_{j})|\leq\epsilon.$ (7) $$

###### Remark .

We want $\epsilon$ in Definition 1 to be close to 0 0 i.e., $\mu_{1},\ldots,\mu_{K}$ are on the $SW_{p}$-sphere with the center $\mu$. However, for a too-small value of $\epsilon$, there might not exist a solution $\mu$.

Duality objective. For admissible $\epsilon>0$, there exist a Lagrange multiplier $\lambda$ such that we have the dual form

$$ $\displaystyle\mathcal{L}(\mu,\lambda)=\frac{1}{K}\sum_{k=1}^{K}\text{SW}_{p}^{ p}(\mu,\mu_{k})+\frac{2\lambda}{(K-1)K}\sum_{i=1}^{K}\sum_{j=i+1}^{K}|\text{SW }_{p}^{p}(\mu,\mu_{i})-\text{SW}_{p}^{p}(\mu,\mu_{j})|-\lambda\epsilon.$ (8) $$

Computational challenges. Firstly, MFSWB in Definition 1 requires admissible $\epsilon>0$ to guarantee the existence of the barycenter $\mu$. In practice, it is unknown if a value of $\epsilon$ satisfies such a property. Secondly, given a $\epsilon$, it is not trivial to obtain the optimal Lagrange multiplier $\lambda^{\star}$ in Equation equation 8 to minimize the duality gap which can be non-zero (weak-duality). Thirdly, using directly the dual objective Equation equation 8 requires hyperparameter tunning for $\lambda$ and might not provide a good landscape for optimization. Moreover, we cannot obtain an unbiased gradient estimate of $\phi$ in the case of the parametric barycenter $\mu_{\phi}$. In greater detail, the Monte Carlo estimation for the absolute distance between two SW distances is biased. Finally, Equation equation 8 has a quadratic time complexity and space complexity in terms of the number of marginals i.e., $\mathcal{O}(K^{2})$.

###### Definition 1 .

###### Remark .

### 3.2 Surrogate Definitions

Since it is not convenient to use the formal MFSWB in applications, we propose two surrogate definitions of MFSWB that are free of hyperparameters and computationally friendly.

First Surrogate Definition. Motivated by Fair PCA [], we propose a practical surrogate MFSWB problem that is hyperparameter-free.

###### Definition 2 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, the surrogate marginal fairness sliced Wasserstein barycenter (s-MFSWB) problem is defined as:

$$ $\displaystyle\min_{\mu}\mathcal{SF}(\mu;\mu_{1:K});\quad\mathcal{SF}(\mu;\mu_{ 1:K})=\max_{k\in\{1,\ldots,K\}}SW_{p}^{p}(\mu,\mu_{k}).$ (9) $$

The s-MFSWB problem tries to minimize the maximal distance from the barycenter to the marginals. Therefore, it can minimize indirectly the overall distances between the barycenter to the marginals and implicitly make the distances to marginals approximately the same.

Gradient estimator. Let $\mu_{\phi}$ be paramterized by $\phi\in\Phi$, and $\mathcal{F}(\phi,k)=SW_{p}^{p}(\mu_{\phi},\mu_{k})$, we would like to compute $\nabla_{\phi}\max_{k\in\{1,\ldots,K\}}\mathcal{F}(\phi,k)$. By Danskin’s envelope theorem [], we have $\nabla_{\phi}\max_{k\in\{1,\ldots,K\}}\mathcal{F}(\phi,k)=\nabla_{\phi}
\mathcal{F}(\phi,k^{\star})=\nabla_{\phi}\text{SW}_{p}^{p}(\mu_{\phi},\mu_{k^{
\star}})$ for $k^{\star}=\operatorname*{arg\,max}_{k\in\{1,\ldots,K\}}\mathcal{F}(\phi,k)$. Nevertheless, $k^{\star}$ is intractable due to the intractablity of $SW_{p}^{p}(\mu_{\phi},\mu_{k})$ for $k=1,\ldots,K$. Hence, we can form the estimation $\hat{k}^{\star}=\operatorname*{arg\,max}_{k\in\{1,\ldots,K\}}\widehat{SW}_{p}^
{p}(\mu_{\phi},\mu_{k};L)$ where $\widehat{SW}_{p}^{p}(\mu_{\phi},\mu_{k};L)=\frac{1}{L}\sum_{l=1}^{L}\text{W}_{
p}^{p}(\theta_{l}\sharp\mu_{\phi},\theta_{l}\sharp\mu_{k})$ with $\theta_{1},\ldots,\theta_{L}\overset{i.i.d}{\sim}\mathcal{U}(\mathbb{S}^{d-1})$. Then, we can estimate $\nabla_{\phi}\text{SW}_{p}^{p}(\mu_{\phi},\mu_{\hat{k}^{\star}})$ as in Equation 3. We refer the reader to Algorithm 2 in Appendix B for the gradient estimation and optimization procedure. The downside of this estimator is that it is biased.

Second Surrogate Definition. To address the biased gradient issue of the first surrogate problem, we propose the second surrogate MFSWB problem.

###### Definition 3 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, the unbiased surrogate marginal fairness sliced Wasserstein barycenter (us-MFSWB) problem is defined as:

$$ $\displaystyle\min_{\mu}\mathcal{USF}(\mu;\mu_{1:K});\quad\mathcal{USF}(\mu;\mu _{1:K})=\mathbb{E}_{\theta\sim\mathcal{U}(\mathbb{S}^{d-1})}\left[\max_{k\in\{ 1,\ldots,K\}}W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})\right].$ (10) $$

In contrast to s-MFSWB which minimizes the maximal SW distance among marginals, us-MFSWB minimizes the expected value of the maximal one-dimensional Wasserstein distance among marginals. By considering fairness on one-dimensional projections, us-MFSWB can yield an unbiased gradient estimate which is the reason why it is named as unbiased s-MFSWB.

Gradient estimator. Let $\mu_{\phi}$ be paramterized by $\phi\in\Phi$, and $\mathcal{F}(\theta,\phi,k)=W_{p}^{p}(\theta\sharp\mu_{\phi},\theta\sharp\mu_{k})$, we would like to compute $\nabla_{\phi}\mathbb{E}_{\theta\sim\mathbb{S}^{d-1}}[\max_{k\in\{1,\ldots,K\}}
\mathcal{F}(\theta,\phi,k)]$ which is equivalent to $\mathbb{E}_{\theta\sim\mathbb{S}^{d-1}}[\nabla_{\phi}\max_{k\in\{1,\ldots,K\}}
\mathcal{F}(\theta,\phi,k)]$ due to the Leibniz’s rule. By Danskin’s envelope theorem, we have $\nabla_{\phi}\max_{k\in\{1,\ldots,K\}}\mathcal{F}(\theta,\phi,k)=\nabla_{\phi}
\mathcal{F}(\theta,\phi,k^{\star})=\nabla_{\phi}\text{W}_{p}^{p}(\theta\sharp
\mu_{\phi},\theta\sharp\mu_{k^{\star}})$ for $k^{\star}_{\theta}=\operatorname*{arg\,max}_{k\in\{1,\ldots,K\}}\mathcal{F}(
\theta,\phi,k)$ where we can estimate $\nabla_{\phi}\text{W}_{p}^{p}(\theta\sharp\mu_{\phi},\theta\sharp\mu_{k^{\star
}_{\theta}})$ can be computed as in Equation 4- 5. Overall, with $\theta_{1},\ldots,\theta_{L}\overset{i.i.d}{\sim}\mathcal{U}(\mathbb{S}^{d-1})$, we can form the final estimation $\frac{1}{L}\sum_{l=1}^{L}\nabla_{\phi}\text{W}_{p}^{p}(\theta_{l}\sharp\mu_{
\phi},\theta_{l}\sharp\mu_{k^{\star}_{\theta_{l}}})$ which is an unbiased estimate. We refer the reader to Algorithm 3 in Appendix B for the gradient estimation and optimization procedure.

###### Proposition 1 .

Given $K\geq 2$ marginals $\mu_{1:K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, we have $\mathcal{SF}(\mu;\mu_{1:K})\leq\mathcal{USF}(\mu;\mu_{1:K})$.

Proof of Proposition 1 is given in Appendix A.1. From the proposition, we see that minimizing the objective of us-MFSWB also reduces the objective of s-MFSWB implicitly.

###### Proposition 2 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, $\theta_{1},\ldots,\theta_{L}\overset{i.i.d}{\sim}\mathcal{U}(\mathbb{S}^{d-1})$, we have:

$$ $\displaystyle\mathbb{E}\left|\nabla_{\phi}\frac{1}{L}\sum_{l=1}^{L}\text{W}_{p }^{p}(\theta_{l}\sharp\mu_{\phi},\theta_{l}\sharp\mu_{k^{\star}_{\theta}})- \nabla_{\phi}\mathcal{USF}(\mu_{\phi};\mu_{1:K})\right|\leq\frac{1}{\sqrt{L}} \text{Var}\left[\nabla_{\phi}\text{W}_{p}^{p}(\theta\sharp\mu_{\phi},\theta \sharp\mu_{k^{\star}_{\theta}})\right]^{\frac{1}{2}},$ (11) $$

where $k^{\star}_{\theta}=\operatorname*{arg\,max}_{k\in\{1,\ldots,K\}}W_{p}^{p}(
\theta\sharp\mu_{\phi},\theta\sharp\mu_{k})$; and the expectation and variance are under the random projecting direction $\theta\sim\mathcal{U}(\mathbb{S}^{d-1})$

Proof of Proposition 2 is given in Appendix A.2. From the proposition, we know that the approximation error of the gradient estimator of us-MFSWB reduces at the order of $\mathcal{O}(L^{-1/2})$. Therefore, increasing $L$ leads to a better gradient approximation. The approximation could be further improved via Quasi-Monte Carlo methods [].

Third Surrogate Definition. The us-MFSWB in Definition 3 utilizes the uniform distribution as the slicing distribution, which is empirically shown to be non-optimal in statistical estimation []. Following the slicing distribution selection approach in [], we propose the third surrogate with a new slicing distribution that focuses on unfair projecting directions.

Marginal fairness energy-based slicing distribution. Since we want to encourage marginal fairness, it is natural to construct the slicing distribution based on fairness energy.

###### Definition 4 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, the marginal fairness energy-based slicing distribution $\sigma(\theta;\mu,\mu_{1:K})\in\mathcal{P}(\mathbb{S}^{d-1})$ is defined with the density function as follow:

$$ $\displaystyle f_{\sigma}(\theta;\mu,\mu_{1:K})\propto\exp\left(\max_{k\in\{1, \ldots,K\}}W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})\right),$ (12) $$

We see that the marginal fairness energy-based slicing distribution in Definition 4 put more mass to a projecting direction $\theta$ that has the larger maximal one-dimensional Wasserstein distance to marginals. Therefore, it will penalize more marginally unfair projecting directions.

Energy-based surrogate MFSWB. From the new proposed slicing distribution, we can define a new surrogate MFSWB problem, named energy-based surrogate MFSWB.

###### Definition 5 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, the energy-based surrogate marginal fairness sliced Wasserstein barycenter (us-MFSWB) problem is defined as:

$$ $\displaystyle\min_{\mu}\mathcal{ESF}(\mu;\mu_{1:K});\quad\mathcal{ESF}(\mu;\mu _{1:K})=\mathbb{E}_{\theta\sim\sigma(\theta;\mu,\mu_{1:K})}\left[\max_{k\in\{1 ,\ldots,K\}}W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})\right].$ (13) $$

Similar to the us-MFSWB, es-MFSWB utilizes the implicit one-dimensional marginal fairness. Nevertheless, es-MFSWB utilizes the marginal fairness energy-based slicing distribution to reweight the importance of each projecting direction instead of considering them equally.

###### Proposition 3 .

Given $K\geq 2$ marginals $\mu_{1:K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, we have $\mathcal{USF}(\mu;\mu_{1:K})\leq\mathcal{ESF}(\mu;\mu_{1:K})$.

Proof of Proposition 3 is given in Appendix A.3. From the proposition, we see that minimizing the objective of es-MFSWB reduces the objective of us-MFSWB implicitly which also decreases the objective of s-MFSWB (Proposition 1).

Gradient estimator. Let $\mu_{\phi}$ be parameterized by $\phi\in\Phi$, we want to estimate $\nabla_{\phi}\mathcal{ESF}(\mu_{\phi};\mu_{1:K})$. Since the slicing distribution is unnormalized, we use importance sampling to form an estimation.
With $\theta_{1},\ldots,\theta_{L}\overset{i.i.d}{\sim}\mathcal{U}(\mathbb{S}^{d-1})$, we can form the importance sampling stochastic gradient estimation:

$$ $\displaystyle\hat{\nabla}_{\phi}\mathcal{ESF}(\mu_{\phi};\mu_{1:K},L)=\frac{1} {L}\sum_{l=1}^{L}\left[\nabla_{\phi}\left(W_{p}^{p}(\theta_{l}\sharp\mu,\theta _{l}\sharp\mu_{k^{\star}_{\theta_{l}}})\frac{\exp\left(W_{p}^{p}(\theta_{l} \sharp\mu,\theta_{l}\sharp\mu_{k^{\star}_{\theta_{l}}})\right)}{\frac{1}{L} \sum_{i=1}^{L}\left[\exp\left(W_{p}^{p}(\theta_{i}\sharp\mu,\theta_{i}\sharp \mu_{k^{\star}_{\theta_{i}}})\right)\right]}\right)\right],$ $$

which can be further derived by using the chain rule and previously discussed techniques. It is worth noting that the above estimation is only asymptotically unbiased. We refer the reader to Algorithm 4 in Appendix B for the gradient estimation and optimization procedure.

Computational complexities of proposed surrogates. For the number of marginals $K$, the three proposed surrogates have a linear time complexity and space complexity i.e., $\mathcal{O}(K)$ which is the same as the conventional SWB and is better than $\mathcal{O}(K^{2})$ of the formal MFSWB. For the number of projections $L$, the number of supports $n$, and the number of dimensions $d$, the proposed surrogates have the time complexity of $\mathcal{O}(Ln(\log n+d))$ and the space complexity of $\mathcal{O}(L(n+d))$ which are similar to the formal MFSWB and SWB.

###### Definition 2 .

###### Definition 3 .

###### Proposition 1 .

###### Proposition 2 .

###### Definition 4 .

###### Definition 5 .

###### Proposition 3 .

### 3.3 Sliced multi-marginal Wasserstein distance with the maximal ground metric

To shed some light on the proposed substrates, we connect them to a special variant of Sliced multi-marginal Wasserstein (SMW) (see Equation 6) i.e., SMW with the maximal ground metric $c(\theta^{\top}x_{1},\ldots,\theta^{\top}x_{K})=\max_{i\in\{1,\ldots,K\},j\in
\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j}|$. We first show that SMW with the maximal ground metric is a generalized metric on the space of probability measures.

###### Proposition 4 .

Sliced multi-marginal Wasserstein distance with the maximal ground metric is a generalized metric i.e., it satisfies non-negativity, marginal exchangeability, generalized triangle inequality, and identity of identity of indiscernibles.

Proof of Proposition 4 is given in Appendix A.4. It is worth noting that SMW with the maximal ground metric has never been defined before. Since our work focuses on the MFSWB problem, we will leave the careful investigation of this variant of SMW to future work.

###### Proposition 5 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, the maximal ground metric $c(\theta^{\top}x_{1},\ldots,\theta^{\top}x_{K})=\max_{i\in\{1,\ldots,K\},j\in
\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j}|$, we have:

$$ $\displaystyle\min_{\mu_{1}}\mathcal{USF}(\mu_{1};\mu_{2:K})\leq\min_{\mu_{1}} SMW_{p}^{p}(\mu_{1},\mu_{2},\ldots,\mu_{K};c).$ (14) $$

Proof of Proposition 5 is given in Appendix A.5 and the inequality holds when changing $\mu_{1}$ to any $\mu_{i}$ with $i=2,\ldots,K$. Combining Proposition 1, we have the corollary of $\min_{\mu_{1}}\mathcal{SF}(\mu_{1};\mu_{2:K})\leq\min_{\mu_{1}}SMW_{p}^{p}(\mu
_{1},\mu_{2},\ldots,\mu_{K};c)$. From the proposition, we see that minimizing the us-MFSWB is equivalent to minimizing a lower bound of SMW with the maximal ground metric. Therefore, this proposition implies the us-MFSWB could try to minimize the multi-marginal distance. Moreover, this proposition can help to understand the proposed surrogates through the gradient flow of SMW. We can further extend the proposition to show the minimizing es-MFSWB objective is the same as minimizing a lower bound of energy-based SMW with the maximal ground metric, a new special variant of SMW. We refer the reader to Propositon 6 in Appendix B for more detail.

Figure: Figure 2: Barycenters from USWB, MFSWB with $\lambda=1$, s-MFSWB, us-MFSWB, and es-MFSWB along gradient iterations with the corresponding F-metric and W-metric.
Refer to caption: x4.png

###### Proposition 4 .

###### Proposition 5 .

## 4 Experiments

In this section, we compare the barycenter found by our proposed surrogate problems i.e., s-MFSWB, us-MFSWB, and es-MFSWB with the barycenter found by USWB and the formal MFSWB. For evaluation, we use two metrics i.e., the F-metric (F) and the W-metric (W) which are defined as follows:
$F=\frac{2}{K(K-1)}\sum_{i=1}^{K}\sum_{j=i+1}^{K}|W_{p}^{p}(\mu,\mu_{i})-W_{p}^
{p}(\mu,\mu_{j})|,\quad W=\frac{1}{K}\sum_{i=1}^{K}W_{p}^{p}(\mu,\mu_{i}),$
where $\mu$ is the barycenter, $\mu_{1},\ldots,\mu_{K}$ are the given marginals, and $W_{p}^{p}$ is the Wasserstein distance [] of the order $p$. Here, the F-metric represents the marginal fairness degree of the barycenter and the W-metric represents the centerness of the barycenter. For all following experiments, we use $p=2$ for the Wasserstein distance and barycenter problems.

### 4.1 Barycenter of Gaussians

We first start with a simple simulation with 4 marginals which are empirical distributions with 100 i.i.d samples from 4 Gaussian distributions i.e., $\mathcal{N}((0,0),I)$, $\mathcal{N}((20,0),I)$, $\mathcal{N}((18,8),I)$, and $\mathcal{N}((18,-8),I)$. We then find the barycenter which is represented as an empirical distribution with 100 supports initialized by sampling i.i.d from $\mathcal{N}((0,-5),I)$. We use stochastic gradient descent with 50000 iterations of learning rate 0.01, the number of projections 100. We show the visualization of the found barycenters with the corresponding F-metric and W-metric by using USWB, s-MFSWB, us-MFSWB, and es-MFSWB at iterations 0, 1000, 5000, and 50000 in Figure 2. We observe that the USWB does not lead to a marginal fairness barycenter. The three proposed surrogate problems help to find a better barycenter faster in both two metrics than USWB. At convergence i.e., iteration 50000, we see that USWB does not give a fair barycenter while the three proposed surrogates lead to a more fair barycenter. Among the proposed surrogates, es-MFSWB gives the most marginal fairness barycenter with a competitive centerness. The formal MFSWB (dual form with $\lambda=1$) leads to the most fair barycenter. However, the performance of the formal MFSWB is quite sensitive to $\lambda$. We show the visualization for $\lambda=0.1$ and $\lambda=10$ in Figure 5 in Appendix D.

### 4.2 3D Point-cloud Averaging

Figure: Figure 3: Averaging point-clouds with USWB, MFSWB ($\lambda=1$), s-MFSWB, us-MFSWB, and es-MFSWB.
Refer to caption: extracted/2405.07482v1/images/shape3.png

**Table 1: F-metric and W-metric along iterations in point-cloud averaging application.**
| Method | Iteration 0 | Iteration 1000 | Iteration 5000 | Iteration 10000 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| F ($\downarrow$) | W ($\downarrow$) | F ($\downarrow$) | W ($\downarrow$) | F ($\downarrow$) | W ($\downarrow$) | F ($\downarrow$) | W ($\downarrow$) |  |
| USWB | $252.24\pm 0.0$ | $3746.05\pm 0.0$ | $4.89\pm 0.28$ | $85.72\pm 0.18$ | $3.79\pm 0.32$ | $45.37\pm 0.18$ | $1.55\pm 0.48$ | $39.81\pm 0.18$ |
| MFSWB $\lambda=0.1$ | $252.24\pm 0.0$ | $3746.05\pm 0.0$ | $4.76\pm 0.27$ | $84.86\pm 0.17$ | $3.78\pm 0.2$ | $45.2\pm 0.11$ | $1.32\pm 0.22$ | $39.73\pm 0.16$ |
| MFSWB $\lambda=1$ | $252.24\pm 0.0$ | $3746.05\pm 0.0$ | $0.49\pm 0.2$ | $79.08\pm 0.15$ | $3.64\pm 0.26$ | $44.71\pm 0.19$ | $1.03\pm 0.06$ | $39.45\pm 0.18$ |
| MFSWB $\lambda=10$ | $252.24\pm 0.0$ | $3746.05\pm 0.0$ | $4.03\pm 2.43$ | $\mathbf{71.24\pm 0.9}$ | $7.32\pm 2.5$ | $45.21\pm 0.2$ | $4.13\pm 2.48$ | $42.56\pm 0.36$ |
| s-MFSWB | $252.24\pm 0.0$ | $3746.05\pm 0.0$ | $2.52\pm 0.77$ | $81.84\pm 0.14$ | $4.01\pm 0.38$ | $44.9\pm 0.13$ | $1.15\pm 0.09$ | $39.58\pm 0.17$ |
| us-MFSWB | $252.24\pm 0.0$ | $3746.05\pm 0.0$ | $0.3\pm 0.18$ | $78.69\pm 0.17$ | $3.74\pm 0.26$ | $44.38\pm 0.1$ | $0.87\pm 0.18$ | $39.26\pm 0.1$ |
| es-MFSWB | $252.24\pm 0.0$ | $3746.05\pm 0.0$ | $\mathbf{0.2\pm 0.19}$ | $78.1\pm 0.16$ | $\mathbf{3.5\pm 0.29}$ | $\mathbf{44.37\pm 0.08}$ | $\mathbf{0.84\pm 0.22}$ | $\mathbf{39.18\pm 0.08}$ |

We aim to find the mean shape of point-cloud shapes by casting a point cloud $X=\{x_{1},\ldots,x_{n}\}$ into an empirical probability measures $P_{X}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$. We select two point-cloud shapes which consist of 2048 points in ShapeNet Core-55 dataset []. We initialize the barycenter with a spherical point-cloud. We use stochastic gradient descent with 10000 iterations of learning rate 0.01, the number of projections 10. We report the found barycenters for two car shapes in Figure 3 at the final iteration and the corresponding F-metric and W-metric at iterations 0, 1000, 5000, and 10000 in Table 1 from three independent runs. As in the Gaussian simulation, s-MFSWB, us-MFSWB, and es-MFSWB help to reduce the two metrics faster than the USWB. With the slicing distribution selection, es-MFSWB performs the best at every iteration, even better than the formal MFSWB with three choices of $\lambda$ i.e., $0.1,1,10$. We also observe a similar phenomenon for two plane shapes in Figure 6 and Table 3 in Appendix D. We refer the reader to Appendix D for a detailed discussion.

### 4.3 Color Harmonization

Figure: Figure 4: Harmonized images from USWB, MFSWB ($\lambda=1$) s-MFSWB, us-MFSWB, and es-MFSWB.

We want to transform the color palette of a source image, denoted as $X=(x_{1},\ldots,x_{n})$ for $n$ is the number of pixels, to be an exact hybrid between two target images. Similar to the previous point-cloud averaging, we transform the color palette of an image into the empirical probability measure over colors (RGB) i.e., $P_{X}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$. We then minimize barycenter losses i.e., USWB, MFSWB ($\lambda\in\{0.1,1,10\}$), s-MFSWB, us-MFSWB, and es-MFSWB by using stochastic gradient descent with the learning rate $0.0001$ and 20000 iterations. We report both the transformed images and the corresponding F-metric and W-metric in Figure 4. We also report the full results in Figure 7- 9 in Appendix D. As in previous experiments, we see that the three proposed surrogates yield a better barycenter faster than USWB. The proposed es-MFSWB is the best variant among all surrogates since it has the lowest F-metric and W-metric at all iterations. We refer the reader to Figure 10-Figure 13 in Appendix D for additional flowers-images example, where a similar relative comparison happens. For the formal MFSWB, it is worse than es-MFSWB in one setting and better than es-MFSWB in one setting with the right choice of $\lambda$. Therefore, it is more convenient to use us-MFSWB in practice.

### 4.4 Sliced Wasserstein Autoencoder with Class-Fair Representation

Problem. We consider training the sliced Wasserstein autoencoder (SWAE)[] with a class-fairness regularization. In particular, we have the data distributions of $K\geq 1$ classes i.e., $\mu_{k}\in\mathcal{P}(\mathbb{R}^{d})$ for $k=1,\ldots,K$ and we would like to estimate an encoder network $f_{\phi}:\mathbb{R}^{d}\to\mathbb{R}^{h}$ ($\phi\in\Phi$) and a decoder network $g_{\psi}:\mathbb{R}^{h}\to\mathbb{R}^{d}$ ($\psi\in\Psi$ with $\mathbb{R}^{h}$ is a low-dimensional latent space. Given a prior distribution $\mu_{0}\in\mathcal{P}(\mathbb{R}^{h})$, $p\geq 1$, $\kappa_{1}\in\mathbb{R}^{+}$, $\kappa_{2}\in\mathbb{R}^{+}$, and a minibatch size $M\geq 1$, we perform the following optimization problem:

$$ $\displaystyle\min_{\phi,\psi}\mathbb{E}\left[\frac{1}{KM}\sum_{k=1}^{K}\sum_{i =1}^{M}c(X_{ki},g_{\psi}(f_{\phi}(X_{ki}))+\kappa_{1}SW_{p}^{p}(P_{Z},P_{(f_{ \phi}(X_{k}))_{k=1}^{K}})+\kappa_{2}\mathcal{B}(P_{Z};P_{f(X_{1})}:P_{f(X_{K}) })\right],$ $$

where $(X_{1},\ldots,X_{K})\sim\mu_{1}^{\otimes M}\otimes\ldots\otimes\mu_{K}^{
\otimes M}$, $Z\sim\mu_{0}^{\otimes M}$, $c$ is a reconstruction loss, $P_{Z}=\frac{1}{M}\sum_{i=1}^{M}\delta_{Z_{i}}$, $P_{(f_{\phi}(X_{k}))_{k=1}^{K}}=\frac{1}{KM}\sum_{k=1}^{K}\sum_{i=1}^{M}\delta
_{f_{\phi}(X_{ki})}$, $P_{f(X_{k})}=\frac{1}{M}\sum_{i=1}^{M}\delta_{X_{ki}}$ for $k=1,\ldots,K$, and $\mathcal{B}$ denotes a barycenter loss i.e., USWB, MFSWB, s-MFSWB, us-MFSWB, and es-MFSWB. This setting can be seen as an inverse barycenter problem i.e., the barycenter is fixed and the marginals are learnt under some constraints (e.g., the reconstruction loss and the aggregated distribution loss).

Results. We train the autoencoder on MNIST dataset [] ($d=28\times 28$) with $\kappa_{1}=8.0$, $\kappa_{2}=0.5$, 250 epochs, and $\mu_{0}$ is the uniform distribution on 2D ball ($h=2$). We refer the reader to Appendix D for more details of the neural network architectures and the optimizer. Following the training phase, we evaluate the trained autoencoders on the MNIST test set. Similar to previous experiments, we use the metrics F (denoted as $F_{\text{latent}}$) and W (denoted as $W_{\text{latent}}$) in the latent space distributions $f_{\phi}\sharp\mu_{1},\ldots,f_{\phi}\sharp\mu_{K}$ and the barycenter $\mu_{0}$. We use the reconstruction loss (binary cross-entropy, denoted as RL), the Wasserstein-2 distance between the prior and aggregated posterior distribution in latent space $\text{W}_{2,\text{latent}}^{2}:=W_{2}^{2}\left(\mu_{0},\frac{1}{K}\sum_{k=1}^{
K}f_{\phi}\sharp\mu_{k}\right)$, as well as in image space $\text{W}_{2,\text{image}}^{2}:=W_{2}^{2}\left(g_{\psi}\sharp\mu_{0},\frac{1}{K
}\sum_{k=1}^{K}\mu_{k}\right)$. During evaluation, we approximate $\mu_{0}$ by its empirical version of 10000 samples. We report the quantitative result in Table 2, and reconstructed images, generated images, and images of latent codes in Figure 14 in Appendix D. From the results, we see that the proposed surrogate MFSWB yield better scores than USWB in all metrics except the generative score i.e., $\text{W}_{2,\text{image}}^{2}$. The formal MFSWB gives good score in reconstruction loss, $\text{W}_{2,\text{image}}^{2}$, and $\text{W}_{2,\text{latent}}^{2}$, however, its F and W scores are high. The es-MFSWB is the best variant among the proposed surrogates. Compared with the conventional SWAE, using a barycenter loss leads to a more class-fair latent representation, however, it sacrifices the image reconstructive and generative quality.

**Table 2: Results of training SWAE with different regularization losses.**
| Methods | RL ($\downarrow$) | $\text{W}_{2,\text{latent}}^{2}\times 10^{2}$($\downarrow$) | $\text{W}_{2,\text{image}}^{2}$ ($\downarrow$) | $\text{F}_{\text{latent}}\times 10^{2}$ ($\downarrow$) | $\text{W}_{\text{latent}}\times 10^{2}$ ($\downarrow$) |
| --- | --- | --- | --- | --- | --- |
| SWAE | 2.95 | 9.41 | 26.91 | 7.05 | 23.86 |
| USWB | 3.15 | 10.46 | 27.41 | 7.02 | 12.73 |
| MFSWB $\lambda=0.1$ | 3.11 | 8.52 | 26.62 | 8.66 | 20.01 |
| MFSWB $\lambda=1.0$ | 3.12 | 9.71 | 26.92 | 10.15 | 21.16 |
| MFSWB $\lambda=10.0$ | 2.79 | 10.32 | 26.95 | 8.27 | 23.39 |
| s-MFSWB | 3.22 | 10.50 | 28.69 | 1.30 | 13.71 |
| us-MFSWB | 3.05 | 7.79 | 27.81 | 2.29 | 9.11 |
| es-MFSWB | 3.35 | 9.60 | 28.27 | 0.98 | 9.92 |

## 5 Conclusion

We introduced marginal fairness sliced Wasserstein barycenter (MFSWB), a special case of sliced Wasserstein barycenter (SWB) which has approximately the same distance to marginals. We first defined the MFSWB as a constrainted uniform SWB problem. After that, to overcome the computational drawbacks of the original problem, we propose three surrogate definitions of MFSWB which are hyperparameter-free and easy to compute. We discussed the relationship of the proposed surrogate problems and their connection to the sliced Multi-marginal Wasserstein distance with the maximal ground metric. Finally, we conduct simulations with Gaussian and experiments on 3D point-cloud averaging, color harmonization, and sliced Wasserstein autoencoder with class-fairness representation to show the benefits of the proposed surrogate MFSWB definitions.

## Acknowledgements

We would like to thank Joydeep Ghosh for his insightful discussion during the course of this project.

Supplement to “Marginal Fairness Sliced Wasserstein Barycenter"

We present skipped proofs in Appendix A. We then provide some additional materials which are mentioned in the main paper in Appendix B. After that, related works are discussed in Appendix C. We then provide additional experimental results in Appendix D. Finally, we report the used computational devices in Appendix E.

## Appendix A Proofs

### A.1 Proof of Proposition 1

Proof. From Definition 2, we have

$$ $\displaystyle\mathcal{SF}(\mu,\mu_{1:K})$ $\displaystyle=\max_{k\in\{1,\ldots,K\}}SW_{p}^{p}(\mu,\mu_{k})$ $\displaystyle=\max_{k\in\{1,\ldots,K\}}\mathbb{E}_{\theta\sim\mathcal{U}( \mathbb{S}^{d-1})}[W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})]$ $$

Let $k^{\star}=\operatorname*{arg\,max}_{k\in\{1,\ldots,K\}}\mathbb{E}_{\theta\sim
\mathcal{U}(\mathbb{S}^{d-1})}[W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})]$, we have

$$ $\displaystyle\mathcal{SF}(\mu,\mu_{1:K})$ $\displaystyle=\mathbb{E}_{\theta\sim\mathcal{U}(\mathbb{S}^{d-1})}[W_{p}^{p}( \theta\sharp\mu,\theta\sharp\mu_{k^{\star}})]$ $\displaystyle\leq\mathbb{E}_{\theta\sim\mathcal{U}(\mathbb{S}^{d-1})}\left[ \max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})\right]$ $\displaystyle=\mathcal{USF}(\mu,\mu_{1:K}),$ $$

as from Definition 3, which completes the proof.

### A.2 Proof of Proposition 2

Using the Holder’s inequality, we have:

$$ $\displaystyle\mathbb{E}\left|\nabla_{\phi}\frac{1}{L}\sum_{l=1}^{L}\text{W}_{p }^{p}(\theta_{l}\sharp\mu_{\phi},\theta_{l}\sharp\mu_{k^{\star}_{\theta_{l}}}) -\nabla_{\phi}\mathcal{USF}(\mu_{\phi};\mu_{1:K})\right|$ $\displaystyle\leq\left(\mathbb{E}\left|\nabla_{\phi}\frac{1}{L}\sum_{l=1}^{L} \text{W}_{p}^{p}(\theta_{l}\sharp\mu_{\phi},\theta_{l}\sharp\mu_{k^{\star}_{ \theta_{l}}})-\nabla_{\phi}\mathcal{USF}(\mu_{\phi};\mu_{1:K})\right|^{2} \right)^{\frac{1}{2}}$ $\displaystyle=\left(\mathbb{E}\left(\nabla_{\phi}\frac{1}{L}\sum_{l=1}^{L} \text{W}_{p}^{p}(\theta_{l}\sharp\mu_{\phi},\theta_{l}\sharp\mu_{k^{\star}_{ \theta_{l}}})-\nabla_{\phi}\mathbb{E}\left[\text{W}_{p}^{p}(\theta\sharp\mu_{ \phi},\theta\sharp\mu_{k^{\star}_{\theta}})\right]\right)^{2}\right)^{\frac{1} {2}}$ $\displaystyle=\left(\mathbb{E}\left(\frac{1}{L}\sum_{l=1}^{L}\nabla_{\phi} \text{W}_{p}^{p}(\theta_{l}\sharp\mu_{\phi},\theta_{l}\sharp\mu_{k^{\star}_{ \theta_{l}}})-\mathbb{E}\left[\nabla_{\phi}\text{W}_{p}^{p}(\theta\sharp\mu_{ \phi},\theta\sharp\mu_{k^{\star}_{\theta}})\right]\right)^{2}\right)^{\frac{1} {2}}$ $\displaystyle=\left(\text{Var}\left[\frac{1}{L}\sum_{l=1}^{L}\nabla_{\phi} \text{W}_{p}^{p}(\theta_{l}\sharp\mu_{\phi},\theta_{l}\sharp\mu_{k^{\star}_{ \theta_{l}}})\right]\right)^{\frac{1}{2}}$ $\displaystyle=\frac{1}{\sqrt{L}}\text{Var}\left[\nabla_{\phi}\text{W}_{p}^{p}( \theta\sharp\mu_{\phi},\theta\sharp\mu_{k^{\star}_{\theta}})\right]^{\frac{1}{ 2}},$ $$

which completes the proof.

### A.3 Proof of Proposition 3

We first restate the following Lemma from [] and provide the proof for completeness.

###### Lemma 1 .

For any $L\geq 1$, $0\leq a_{1}\leq a_{2}\leq\ldots\leq a_{L}$ and $0<b_{1}\leq b_{2}\leq\ldots\leq b_{L}$, we have:

$$ $\displaystyle\frac{1}{L}(\sum_{i=1}^{L}a_{i})(\sum_{i=1}^{L}b_{i})\leq\sum_{i= 1}^{L}a_{i}b_{i}.$ (15) $$

###### Proof.

For $L=1$, we directly have $a_{i}b_{i}=a_{i}b_{i}$. Assuming that for $L$ the inequality holds i.e., $\frac{1}{L}(\sum_{i=1}^{L}a_{i})(\sum_{i=1}^{L}b_{i})\leq\sum_{i=1}^{L}a_{i}b_
{i}$ which is equivalent to $(\sum_{i=1}^{L}a_{i})(\sum_{i=1}^{L}b_{i})\leq L\sum_{i=1}^{L}a_{i}b_{i}.$
Now, we show that $\frac{1}{L}(\sum_{i=1}^{L}a_{i})(\sum_{i=1}^{L}b_{i})\leq\sum_{i=1}^{L}a_{i}b_
{i}$ i.e., the inequality holds for $L+1$. We have

$$ $\displaystyle(\sum_{i=1}^{L+1}a_{i})(\sum_{i=1}^{L+1}b_{i})$ $\displaystyle=(\sum_{i=1}^{L}a_{i})(\sum_{i=1}^{L}b_{i})+(\sum_{i=1}^{L}a_{i}) b_{L+1}+(\sum_{i=1}^{L}b_{i})a_{L+1}+a_{L+1}b_{L+1}$ $\displaystyle\leq L\sum_{i=1}^{L}a_{i}b_{i}+(\sum_{i=1}^{L}a_{i})b_{L+1}+(\sum _{i=1}^{L}b_{i})a_{L+1}+a_{L+1}b_{L+1}.$ $$

Since $a_{L+1}b_{L+1}+a_{i}b_{i}\geq a_{L+1}b_{i}+b_{L+1}a_{i}$ for all $1\leq i\leq L$ by rearrangement inequality. By taking the sum of these inequalities over $i$ from $1$ to $L$, we obtain:

$$ $\displaystyle(\sum_{i=1}^{L}a_{i})b_{L+1}+(\sum_{i=1}^{L}b_{i})a_{L+1}\leq\sum _{i=1}^{L}a_{i}b_{i}+La_{L+1}b_{L+1}.$ $$

Then, we have

$$ $\displaystyle(\sum_{i=1}^{L+1}a_{i})(\sum_{i=1}^{L+1}b_{i})$ $\displaystyle\leq L\sum_{i=1}^{L}a_{i}b_{i}+(\sum_{i=1}^{L}a_{i})b_{L+1}+(\sum _{i=1}^{L}b_{i})a_{L+1}+a_{L+1}b_{L+1}$ $\displaystyle\leq L\sum_{i=1}^{L}a_{i}b_{i}+\sum_{i=1}^{L}a_{i}b_{i}+La_{L+1}b _{L+1}+a_{L+1}b_{L+1}$ $\displaystyle=(L+1)(\sum_{i=1}^{L+1}a_{i}b_{i}),$ $$

which completes the proof.
∎

Now, we go back to the main inequality which is $\mathcal{USF}(\mu;\mu_{1:K})\leq\mathcal{ESF}(\mu;\mu_{1:K})$. From Definition 5, we have:

$$ $\displaystyle\mathcal{ESF}(\mu;\mu_{1:K})$ $\displaystyle=\mathbb{E}_{\theta\sim\sigma(\theta;\mu,\mu_{1:K})}\left[\max_{k \in\{1,\ldots,K\}}W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})\right]$ $\displaystyle=\mathbb{E}_{\theta\sim\mathcal{U}(\mathbb{S}^{d-1})}\left[\max_{ k\in\{1,\ldots,K\}}W_{p}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})\frac{f_{ \sigma}(\theta;\mu,\mu_{1:K})}{\frac{\Gamma(d/2)}{2\pi^{d/2}}}\right],$ $$

where $f_{\sigma}(\theta;\mu,\mu_{1:K})\propto\exp\left(\max_{k\in\{1,\ldots,K\}}W_{p
}^{p}(\theta\sharp\mu,\theta\sharp\mu_{k})\right)$. Now, we consider a Monte Carlo estimation of $\mathcal{ESF}(\mu;\mu_{1:K})$ by importance sampling:

$$ $\displaystyle\widehat{\mathcal{ESF}}(\mu;\mu_{1:K},L)=\frac{1}{L}\sum_{l=1}^{L }\left[\max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta_{l}\sharp\mu,\theta_{l}\sharp \mu_{k})\frac{\exp\left(\max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta_{l}\sharp\mu ,\theta_{l}\sharp\mu_{k})\right)}{\sum_{i=1}^{L}\exp\left(\max_{k\in\{1,\ldots ,K\}}W_{p}^{p}(\theta_{i}\sharp\mu,\theta_{i}\sharp\mu_{k})\right)}\right],$ $$

where $\theta_{1},\ldots,\theta_{L}\overset{i.i.d}{\sim}\mathcal{U}(\mathbb{S}^{d-1})$. Similarly, we consider a Monte Carlo estimation of $\mathcal{USF}(\mu;\mu_{1:K})$:

$$ $\displaystyle\widehat{\mathcal{USF}}(\mu;\mu_{1:K},L)=\frac{1}{L}\sum_{l=1}^{L }\left[\max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta_{l}\sharp\mu,\theta_{l}\sharp \mu_{k})\right],$ $$

for the same set of $\theta_{1},\ldots,\theta_{L}$. Without losing generality, we assume that $\max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta_{1}\sharp\mu,\theta_{1}\sharp\mu_{k}
)\leq\ldots\leq\max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta_{L}\sharp\mu,\theta_{
L}\sharp\mu_{k})$.

Let $\max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta_{i}\sharp\mu,\theta_{i}\sharp\mu_{k}
)=a_{i}$ and $\exp\left(\max_{k\in\{1,\ldots,K\}}W_{p}^{p}(\theta_{i}\sharp\mu,\theta_{i}
\sharp\mu_{k})\right)=b_{i}$, applying Lemma 1, we have:

$$ $\displaystyle\widehat{\mathcal{USF}}(\mu;\mu_{1:K},L)\leq\widehat{\mathcal{ESF }}(\mu;\mu_{1:K},L)\quad\forall L\geq 1.$ $$

By letting $L\to\infty$ and applying the law of large numbers, we obtain:

$$ $\displaystyle\mathcal{USF}(\mu;\mu_{1:K})\leq\mathcal{ESF}(\mu;\mu_{1:K}),$ $$

which completes the proof.

###### Lemma 1 .

###### Proof.

### A.4 Proof of Proposition 4

We first recall the definition of the SMW with the maximal ground metric:

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int\max _{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j} |^{p}d\pi(x_{1},\ldots,x_{K})\right].$ $$

Non-negativity. Since $\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x
_{j}|^{p}\geq 0$ for any $x_{1},\ldots,x_{K}$ and for any $\theta$, we can obtain the desired property $SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)\geq 0$ which implies $SMW_{p}(\mu_{1},\ldots,\mu_{K};c)\geq 0$.

Marginal Exchangeability. For any permutation $\sigma:[[K]]\to[[K]]$, we have:

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int\max _{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j} |^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $\displaystyle=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{\sigma(1)},\ldots,\mu_{ \sigma(K)})}\int\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}|\theta^{\top}x_{i }-\theta^{\top}x_{j}|^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $\displaystyle=SMW_{p}^{p}(\mu_{\sigma(1)},\ldots,\mu_{\sigma(K)};c).$ $$

Generalized Triangle Inequality. For $\mu\in\mathcal{P}_{p}(\mathbb{R}^{d})$, we have :

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int\max _{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j} |^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $\displaystyle\leq\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int \sum_{k=1}^{K}\max_{i\in\{1,\ldots,K\}\setminus\{{k}\},j\in\{1,\ldots,K\} \setminus\{{k}\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j}|^{p}d\pi(x_{1},\ldots, x_{K})\right]$ $\displaystyle=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\sum_{k= 1}^{K}\int\max_{i\in\{1,\ldots,K\}\setminus\{{k}\},j\in\{1,\ldots,K\}\setminus \{{k}\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j}|^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $\displaystyle=\mathbb{E}\left[\sum_{k=1}^{K}\int\max_{i\in\{1,\ldots,K\} \setminus\{{k}\},j\in\{1,\ldots,K\}\setminus\{{k}\}}|\theta^{\top}x_{i}-\theta ^{\top}x_{j}|^{p}d\pi^{\star}(x_{1},\ldots,x_{k-1},x_{k+1},\ldots x_{K})\right]$ $$

for $\pi^{\star}$ is the optimal multi-marginal transportation plan and $\pi^{\star}(x_{1},\ldots,x_{k-1},x_{k+1},x_{K}\ )$ is the marginal joint distribution by integrating out $x_{k}$. By the gluing lemma [], there exists optimal plans $\pi^{\star}(x_{1},\ldots,x_{k-1},y,x_{k+1},x_{K}\ )$ for any $k\in[[K]]$ and $y$ follows $\mu$. We further have:

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle\leq\mathbb{E}\left[\sum_{k=1}^{K}\int\max\left(\max_{i\in\{1, \ldots,K\}\setminus\{{k}\},j\in\{1,\ldots,K\}\setminus\{{k}\}}|\theta^{\top}x_ {i}-\theta^{\top}x_{j}|^{p},\right.\right.$ $\displaystyle\quad\quad\max_{i\in\{1,\ldots,K\}\setminus\{{k}\}}|\theta^{\top} x_{i}-\theta^{\top}y|^{p}\left.\left.\right)d\pi^{\star}(x_{1},\ldots,x_{k-1}, y,x_{k+1},\ldots x_{K})\right]$ $\displaystyle=\sum_{k=1}^{K}\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu _{k-1},\mu,\mu_{k+1},\ldots,\mu_{K})}\int\max_{i\in\{1,\ldots,K\},j\in\{1, \ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j}|^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $\displaystyle=\sum_{k=1}^{K}SMW_{p}^{p}(\mu_{1},\ldots,\mu_{k-1},\mu,\mu_{k+1} ,\ldots,\mu_{K};c).$ $$

Applying the Minkowski’s inequality, we obtain the desired property:

$$ $\displaystyle SMW_{p}(\mu_{1},\ldots,\mu_{K};c)\leq\sum_{k=1}^{K}SMW_{p}(\mu_{ 1},\ldots,\mu_{k-1},\mu,\mu_{k+1},\ldots,\mu_{K};c).$ $$

Identity of Indiscernibles. From the proof in Appendix A.5, we have:

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle\geq\mathbb{E}\left[\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}W _{p}^{p}(\theta\sharp\mu_{i},\theta\sharp\mu_{j})\right]$ $\displaystyle\geq\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}\mathbb{E}\left[W _{p}^{p}(\theta\sharp\mu_{i},\theta\sharp\mu_{j})\right]$ $\displaystyle=\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}SW_{p}^{p}(\mu_{i}, \mu_{j}).$ $$

Therefore, when $SMW_{p}(\mu_{1},\ldots,\mu_{K};c)=0$, we have $SW_{p}^{p}(\mu_{i},\mu_{j})=0$ which implies $\mu_{i}=\mu_{j}$ for any $i,j\in[[K]]$. As a result, $\mu_{1}=\ldots=\mu_{K}$ from the metricity of the SW distance. For the other direction, it is easy to see that if $\mu_{1}=\ldots\mu_{K}$, we have $SMW_{p}(\mu_{1},\ldots,\mu_{K};c)=0$ based on the definition and the metricity of the Wasserstein distance.

### A.5 Proof of Proposition 5

Given the maximal ground metric $c(\theta^{\top}x_{1},\ldots,\theta^{\top}x_{K})=\max_{i\in\{1,\ldots,K\},j\in
\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j}|$, from Equation 6

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int c( \theta^{\top}x_{1},\ldots,\theta^{\top}x_{K})^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $\displaystyle=\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int\max _{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j} |^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $$

By Jensen inequality i.e., $(x_{1},\ldots,x_{K})\to\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}|\theta^{
\top}x_{i}-\theta^{\top}x_{j}|^{p}$ is a convex function, we have:

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle\geq\mathbb{E}\left[\inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\max_ {i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}\int|\theta^{\top}x_{i}-\theta^{\top}x_ {j}|^{p}d\pi(x_{1},\ldots,x_{K})\right].$ $$

Using max-min inequality, we have:

$$ $\displaystyle SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle\geq\mathbb{E}\left[\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}} \inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int|\theta^{\top}x_{i}-\theta^{\top}x _{j}|^{p}d\pi(x_{1},\ldots,x_{K})\right]$ $\displaystyle\geq\mathbb{E}\left[\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}} \inf_{\pi\in\Pi(\mu_{i},\mu_{j})}\int|\theta^{\top}x_{i}-\theta^{\top}x_{j}|^{ p}d\pi(x_{i},x_{j})\right]$ $\displaystyle=\mathbb{E}\left[\max_{i\in\{1,\ldots,K\},j\in\{1,\ldots,K\}}W_{p }^{p}(\theta\sharp\mu_{i},\theta\sharp\mu_{j})\right].$ $$

Therefore, minimizing two sides with respect to $\mu_{1}$, we have:

$$ $\displaystyle\min_{\mu_{1}}SMW_{p}^{p}(\mu_{1},\ldots,\mu_{K};c)$ $\displaystyle\geq\min_{\mu_{1}}\mathbb{E}\left[\max_{i\in\{1,\ldots,K\},j\in\{ 1,\ldots,K\}}W_{p}^{p}(\theta\sharp\mu_{i},\theta\sharp\mu_{j})\right]$ $\displaystyle\geq\min_{\mu_{1}}\mathbb{E}\left[\max_{i\in\{2,\ldots,K\}}W_{p}^ {p}(\theta\sharp\mu_{1},\theta\sharp\mu_{i})\right]$ $\displaystyle=\min_{\mu_{1}}\mathcal{USF}(\mu_{1};\mu_{2:K}),$ $$

which completes the proof.

## Appendix B Additional Materials

Algorithms. As mentioned in the main paper, we present the computational algorithm for SWB in Algorithm 1, for s-MFSWB in Algorithm 2, for us-MFSWB in Algorithm 3, and for es-MFSWB in Algorithm 4.

Figure: Algorithm 1 Computational algorithm of the SWB problem

Figure: Algorithm 2 Computational algorithm of the s-MFSWB problem

Figure: Algorithm 3 Computational algorithm of the us-MFSWB problem

Figure: Algorithm 4 Computational algorithm of the es-MFSWB problem

Energy-based sliced Multi-marginal Wasserstein. As shown in Proposition 5, us-MFSWB is equivalent to minimizing a lower bound of SMW with the maximal ground metric. We now show that es-MFSWB is also equivalent to minimizing a lower bound of a variant of SMW i.e., Energy-based sliced Multi-marginal Wasserstein with the maximal ground metric. We refer the reader to Proposition 6 for a detailed definition. The proof of Proposition 6 is similar to the proof of Proposition 5 in Appendix A.5.

###### Proposition 6 .

Given $K\geq 2$ marginals $\mu_{1},\ldots,\mu_{K}\in\mathcal{P}_{p}(\mathbb{R}^{d})$, the maximal ground metric $c(\theta^{\top}x_{1},\ldots,\theta^{\top}x_{K})=\max_{i\in\{1,\ldots,K\},j\in
\{1,\ldots,K\}}|\theta^{\top}x_{i}-\theta^{\top}x_{j}|$, we have:

$$ $\displaystyle\min_{\mu_{1}}\mathcal{ESF}(\mu_{1};\mu_{2:K})\leq\min_{\mu_{1}} ESMW_{p}^{p}(\mu_{1},\mu_{2},\ldots,\mu_{K};c),$ (16) $$

where

$$ $\displaystyle ESMW_{p}^{p}(\mu_{1},\mu_{2},\ldots,\mu_{K};c)=\mathbb{E}\left[ \inf_{\pi\in\Pi(\mu_{1},\ldots,\mu_{K})}\int c(\theta^{\top}x_{1},\ldots, \theta^{\top}x_{K})^{p}d\pi(x_{1},\ldots,x_{K})\right],$ $$

and the expectation is with respect to $\sigma(\theta)$ i.e.,

$$ $f_{\sigma}(\theta;\mu_{1},\mu_{2:K})\propto\exp\left(\max_{k\in\{2,\ldots,K\}} W_{p}^{p}(\theta\sharp\mu_{1},\theta\sharp\mu_{k})\right).$ $$

###### Proposition 6 .

## Appendix C Related Works

Demographic fairness with Wasserstein Barycenter. A connection between fair regression and one-dimensional Wasserstein barycenter is established by deriving the expression for the optimal function minimizing squared risk under Demographic Parity constraints []. Similarly, Demographic Parity fair classification is connected to one-dimensional Wasserstein-1 distance barycenter in []. The work [] extends the Demographic Parity constraint to multi-task problems for regression and classification and connects them to the one-dimensional Wasserstein-2 distance barycenters. A method to augment the input so that predictability of the protected attribute is impossible, by using Wasserstein-2 distance Barycenters to repair the data is proposed in [].
A general approach for using one-dimensional Wasserstein-1 distance barycenter to obtain Demographic Parity in classification and regression is proposed in [].
Overall, all discussed works define fairness in terms of Demographic Parity constraints in applications with a response variable (classification and regression) in one dimension. In contrast, we focus on marginal fairness barycenter i.e., using a set of measures only, in any dimensions.

Other possible applications. Wasserstein barycenter has been used to cluster measures in []. In particular, a K-mean algorithm for measures is proposed with Wasserstein barycenter as the averaging operator. Therefore, our MFSWB can be directly used to enforce the fairness for averaging inside each cluster. The proposed MFSWB can be also used to average meshes by changing the SW to H2SW which is proposed in [].

## Appendix D Additional Experiments

Figure: Figure 5: Barycenters from MFSWB with $\lambda=0.1$ and $\lambda=10$ along gradient iterations with the corresponding F-metric and W-metric.
Refer to caption: x26.png

Gaussians barycenter with the formal MFSWB. We present the result of finding barycenters of Gaussian distributions with MFSWB $\lambda=0.1$ and $\lambda=10$ in Figure 5.

Figure: Figure 6: Averaging point-clouds with USWB, MFSWB ($\lambda=1$), s-MFSWB, us-MFSWB, and es-MFSWB.
Refer to caption: extracted/2405.07482v1/images/shape2.png

**Table 3: F-metric and W-metric along iterations in point-cloud averaging application.**
| Method | Iteration 0 | Epoch 1000 | Epoch 5000 | Epoch 10000 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| F ($\downarrow$) | W ($\downarrow$) | F ($\downarrow$) | W ($\downarrow$) | F ($\downarrow$) | W ($\downarrow$) | F ($\downarrow$) | W ($\downarrow$) |  |
| USWB | $746.67\pm 0.0$ | $4814.71\pm 0.0$ | $35.22\pm 1.04$ | $161.11\pm 0.54$ | $7.82\pm 0.26$ | $109.82\pm 0.28$ | $11.08\pm 0.06$ | $108.52\pm 0.17$ |
| MFSWB $\lambda=0.1$ | $746.67\pm 0.0$ | $4814.71\pm 0.0$ | $35.15\pm 0.36$ | $159.84\pm 0.55$ | $4.95\pm 0.23$ | $109.14\pm 0.33$ | $6.95\pm 0.8$ | $107.83\pm 0.16$ |
| MFSWB $\lambda=1$ | $746.67\pm 0.0$ | $4814.71\pm 0.0$ | $33.21\pm 2.72$ | $151.24\pm 0.64$ | $2.54\pm 1.5$ | $109.66\pm 0.26$ | $4.66\pm 2.1$ | $\mathbf{108.1\pm 0.05}$ |
| MFSWB $\lambda=10$ | $746.67\pm 0.0$ | $4814.71\pm 0.0$ | $34.03\pm 22.6$ | $158.66\pm 1.39$ | $29.19\pm 14.29$ | $122.66\pm 0.88$ | $20.55\pm 13.57$ | $123.65\pm 1.52$ |
| s-MFSWB | $746.67\pm 0.0$ | $4814.71\pm 0.0$ | $36.23\pm 1.88$ | $154.4\pm 0.67$ | $\mathbf{0.66\pm 0.44}$ | $\mathbf{109.17\pm 0.34}$ | $2.54\pm 2.06$ | $107.57\pm 0.19$ |
| us-MFSWB | $746.67\pm 0.0$ | $4814.71\pm 0.0$ | $28.65\pm 1.37$ | $144.27\pm 0.65$ | $1.02\pm 0.8$ | $109.67\pm 0.1$ | $\mathbf{1.35\pm 0.77}$ | $108.2\pm 0.19$ |
| es-MFSWB | $746.67\pm 0.0$ | $4814.71\pm 0.0$ | $\mathbf{28.05\pm 1.16}$ | $\mathbf{143.24\pm 0.76}$ | $0.99\pm 0.32$ | $109.68\pm 0.14$ | $1.36\pm 0.62$ | $108.28\pm 0.07$ |

Point-cloud averaging. We report the averaging results of two point-clouds of plane shapes n Figure 6 and the corresponding F-metrics and W-metric along iterations in Table 3. We see that the proposed surrogates achieve better F-metric and W-metric than the USWB. In this case, us-MFSWB gives the best F-metric at the final epoch, however, es-MFSWB also gives a comparable performance and performs better at earlier epochs. For the formal MFSWB, it does not perform well with the chosen set of $\lambda$.

Figure: Figure 7: Harmonized images from USWB, MFSWB ($\lambda=1$), s-MFSWB, us-MFSWB, and es-MFSWB at iteration 5000.

Figure: Figure 8: Harmonized images from USWB, MFSWB ($\lambda=1$), s-MFSWB, us-MFSWB, and es-MFSWB at iteration 10000.

Figure: Figure 9: Harmonized images from MFSWB with $\lambda=0.1$ and $\lambda=10$ at iterations 5000, 10000, and 20000.

Figure: Figure 10: Harmonized images from USWB, MFSWB ($\lambda=1$) s-MFSWB, us-MFSWB, and es-MFSWB at iteration 5000.

Figure: Figure 11: Harmonized images from USWB, MFSWB ($\lambda=1$), s-MFSWB, us-MFSWB, and es-MFSWB at iteration 10000.

Figure: Figure 12: Harmonized images from USWB, MFSWB ($\lambda=1$), s-MFSWB, us-MFSWB, and es-MFSWB at iterations 20000.

Figure: Figure 13: Color harmonized images from MFSWB with $\lambda=0.1$ and $\lambda=10$ at iterations 5000, 10000, and 20000.

Color Harmonization. We first present the harmonized images of different methods including USWB, MFSWB $(\lambda=1)$, s-MFSWB, us-MFSWB, and es-MFSWB at iteration 5000 and 10000 for the demonstrated images in the main text in Figure 7-Figure 8. Moreover, we report the results of MFSWB $(\lambda=0.1,10)$ at iteration 5000, 10000, and 20000 in Figure 9. Similarly, we repeat the same experiments with flower images in Figure 10- 13. Overall, we see that es-MFSWB helps to reduce both F-metric and W-metric faster than USWB and other surrogates. For the formal MFSWB, the performance depends significantly on the choice of $\lambda$.

Sliced Wasserstein autoencoder with class-fairness representation. We use the RMSprop optimizer with learning rate $0.01$, alpha=$0.99$, eps=$1e-8$. As mentioned in the main text, we report the used neural network architectures:

We report some randomly selected reconstructed images, some randomly generated images, and the test latent codes of trained autoencoders in Figure 14. Overall, we observe that the qualitative results are consistent with the quantitive results in Table 2. From the latent spaces, we see that the proposed surrogates help to make the codes of classes have approximately the same structure which do appear in the conventional SWAE’s latent codes.

Figure: Figure 14: Reconstructed images, generated images and latent space of all methods.

## Appendix E Computational Devices

For the Gaussian simulation, point-cloud averaging, and color harmonization, we use a HP Omen 25L desktop for conducting experiments. Additionally, for the Sliced Wasserstein Autoencoder with class-fair representation experiment, we employ the NVIDIA Tesla V100 GPU.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Did you have a typo in line 506? Should $P_{f(X_k)}$ be equal to $\frac{1}{M} \sum_{i=1}^M \delta_{f(X_{ki})}$?
soundness: 3
strengths: This paper addresses an important problem that has been overlooked. In traditional SWB, the weights between barycenter and marginals are uniform. However, this uniform weights actually leads to non-uniform distances between the barycenter and the marginals. The authors have a simple experiment demonstrating that in the paper. To ensure equal distances, they propose the MFSWB to address this problem. They define MFSWB such that the average distance to be within an $\epsilon$.

The proposed MFSWB is hard to solve. The authors proposed three surrogate definitions for MFSWB. The first one, called s-MFSWB, aims to minimize the maximal distance from the barycenter to the marginals. Yet, the s-MFSWB is biased. To rectify this, they proposed the second unbiased surrogate, us-MFSWB. They show that us-MFSWB's objective is an upper bound of s-MFSWB, meaning that minimizing us-MFSWB implicitly minimizes s-MFSWB. Furthermore, They also show that the approximation error of the gradient estimator of us-MFSWB reduces at the order of $\mathcal{O}(L^{−1/2})$. Since s-MFSWB and us-MFSWB both use the uniform distribution as the slicing distribution, which can empirically result in non-optimal in statistical estimation, the authors proposed the energy-based surrogate MFSWB (es-MFSWB) which considers the importance of the sampling directions. 

The authors conducted extensive experiments evaluating the proposed methods and comparing the methods against the conventional SWB. The results are convincing. Using the synthetic Gaussian data, they showed that SWB produced the barycenter with non-uniform distances from the marginals, whereas the proposed methods yielded barycenters with approximately equal distances from the marginals. The authors also performed experiments on the 3D points clouds, Color Harmonization, and generative models with both quantitative results and qualitative results showing the effectiveness of the proposed method. The proposed es-MFSWB appears to be the best among all three proposed methods. 

The authors shows the potential of using the proposed method in learning conditional generative models as an inverse barycenter problem.
summary: The traditional Sliced Wasserstein Barycenter (SWB) employs uniform weights between the barycenter and its marginals, leading to imbalances in the distances between them. This paper introduces the Marginal Fairness Sliced Wasserstein Barycenter (MFSWB), which aims to equalize these distances, ensuring a fairer representation. The authors propose three algorithms to address the MFSWB formulation and conduct experiments on synthetic Gaussian distributions, 3D point clouds, color harmonization, and generative models. The results demonstrate that the proposed method is highly effective, validating the authors' claims.
weaknesses: I didn't really catch any weakness of this paper. The paper is complete with clear motivation, effective solutions and sufficient experimental results.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can the authors elaborate on the significance of MFSWB? The description in current version is somewhat high-level.
2. Can the authors provide a clearer description about the formulation in line 148?
3. What does the symbol $\sharp$ in Eq.(1) denote?
4. What does the $x_i(y_i)$ in line 144 denote?
5. What does the $\pi$ in line 144 denote?
6. What modification to the methodology presented in this paper are necessary to replace SWD with GSWD [1] or ASWD [2]?

[1] Kolouri, Soheil, et al. "Generalized sliced wasserstein distances." Advances in neural information processing systems 32 (2019).

[2] Chen, Xiongjie, Yongxin Yang, and Yunpeng Li. "Augmented sliced Wasserstein distances." arXiv preprint arXiv:2006.08812 (2020).
soundness: 3
strengths: 1. The paper is well-organized and the writing is good.
2. The paper has solid theoretical proofs.
3. The story line is easy to follow.
summary: This paper aims at achieving marginal fairness sliced wasserstein barycenter (MFSWB). To tackle the problem, the authors define the MFSWB as a constraint SWB problem. Due to the computational disadvantages of the formal definition, the authors propose two hyperparameter-free and computationally tractable surrogate MFSWB problems that implicitly minimize the distances to marginals and encourage marginal fairness at the same time. To further improve the efficiency, the authors perform slicing distribution selection and obtain the third surrogate definition by introducing a new slicing distribution that focuses more on marginally unfair projecting directions. The experimental results demonstrate the favorable performance of the proposed MFSWB performance.
weaknesses: Please refer to questions.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - What does USWB stand for in the paper? 
- Can the method extended to settings with more than two images or point cloud?
soundness: 3
strengths: - The presented problem is novel and reasonable. If fairness of barycenter can be achieved w.r.t. marginals, then varous future applicaiton can be conducted e.g. ensuring fairness of ML algorithms. 
- The paper organization, problem definition and writing is good and easy to follow. If audience has some Sliced Wasserstein knowledge, then it will be easy to follow the idea. 
- The Lagaragian formulation as well as three surrogates are proposed to deal with the Fairness barycenter estimation with unbiased gradient estimator.   
The space and time complexity analysis show that the proposed e-MFSWB has good space and time complexity compared with conventional SWB.
summary: This paper presents the marginal fairness sliced Wasserstein barycenter (MFSWB) problem. The goal of MFSWB is to achieve fairness of barycenter to marginals. To solve the problem, the Lagaragian format and three surrogate formulations are proposed, from biased estimation to unbiased estimation. Both time and space complexities are analyzed. 

For the experiments, Gaussian toy data is utilized for verification. Then, point cloud, color hamonization, and SWAE experiments are conducted on more realistic data.   
Both fairness metirc and barycenter metrics are reported.
weaknesses: - Most of the experiments conducted are simple or on simple dataset, such as Gaussian, Simple Color Hamonization of two images or point cloud of only two shapes, AE on simple MNIST dataset.   
I am wondering the scability and real application of the proposed method. For example, if more than two (e.g. >5) images or point clouds are adopted, what might be the results?   
Or are there more realistic potential appplications of the proposed method? If related work or potential applications are discussed, it would be better. For example, CelebA dataset results or Fairness generation [R1]. 
- USWB first appears in Figure 1 and used in several places. But it is not defined. 
- I can understand the F-metric is significantly improved with the proposed solution. But for W-metric why it is also improved. Are there feasible explanation or insights? Can you provide an explanation or discussion on why the W-metric improves along with the F-metric? This would be valuableto better understand the proposed method. 

[R1] Fair generative modeling via weak supervision.

–––––––––––––––
After Rebuttal: 
The authors' rebuttal addressed those questions comprehensively, I would like to increase my overall score to 8.

Meta-review: The traditional Sliced Wasserstein Barycenter (SWB) employs uniform weights between the barycenter and its marginals, often resulting in imbalances in the distances between them. This paper introduces the Marginal Fairness Sliced Wasserstein Barycenter (MFSWB), designed to equalize these distances and ensure a more balanced representation. To address the MFSWB formulation, the authors propose three algorithms and evaluate their effectiveness through experiments on synthetic Gaussian distributions, 3D point clouds, color harmonization, and generative models. The results demonstrate the high efficacy of the proposed method, providing strong validation for the authors' claims.  The paper is well-written and well-organized with sufficient theoretical and empirical justification, making a valuable contribution to the community.  The rebuttal addressed the concerns raised by the reviewers in their initial assessments. This led to a consensus among all three reviewers, who agreed to accept the paper based on its innovation and empirical contributions. The AC concurs with the reviewers and recommends accepting the paper.

Decision: Accept (Spotlight)