Federated Wasserstein Distance
Authors: Alain Rakotomamonjy
Kimia Nadjahi
Liva Ralaivola
Venue: ICLR 2024 poster (2024)
Keywords: Wasserstein distance
Federated Learning ; Triangle inequality
Primary area: general machine learning (i.e., none of the above)
TL;DR: We propose a federated manner of computing Wasserstein distance by exploiting its metric property and its geodesics

Abstract: We introduce a principled way of computing the Wasserstein distance between two distributions in a federated manner.
Namely, we show how to estimate the Wasserstein distance between two samples stored and
 kept on different devices/clients whilst a central entity/server orchestrates the computations 
 (again, without having access to the samples). To achieve this feat, we take advantage of the geometric 
properties of the Wasserstein distance -- in particular, the triangle inequality -- 
 and that of the associated {\em geodesics}: our algorithm, FedWad (for Federated Wasserstein Distance), iteratively approximates 
 the Wasserstein distance by manipulating and exchanging distributions from the
  space of geodesics in lieu of the input samples. 
  In addition to establishing the convergence properties of FedWad,
   we provide empirical results on federated coresets and federate 
   optimal transport dataset distance, that we respectively exploit for
   building a novel federated model and for boosting performance of popular federated learning algorithms.

Published as a conference paper at ICLR 2024
FEDERATED WASSERSTEIN DISTANCE
Alain Rakotomamonjy
Kimia Nadjahi CSAIL, MIT Boston, MA knadjahi @ mit.edu
Liva Ralaivola
Criteo AI Lab
Criteo Al Lab
Paris, France
Paris, France
Tralivola@crrito..co
ABSTRACT
We introduce principled way of computing the Wasserstein distance between two distributions in a federated manner. Namely, we show how to estimate the Wasser- stein distance between two samples stored and kept different devices/clients whilst central entity/server orchestrates the computations (again, without having access to the samples). To achieve this feat, we take advantage of the geomet- ric properties of the Wasserstein distance in particular, the triangle inequality and that of the associated geodesics: our algorithm, FedWaD (for Federated Wasserstein Distance), iteratively approximates the Wasserstein distance by ma- nipulating and exchanging distributions from the space of geodesics lieu of the input samples. In addition to establishing the convergence properties of FedWaD, we provide empirical results on federated coresets and federate optimal transport dataset distance, that respectively exploit for building novel federated model and for boosting performance popular federated learning algorithms.
INTRODUCTION
Context. Federated Learning (FL) is form of distributed machine learning (ML) dedicated to train global model from data stored on local devices/clients, while ensuring these clients never share their data (Kairouz et al.] 2021: Wang et al.) 2021). FL provides elegant and convenient solutions concerns in data privacy, computational and storage costs of centralized training, and makes possible to take advantage of large amounts of data stored on local devices. A typical FL approach to learn parameterized global model alternate between the two following steps: update local versions of the global model using local data, and ii) send and aggregate the parameters of the local models on central server (McMahan et al. 2017) update the global model.
Problem. In some practical situations, the goal is not to learn prediction model, but rather to compute a certain quantity from the data stored on the clients. For instance, one's goal may be to compute, in federated way, some prototypes of client's 's data, that can be leveraged for federated clustering or for classification models (Gribonval et al. 202T Phillips] 2016 Munteanu et al. 2018 Agarwal et al. (2005). In another learning scenarios where data scarce, may want look for similarity between datasets in order to evaluate dataset heterogeneity over clients and leverage this information improve the performance of federated learning algorithms. In this work, we address the problem of computing, a federated way, the Wasserstein distance between two distributions μ and y when samples from each distribution are stored on local devices. A solution to this problem will be useful in the aforementioned situations, where the Wasserstein distance is used similarity measure between two datasets and is the key tool for computing some coresets of the data distribution or cluster prototypes. We provide a solution to this problem which hinges on the geometry of the Wasserstein distance and more specifically, its geodesics. We leverage the property that for any element E* of the geodesic between two distributions μ and the following equality holds, Wp(u, v) = Wp( (++++ + WW where W, denotes the p- Wassersstiin distance. This property is especially useful compute (p,,,, v)) in federated manner, leading novel theoretically-jjsttiice procedure coined FedMaa, for Federated Wasserstein Distance.
Contribution:
FedWaD. The principle of FedWaD is to iteratively approximate which, in terms of traditional FL, can be interpreted as the global model. At iteration k; our procedure consists in i) computing, on the clients, distributions &h and ek from the geodesies between the current

---

