Local Bures-Wasserstein Transport: A Practical and Fast Mapping Approximation

TL;DR

This paper introduces a fast, practical method for approximating optimal transport maps using local Bures-Wasserstein transport, significantly reducing computation time and complexity compared to kernel-based approaches.

Contribution

It proposes a novel, efficient approach to learn approximate transport maps and Wasserstein barycenters leveraging Gaussian closed-form solutions, enabling faster and more scalable computations.

Findings

80x faster overall runtime compared to existing methods

Requires fewer components to recover barycenter support

Generalizes well to out-of-sample data

Abstract

Optimal transport (OT)-based methods have a wide range of applications and have attracted a tremendous amount of attention in recent years. However, most of the computational approaches of OT do not learn the underlying transport map. Although some algorithms have been proposed to learn this map, they rely on kernel-based methods, which makes them prohibitively slow when the number of samples increases. Here, we propose a way to learn an approximate transport map and a parametric approximation of the Wasserstein barycenter. We build an approximated transport mapping by leveraging the closed-form of Gaussian (Bures-Wasserstein) transport; we compute local transport plans between matched pairs of the Gaussian components of each density. The learned map generalizes to out-of-sample examples. We provide experimental results on simulated and real data, comparing our proposed method with other mapping estimation algorithms. Preliminary experiments suggest that our proposed method is not only faster, with a factor 80 overall running time, but it also requires fewer components than state-of-the-art methods to recover the support of the barycenter. From a practical standpoint, it is straightforward to implement and can be used with a conventional machine learning pipeline.

Equations

$$W^2(\mu ,\nu )=\underset{\mathbf{T}\in \Pi (\mu ,\nu )}{\min }\mathrm{Tr}({\mathbf{T}}^{\mathsf{T}}\mathbf{C}),$$

$$W_2^2(\mu ,\nu )=\Vert {\mathbf{m}}_1-{\mathbf{m}}_2{\Vert }_2^2+\mathrm{Tr}({\mathbf{\Sigma }}_1+{\mathbf{\Sigma }}_2)-2\mathrm{Tr}\left[{\left({\mathbf{\Sigma }}_1^{1/2}{\mathbf{\Sigma }}_2{\mathbf{\Sigma }}_1^{1/2}\right)}^{1/2}\right],$$

$${\mathbf{T}}_{{\mathbf{\Sigma }}_1}^{{\mathbf{\Sigma }}_2}={\mathbf{\Sigma }}_1^{-1/2}{\left({\mathbf{\Sigma }}_1^{1/2}{\mathbf{\Sigma }}_2{\mathbf{\Sigma }}_1^{1/2}\right)}^{1/2}{\mathbf{\Sigma }}_1^{-1/2},$$

$$\mathbf{T}(\mathbf{x})={\mathbf{m}}_2+{\mathbf{T}}_{{\mathbf{\Sigma }}_1}^{{\mathbf{\Sigma }}_2}(\mathbf{x}-{\mathbf{m}}_1).$$

$$\underset{\bar{\mathbf{m}}\in {\mathbb{R}}^p,\bar{\mathbf{\Sigma }}\in {\mathbb{S}}_{+}}{\min }\sum _{i=1}^l{\mathbf{\lambda }}_i{W}_2^2(\mathcal{N}(\bar{\mathbf{m}},\bar{\mathbf{\Sigma }}),\mathcal{N}({\mathbf{m}}_i,{\mathbf{\Sigma }}_i)).$$

$$\bar{\mathbf{\Sigma }}=\sum _{i=1}^l{\mathbf{\lambda }}_i{\left({\bar{\mathbf{\Sigma }}}^{1/2}{\mathbf{\Sigma }}_i{\bar{\mathbf{\Sigma }}}^{1/2}\right)}^{1/2},$$

$${\mathbb{P}}_k[\text{x}]=\sum _{j=1}^k{\mathbf{\pi }}_j^{\mu }\varphi (\mathbf{x};{\mathbf{m}}_j^{\mu },{\mathbf{\Sigma }}_j^{\mu }),$$

$$\delta (\mathbf{x})={\mathrm{argmax}}_{j\in [k]}{\mathbf{\gamma }}_j(\mathbf{x}).$$

$$\begin{array}{rl}{\hat{\mathbf{T}}}_a^b(\mathbf{x})& =\sum _{(i,j)\in [k{]}^2}{\mathbf{A}}_{i,j}\left[{\mathbf{m}}_j^b+\Gamma (\mathbf{x},i)\right],\\ \text{where }\Gamma (\mathbf{x},i)& =\{\begin{array}{ll}{\mathbf{T}}_{{\mathbf{\Sigma }}_i^a}^{{\mathbf{\Sigma }}_j^b}(\mathbf{x}-{\mathbf{m}}_i^a)& {\delta }^a(\mathbf{x})=i,\\ 0& \text{otherwise},\end{array}\end{array}$$

$$D\left[\mathbb{P}[\mathbf{y}|{\hat{\mathbf{T}}}_a^c({\mathbf{X}}^a);\theta ],\mathbb{P}[\mathbf{y}|{\hat{\mathbf{T}}}_b^c({\mathbf{X}}^b);\theta ]\right]\le ϵ,$$

Full text

1 Introduction

In recent years, Optimal Transport-based algorithms (OT)[24] are increasingly attracting the machine learning community. OT leverages useful information about the nature of a problem by encoding the geometry of the underlying space through a cost metric. As OT measures distances between probability distributions, it is generally used as an optimization loss [1, 12, 14, 29]. OT aims at finding a probabilistic matrix that couples two distributions; this matrix is also called a transportation matrix. However, we can only use this matrix with the data used to learn it. Thus, one has to recompute an OT problem to handle new samples. Hence, it is prohibitive for some applications where the OT problem is not used as a distance to be minimized, for instance, transport out-of-sample data into a fixed distribution [7]. Another critical application is computing a weighted mean in a Wasserstein space. This weighted mean is known as the Wasserstein barycenter, which has numerous applications in unsupervised learning. It is used to represent the input data [26], to cluster densities [31], to analyze shapes [25], to build fair classifiers [10], among other applications.

To handle unseen data, [23] proposes to learn a transport map. However, this approach relies on kernel-based methods, which makes it slow when the number of samples increases. To approximate the barycenter, [5] proposes to decompose the input data into sums of radial basis functions and perform partial transport between matched pairs of functions coming from each distribution. However, it is not clear what is the influence of noise in this approximation as it handles the positive and negative parts of the data independently. Additionally, the extension of this approach to out-of-sample data is not explored.

Contributions.

We propose to use a local approximation of the Gaussian (Bures-Wasserstein) transport to speed up learning mappings and barycenters that can be applied to unseen data. This approach assumes that the densities involved in the transport problem are highly concentrated in a few clusters/components. This idea is reminiscent of the computation of barycenters by partial transport [5]. The key point is to express each density as a mixture of Gaussians; contrary to [5], it is no longer necessary to split the input data into its positive and negative parts. Then, we can use the means and covariances of these Gaussian functions to approximate the transport map as well as the barycenter. To do so, we use these means as a new spatial reference. Then, we match pairs of means of different distributions and solve OT problems for each matched pair.

The proposed method unifies the learning of the barycenter and transport map. We provide some empirical evidence for the usefulness of our approach in two tasks: i) interpolating of clouds of points, and ii) building fair classifiers. Through experiments on public datasets, we show the improvements brought by our scheme with respect to standard mapping estimation methods. Our approach can be easily implemented and combined with conventional machine learning pipelines.

Notation.