Published as a conference paper at ICLR 2024
approximation of 5" and the two secluded distributions μ and v < and e playing the role of the local versions of the global model, and ii) aggregating them on the global model to update 8*.
Organization of the paper. Section formalizes the problem we address, and provides the necessary technical background to devise our algorithm FedWaD. Section devoted to the depiction of FedWaD, pathways to speed-up its executions, and a theoretical justification that FedWaD is guaranteed to converge to the desired quantity. In Section 3 FedWaD on different use-cases (Wasserstein coresets and Optimal Transport Dataset distance) which rely on the computation of the Wasserstein distance. We unveil how these problems can be solved in our FL setting and demonstrates the remarkable versatility of our approach. In particular, we expose the impact of federated coresets. By learning single global model on the server based on the coreset, our method can outperform personalized FL. models. In addition, our ability to compute inter-device dataset distances significantly helps amplify performances of popular federated learning algorithms, such as FedAvg, FedRep, and FedPer. We achieve this by clustering clients and harnessing the power of reduced dataset heterogeneity.
we conduct empirical analysis of
RELATED WORKS AND BACKGROUND
2.1 WASSERSTEIN DISTANCE AND GEODESICS
Throughout, denote by (X) the of probability measures in X. Letp 2 1 and define P, (X) the subset of measures in (X) with finite p-moment, i.e., P Px(X)= = {" E (X)) : < oo). where Mp(n)= = fx (x, ()dr)(a)) and dx is metric on X often referred
the ground cost. For
(X) and v/ E (Y),,, vI(,,v) c (X Y) is the collection of probability measures or couplings on X x Y defined as
(( { ( XAC
= ((A) and (X x B) = xBB).
The p- Wasserstein distance Wp, v) between the measures ja and assuued to be defined over the same ground space, i.e. X Y- is defined
1/2
Wp( =
inf (Ell().))) xxx
(1)
It is proven that the infimum in ①1 is attained (Peyré et al. 2019) and any probability which realizes the minimum is an optimal transport plan. In the discrete case, we denote the two marginal measures as y Σ agdg, and = Em= b.óag. with a;, b, 2 0 and Σ a; > by 1. The Kantorovitch relaxation of (I) seeks for transportation coupling P that solves the problem
1/P
W)
min (C.P)
(Pell(.bb)
where C = (d& (zi, (xi,, Raxm is the matrix of all pairwise costs, and II(a, b) = (P
I R
P1 = a, P = b) is the transportation polytope (i.e. the set of all transportation plans) between the distributions and b.
Property (Peyré et al. 2019 For any p 2 1, Wp is metric Px(X). As such it satisfies the triangle inequality:
VLU (AMEEO,(X)) Wp(,,D) S Wp((, + W))
might be convenient consider geodesics as structuring tools of metric spaces.
Definition 1 (Geodesics, Ambrosio et al. (2005)). Let (X, d) be metric space. A constant speed geodesic x [0. 1] + X between To, 41 E X is a continuous curve such that Vs, E [0, 1]. d(x(s), (())
1| 2()0,
Property 2 (Interpolating point, Ambrosio et al. (2005). Any point v, from constant speed geodesic (x(t))nejo,) is an interpolating point and verifies, d(zo. I1) = d(zo, xx) ) + d(z,, 21), f.e. the triangle inequality becomes equality.
2

---

Published as a conference paper at ICLR 2024
5o *
At first iteration the current estimation Ço of 5" is sent each client order com- pute two interpolating measures Ey and &u, which are sent back to the server.
cla
4
Server then computes an interpolating mea- sure between Eμ and & to obtain the next iterate of geodesic element <²
The process is repeated until convergence to obtain 5 and we define Wo(,, Wy(m, <*) + W,(u,
300
CR
50
Figure 1: The Wasserstein distance between μ and which are on their respective clients can be computed as v) = ()( (t, where 5* an element on the geodesic between j and FedWaD seeks estimating ** with g using an iterative algorithm and plugs in this estimation to obtain Wp(ss, v). Iterates of &, are computed on the server and sent to clients in order to compute measures f and E that are on geodesics of p and & and v and & respectively.
These definitions and properties carry over the case of the Wasserstein distance:
Definition 2 (Wasserstein Geodesics, Interpolating measure, Ambrosio et al. (2005): Kolouri et al. (2117)) Let Po- #r E (X)) with C. Rd compact, convex and equipped with Wp. Let y E [[ p1) be optimal transport plan. For t [0, L]. let μl = #7 where (I( y) = t)z ty, i.e. Pt is the forward measure of under the map Tr. Then, the curve J (H4))e is constant speed geodesic between po and μ1; call Wasserstein geodesics between pi and μl1
Any point μr of the geodesics is an interpolating measure between po and μM and, as expected:
Wy m1) Wp ue) + Wp(pe,
In the discrete case, and for a fixed one can obtain such interpolating measure 14 given the optimal transport map P solution of Equation ( as follows (Peyré et al. 2019 Remark 7.1):
Pr Σ
(5)
where P the (i.j)-th entry of P*; as interpolating measure, 14, obviously complies with (4)
2.2 PROBLEM STATEMENT
Our goal is to compute the Wasserstein distance between two data distributions μ and global server with the constraint that μ and v are distributed on two different clients which do not share any data samples to the server. From mathematical point of view, our objective is estimate an element g* on the geodesic of I and without having access to them by leveraging two other elements &p and & the geodesics of y and &* and and g respectively.
2.3 RELATED WORKS
Our work touches the specific question of learning/aproxximating distance between distributions whose samples are secluded on isolated clients. As far as we are aware of, this is a problem that has never been investigated before and there are only few works that we see closely connected to ours. Some problems have addressed the objective of retrieving nearest neighbours of one vector in a federated manner. For instance, Liu et al. (202T) consider to exchange encrypted versions of the dataset on client to the central server and Schoppmann et al. (2018) consider the exchange of differentially private statistics about the client dataset. Zhang et al. (2023) propose federated approximate k-nearest approach based on a specific spatial data federation. Compared to these works that compute distances in federated manner, we address the case of distances distributions without any specific encryption of the data and we exploit the properties of the Wasserstein distances and its geodesics, which have been overlooked in the mentioned works. If these properties have been
3

---

Published as a conference paper at ICLR 2024
relied upon as key tool in some computer vision applications (Bauer et al. 2015 Maas al. 2017) and trajectory inference (Huguet et al. 2022), they have not been employed as a privacy preseevig tool.
3 COMPUTING THE FEDERATED WASSERSTEIN DISTANCE
In this section, we develop methodology to compute, on a global server, the Wasserstein distance between two distributions f and v, stored on two different clients which not share this information to the server. Our approach leverages the topology induced by the Wasserstein distance in the space of probability measures, and more precisely, the geodesics.
Outline of our methodology. A key property is that Wp is a metric, thus satisfies the triangle inequality: for any μ. U E Px()),
Wp(u,D) I Wp(m:)) {) + Wo(5,
(6)
with equality if and only i E", where 5" is an interpolating measure. Consequently, one can compute (n v) by computing (p, ( and W,(E", and adding these terms. This result is useful in the federated setting and inspires our methodology. as described hereafter. The global server computes 5" and communicate e* to the two clients. The clients respectively compute (p, {*) and Wy(5*, v), then send these to the global server. Finally, the global server adds the two received terms to return Wp(ss, 2/).
The main bottleneck of this procedure is that the global server needs to compute 5 (which by definition, depends on J.. 1/) while not having access to μu, v (which are stored on two clients). We then propose simple workaround to overcome this challenge, based on an additional application of the triangle inequality: for any & E (X),,
(7) where Eu and Eu are interpolating measures respectively between μ and & and & and v. Hence, computing g* can be done through intermediate measures En and fu. to ensure that jr, v stay on their respective clients. To this end, we develop an optimization procedure which essentially consists in iteratively estimating interpolating measure g(A) between μ and
Wp( S (p(m,E) + Wp(C,v) = (,(u, + (P {) + Wp(E,E + WP(Eu, v),
the server, by using (k) and <(h) which were computed and communicated by the clients. More precisely, the objective to minimize 7 over & as follows: at iteration k, the clients receive current iterate c(k-1) and compute < and <!*) (as interpolating measures between and (kk-1), and between g(k-1) and respectively). By the triangle inequality,
Wp(,,V) S Wp(,, e(4)) + (p(E ((k-1)) + (o(g(k-1),g!" + Wp(s! ,)),
(8)
therefore, the clients then send <(*²) and e! to the server, which turn, computes the next iterate {(k) by minimizing the left-hand side term of (8), i.e.
(()) argmin Wp ( ,E) + Wp(E,
(9)
Our methodology is illustrated in Figure 1 and summarized in Algorithm I can be applied to continuous measures as long as an interpolating measure between two distributions can be computed in closed form. Regarding communication cost, each iteration, the communication cost involves the transfer between the server and the clients of four interpolating measures: ((k-1) (twice), (k) ((k) Hence, if the support size of g(k-1) is S, the communication cost is in O(4SKd), with d the data dimension and K the number of iterations.
So
Reducing the computational complexity. In terms of computational complexity, we need to compute three OT plans per iteration which single cost, based on the network simplex is O((n m)nmllog(n + m)). More importantly, consider that μ and v are discrete measures, then, any interpolating measure between μ and is supported on most on + + 1 points. Hence, even if the size of the support of ((0) is small, but is large, the support of the next interpolating measures may get larger and larger, and this can yield an important computational overhead when computing (, ((k)) and (K(E(k),
4

---

Published as a conference paper at ICLR 2024
To reduce this complexity, resort approx- imations of the interpolating measures which goal is to fix the support size of the interpolating measures to a small number S. The solution we consider is approximate the McCann's 's inter- polation equation which formalizes geodesics & given an optimal transport map between two distributions,say, & and 5", based on the equation & = ((1 t) + + Peyré et al. )0119) Using barycentric mapping approximation of the map T (Courty et al., 2018), we propose to approximate the interpolating measures & as
Algorithm 1 FedWaD
Input:
and v. initialisation of ( function Interp- Meas that computes an interpolating measure be- tween two measures using Equation (5) or Equa- tion (10) for any 0 < <t<l.
1:
for l = 1 to K do
2:
Send c(k-1) to clients
3:
// Compute on clients with optional return of dis-
lances
4: E. ..W,W, [10)] < Interppeas(s,, ((k-1))
5: E WWpp( (,b((,,,) IInterpoas((, ((-2))
6: // Send E and El to server
Σ 8(1 (m (10) 1=1
&. =
7: g()) + -ItterpMeas((f" (((t)
8: end for
9: // Send (i ((k)), ()((((k), v) to server
where P is the optimal transportation plan be- tween & and 5' z, and x are the samples from these distributions and X' is the matrix of sam- ples from 5". Note that by choosing the appro-
Output: return d,, on ((k)) + W,, (c() v)
10: W,(H,,) = (
priate formulation of the equation, the support size of this interpolating measure can be chosen as the one of & or g'. In practice, we always opt for the choice that leads to the smallest support of the interpolating measure. Hence, if the support size of g()) is S, we have the guarantee that the support of g(k) is S for all k. Then, for computing Wp (1, (0(4g((() (+)) using approximated interpolating measures, it costs ((3 * (Sn² + loog)n + 5)) at each iteration and if S and the number of iterations K are small enough, the approach we propose is even competitive compared to exact OT. Our experiments reported later that for larger number of samples (2 5000), our approach is as fast as exact optimal transport and less prone to numerical errors.
Mitigating privacy issues. As for many FL algorithms, we do not provide or have a formal guarantee of privacy. However, we have components of the algorithm that helps mitigate risks of privacy leak. First, the interpolating measures can be computed for a randomized value of t; second, distances are not communicated to the server until the last iteration, and finally the use of the approximated interpolating measures in Equation (TO) helps in obfuscation since interpolating measure supports depend on the transport plan which is not reveal to the server. If formal differential privacy guarantee is required, one need to incorporate an (adapted) differentially private version of the Wasserstein distance (Le Tien et al. 2019 Goldfeld & Greenewald 2020).
Theoretical guarantees. We discuss in this section some theoretical properties of the components of FedWaD. At first, we show that the approximated interpolating measure is tight in the sense that there exists some situations where the resulting approximation exact.
Theorem 1. Consider two discrete distributions μ and with the same number of samples n and uniform weights, then for any t, the approximated interpolating measure, between and v given by Equation (10) is equal to the exact one Equation (5).
Proof is given in Appendix A. In practice, this property does not have much impact, but it ensures us about the soundness of the approach. In the next theorem, we prove that Algorithm iss theoretically justified, in the sense that its output converges (p,)v))
Theorem 2.
Let i and v be two measures in P PX),), ((k)) <l()) and g(k) be the interpolating measures computed at iteration k as defined in Algorithm 1 Denote as
(k) Wplu,, e(12) + W,(e!") g(t)) () <(hh) Wp(E ,
Then the sequence (A())) is non-increusing and converges to Wn,
We provide hereafter a sketch of the proof, and refer to Appendix B for full details. First, we show
(k+1) ck++)) and c(k+1) based on geodesics (a minimizer of the triangle inequality). Then, show that the sequence (A(1)) ) is
that the sequence (A(K)) ))& is non-increasing, as we iteratively update & S2
5

---

Published as a conference paper at ICLR 2024
Samp relto
Sampl NiS o1
Susport
Suppan 301 20
Fidhladie
Fazikade
ND (2.33
Fedhise sciti
Foedleo Ct
Techiapie-0 200
Fecefled (L:2
Fediiii
WBrkk:
Tediidde [2:3]
fadilade
Number campies.
en
kunber sanplesl
Figure 2: Analysis the different Wasserstein distance computation methods (most-left panels) for varying support size of the approximated FedWaD and (most-right panels) for varying sample ratio in the two distributions and fixed support size. For each couple of panels, for increasing number of samples, we report the running time and the relative error of the Wasserstein distance (WD), our exact Fed (FedWad-e) and our approximate FedWaD (FedWad-a) with support size of 2, 10 and 100. For the most-right panels, we have the support size of the interpolating measure to 10. For sample ratio (1:3), the first distribution has number of samples N and the second ones N/3.
bounded below by Wp(,,v) We conclude the proof by proving that the sequence (A(*))& converges (,,,
In the next theorem, we show that when and v are Gaussians then we can recover some nicer properties our algorithm and provide a convergence rate (proof in Appendix C).
Theorem 3. Assume that μ., v and g(0) are three Gaussian distributions with the same covariance matrix E ie µ ~ (mm,, Σ)), v ~ (m,, M and g(0) ~ ~ Nmmgoo E). Further assume that we are not the trivial case where my my, and mg are aligned. Applying our Algorithm 7 with 0.5 and the squared Euclidean cost, we have the following properties:
1. all interpolating measures </ (kk ance matrix E,
.50 , ((0) are Gaussian distributions with the same covari-
2. for any k 2 L, W2(H,V) || m2ll2 2||m m mg ll2 (W ,0 e
3. Wa(g(k), (
4. (2(,,,(()) + (2(g((), v) - W2(u,,) v) S g())
Interestingly, this theorem also says that in this specific case, only one iteration is needed to recover W2(p.v)
4 EXPERIMENTS
This section presents numerical applications, where FedWaD can successfully be used and show how it can boost performances of federated learning algorithms. The code for reproducing part of the results is available at /ittps::/// htt:://gitttht com/ arakot fedwad and built on top of the Python Optimal Transport library (Flamary et al. 2021). Full details provided in Appendix D.
Toy analysis. We illustrate the evolution of interpolating measures using FedWaD for cal- culating the Wasserstein distance between two Gaussian distributions. We sample 200 points from two 2D Gaussian distributions with differ- ent means and the same covariance matrix. We compute the interpolating measure at = 0.5 using both the analytical formula (5) and the approximation (T)). Figure (left panel) shows (left
Figure (left) Evolution of the interpolating mea- sure {())_ in blue (right) the estimated Wasser- stein distance between two Gaussian distributions i and v.
how the interpolating measure evolves across iterations. We also observe, in Figure 3 (right panel), that the enror on the true Wasserstein dis- tance for the approximated interpolating measure reaches 10-3, while for the exact interpolating

---

Published as a conference paper at ICLR 2024
measure, it drops to minimum of 10-4 before increasing. This discrepancy occurs the support size of the interpolating measure expands across iterations leading to numerical errors when com- puting the optimal transport plan between g(k) and ((k) E Hence, using the approximation Equation (10) is a more robust alternative to exact computation Equation (5).
We also examine computational complexity and approximation errors for both methods as we increase sample sizes in the distributions, as displayed in Figure 2 Key findings include: The approximated interpolating measure significantly improves computational efficiency, being at least 10 times faster with sample size exceeding 100, especially with smaller support sizes. It also achieves a similar relative approximation error as FedWaD using the exact interpolating measure and true non-federated Wasserstein distance. Importantly, demonstrates greater robustness with larger sample sizes compared to true Wasserstein distance for such small dimensional problem.
Wasserstein coreset and application to federated learning. summarizing data into fewer representative samples is routinely done to deal with large datasets. The notion of coreset has been relevant to extract such samples and admit several formulations (Phillips) 2016; MMnteaau et al., 2018). In this experiment, we show that Wasser- stein coresets (Claici et al.] 2018) can be computed in federated way via F FedWaD. Formally, given a dataset described by the distribution μ., the Wasserstein coreset aims at finding the empirical distribution that minimizes
In many ML applications.
Foflos We
2
7
5
2
3
5
0
7
/ 9
5
9
9
3
3
300 19 Fodtt
7
5
3
9
5
4
6
6
5
1
/
0
7
9 9
2
4
5
0
7
0
minzz..-..2 2x Wo Σ >i=10,; 1z): We solve this prob- lem in the following federated setting: we assume that either the samples drawn from μ are stored on an unique client or distributed across different clients, and the objec- tive is to learn the coreset samples {a[} the server. In our setting, we can compute the federated Wasserstein dis- tances between the current coreset and some subsamples of all active client datasets, then update the coreset given the aggregated gradients of these distances with respect to the coreset support. We sampled 20000 examples ran- domly from the MNIST dataset, and dispatched them at random on 100 clients. We compare the results we obtained with FedWaD with those obtained with exact non-federated Wasserstein distance The results are shown Figure We can note that when classes are almost equally spread across clients (with K 8 different classes per client), F FedWaD is able to capture the 10 modes of the dataset. However, as the diversity in classes between clients increases, FedWaD has more difficulty capture all the modes of the dataset. Nonetheless, observe that the exact Wasserstein distance is not able to recover those modes either. We can thus conjecture that this failure is likely due to the coreset approach itself, rather than to the approximated distance returned by FedWaD. We also note that the support size of the interpolating measure has less impact on the coreset. We believe this is very interesting result, as it shows that FedWaD can provide useful gradient to the problem even with poorer estimation of the distance.
Figure Examples of the 10 coreset we obtained, with for each panel (top-row) the exact Wasserstein and (hottow-row) FedWaD for the MNTST dataset. Differ- ent panels correspond to different num- ber of classes K cach client: (top) K 8, (middle) K 2, (bottom) sup- port of the interpolating measure varying from 10 to 100.
also
Federated coreset classification model Those federated coresets can also be used for classifi- cation tasks. As such, we have learned coresets for each client, and used all the coresets from all clients as the examples for one-nearest neighbor global classifier shared to all clients. Note that since a coreset computation is an unsupervised task, we have assigned to each element of a coreset the label of the closest element in the client dataset. For this task, we have used the KNIST dataset which has been autoencoded in order to reduce its dimensionality. Half of the training samples have been used for learning the autoencoder and the other half for the classification task. Those samples and the test samples of dataset have been distributed across clients while ensuring that each client has samples from only classes. We have then computed the accuracy of this federated classifier for varying number of clients and number of coresets and compared its perfor- mance to the ones of FedRep (Collins et al., 2021) and FedPer (Arivazhagan
al. 2019). Results are reported in Figure 5 We can see that our simple approach is highly competitive with these personalized FL approaches, and even outperforms them when the number of users becomes large.
7

---

Published as a conference paper at ICLR 2024
Doact CTBO
Facerabed OTAD 23 ap 2003 ago
Tazerbted OTSC 503 sp 20-3000 mgo
Pacernted 0730 22 7003 ago
315
1838
1533
1476
821
2444
2365
1832
793
2373
2327
1785
M00ST
314
2039
1552
1536
1220
1516
PART RADENNORNT
Clsh RACHIMMNNTT
1804
1522
1748
676
1381
2336
2474
1670
1954
2208
2341
1612
3885
1351
1847
680
1400
1500
Figure 6: Comparison of the matrix of distances between digits datasets computed by FedWaD and the true OTDD distance between the same datasets. (left) OTDD, (middle-left) FedWaD with 20 epochs and 1000 support points, (middle-right) FedWaD with 500 epochs and 1000 support points, (right) PedWaD with 20 epochs and 5000 support points
Geometric dataset distances via federated Wasserstein distance. Our goal is to improve on the seminal algo- rithm of Alvvrrez-Melis & Fusi 2020) that seeks at comput- ing distance between two datasets D and D' using optimal transport. We want to make it federated. This extension will pave the way better federated learning algorithms for transfer learning and domain adaptation or can simply be used for boosting federated learning algorithms, as we illustrate next. Alvarez-Melis & Fusi| (2020) considers a Wasserstein distance with a ground metric that mixes distances between features and tractable distance between caas-cooniiional distributions, For our extension, we will the same ground metric, but will compute the Wasserstein distance using FedWaD. Details are provided in Appendix D.5.
Accuracy (9%)
Fed
FedRep
coresets/client
10 coresetstclient
coresetstelient
Number of clients
Nearest neighbor classifier based the coresets learnt from each client for varying number of clients and number of coresets per clients We have compared to the performance of two per- sonalized FL algorithms.
Figure 5:
We replicated the experiments of Alvarez-MMMl & Fusi (2020) on the dataset selection for transfer learning: given a source dataset, the goal is to find target one which is the most similar to the source. We considered four real datasets, namely MNIST, KMNIST, USPS and F FashionMNIST and we have computed all the pairwise distance between 5000 randomly selected examples from each dataset using the original OTDD of Alvarez- & Fusi (2000) and our FedWaD approach. For FedWaD, we chose the support size of the interpolating measure to be 1000 and 5000 and the number of epochs to 20 and 500. Results, averaged over 5 random draw of the samples, are depicted in Figure 6 We can see that distance matrices produced by FedWaD are semantically similar to the ones for OTDD distance, which means that order relations pairs of datasets the USPS row). More importantly, running more epochs leads to slightly better approximation of the OTDD distance, but the exact order relations are already uncovered using only 20 epochs in FedWaD. Detailed ablation studies on these parameters are provided in Appendix D.6.
well-preserved for most pairwise distances (except only for two
Boosting FL methods One of the challenges in FL the heterogeneity of the data distribution among clients. This heterogeneity is usually due to shift in class-conditions distributions or to a label shift (some classes being absent on client). As such, we propose to investigate simple approach that allows address dataset heterogeneity (in terms of distributions) among clients, by leveraging on our ability to compute distance between datasets in a federated way.
Our proposal involves computing pairwise dataset distances between clients, clustering them based on their (di)-similarities using a spectral clustering algorithm (Von Luxburg 2007), and using this clustering knowledge to enhance existing federated learning algorithms. In approach, we run the FL algorithm for each of the K clusters of clients instead of all clients to avoid information exchange between clients with diverse datasets. For example, for FedAvg, this means learning
8

---

Published as a conference paper at ICLR 2024
Table I: MNIST/CIFARIO Average performances over 5 trials of three FL algorithms, FedAvg, FedRep and FedPer. For each algorithm we compare the vanilla performance with the ones obtained after clustering the clients using the FedWaD OTDD distance and three different setting of the spectral clustering algorithm (details in Appendix) and for support size of 10. The number of clients varies from 20 100. Bolded number indicate the best performing approach (and clustering parameters).
Strong
Clusteingg
Closteing
Vanilla
Affinity
Spane G. (3)
Sparse G. (5)
Vanilla
Affinity
Spurs G. (3)
Spane G. (5)
FodAvg
20
16.3
3.8
99.5 2:
99.5 : 0.0
91.5 10.3
15.1 6.6
71.3 1: 7,3
59.5 4: 3.0
57.0
40
19.1
9.0
99.2 2:
91.1 1: 6.5
94.5
42.5 = 10.5
70.8 = 13.5
60.0 -: 3.7
58.1 4
29.2
7.7
98.9 1:
95.9
98.4 0.8
52.6 3,9
64.4 1:
76.3 1:
61.9
FedRRp
20
31.1
8.1
99.1 1:
99.1 1 0.0
96.2 1:
75.6 9.3
87.5 1:
$1.4 1:
85.3 4:
40
10.4
98.9 1:
93.3
88.0
78.9
3,9
98.6
98.5
86.0 4.8
91.6
86.3
FedPor
20
99.5
99.3
2.4
93.0
93.8 +
40
Average Uplitt
26.4 27.5
24.4 26.5
24.4 +: 25.6
12.7 14.6
8.7 + 12.7
7:2 11.4
CIFRRIO
FedAog
20
42.5
15.9
27.6
FedRep
$1.0
79.4
75.2
75.4
77.9
91.4
77.8
70.0
FedPer
20
2.3
91.0
87.2
87.8 0.9
$8.7
1.7
92.3
89.8
2.0
90.1
10
0.8
87.2
$3.7
2.5
84.3
S8.1
0.7
98.8
86.0
$1.9
82.2
0.4
85.1
50.3 -
2.0
80.9
85.1
0.6
94.0
$2.0
83.0
Average Uplin
18.8
global model for each cluster of clients, resulting in K global models. For personalized models like FedRep (Collins et al. 2021), or FedPer (Arivazhagan et al., (2019.. we run the personalized algorithm clustered client, ensure information
on each cluster of clients. By running FL algorithms exchange only between similar clients and improves the overall performance of federated learning algorithms reducing the statistical dataset heterogeneity among clients.
We have run experiments on MNIST and CIFARIO in which client datasets hold clear cluster structure. We have also run experiments where there is no cluster structure in which clients are randomly assigned a pair of classes. In practice, we used the code of FedRep Collins et al. (2021) for the FedAvg, FedRep and FedPer and the spectral clustering method of scikit-learn (Pedregosa et al., 2011) (details are in Appendix D.7). Results are reported in Table (with details in Appendix D.7). can see that when there is clear clustering structure among the clients, FedWaD is able to recover it and always improve the performance of the original federated learning algorithms. Depending on the algorithm, the improvement can be highly significant. For instance, for FedRep, the performance without
can be improved by points for CIFAR10 10 and up to 29 for MNIST. Interestingly, clear clustering structure, FedWaD is able to almost always improve the performance of all federated learning al algorithms (except for some specific cases of FedPer). Again for FedRep, the performance uplift can reach 19 points for CIFAR 10 and 36 for MNIST. In terms of clustering, the "affinity" parameter of the spectral clustering algorithm seems to the most efficient and robust one.
5 CONCLUSION
In this paper, we presented a principled approach for computing the Wasserstein distance between two distributions in a federated manner. Our proposed algorithm, called FedWaD, leverages the geometric properties the Wasserstein distance and associated geodesics to estimate the distance while respecting the privacy of the samples stored on different devices. We established the conver- gence properties of F FedWaD and provided empirical evidence of its practical effectiveness through simulations various problems, including dataset distance and coreset computation. Our approach shows potential applications in the fields of machine learning and grivacy-preserte data analysis, where computing distances for distributed data is fundamental task.

---

Published as a conference paper at ICLR 2024
REFERENCES
Pankaj K Agarwal, Sariel Har-Peled, Kasturi R Varadarajan, et al. Geometric approximation via coresets. Combinatorial and computational geometry, 52(1):1-30, 2005.
Martial Aguch and Guillaume Carlier. Barycenters the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-924, 2011.
David Livare--Meli and Nicolo Fusi. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems, 33:21428 22439,,, 2020.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Feder- ated learning with personalization layers. arXiv preprint 1r2.0:19/2.00818, 2019.
Martin Bauer, Sarang Joshi, and Klas Modin. Diffeomorphic density matching by optimal information transport. SIAM Journal Imaging Sciences, 8(3):1718-177) 2015.
Sehastian Claici, Aude Genevay, and Justin Solomon. Wasserstein measure coresets. arXiv preprint arXiv:/805.074/2, 2018.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting Shared Represen- tations for Personalized Federated Learning. In International Conference on Machine Learning, pp. 2089-20999 2021.
Nicolas Courty, Remi Flamary, and Melanie Ducoffe. Learning wasserstein embeddings. In Interna- lional Conference on Learning Representations (ICLR), 2018
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalic T.H. Gayraud. Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22((8):1-8, 2021. URLpttp: / // wttp://..rrrgppps///.22... org/ wwppps/w22eeppppppp 200222/224455 html
Nicolas Fournier and Arnaud Guillin. On the rate of convergence in wasserstein distance of the empirical measure. Probability theory and related fields, 162(3-4):707-738 2015.
Ziv Goldfeld and Kristjan Greenewald. Gaussian smoothed optimal transport: Metric structure and statistical efficiency. In International Conference on Artificial Intelligence and Statistics, pp- 3327-3337. PMLR, 2020.
Remi Gribonval, Antoine Chatalic, Nicolas Keriven, Vincent Schellekens, Laurent Jacques, and Philip Schniter. Sketching data sets for large-scale learning: Keeping only what you need. IEEE Signal Processing Magazine, 38(5):12-36, 2021.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy, Manifold interpolating optimal-transpe flows for trajectory inference. Advances in Neural Information Processing Systems, 35:29705-29718. 2022.
Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, K. A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G.L. D'Oliveira, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konecny, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateck Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramer, Praneeth Vepakomma, Jianyu Wang. Li Xiong, Zheng Xu. Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and Open Problems in Federated Learning. Foundations and Trends in Machine Learning, 14(1-2):1-210 2021.
10

---

Published as a conference paper at ICLR 2024
Soheil Kolouri, Se Rim Park, Matthew Thorpe, Dejan Slepcev, and Gustavo K Rohde. Optimal mass transport: Signal processing and machine-learning applications. IEEE signal processing magazine, 34(4):43-59,22 2017.
Nam Lê Tien, Amaury Habrard, and Marc Sebban. Differentially private optimal transport: Applica- tion to domain adaptation. In IJCAI, pp. 2852-2858, 2019.
Zhaorong Liu, Leye Wang. and Kai Chen. Secure efficient federated knn for recommendation systems. In Advances in Natural Computation, Fuzzy Systems and Knowledge Discovery, pp. 1808-1819. Springer, 2021.
Jan Maas, Martin Rumpf, and Stefan Simon. Transport based image morphing with intensity modulation. In Scale Space and Variational Methods in Computer Vision: 6th International Conference, SSVM 2017, Kolding, Denmark, June 4-8, 2017, Proceedings, pp. 563-577. Springer, 2017.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-effie learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1282. PMLR, 2017.
Alexander Muntcanu, Chris Schwiegelshohn, Christian Sohler, and David Woodruff. On coresets for logistic regression. Advances Neural Information Processing Systems, 31, 2018.
Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-lcarn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830 2011.
Gabriel Peyré, Marco Cuturi, et Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 116--6):3556077 2019.
Jeff M Phillips. Coresets and sketches. arXiv preprint arXiv:1601.00617., 2016.
Phillipp Schoppmann, Adrià Gascón, and Borja Balle. Private nearest neighbors classification in federated databases. IACR Cryptol. ePrint Arch., 2018:289, 2018.
Ulrike Von Luxburg. A tutorial on spectral clustering. Statistics and computing, 17:395-416, 2007.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H. Brendan McMahan, Blaise y Arcas, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, Suhas Diggavi, Hubert Eichner, Advait Gadhikar, Zachary Garrett, Antonious M. Girgis, Filip Hanzely, Andrew Hard, Chaoyang He, Samuel Horvath, Zhouyuan Huo, Alex Ingerman, Martin Jaggi. Tara Javidi, Peter Kairouz, Satyen Kale, Sai Praneeth Karimireddy. Jakub Konecny, Sanmi Koyejo, Tian Li, Luyang Liu, Mehryar Mohri, Hang Qi, Sashank J. Reddi, Peter Richtarik, Karan Singhal, Virginia Smith, Mahdi Soltanolkotabi, Weikang Song, Ananda Theertha Suresh, Sebastian U. Stich, Ameet Talwalkar, Hongyi Wang, Blake Woodworth, Shanshan Wu, Felix X. Yu, Honglin Yuan, Manzil Zaheer, Mi Zhang, Tong Zhang, Chunxiang Zheng, Chen Zhu, and Wennan Zhu. A Field 2021.
Guide to Federated Optimization. arXiv preprint
Kaining Zhang, Yongxin Tong, Yexuan Shi, Yuxiang Zeng, Yi Xu, Lei Chen, Zimu Zhou, Ke Xu, Weifeng Lv, and Zhiming Zheng. Approximate k-nearest neighbor query over spatial data federa- tion. In Database Systems for Advanced Applications: 28th International Conference, DASTAA 2023, Tianjin, China, April 17-20, 2023, Proceedings, Part 1, pp. 351-368. Springer, 2023.
11

---

Published as a conference paper at ICLR 2024
Appendices
A PROPERTY OF THE APPROXIMATING INTERPOLATING MEASURE
Theorem 1. Assume that μ and g(k) are two discrete distributions with the same number of samples and uniform weights. Then for any t, the approximating interpolating measure given by equation Equation (10) is equal to the exact one Equation 5
Proof. Remind that the approximating interpolating measure is defined as
εe
(11)
whereas the exact interpolating measure is defined as
TL,
P
P
(12)
3,
where P is the optimal transportation plan between and f'. 2; and x/j are the samples from these distributions and is the matrix of samples from 5'. Because μ and (kk) have the same number of samples and uniform weights, P* is weighted (by 1/n) permutation matrix Peyré al. (2019).
Let us denote by a the permutation associated to P Then, for the approximation, we have
8 (P-)))
8(1 =)=t))=+t=
6(1-t)xc+tx
n
Ft
where the last equality comes from the fact that for each row i, Pr,j is non-zero only for (i) column and P (s) 1/n.
B PROOF OF THEOREM 2
Theorem 2. Let μ and be two measures in Px(X). For k. N, let (k) interpolating measures computed at iteration le as defined in Algorithm
cl and g(t) be
[ Define
(k) = Wp(,,e) + Wo(e(h) ((0)) + (,(g((),,5))
Wp(e(+),
Then, the sequence (A())) is non-increasing and converges W,(,nv)
Proof. First, remind that <(x)) and ( the interpolating measures between l and g(k-1) and between g(k-1) and e respectively, as defined in Algorithm I Likewise, <( <<<-1) and ((++1) are interpolating measures between μ and (()) and between (R) and respectively. Hence, we have
Wp(u,g(k+1)) + (T( ((()) < W, (()) Wp(c ((()))
and (((x) W, (E((t))) - WWpl ((4)) + Wr(E!t) g(()))))((())))
These two inequalities lead to,
Wyln c(k+1)) ((.((((+)),(()))) Wo, (,, + ((
12

---

Published as a conference paper at ICLR 2024
I Ep W (( + Wolu, = + Wr Wp(EL ((())
Besides, since ( 1) is an interpolating measure between E/ (((=1) and &o el we have
Wp(esh g( + Wp ,, - Wp(st ; + W elett)
and
1 S (()
Hence, the sequence (A(k)) is non-increasing. Additionally, by the triangle inequality, we have for any & E N,
Wy(,, < (A))
We conclude by using the monotone convergence theorem: since (A(k)) is non-increasing and bounded sequence below, then it converges to its infimum.
We now justify why the limit of (A(#)) is Wp(,,, v). At convergence, we have reached a stationary point in the (A (k)),
lim A Wy(H, +
2((0) c(o))) + Wipel ((0))) + Wo(g(0o),
Wi(e!
and there are an infinite number of triplets (x)) c(oc) . g(oo)) that allow to reach this value A by the nature of the algorithm. By definition, c(00) and glo)) are interpolating measures between j and ((0o) and between g(xe) and respectively. At convergence, (5/0=, _,(0), g(oe)) are fixed points of the algorithm, and we show here that g(00) is an interpolating measure of μ and in addition to be
an interpolating measure of <((²) and g!>)) For any (0x)) </x²) can be chosen as any interpolating
measure between and g(xo), The same reasoning holds for g!x²) and v. Then since g(ne) is an interpolating measure of c(20) and g(<x) and μr. and v possible choices of interpolating measures, yields that g( is indeed an interpolating measure of μ and v. Hence, we have
lim (k) Wp(u, <(00)) + Wp(e , c((o)) + Wp(5 ,00))) + Wpp(( k-++00
= (p(H, + Wp(e(oo), .v) Wp(u,e)
where the first equality results from the fact that gl and g000) are interpolating measures between ja and g(oo) and between g(>x) and respectively and the second equality is obtained from the fact that
g(s>) is also an interpolating measure between μ and as belonging to the geodesic between μ and
v.
C CONVERGENCE RATE OF THE ALGORITHH FOR GAUSSIAN DISTRIBUTIONS WITH SAME COVARIANCE
In this section, we show that when fa and are Gaussians after iteration, we infer W (p,, na) and the sequence of iterates (c()) obtained for = 0.5 converges to the &* the interpolating measure between j and v for / = 0.5
Theorem 3. Assume that μd, v and c(0) are three Gaussian distributions with the same covariance matrix Σ ie μ ~ Nmmp. Σ), v N(m,, Σ) and g()) ~ ~((mo E). Further assume that not in the trivial case where m.. m,, and meis are aligned. Applying our algorithm Algorithm with 0.5 and the squared Euclidean cost, we have the following properties:
1. all interpolating measures E/ c(k) g(k) are isotropic Gaussian distributions with the same covariance matrix
2. for any & &21 1, W2(H,v) = -Imm, m2ll2 = 2||m m m.
3. Wa(g()), -)22(E(-1)E))
13

---

Published as a conference paper at ICLR 2024
4. (2(1,,((0)) + Wa(5(),,) W2(m.2) 2 M2(5(0)
Proof. The first point comes from the fact that Wasserstein barycenter of Gaussians are Gaussians Agueh Carlier (20TT); Peyré et al. (2019). For isotropic Gaussians with same covariance, the covaniance matrice of the barycenter remains unchanged while the mean is the barycenter mean. So, our case, the interpolating measure with 0.5 i.e the uniform barycenter of two measures, say j
and g(k-1), is g ~ N ((m,(.) 2), where g = 1 (m, + mę(+-1)). ). The consequence of this first point of the theorem is that since are going deal with same covariance Gaussian distributions, then the Wasserstein distance between any pair of measures involved in our algorithm only depends on the Euclidean distance of their means and we will interchangeably the Euclidean distance and the Wasserstein distance.
The second point is proven by using geometrical arguments in the plane (P) in which the three points, for & 2 1, m2, my, Mg(k-1) lie (note that based on our assumption, this plane always exists). By definition of elAi and ((k) and given the above point, we have
1 mm + ma 12) and m
mm + Me 1))
(M) 2
2
By using the intercept theorem, since = 1.4 in the plane (P), the segment [m m000 m is parallel to the segment [m, my] and we have
1 m (M) m 1)ll2
m g 1yll2
m m ll2
2 m me(k-1) l2
m, 1g(k-1)
m, m
which gives us the second point.
For the third point, we are going to consider geometrical arguments similar as above. How- ever, we are going to first show that for given k, the mid point, denoted as g(A),, between g(k-1) and E* is also (R) defined by our algorithm.
By defimition, 5* is the mid point interpolating measure between μ and v, whose mean is mm + m.). Since (A) and <(())) are respectively the mid point measure between c(k-1) and 5* and and g(k-1), can apply the intercept theorem in the appropriate plane and get
- 1
<(<²) ( 55
W2
Using a similar reasoning using v. we get
W2((()) <())) 1 Wa(k,
Summing these two equations, obtain
1 W2(m)))) + W2(v,))) = WW2(H, v) = W2(e(()), ((()))
W2( -(()) + (2((()) ((+)) 2
where the second equality comes from the fact that 8* is an interpolant measure of ju and v, while the last equality comes from the second point of the theorem.
Hence, since we have (z(E((), g((())
((0))
(E(K)) ( W2(C!() ,<(+)), it also mean than
1 (()) E arg min W2(5((1) .E) W2(5, (())
2
and ((k) is also the midpoint interpolating measure between (*) and e!
Then, applying the intercept theorem with 1( (((), E* ,t, and E( we obtain the desired result
1
W2(e(( 11) 1),g)))))=))) = (2(g(()e5)
2
14

---

Published as a conference paper at ICLR 2024
{(4)
(A) Sj
(())
5*
Figure 7: Illustration of the geometrical interpretation of the algorithm and its convergence proof for Gaussian distributions with same covariance, based on the intercept theorem.
Finally, given the above, it simple to show the convergence rate of the algorithm using simple triangle inequalities.
(2(u,(()) + (2(g((), WV((,,)) s Wo(ugEt) + Wa(E", ((0) W2( + W2(5", ,)) M2(u,
2W2(5(k), 55)
1
2 Wo
D ADDITIONAL EXPERIMENTS
D.1 TOY ANALYSIS THE IMPACT OF APPROXIMATING THE INTERPOLATING MEASURE
We propose to analyze in this section the benefits and disadvantages of approximating the interpolating measure instead of using the exact one as given in Equation (5). For this purpose, we compare the running time and the accuracy of the exact Wasserstein distance, our exact FedWaD, and our approximate FedWaD for estimating the Wasserstein distance between two Gaussians distributions. The Gaussians are different means but same covariances so that the true Wasserstein distance is known and equal to the Euclidean distance between the means. have considered two different settings (d and d 50) of Gaussians dimensionality. For the first case (d = 2). detail the results presented in the main paper. Note that when the dimensionality of the Gaussians are set to 50, do not expect the Wasserstein distance nor F FedWaD to provide good estimation of the closed form distance between these distributions, due to the curse of dimensionality the Wasserstein distance (Fournier & Guillin 2015)
As default parameter for our approximate FedWaD, we considered 20 iterations and a support of size 10, then varied the number of samples from 10 to 10000. We have run experiments in different settings
we analyzed the impact of sample ratio between the two distributions, as this may impact the support size of the approximating interpolating measure accross FedWaD iterations.
we made varying the support size of the approximating interpolating measure at fixed sample ratio.
Results have been averaged over 10 runs.
Analyzing the impact of sample ratio Given the setting with uniform weights, when the sample ratio 1, the optimal plan is theoretically a scaled permutation matrix. Hence, the support size of the exact interpolating measure is expected, in theory, be fixed and equal to N. When the ratio of samples different to 1, the support size of the exact interpolating may increase at each iteration of the algorithm and leads larger running time. Figure left panel- shows the running time of
15

---

Published as a conference paper at ICLR 2024
Support size 10
Support size 10
Relative Approximation Error
WD (1:3)
12
30t
WD (1:1)
100
30t
FedWad-e (1:3)
Running time
WD (1:3)
FedWad-a (1:1)
WD (1:1)
FedWad-a (1:3)
10-1
FedWad-e (1:3)
10-
FedWad-a (1:1)
10-
FedWad-e (1:1)
20-1
FedWad-a (1:3)
FedWad-a (1:1)
101
10
10a
10
101
19
19'
10
Number of samples N
Number of samples N
Support size 10
Support size 10
101
10-
Relative Approximation Error
10a
Running time (s)
107
WD (1:3)
WD (1:3)
109
WD WD (1:1)
10t
WD (1:1)
10-3
FedWad-e (1:3)
FedWad-e (1:3)
FedWad-e (1:1)
FedWad-e (1:1)
FedWad-a (1:3)
FedWad-a (1:3)
10-1
FedWad-a (1:1)
FedWad-a (1:1)
10
10:4
Number of samples N
Number of samples N
Figure 8: For different sample ratios, (1:3) or (1:1), in the two distributions report the performance of the different models. For our approximated FedWaD, we have set the support size to 10. (top) d 2 (bottom) d 50. (left) running time. (right) relative error.
all compared methods as well as their relative error right panel compared to the true Wasserstein distance. We note that for 2d Gaussians, both the Wasserstein distance and our approximated FedWaD with support size of 10 the running time is increasing with natural computational overhead for the 1:1 sample ratio (as we have more samples). For the exact FedWaD, the behavior is different. the running time for the 1:3 sample ratio is larger than the 1:1. This due the optimal transportation plan P* not being exactly a scaled permutation matrix. As result, the support size of the interpolating measure increases with the number of samples, leading computational overhead for the method. For 50d Gaussians, the differences in rumning time between the different sample ratio are negligible.
In the case of 2d Gaussians (top row), For the relative error, for N < 1000, we that all methods achieve similar errors. Numerical errors start to appear for exact FedWaD and the Wasserstein distance for respectively N 1000 and N 2 5000 depending on the sample ratio. Interestingly, the approximated FedWaD is robust to large number of samples and achieves similar errors for small number of samples. For higher dimensions (bottom row), all the methods are not able to provide accurate estimation of the Wasserstein distance and with the worst relative eπror for the approximated FedWaD with a support size of 10. Nonetheless, we want to emphasize that despite this lack of accuracy, the approximated FedWaD be useful in high-dimension problems have shown for the other experiments.
Analyzing the support size of approximated interpolating measure Figure9sshvee ssows the running 3 and when the
time and the relative error of the different methods for sample ratio support sizes of the approximating interpolating measure are 2. or 100. We clearly remark the computational cost of a larger support size with a benefit terms of approximation error appearing mostly when N 1000 and for small dimension problems (top row). For higher dimension problems (bottom row), we see again the benefit on running time the approximated approach, yet with larger approximation error.
16

---

Published as a conference paper at ICLR 2024
Sample ratio of 1:3
Sample ratio of 1:3
Relative Approximation Error
WD
WD
12
30t
FedWad-e
FedWad-e
10
FedWad-a-100
30t
FedWad-a 100
Running time
FedWad-a-10
FedWad-a 10
FedWad-a-2
FedWad-a
10-1
10-
10
20-1
101
10
10a
10
101
18
19'
10
Number of samples N
Number of samples N
Sample ratio of 1:3
Sample ratio of 1:3
101
10.
Relative Approximation Error
WD
WD
10.
FedWad-e
FedWad-e
Running time (s)
FedWad-a-100
FedWad-a 100
FedWad-a-10
FedWad-a 10
109
FedWad-a-2
10t
FedWad-a
10-3
10-1
10
10
Number of samples N
Number of samples N
Figure 9: For increasing number of samples, report (top) d (bottom) d = 50. (left) Running time of the Wasserstein distance, our exact FedWaD and our approximate FedWaD. (right) the relative enror of the different models: the computed Wasserstein distance, exact FedWaD and the approximated FedWad with support size of 10 and 100. The first distribution has number of samples N and the second ones N/3.
7.75
((0)
FedWad
True Wasserstein distance
2-
7.25
7.00
6.50
5 ((00)
6.25
14
5.75
Figure 10: We illustrate here how algorithm behaves when the distributions continuous. (left) we plot the distributions and as well as the interpolating measure (A) (right) we plot the evolution of the Wasserstein distance between y and u as computed by FedWad.
D.2 TOY ANALYSIS : CONTINUOUS DISTRIBUTIONS
Our algorithm can applied to continuous distributions as long as possible to compute an element of the geodesic between the two distributions. For multivariate Gaussian distributions, the transport map exists and elements of the geodesics are well-defined. However, fooed--omm of the mean and the covariance matrix of interpolating measures available except when the covariance
17

---

Published as a conference paper at ICLR 2024
Error ((Approx WD True WD
Gaussians
107
Moons
10f
20-2
10-1
Iterations
Figure 11: Example of convergence FedWad when computing the distance between two Gaussian distributions and between two moon-shaped distributions.
matrices between and are jointly diagonalizable. Hence, as example, we have applied our algorithm for two continuous Gaussians distributions μ ~ N (m,, ,,,) and v ~ N (m,,, En)) where En and Σn are diagonal matrices. c(0) is also defined as diagonal Gaussian distribution. Mean and covariance of an interpolating measures for t 0.5 are computed as follows (say between μ and v/):
1
2 m + m,
E = 2
1 2\ + x1/2)²
Figure ToreeToOshen shows an example of the evolution of the Wasserstein distance between μ and v as computed by our algorithm well as the interpolating measure (( for different values of k. We the Wasserstein distance converges to the true Wasserstein distance between jr and v in about 10 iterations confirming the linear convergence rate.
that
D.3 COMPARING CONVERGENCE RATE
In order to gain an insight about the convergence rate of our algorithm for non non- Gaussian distributions, we have compared how fast FedWaD converges the true Wasserstein distance when comparing two 2D Gaussians and when comparing two 2D moon-shaped distributions. We have considered 200 samples per distribution and computed the exact FedWal using 10 iterations. Figure TT shows the evolution of the Wasserstein distance between the two distributions a function of the number of iterations. We can see that the convergence rate for the two moon-shaped distributions is slower than the ones of the Gaussians, which is about 1.5 order of magnitude, and it tends to decrease as iterations increase.
D.4 DETAILS ON CORESET AND ADDITONAL RESULTS
Experimental setting We sampled 20000 examples randomly from the MNIST dataset, and dis- patched them at random on 100 clients but such that only a subset K of the 10 classes is present on each client. We learn 10 coresets over 1000 epochs and each epoch, random clients are available and can be used for computing F FedWaD. For FedWaD, the support size of the interpolating measure has been set either 10 or 100 and the number of iteration in FedWaD to 20.
assume that only 10
We have reproduced in here the same MNIST experiment (which results are reproduced in Figure 12 on coreset for the F Fashi SN dataset, and we can notice, in Figure 13 that we obtain similar results as for the MNIST dataset. When the number of shared classes K large enough, the coreset is not able to capture the different modes in the dataset. And again, we remark that the support size of the approximate interpolating measure has few impacts on the result. For both datasets, the loss landscape of the coreset learning reveals that our FedWaD-based approaches yield to worse minimum than the exact Wasserstein distance, which is mostly due to the interpolating measure approximation. Figure 14 plots the performance of nearest neighbor classifier based on the coresets
18

---

Published as a conference paper at ICLR 2024
FedVadWD Fedivdd WD
3
7
5
6
4
9
6
O
19P
Feddted 10 Feciood 160
4
3
6
/
1
9
5
3
9
4
8
3
6
3
9
9
lose
/
5
3
9
8
9
5
7
3
100 10
6
2
3
5
6
0
7 9 9
2
4
3
5
0
7
6
0
Reeations
Figure 12: Examples of the 10 coreset obtained with for each panel (lop-row) the exact Wasserstein and (botlow-row) FedWaD for the MNIST dataset. Different panels correspond to different number of classes K on each client: (top) K 8, (middle) K 2, (bottom) support of the interpolating measure varying from 10 to 100, As class diversity on each client increases, the coreset is less effective capturing the 10 modes of the dataset
291
FecdWed 10c
LO05
200
Itorations
Figure 13: Examples of the 10 coreset obtained with for each panel (top-row) the exact Wasserstein and (bottow-row), our FedWaD for the FashionMNIST dataset. Different panels correspond to different number of classes on each elient: (top) K 8, (middle) K 2, (bottom) support of the interpolating measure for K == 8.
learnt from each client for varying number clients. Results show that coreset-based approaches competitive, especially for high number of clients, with personalized FL algorithms, which are known to be the best performing FL algorithms practice.
D.5 DETAILS ON FEDERATED OTDD EXPERIMENTS
Geometric dataset distances via federated Wasserstein distance. Transfer learning and domain adaptation are important ML paradigms, which aim transferring knowledge across similar domains. The main underlying concept in these approaches is the notion of distance or similarity between datasets. Transferring knowledge between comparable domains typically simpler than between dis- tant ones. In certain applications, it is relevant to find datasets from which one can transfer knowledge from without disclosing the target dataset. This may be case, for instance, in applications with low-resoures clients storing sensitive data. In this case, the practitioner may want to find dataset similar enough to the client's 's dataset, in order to transfer knowledge from it. In practice, a server would train classifier a dataset that is similar to the client dataset, and the client would then use this classifier to perform inference on its own data.
In that context, goal is propose distance between datasets that can be computed in federated way based on FedwaD. We leverage the distance proposed in Alvarez- Melis & Fusi] (2020), which is based on the Wasserstein distance between two labeled datasets D and D'. The ground metric is defined by.
do((r, y), (x/,1/)) & (d(a, x') + W2 2/2
(13)
19

---

Published as a conference paper at ICLR 2024
Accuracy (%)
60-
FedPer
FedRep
coresets/client
10 coresets/client
20 coresets/client
250
300
730
15t0
1250
1550
1750
2000
Number of clients
Figure 14: FashionMNIST performance of nearest neighbor classifier based on the coresets learnt from each client for varying number of clients and number of coresets per clients We have compared to the performance of two personalized FL algorithms.
where d is a distance between two features a and x', and a, is the class-conditional distribution of given y. In order reduce computational complexity, Alvarez- & Fusi (2020) assume the cladsicon are Gaussian, so that W2 boils down the 2- 2-Bures Wasserstein distance, which is available in closed form:
W (a ogy) = || - 1 + E En 11:
(14)
where my and Mg denote the mean and covariance ag..
FedWaD needs vectorial representations of the data compute intermediate measures. The Bures- Wasserstein distance allows us to conveniently represent a, as the concatenation the mean my and vectorized covariance Σy. Hence, we compute the distance between two datasets D and D' by augmenting each example from those datasets with the corresponding class-conditional mean and vectorized covariance, and using the & norm the ground metric in the Wasserstein distance. One can eventually reduce the dimension the augmented representation by considering only the diagonal of the covariance matrix.
D.6 FEDERATED OTDD ANALYSIS
To evaluate our procedure, we replicated the experiments Alvarez-Mllls & Fusi (0220) on the dataset selection for transfer learning: given a source dataset, the goal is to find target one which is most similar to the source. and Fashi oNMNIST. We first analyze the impact of two hyperparameters, the number of epochs and the number support points in the interpolating measure, on the distance computation between 5000 samples from MNT ST and KMNT ST, Figure 15 shows the evolution of the distance between MNIST and KMN INNSS ST as well as the running time for varying values of hyperparameters. The number of epochs has very small impact on the distance and using 10 epochs suffices to get reasonably accurate approximation of the distance. On the other hand, the number of support point seems more critical, and need at least 5000 support points obtain a very accurate approximation, although we have nice linear convergence of the distance with respect to support size.
considered four real datasets, namely MNIST, KMNIST, USPS
We also analyzed the impact of the dataset size on the distance computation and running time: Figure 1olshows the evolution of the distance and the running time with respect to the the sample size in the two distributions. We note that the order relation is preserved between the two distances for all possible range of sample size. Another interesting observation is that as long as the sample size is smaller than the support size of the interpolating measure, FedWaD provides an accurate estimation of the distance. When the sample size is larger then the distance is overestimated. This is due to less accurate estimation of an exact interpolating measure (which is supported on 2n + 1 points). Regarding computational efficiency, measure, the running time increases at the same rate as the sample size, whereas for larger support size, the running time increases 10-fold for an 100-fold increase in sample size.
observe that for small support size of the interpolating
20

---

Published as a conference paper at ICLR 2024
2500
Time Exatt Time
Distance Exact Distarice
0010
Distance
250-1
CTOO Distance
2750
Time (5)
2250
32
2000
1750
1750
1540
Epachs FedWWad
Number support
Figure 15: FedWaD and OTDD distances on MNIST-KMNIST and its running time against (left) the number of epochs and (right) the number of support points in the interpolating measure. For each plot, the left and right y-axis report the distance and the running time respectively.
3500
207
Exact MNIST-KMNIST
5009
Exact MNIST-USPS
307
0530
Fed MNIS T-KNNIST [5|=1000
Fed MNIST-USPS 15/=1000
300
6000
15|-10
20:00
Fed MNIIT-USEE T-USPS [S|=10
3.0
Exact MNIST-IMMIR
200
Fed MNISTKMMIST
2530
10"
Exact MNIST-USPS 1000
Fed MNIS MNIST-USPS 19/-1000
2009
Fed MNISTIMNIST
Fed MNNIT-USEE 19|=10
15a
101
10
15"
10"
Sample Size
Sample Size
Figure 16: (left) Distance and (right) running time against the dataset size for the MNIST-KMNIST an MNIST-USPS distances, for varying number of support points |S]
D.7 BOOSTING FL. METHODS
We provide here more detailed results about our experiments on boosting FL methods. Figure 17 shows the distance matrices obtained for MNIST and CIFARIO when the number of clients is 20 for different structures on the clients datasets. We can clearly see the cluster structure on the MNIST dataset when it exists, but when there is no structure, the distance matrix is more uniform yet show some variations For CIFAR10, 10, no clear structure is visible on the distance matrix as the dataset is more complex. Nonetheless, our experiments on boosting FL methods show that even in this case, clustering lients can help improve the performances of federated learning algorithms.
Those distance matrices are the one we use as the input of the spectral clustering algorithm. We used the spectral clustering algorithm of scikit-learn (Pedregosa et al. 2011) with the following setting::
we denoted as "affinity", the setting in which the distance matrix, after rescaling, used as affinity matrix, where larger values indicate greater similarity between instances. (see affinity parameter set 'precomputed' scikit-learn)
we denote as Sparse G. (3) and Sparse (5) the setting in which the distance matrix is interpreted as sparse graph of distances, and construct binary affinity matrix from the (3 or 5) nearest neighbors of each instance. matrix is computed
Details on the cluster structure We have built this cluster structure on the client datasets by assign- ing each client one pair of classes among the following 5 ones [(0, 1), (2, (4, 5), (6, 7), (8, 9)]. When the number of clients in equal to 10, each cluster is composed of clients. For a larger number of clients, each cluster is of random size with minimum of 2 clients.
21

---

Published as a conference paper at ICLR 2024
3
5
-6
8
9
-11
12
14
-15
16
17
-18
19
1
.
4
5
100
11
12
133
-14
15
16
-
18
-19
0
3
8
8
11
11
12
12
14
14
15
16
17
17
3
8
110
11
12
13
14
15
16
17
18
19
3
6
8
9
110
12
13
14
1.5
16
17
18
19
2
4
4
5
6
7
9
10
10
12
13
13-
15
15
16
18
18
19
Figure 17: (left) MNIST and (right) CIFARIO distance matrices for 20 clients computed using our Federated OTDD. On the top row, we have imposed cluster structure on client datasets while on the bottom row, there is no specific sructure. We can note that this structure clearly visible on the MNIST dataset but less on CIFARIO. Eventhough, clustering clients will help improve federated learning algorithm performances.
Practical algorithmic details In practice, we used the code of FedRep Collins et al. (2021) for the FedAvg, FedRep and FedPer and the spectral clustering method of scikit-learn Pedregosa al.) (2011). The federated OT distance dataset has been computed on the original data space while for CIFAR1O, we have worked on the 784-dimensional code obtained from an (untrained) randomly initialized autoencoder. We have also considered the case where the there is no specific clustering structure on clients as they randomly select pair of classes among the 10 ones.
Extra results Performance results federated learning reported below for different settings.
Table and Table 3 show the results for MNIST respectively with and without client structure. Table and Table report similar results for CIFARIO.
22

---

Published as a conference paper at ICLR 2024
Clustered
Affinity
Sparse G. (3)
Sparse G. (S)
Vanilla
10
100
10
100
10
100
FedAvg
10
19.6 1 0.9
99.6 1 0.0
99.6 19.6.100 0.0
90.5 + 8.7
91.8 1 9.6
84.5 1 8.0
85.2 5.7
20
26.3 + 3.8
99.5 0.0
99.5 0.0
99.5 0.0
99.5 = 0.0
91.5 10.3
96.5 6.0
40
39.1 + 9.0
99.2 + 0.1
99.2 + 0.1
91.1 + 6.5
99.2 + 0.1
9.4
99.2 = 0.1
+
98.9
98.9 0.0
95.9 4.6
=
0.8
98.9 0.0
FedRep
71.6 10.5
+
+ 0.1
94.3 7.7
99.0 +
+
90.5
20
81.1 + 8.1
99.1 + 0.0
99.1 + 0.1
99.1 + 0.0
99.1 0.0
1.3
99.0 + 0.2
40
88.8
10.4
0.1
98.9 0.0
93.
99.0 0.1
9.77
99.0 0.1
100
93.0 + 3,9
98.6 0.1
98.6 0.1
0.14
98.4 0.1
98.5
98.5 0.1
FedPer
10
86.7
4.3
99.6 1 0.0
99.6 1 0.0
9.1
99.6 - 0.1
9.
98.9 1.0
20
94.3 4.3
99.5 + 0.0
0.0
0.0
9 0.0
99.3
99.5 0.0
40
94.7 7.6
99.2 0.1
99.2 99.2 0.1
99.1 0.2
99.2 + 0.1
97.9
99.2 0.1
100
98.1 + 0.1
98.9 + 0.0
98.9 + 0.0
98.8 + 0.2
98.8 + 0.1
98.9 + 0.0
98.9 0.0
Average Uplift
29.8 + 28.4
29.8 + 28.4
27.2 :: 26.6
29.0 27.2
26.7 +: 25.2
27.6 + 26.3
Table 2: MNIST Average performances over 5 trials of three FL algorithms: FedAvg, FedRep and FedPer. For each algorithm we compare the vanilla performance with the ones obtained after clustering the clients using the FedOTDD distance, using three different parameters of the spectral clustering algorithm and for support size of 10 and 100. The number clients varies from to 100. For this table, datasets from clients do have clear cluster structure
Clustered
Affinity
Sparse G. (3)
Sparse G. (S) 100
Vanilla
10
100
10
100
10
FedAvg
10
20.2 + 0.6
81.0 + 4.2
81.3 + 4.5
78.0 + 6.0
77.7 + 6.6
71.5 5.1
72.0 + 6.0
20
25.1 + 6.6
71.3 7.3
72.0 4.3
59.5 3.0
5.7
57.0 + 4.4
60.5 + 2.3
40
42.5 + 10.5
70.8 + 13.5
70.3 + 13.3
60.0 + 3.7
59.5 + 10.6
58.1 + 6,3
56.9 + 6.1
100
52.6 +
9.6
60.4 = 11.3
76.3 5.4
08.26.11
67.9
65.4 =.7
FedRep
54.3 11.2
90.1 6.7
90.1 7.5
92.1 4.2
1
94.0 3.1
20
75.6 9,3
87.5 1 4.5
2.6
81.4 1 8.6
85.1 1 6.3
85.3 7.3
87.1 5.5
40
78.0 +
1.0+44. 4.3
4.84
78.9 + 7.9
74.9 8.7
5.6
79.6 - 5.7
100
86.0 + 4.8
91.6 + 3.1
9.77
89.1 5.0
84.5 + 2.9
86.3 4.3
84.9 + 3.6
FedPer
10
82.0
10.1
98.4 1.4
96.5
96.4 3.5
96.5 3.6
98.5 1.4
98.3 1.3
20
90.5 2.4
92.7 1.715.
0.5
93.0 - 4.3
96.2 - 3.0
93.8
94.5 - 2.5
40
92.3 + 1.3
90.2 4.7
91.0
87.7 4.1.
87.0 3.7
89.2 09.2::2.. 2.3
87.5 5.4
100
96.6 1 0.9
96.6 1 1.6
96.4 2.0
92.1 4 3.3
93.0 +: 2.3
90.2 1 4.9
86.9 1.7
Average Uplift
18.9 1 18.9
18.3 1. 19.2
15.7 1 18.6
14.8 18.6
14.1 1: 17.1
14.3 1 18.1
Table 3: MNIST Average performances over 5 trials of three FL algorithms: FedAvg, FedRep and FedPer. For each algorithm we compare the vanilla performance with the ones obtained after clustering the clients using the FedOTDD distance, using three different parameters of the spectral clustering algorithm and for support size of 10 and 100. The number clients varies from to 100. For this table, datasets from clients do not have a clear cluster structure
23

---

Published as a conference paper at ICLR 2024
Clustered
Affinity
Sparse G. (3)
Sparse G. (5)
Vanilla
10
100
10
100
10
100
FedAvg
10
17.6 + 1.1
79.1 + 6.3
78.6 1 6.0
61.6 1 2.6
69.5 + 5.1
72.2 1 9.4
72.3 1 6.0
20
22.0 2.6
75.1 6.2
66.9 9.1
42.6 + 4.5
52.4 + 17.0
52.2 - 8.8
56.2 13.6
40
26.1 7.1
65.9 = 7.1
70.1 + 5.7
36.7 + 18.3
46.2 + 15.7
48.8 + 8.3
49.9 + 12.1
26.4 4.3
68.0 5.1
=
37.4 11.4
44.9 + 13.0
39.8 8.0
43.1 10.4
Fedrep
82.4 2:44222
91.1 1.2
90.7 + 1.2
90.3
89.7 2.3
90.0 + 1.1
20
81.8 + 1.8
88.1 + 2.0
85.9 1.4
84.4 0.5
86.0 2.1
85.3
86.8 + 1.4
40
80.3 0.8
83.7 = 2.0
86.2 0.9
81.0 2.1
82.1 8.76 2.5
100
75.0 0.9
79.4 = 2.3
78.5 = 1.7
75.2
76.3 1.6
75.4 1.5
76.9 1.1
FedPer
10
82.1
2.3
93.2 + 1.1
93.0 + 0.8
91.7 0.5
93.0 0.8
92.3 2.0
92.7 1 1.0
20
85.4 2.3
91.0 + 1.9
89.1 + 1.8
87.2 0.5
88.7 + 2.5
87.8
89.5 1 1,9
40
85.9
+ 0.8
87.2 2.2
89.7 + 1.4
82.7
85.4 2.7
84.3
84.9 1.6
100
82.2 + 0.4
85.1 + 1.8
83.4 + 2.7
80.3 1.3::22.. 2.0
81.3 1.8
80.9 + 1.7
82.5
Average Uplift
20.0 + 21.3
19.4 + 20.8
8.6 +: 12.5
12.4 + 15.1
11.9 + 16.0
13.3 + 16.1
Table 4: CIFAR10 Average performances over 5 trials of three FL algorithms: FedAvg, FedRep and FedPer. For each algorithm we compare the vanilla performance with the ones obtained after clustering the clients using the FedOTDD distance, using three different parameters of the spectral clustering algorithm and for support size of 10 and 100. The number clients varies from to 100. For this table, datasets from clients do have cluster structure
Clustered
Affinity
Sparse G. (3)
Sparse G. (5)
Vanilla
10
100
10
100
10
100
FedAvg
10
18.1 0.7
71.3 + 7.3
71.0 + 3,4
72.7 6.2
72.6 4.1
76.6 + 2.6
72.4 1.6
20
23.5 + 6.9
71.4 + 9.7
71.2 7.9
42.5 4.7
47.8 + 4.8
49.7 4.7
44.4 + 8.1
40
26.6
+ 5.1
73.4 + 15.9
71.1 +1.1::1 15.0
36,3 4.5
30.9 7.1 7.1
32.3 + 11.6
30.3 + 4.6
100
27.5 = 2.0
54.6 10.1
54.6 10.2
29.8 2.1
29.0
28.3 5.6
FedRep
83.6
90.3 3.1
2.4
91.2 1.6
91.1
91.1 2.7
20
85.3 2.0
90.7 2.5
91.5 2.6
87.9 2.0
18.4
8.4
88.6 1 1.8
40
84.1 0.8
93.6 - 2.9
93.3 +
84.8 1.7
8..4
0.5
85.3 + 1.2
100
77.9 + 1.4
91.4 1.44200 2.0
91.6 + 1.9
77.8 + 1.7
78.0 2.4
79.0
79.4 1.7
FedPer
10
83.1
2.1
92.6 = 2.2
92.7 = 1.4
93.0 1.4
93.1 1.5
93.0
93.1 + 1.3
20
88.7 + 1.7
92.3 - 1.8
92.7 - 2,4
89.8 2.0
90.2 2 1.8
90.1 1.5
90.0 + 1.2
40
88.1
+ 0.7
94.8 2.6
94.6 + 2.5
86.0 2.3
86.5 06.5::00 0.7
84.9 3.3
85.7 1.4
100
85.1
0.6
94.0 + 1.4
94.1 1 1.3
82.0 2.4
82.3 1 2.2
83.0
83.6 1 1.6
Average Uplift
19.9 1 18.0
19.7 1 17.5
8.3 1 15.2
8.6 1: 15.4
9.1 1 16.6
8.4 15.1
Table 5; CIFAR10 Average performances over 5 trials of three FL algorithms: FedAvg, FedRep and FedPer. For each algorithm we compare the vanilla performance with the ones obtained after clustering the clients using the FedOTDD distance, using three different parameters of the spectral clustering algorithm and for support size of 10 and 100. The number clients varies from to 100. For this table, datasets from clients do not have a clear cluster structure
24


Reviews:
[{"review_id": "eFZUqcy1aN", "reviewer": "Reviewer_FgAh", "summary": "The paper proposed a novel algorithm for computing the widely used Wasserstein distance (WD) that is in line with federated learning (FL). In particular, as WD is used in a wide range of tasks where privacy is of major concern, FL is a favorable paradigm to conduct its computation in practice. To this end, the paper proposed to utilize the structure of WD and to communicate between server and each client individually. For each iteration of the proposed algorithm, a measure is sent from server to a client, and a barycenter between this and the clients data is computed and sent back. The paper provides theoretical guarantees that the resulting aggregated distance converges to the actual WD in the limit, and in some cases even in 1 step. Computational complexity of the proposed algorithm is also discussed, with certain reduction possible. Various experiments are given to illustrate the performance of the proposed algorithm compared to the vanilla WD, showing that the performance is not compromised.", "questions": "The main question is: under an FL setting, is it in line to essentially reveal raw data to the server? As I remarked above, the server essentially have all of the raw data once it receive both responses, so it is unclear why it cannot just compute WD thereafter. Please see above (section Weaknesses) for more details.\n\nCertainly some amount of information has to be revealed to the server in FL, as client is communicating with the server with a locally trained object. Thus one possible way to address this is to perhaps quantify the privacy guarantee in the framework. I'm happy to raise score once this is fully addressed.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is overall well presented and the writing is clear. The contribution is original to the knowledge of the reviewer.\nSome main strengths:\n1: The main problem the paper, i.e. FL for WD, is novel and deserves attention. As WD heavily depends on both data sets at the same time, how to minimize the exposure of the raw data is of theoretical importance. The paper proposes a framework that is in line with FL.\n2. The paper covers most major aspects of practicality of FL for WD, including both theoretical and computational properties.", "weaknesses": "Although the idea of FL for WD is indeed interesting and deserves attention, I have a major concern on whether the proposed algorithm is really compliant with FL principles. Specifically, for each iteration, the server sends a distribution $\\xi$ to client with data $\\mu$, then a $t$-barycenter $\\xi_\\mu$ is sent back to the server. By structure of WD geodesic, with knowledge of $\\xi$ and $\\xi_\\mu$ it is already very immediate to reconstruct $\\mu$. For instance, if $\\xi = \\frac{1}{n}\\sum_i \\delta_{x_i}$ and $\\mu = \\frac{1}{n}\\sum_i \\delta_{y_i}$ are both $n$ points uniformly distributed, with optimal correspondence being $T:x_i\\to y_i$, then $\\xi_\\mu = \\frac{1}{n}\\sum_i \\delta_{(1-t)x_i+ty_i}$, and the optimal correspondence between $\\xi$ and $\\xi_\\mu$ will be $T':y_i\\to (1-t)x_i+ty_i$. Thus by computing $T'$ between $\\xi$ and $\\xi_\\mu$ on server, the server gains full information of the raw data. Surely $t$ is needed, but can also be inferred from the ratios of WDs between $\\xi$ and $\\mu$, and between $\\xi$ and $\\xi_\\mu$ (to be fair the paper claimed that WD is not necessary to report to the server). The reviewer does not claim full expertise in FL, but does believe that this not fully in line with the FL paradigm, which seeks to minimize the revealing of data from the client to the server. It would be great if this point can be addressed by the authors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MeMbll64uP", "reviewer": "Reviewer_vrdV", "summary": "The authors propose an algorithm for a federated computation of the wasserstein distance, where the datasets are held by different agents. Intuitively, this algorithm finds a fixed point $(\\xi,\\xi^\\mu,\\xi^\\nu)$ of the following system of equations:\n$$\\xi^\\mu=BC(\\xi,\\mu)$$\n$$\\xi^\\nu=BC(\\xi,\\nu)$$\n$$\\xi=BC(\\xi^\\mu,\\xi^\\nu)$$\nwhere $BC$ is the operation of computing a Wasserstein barycenter. The authors show communication cost, convergence and a theoretical analysis of some special cases. The proposed algorithm is then applied to a variety of problems.", "questions": "I have two minor comments/questions:\n- In Algorithm 1, you use $d_{\\mu,\\xi^{(k)}}$. While I can infer from the context what it probably means, this symbol has not been used before. I recommend defining it somewhere.\n- In the experiments, for the Wasserstein Coreset it does not seem straight forward how your algorithm can be applied to this problem. Could you please explain it to me which agents hold which distributions and where your algorithm enters the computation?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I think the authors propose an algorithm which could be valuable in federated learning and other data-science applications in the federated setting. The paper is well written and the concepts and development is easy to follow. I especially found Figure 1 to be very insightful. The Theorems support the claims and are valuable. While the setting of Theorem 3 is arguably a bit restrictive, the result is nonetheless very interesting. Specifically the fact that you can compute the Wasserstein distance in only one communication round!", "weaknesses": "Two main desiderata of federated learning are privacy and low communcation cost.\nWhile the problem of communication cost is addressed with (10), the problem of privacy remains largely unanswered. \nIf the authors convincingly address the problem of privacy leak, I am open to change my recommendation.\nIf I correctly understand the reasoning at the bottom of page 4, the authors propose to randomly select $t$ such that the server cannot easily infer $d_{\\mu,\\xi^{(k)}}$. While I agree that this is improves the privacy, I think that the server can, after just two communication rounds, infer $\\mu$ and $\\nu$. Consider the simplified setting of theorem 3. \nIf the server knows $(m_{\\xi^{(k)}}, m_{\\xi^{(k)}\\mu})$ and $(m_{\\xi^{(k+1)}}, m_{\\xi^{(k+1)}\\mu})$, she knows that $\\mu$ must line on the line through $(m_{\\xi^{(k)}}, m_{\\xi^{(k)}\\mu})$ and on the line through $(m_{\\xi^{(k+1)}}, m_{\\xi^{(k+1)}\\mu})$. These lines must cross at $m_\\mu$ by construction. (Here, $m_{\\xi^{(k)}\\mu}$ denotes the mean of $\\xi^{(k)}_\\mu$).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bGpEQfPzMG", "reviewer": "Reviewer_71t9", "summary": "The paper motivates the Federated Wasserstein distance that would be useful in situations where we what to know the distance between two distributions while the samples of the two distributions are not shared. The authors propose a converging algorithm that can compute such a Federated Wasserstein distance.", "questions": "1. Do we need the two distributions to be discrete and the samples of the two distributions to be the same? Theorem 1 seems to be less general than Theorem 2. If so, the motivation of Theorem 1 should be emphasized and highlighted.\n\n2. I think an experiment where the datasets are unbalanced would be super helpful; say what if a dataset is always a certain times of the other dataset.\n\n3. An experiment that compares the Gaussian & non-Gaussian assumptions to show the convergence speed could be better.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I found the problem is well motivated. The theoretical treatment is nice and sound. The numerical experiments are convincing.", "weaknesses": "The presentation of the manuscript could be better.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x2h33CwqmJ", "reviewer": "Reviewer_xDsE", "summary": "The paper proposes Federated Wasserstein Distance, an algorithm to compute the Wasserstein distance between two measures in the setting of federated learning. Due to the setting, a measure on a client cannot be transferred to the server or the other client to compute the distance. Therefore, the paper utilizes the geometry of Wasserstein distance i.e., triangle inequality and interpolating measure to create transferable measures. To ensure privacy of two interested measures $\\mu$, $\\nu$, three measures are created $\\xi$, $\\xi_\\mu$, $\\xi_\\nu$ where $\\xi_\\mu$ is the interpolating measure between $\\xi$ and $\\mu$,  $\\xi_\\nu$ is the interpolating measure between $\\xi$ and $\\nu$, and $\\xi$ is  the interpolating measure between $\\xi_\\mu$ and $\\xi_n$. For a fast computation, McCann’s interpolation is used for constructing the interpolating measure. Experiments on toy examples (Gaussian data) show that FedWD can approximate the Wasserstein distance with a small error. In addition, FedWD is used as the approximation for Wasserstein in coreset classification model, Geometric dataset distances, and  Boosting FL methods. Overall, FedWD shows a \"similar\" transportation cost as the exact Wasserstein distance.", "questions": "* Could FedWd be extended to continuous cases?\n* Could FedWd yield any approximations for the gradient?\n* For some ground costs that are not metric e.g., in domain adaptation application, the corresponding Wasserstein cost is not metric hence there is no triangle inequality. In this case, could we have any solutions?\n* How does $t$ affect the performance?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper tackles a new practice setting for computing Wasserstein distance i.e., federated learning. This is the first work that considers computing Wasserstein distance in this setting.\n* The usage of the interpolating measures is the key in this setup which is a novel contribution. \n* There are some (asymptotically) theoretical guarantees for the approximation from FedWD.\n* Experiments are conducted on various applications of Wasserstein distance to show a wide range of applications.", "weaknesses": "* Despite interpolating measures being well-defined for general measures, the paper only considers discrete measures setting. In this setting, interpolating measures can be computed in closed form (after knowing the transportation plan). In the continuous setting, this could be not doable (it seems that a transport map or a parametric measure must be used as the proxy in this setting). Considering the continuous setting e.g., with square L2 cost, could make the paper much stronger.\n* The gradient approximation scheme should be discussed. For example, If an application wants to estimate the gradient with respect to the support of a measure of the Wasserstein distance, could FedWd yield any approximations for the gradient?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)