Column vectors are denoted as bold lower-case, e.g, $\mathbf{a}$. Matrices are written using bold capital letters, e.g., $\mathbf{A}$. Let $\|\cdot\|_{F}$ be the Frobenius norm of a matrix, $\mathop{\rm Tr}(\cdot)$ is the trace of a matrix, and $\mathbf{A}^{\mathsf{T}}$ is the transpose of $\mathbf{A}$. We denote the $(m-1)$-dimensional probability simplex by $\Delta^{m}\coloneqq\{\mathbf{y}\in\mathbb{R}_{+}^{m}:\|\mathbf{y}\|_{1}=1\}$. Let $n$ be the number of data points and $[n]$ denotes $\{1,\ldots,n\}$.

2 Background: Gaussian optimal transport

We will restrict ourselves to the discrete setting (i.e., histograms); thus, $\mu\in\Delta^{p}$ and $\nu\in\Delta^{q}$. The cost matrix $\mathbf{C}\in\mathbb{R}^{p\times q}$ represents the cost function (e.g., the squared Euclidean distance), which contains the cost of transportation between any two locations in the discretized grid. The OT distance between two histograms $W(\mu,\nu)$ is the solution of the discretized Monge-Kantorovich problem [30]:

[TABLE]

where $\mathbf{T}:\mathbb{R}^{p}\rightarrow\mathbb{R}^{q}$ is the transport plan, $\Pi(\mu,\nu)$ denotes the set of admissible couplings between $\mu$ and $\nu$, that is, the set of matrices with rows summing to $\mu$ and columns to $\nu$.

However, computing Eq.(1) for other distributions than Gaussians requires to solve a linear programming problem that grows at least $O(p^{3}\log p)$ [24] for histograms of dimension $p$. We can constrain the entropy of the transport plan $\mathbf{T}$ to alleviate this issue [8].

Bures-Wasserstein transport.

The OT distance in Eq.(1) can be calculated in closed-form when $\mu$ and $\nu$ are Gaussians measures, as demonstrated by [15] and discussed by [28]. This corresponds to the Bures-Wasserstein distance [3] between $\mu\sim\mathcal{N}(\mathbf{m}_{1},\mathbf{\Sigma}_{1})$ and $\nu\sim\mathcal{N}(\mathbf{m}_{2},\mathbf{\Sigma}_{2})$:

[TABLE]

where $\mathbf{m}_{i}\in\mathbb{R}^{p}$ and $\mathbf{\Sigma}_{i}\in\mathbb{R}_{+}^{p\times p}$ for $i\in[2]$ are the mean vectors and the covariance matrices, respectively. $W_{2}^{2}(\cdot,\cdot)$ uses the squared Euclidean distance as transport cost. For zero-mean Gaussian random variables, the Wasserstein distance reduces to Frobenius of the covariance roots [18], and the transport map is:

[TABLE]

where $\mathbf{T}_{\mathbf{\Sigma}_{1}}^{\mathbf{\Sigma}_{2}}$ is the transport map from $\mathbf{\Sigma}_{1}$ to $\mathbf{\Sigma}_{2}$. In general, the Gaussian transport $\mathbf{T}:\mathbb{R}^{p}\rightarrow\mathbb{R}^{p}$ is

[TABLE]

Bures-Wasserstein barycenter.

The Wasserstein barycenter problem attempts to summarize a collection of probability distributions by taking their weighted average with respect to the Wasserstein metric. In particular, the Wasserstein barycenter of $l$ Normal distributions $\mathcal{N}(\mathbf{m}_{0},\mathbf{\Sigma}_{0}),\ldots,\mathcal{N}(\mathbf{m}_{l},\mathbf{\Sigma}_{l})$, with weights $\boldsymbol{\lambda}\in\Delta^{l}$ is the solution of the minimization problem

[TABLE]

The barycenter of Gaussians is still Gaussian, and its mean $\bar{\mathbf{m}}$ can be computed directly as $\bar{\mathbf{m}}=\sum_{i=0}^{l}\boldsymbol{\lambda}_{i}\,\mathbf{m}_{i}$. However, its covariance $\bar{\mathbf{\Sigma}}$ is the solution of the following matrix equation

[TABLE]

which has not a closed-form solution for $l>2$.

The algorithm 1[18] finds a fixed point of Eq.(6). It initializes the barycenter with a positive definite matrix. Then, it lifts all observation to the tangent space at the initial guess via log map and performs a linear averaging. This average is then retracted onto the manifold via the exponential map, providing a new guess. The algorithm iterates until it reaches the stop condition (e.g., number of iterations).

3 The model

We are interested in a fast approximation of a transport map that generalizes to out-of-samples data. We assume the data lives in a low-dimensional feature space, and a few Gaussians concentrate the samples. Thus, we use mixtures of Gaussians to approximate independently the densities $\mu$ and $\nu$ involved in Eq.(1). Then, we compute the pairwise transport plans (see Eq.4) at the component level. We refer to this approach as local Bures-Wasserstein (L-BW) transport.

Gaussian Mixture Model (GMM).

We model the density $\mu$ as a weighted average of $k$ Gaussian functions, where each Gaussian or component represents a subpopulation of the data. The parameters of each component $i\in[k]$ are the mean vector $\mathbf{m}_{i}^{\mu}\in\mathbb{R}^{p}$ and covariance matrix $\mathbf{\Sigma}_{i}^{\mu}\in\mathbb{R}_{+}^{p\times p}$.

[TABLE]

where $\boldsymbol{\pi}^{\mu}\in\Delta^{k}$ is the mixing proportions, and $\phi(\cdot;\,\mathbf{m}_{j},\mathbf{\Sigma}_{j})$ denotes a $p$-dimensional Normal density with mean vector $\mathbf{m}_{j}$ and covariance matrix $\mathbf{\Sigma}_{j}$. We usually use the Expectation-Maximization (EM) algorithm to fit a GMM [13]. This algorithm alternates between estimating the probability to assign a sample $\mathbf{x}$ to component $j$, $\boldsymbol{\gamma}_{j}(\mathbf{x})\in\Delta^{k}$, and calculating the parameters of model. We denote $\delta(\mathbf{x})$ the (hard) assignment of sample $\mathbf{x}$ to a component,

[TABLE]

3.1 Leveraging Bures-Wasserstein transport.

For simplicity, we describe our approach for two densities/groups. We depict how to extend it to more groups at the end of this section. Let $g\in\{a,b\}$ be the group indicator, and $\mathbf{X}^{g}\in\mathbb{R}^{|q|\times p}$ be the set of samples $\mathbf{x}$ that belong to the group $g$, where $|g|$ denotes the number of samples of group $g$.

Getting the parameters.

We aim to leverage the closed-form of the Bures-Wasserstein transport (see Eq.(4)) to learn an approximated transport map. To do so, for each $g\in\{a,b\}$, we use a $k$ mixture of Gaussians to approximate the density of $\mathbf{X}^{g}$. We use the means $\{\mathbf{m}_{i}^{g}\}_{i=1}^{k}$ as a rough approximation of the geometry. Then, we solve $k$ pairwise Bures-Wasserstein OT problems. We select pairs because the OT of Gaussians is still Gaussian, removing important geometrical information.

To select which components from each group to pair, we solve a linear allocation problem [6] of their mean vectors. In particular, we use the Hungarian algorithm [19] to match these values. This algorithm has a time complexity $O(k^{3})$, where $k\ll\min(|a|,|b|)$, and solves an OT problem for uniform distributions of the same size [24]. We set the transport cost to $\mathbf{C}_{i,j}=\|\mathbf{m}_{i}^{a}-\mathbf{m}_{j}^{b}\|^{2},\,\mathbf{C}\in\mathbb{R}^{k\times k}$. We denote $\mathbf{A}\in[0,1]^{k\times k}$ a matrix that encodes the pairwise matching of components. Thus, for $(i,j)\in[k]^{2}$, $\mathbf{A}_{i,j}=1$ if $\mathbf{m}_{i}^{a}$ match $\mathbf{m}_{j}^{b}$ and [math] otherwise.

From now on, we rely on $\mathbf{A}$, $\left\{(\mathbf{m}_{l}^{g},\mathbf{\Sigma}_{l}^{g})\right\}_{l=1}^{k}$ and $\delta(\cdot)^{g}$ to approximate the barycenter and the transport map. The algorithm 2 presents a pseudo-code of learning these parameters.

We can extend this approach for a number $s$ of groups by sequentially computing one reference group and the other populations (a.k.a. one-vs-rest approach). This has a time complexity of $O(s\,k^{3})$.

Approximate a transport map.

To transport sample $\mathbf{x}$ from a distribution $a$ to another $b$, we use the algorithm 2 and apply Eq.(4) for each matched pair of components, as follows

[TABLE]

where $\hat{\mathbf{T}}_{a}^{b}(\mathbf{x})$ denotes the approximated transport map of $\mathbf{x}$ from the population $a$ to population $b$. In a nutshell, our algorithm performs pairwise transport between components of two distributions, $a$ and $b$. First, it uses $\delta^{a}(\mathbf{x})$ to assign a sample $\mathbf{x}$ to a $i$ subpopulation of $a$, which has parameters $\left(\mathbf{m}_{i}^{a},\mathbf{\Sigma}_{i}^{a}\right)$. Then, it matches the $i$-Gaussian of group $a$ with a $j$-component of group $b$, parametrized by $\left(\mathbf{m}_{j}^{b},\mathbf{\Sigma}_{j}^{b}\right)$. Next, it computes the Bures-Wasserstein transport of these two Gaussians. Finally, It repeats for all pairs of matched components111 We can use our approach with Spark [33] for large-scale data applications; the GMM is part of SparkML. .

Approximate a barycenter.

Let $c$ be the barycenter of the distributions $a$ and $b$. To approximate this barycenter, we use algorithm 2 to compute the parameters of the L-BW for $a$ and $b$. Then, we calculate the Bures-Wasserstein barycenter of every $(i,j)$-pair that satisfies the match condition. Thus, the approximated barycenter $c$ is in the space of $k$ Gaussian mixtures with parameters $\{(\bar{\mathbf{m}}_{h},\bar{\mathbf{\Sigma}}_{h})\}_{h=1}^{k}$. The algorithm 3 presents the pseudo-code of the approximate barycenter by L-BW222 Note that using one-vs-rest allows us to extend this method to more than two distributions.. Fig.1 shows an illustration of this procedure.

4 Related work

Mapping Estimation (ME).

To extend OT-based methods to unseen data, [23] proposes an algorithm to learn the coupling and an approximation of the transport map. This optimization problem learns a transformation regularized by a transport map. However, ME relies on kernel-based methods, which usually have a time complexity of $O(n^{3})$ [2] for $n$ samples. Additionally, it requires tuning the parameters associated with each kernel and the regularization parameter.

Practical computation of the barycenter.

[5] proposes a first approach to approximate the barycenter using signal decomposition and partial transport for multiple frequency bands. However, it requires the nontrivial task of selecting the frequency bands. A second approach by [9] proposes a fast algorithm to compute the free support Wasserstein barycenter. This algorithm optimizes the locations of the barycenter and not the weights associated to the discrete measure. It uses $k$ atoms or points to constrain the support of the barycenter. It is similar to our method, as both rely on the approximation of the geometry by $k$ atoms/components. However, our approach requires fewer atoms because the covariance matrices of the Gaussian mixtures encode more geometrical information.

5 Applications

We conduct a series of experiments to highlight the versatility of Local Bures-Wasserstein (L-BW) transport. We consider two tasks: i) computing the barycenter of clouds of points and ii) building fair classifiers. Both problems use an approximation of the barycenter and the transport map.

We denote $\hat{\mathbf{T}}_{a}^{c}(\cdot)$ the learned transport map from population $a$ to $c$, where $c$ is the learned barycenter of densities $a$ and $b$. We test the learned transport map on held-out data, in a cross-validation scheme splitting the data ten times at random. Each time, we learn the barycenter and the transport map on $70\%$ of the data and apply it to the other $30\%$. Besides L-BW, we also investigate Mapping Estimation (ME) [5] with both, linear and Gaussian kernels. For the ME, we first compute the free support Wasserstein barycenter [9], and then we use ME to learn the mapping to this barycenter.

Henceforth, we refer to $k$ as the number of atoms. It represents, both, the number of components in L-BW and the number of points in the free support Wasserstein barycenter used in ME.

Technical aspects.

We used standard data processing and classification methods implemented in scikit-learn [22]. We used the POT library [11] for the convolutional Wasserstein barycenter and Mapping Estimation. We used Scipy [16] for the Hungarian algorithm. Our method is available333Link to the Github repository.

5.1 Shape interpolation of clouds of points

A straightforward application of Wasserstein barycenter is the interpolation of shapes, as it relies on performing optimal transportation over geometric domains. We vary the number of atoms $k\in\{2,\ldots,100\}$ and evaluate the performance of the approximated transport map and barycenter with two measures: i) support recovery and ii) computation time. For the ME, we set the regularization parameters of the linear and Gaussian kernels to $\eta=10^{-3}$ and $\sigma=1$, respectively444 These are the parameters by default in the POT package [11]..

Synthetic dataset. We generated a synthetic dataset that consists of three binary $200\times 200$ images: cat, rabbit, and tooth (see Fig.6 in supp. materials). We put them in a $1200\times 1200$ square grid. We sample $\approx 10\,000$ points at random inside every silhouette. We aim to calculate the barycenter of these images given some weights $\boldsymbol{\lambda}\in\Delta^{3}$.

Quality assessment.

To compare the quality of the approximated barycenter, we measure its agreement proportion with ground-truth. The ground-truth is in itself a silhouette, and we build it using the convolutional 2-D Wasserstein barycenter algorithm [27]555 To use this algorithm, the images are normalized to sum to $1$ and concatenated in a third dimension., followed by binarization by thresholding. We set the threshold via the Otsu’s method [21].

Fig.2a displays the support recovery of intermediate silhouettes as a function of the number of atoms. ME-based methods require five times more atoms than L-BW to improve the quality of their resulting images. L-BW has higher percentage of agreement with ground truth than ME with both linear and Gaussian kernels. It seems to reach a plateau at $k=20$, which corresponds to the best representation according to the Akaike information criterion (see Fig.7 in supp. materials).

Fig.3 shows the barycenters found by various algorithms for $k=20$ on the synthetic dataset. Fig.3a displays our ground-truth. ME-Gaussian generates barycenters that reproduce the shapes, yet it concentrates the transported samples in a few atoms. In contrast, ME-Linear does not concentrate the transported samples; however, it fails to preserve the silhouette. L-BW reproduces the transitions between inputs, and creates plausible intermediate silhouettes.

Computation time.

We measured the running time of various approaches, including both the training time and the transport time.

Fig.2b gives computation time for different methods. L-BW displays overall the lowest computation time for all atoms; on average $80$ times faster than ME-Linear. For $k<40$, ME-Linear is ten times faster than ME-Gaussian; when the number of atoms increases, this factor drops to approximately 1.5.

5.2 Demographic-parity fairness

Another application of OT is to build fair classifiers. We rely on demographic-parity, which requires that predictions of an estimator cannot reveal information about the protected class any better (up to a tolerance) than random guessing [20, 32].

We consider mixing distributions of two groups (e.g., Female and Male) via OT to produce naive fair classifiers. We modify the input data $\mathbf{X}^{g}$ by transporting each sample $\mathbf{x}$ to the barycenter $c$ of the two populations, $g\in\{a,b\}$. This is proposed in [10] for Gaussian random variables. We extend this idea to more complex distributions via L-BW, as follows

[TABLE]

where $\theta$ denotes the parameter of the classifier, $\mathbf{y}$ is the binary classification target, and $D(\mu,\nu)$ is a similarity measure of two distributions $\mu$ and $\nu$ (e.g., Kullback-Lieber, or Wasserstein). $\epsilon$ is the desired tolerance. This transformation of the input data is oblivious to the target; hence we do not have precise control of the tolerance $\epsilon$.

Fairness experiment.

We present an empirical study on five real-world datasets. In all classification tasks, we collect statistics concerning prediction and fairness score measure. We use nested cross-validation to set hyperparameters. Our preprocessing pipeline comprises standardization and dimensionality reduction via Principal Component Analysis (PCA). We use one-hot-encoding to handle categorical variables.

We use the Demographic-Parity-$\gamma$ score [20] to measure fairness. It is the maximum distance of the Receiver Operating Characteristic (ROC) curve for the group membership label from the diagonal, where zero represents an absolute lack of predictability (see Fig.10 in supp. materials). Thus, the lower the fairness score, the better. For prediction performance, we use the area under the ROC curve (AUC). This score upper bounded by $1$ and higher values mean better predictions.

Setting hyperparameters.

We vary the number of principal components $d\in\{2,4,\ldots,30\}$, and the number of atoms $k\in\{1,\ldots,10\}$ for the L-BW. We set the parameter to regularize the covariance of the GMMs to $10^{-6}$. To compute the free support Wasserstein barycenter, we set the grid $k\in\{10,20,50,100,200,300\}$. We set grid for the parameters of ME to $\eta\in\{10^{3},10^{2},10^{1},1,10\}$ and $\sigma\in\{1,10,10^{2},10^{3}\}$ for the linear and Gaussian kernels, respectively. We set weights of the barycenter to $\boldsymbol{\lambda}=[0.5,0.5]$, the new distribution should not benefit any of two populations.

Datasets.

We investigate our approach in several binary classification problems based on four datasets from the UCI Machine Learning Repository[4]. See Fig.11 in supp. materials a visualization of the first two principal components of these datasets conditioned on the sensitive groups.

Wine quality666https://archive.ics.uci.edu/ml/datasets/wine+quality: it contains $12$ attributes of $4\,898$ wines. Professional taste-tester provided a ranking out of $10$. The task is to predict if a wine has a rating higher or equal than $6$. The protected groups are color-based (White and Red). All explanatory variables are continuous.

German Credit Data777 http://archive.ics.uci.edu/ml/datasets/statlog+(german+credit+data): it contains $20$ features ($7$ numerical and $13$ categorical) related to the economic situation of $1\,000$ German applicants for loans. The aim is to predict whether or not an applicant is going to default the credit loan. The protected groups are gender-based (Female and Male).

Adult Income888https://archive.ics.uci.edu/ml/datasets/adult: it contains $14$ features concerning demographic characteristics of $45\,222$ individuals. The task is to predict whether or not a person has more than $50\,000$ as an income per year. The protected groups are education-based (Higher and Lower education) and gender-based (Male and Female).

Taiwan credit999https://archive.ics.uci.edu/ml/datasets/default+of+credit+card+clients: it contains the credit information of $30\,000$ individuals $24$ features represent this information, where $9$ of them are categorical. The task is to predict customers if default payments.

COMPAS [17]: it is a collection of $10\,000$ criminal defendants screened in Broward County, Florida U.S. in $2013$ and $2014$. The features are demographics and criminal records of offenders. The task is to score an individual’s likelihood of reoffending. The protected groups are race-based (Black and White) and gender-based (Male and Female).

Fairness on a time budget.

We examine the impact of mixing groups on prediction time for various approximate transport map methods. We are interested in the total computation time needed to learn a model: the cost of computing the approximate representation (i.e., fitting a barycenter and a transport mapping) and of training the classifier.

Fig.4 shows the prediction accuracy as a function of the computation time for training a logistic regression with $\ell_{2}$-penalty, and it displays the mean the across folds and the percentiles $25$ and $75$. In both datasets, L-BW requires less time to converge to an estimator with a lower fairness score and displays the same predictive performance as ME-Gaussian. ME-Linear has the best classification performance in the income dataset; however, it has a more significant unfair behavior than L-BW.

Fair classifiers.

We analyze the impact of mixing groups by L-BW in various conventional classifiers: Logistic regression ($\ell_{2}$ penalty), Gradient Boosting Trees ($100$ trees), and Random Forest ($100$ trees). See Fig.13 in supp. materials for more classifiers.

Fig.5 shows the performance classifiers after mixing groups with L-BW. The performance is relative to the mean performance of the raw classifier (per dataset). The fairness score (DP-$\gamma$) of Random-Forest does not display any improvement, whereas the other classifiers present significantly better results ($\text{p}_{\text{value}}<10^{-8}$, paired Wilcoxon rank test). L-BW decreases the predictive performance of Gradient Boosting Trees to the raw classifier. This drop in accuracy is not significant for other classifiers.

6 Discussion

We proposed a fast method to compute an approximation of the barycenter and transport mapping, which can be used in out-of-sample data. Our method directly exploits highly concentrated samples in low dimensions to approximate the Wasserstein barycenter as well as the transport map. First, it approximates the training data with Gaussian mixtures. For each distribution, we fix the means of these Gaussians as the new geometrical reference. Then, we select the closest means between densities. Finally, we perform pointwise Gaussian transport between matched components. Additionally, we propose an extension of our approach to more than two densities, which has linear complexity in the number of densities. From a machine learning perspective, we can directly project held-out-data into the approximate barycenter, which makes it possible to use in standard machine learning pipelines.

We demonstrated the effectiveness of our approximation on two problems: shape interpolation of clouds of points and naive demographic-parity fair classifiers. In both problems, our method displays the fastest computation time, outperforming the ME. L-BW extends [10] to densities more complex than Gaussians. It requires fewer components to approximate the barycenter and, for most classifiers, it produces naive fair estimators while maintaining prediction accuracy. Our approach is similar to [9], as it relies on $k$ Gaussian components (atoms). Yet, our approach also takes the covariance of subpopulations into account. Experimental results confirm that this additional information benefits the support recovery of the barycenter.

However, the current framework has the following limitations: i) it assumes that the few clusters represent the densities involved, which limits its applications to low/medium dimensionality settings; ii) it assumes that the number of atoms required to approximate each distribution included in the barycenter calculation is the same, which can be problematic for some data; iii) the theoretical properties of the L-BW are out of the scope of this paper. Addressing the above limitations are new directions for future work.

7 Applications: Additional analysis

7.1 Shape interpolation of a cloud of points experiment

7.2 Statistical-parity fairness experiment

Fig.13 shows the impact of mixing groups by L-BW in various conventional classifiers: Logistic regression ($\ell_{2}$ penalty), Naive Bayes (Bernoulli), k-Nearest Neighbors (k-NN), Linear Discriminant Analysis (LDA), Gradient Boosting Trees ($100$ trees), Random Forest ($100$ trees), and Multi Layer Perceptron (MLP, $100$ hidden units and ReLu as activation function). The fairness score (DP-$\gamma$) of K-NN and Random-Forest does not display any improvement, whereas the other classifiers present significant results ($\text{p}_{\text{value}}<10^{-8}$, paired Wilcoxon rank test). L-BW reduces the predictive performance of Gradient Boosting Trees and Linear Discriminant Analysis. Contrary to other classifiers, where the difference to the raw classifier is not significant.