ON THE GENERALIZATION OF WASSERSTEIN ROBUST FEDERATED LEARNING
Authors: Long Tan Le
Josh Nguyen
Canh T. Dinh
Nguyen Hoang Tran
Venue: ICLR 2022 Submitted (2022)
Keywords: Federated Learning
Robust Optimization
Adversarial Training

Abstract: In Federated learning (FL), participating clients typically possess non-i.i.d. data, posing a significant challenge to generalization to unseen distributions. To address this, we propose a Wasserstein distributionally robust optimization scheme called WAFL. Leveraging its duality, we frame WAFL as an empirical surrogate risk minimization problem, and solve it using a novel local SGD-based algorithm with convergence guarantees. We show that the robustness of WAFL is more general than related approaches, and the generalization bound is robust to all adversarial distributions inside the Wasserstein ball (ambiguity set). Since the center location and radius of the Wasserstein ball can be suitably modified, WAFL shows its applicability not only in robustness but also in domain adaptation. Through empirical evaluation, we demonstrate that WAFL generalizes better than the vanilla FedAvg in non-i.i.d. settings, and is more robust than other related methods in distribution shift settings. Further, using benchmark datasets we show that WAFL is capable of generalizing to unseen target domains.

Under review as a conference paper at ICLR 2022
ON THE GENERALIZATION OF WASSERSTEIN ROBUST FEDERATED LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
In Federated Learning (FL), participating clients typically possess non-i.idd data, posing a significant challenge to generalization to unseen distributions. To address this, we propose a Wasserstein distributionally robust optimization scheme called WAFL. Leveraging its duality, we frame WAFL as an empirical surrogate risk minimization problem, and solve it using a novel local SGD- algorithm with convergence guarantees. We show that the robustness of WAFL is more general than related approaches, and the generalization bound is robust to all adversarial distributions inside the Wasserstein ball (ambiguity set). Since the center loca- tion and radius of the Wasserstein ball can be suitably modified, WAFL shows its applicability not only in robustness but also in domain adaptation. Through em- pirical evaluation, we demonstrate that WAFL generalizes better than the vanilla FedAvg non-i.id. settings, and is more robust than other related methods in distribution shift settings. Further, using benchmark datasets we show that WAFL is capable of generalizing to unseen target domains.
INTRODUCTION
Federated Learning (FL) (Konečný al.) 2016 McMahan al. 2017) has emerged cutting- edge technique in distributed and srivacy-preserving machine learning. The nature of highly non- independent and identically distributed (non-i.i.d.) data in clients' devices poses an important chal- lenge FL commonly called statistical heterogeneity. The global model trained on this data using the de facto FedAvg algorithm (McMahan et al.) (017) has been shown to generalize poorly to individual clients' data, and further unseen distributions new clients they enter the network.
Several solutions to data heterogeneity have been proposed. Personalized FL (Fallah et al. 2020 Deng et al., 2020a| Dinh et al.) 2021 Li et al. 2021 and multi-task FL (Smith al. 2018) are client-adaptive approaches, where a personalized model is adapted to each client. From another perspective, distributionally robust FL trains model using worst-case objective over ambiguity set (Mohri et al.) 2019; Du et al.) 22200) Reisizadeh et al. 2020 Deng et al., 2020b). This approach client-uniform because single global model is judiciously learned deliver uniformly good performance not only for all training clients but also for new/unseen clients with unknown data dis- tributions. is specifically useful when test distributions drift away from the training distributions.
A natural question when designing distributionally robust FL frameworks is generalization: how can minimizing the training eтror also bound the test eиror (generalization bounds)? FL, Mohri et al.) (019) proposed agnostic FI. where a model is designed to be robust against any ambiguity set as convex combination of the clients' distributions. |Reisizadeh et al. (2020) applied the general affine covariate shift - used in the standard adversarial robust training into FL training. In characterizing the generalization bounds, while Mohri et al. (2019) relied on the standard Rademacher complexity, Reisizadeh al. 2220) use the basgin-based technique developed by Bartlett et al. (211)).
In this work, we take a different approach called WAsserstein distributionally robust FL (WAFL for short). The ambiguity set in WAFL is Wasserstein ball of all adversarial distributions in close proximity to the nominal data distribution at the center. Our main contributions are:
We propose distributionally robust optimization problem to address statistical heterogene- ity in FL. By controlling the center and radius of the Wasserstein ball, we show that WAFL is robust to wider range of adversarial distributions than is agnostic adversarial FL..

---

Under review as a conference paper at ICLR 2022
To make WAFL more amenable distributed optimization, transform the original prob- lem into a minimization of the empirical surrogate risk. We propose a local SGD-based al- gorithm to solve this surrogate problem. With additional Lipschitz smoothness conditions, standard techniques can he applied to find the convergence rate for the proposed algorithm.
We show how WAFL's output can reduce the test error by bounding its excess risk. We call this the robust generalization bound as it is applicable to all adversarial distributions inside the Wasserstein ball. By scaling the Wasserstein radius based on local data sizes, we show this bound applicable the true (unknown) data distribution among all clients.
We show WAFL's extra flexibility controlling the location of the Wasserstein ball by adjusting the nominal distribution. This enables applications such as multi-source domain adaptation and robustness to all clients' unknown distributions with minimal Wasserstein radius.
Experimentally, discuss how to control this radius for given center location by fine- tuning a robust hyperparameter. We show that WAFL generalizes better than FedAvg in non-i.i.d. settings and further outperforms existing robust methods in distribution shift set- tings. We finally explore WAFL's capability in transferring knowledge from multi-source domains to related target domains with much less data and/or without labels.
2 RELATED WORK
Federated Learning was introduced in response to three challenges of machine learning at scale: massive data quantities at the edge, communication-eritte
networks of participating devices, and
Divacy-preserving learning without central data storage (Konečný et al., 2016 McMahan et al. 2017]. The de facto federated optimization algorithm FedAvg (McMahan al. 2017) is based on local stochastic gradient descent (SGD) and averaging and is often considered a baseline in FL..
Most challenges of FL categorized into systems heterogeneity and statistical heterogeneity. The former focuses on communication problems such as connection loss and bandwidth minimization. This motivated some prior works to design more communication--ffiiiee methods (Konečný al., 2016 2017] Suresh et al.] 2017). On the other hand, statitical heterogeneity is concerned with clients' non-i.i.d. data, which is the main cause behind aggregating very different models leading to one which does not perform well any data distribution. To address this, many ideas have been introduced, [L et al. (2020) provided much theoretical analysis FI. non-i.i.d. settings. Zhao et al. (2018) proposed an FI. framework which globally shares a small subset of data among clients to train the model with non-i.idd data, /Smith et al. (2018) introduced a multi-task FL framework in which each client individually learns its own data pattern while borrowing information from other clients. Mansour et al. (2020) proposed three approaches to adapt the FL. model to enable personalization, in reponse to distribution shift. Several personalized FL models have also been developed, including
Fallah et al. (2220); Deng et al. (2020a): Dinh et al. 2021 Li et al. (2021).
Wasserstein Distributionally Robust Optimization (WDRO) aims lean robust model against adversarially manipulated data. The unknown data distribution is assumed to lie within a Wasserstein ball centered around the empirical distribution (Kuhn al. 2019). WDRO has received attention as a promising tool for training parametric models, both in centralized and federated learning settings.
In centralized learning, many studies have proposed solutions based on WDRO problems for cer- tain machine learning tasks (Shafieezadeh Abadeh et al. 2015) Gao & Kleywegt 2016 Essthhai
& Kuhn 2017 Chen & Paschalidis 20188] [Sinha et al. [2020] Blanchet et al.) 2019 Shafieezadeh
Abadeh et al. 2019: Gao et al. 2020). For instance, Shaffiezzadhh Abadeh et al. (2015) considered
a robust logistic regression model under the assumption that the probability distributions lie in Wasserstein ball. [Chen & Paschalidis (2018); Blanchet et al. (2019); Gao et al. (2020) leveraged WDRO to recover regularization formulations classification and regression. Gao & Kleywegt 2016) proposed minimizer based
tractable approximation of the local worst-case risk. [Es- fahani & Kuhn (2017) used WDRO to formulate the search for the largest perturbation range as an optimization problem and solve its dual problem. Sinha et al. (2020) introduced a robustness certifi- cate based on a Lagrangian relaxation of the loss function which provably robust against adversarial input distributions within Wasserstein ball centered around the original input distribution.
2

---

Under review as a conference paper at ICLR 2022
In FL. only few works have explored the Wasserstein distance (Reisizadeh et al. 20200] Diamandis et al. 2021). Specially, Reisizadeh et al. (2020) proposed FedRobust based on adversarial robust tramiig enhance robustness. Regarding distributionally robust learning, Deng et al. 2020b) pro- posed DRFA, a communication efficient distributed algorithm. Besides, based on the agnostic FI. framework suggested by Mohri al. (2019), Du al. (2020) introduced AgnosticFair, a two-player adversarial minimax game between the learner and the adversary, to achieve fairness.
3
WASSERSTEIN ROBUST FEDERATED LEARNING
3.1 EXPECTED RISK AND EMPIRICAL RISK MINIMIZATION IN FEDERATED LEARNING
In a federated setting, there are m clients, and each client [m] == {I, ... m) has its data gen- erating distribution P; supported on domain Z; := (X, M)). Consider the parametrized hypothesis class H. = ho 0 E R where each member hg is a mapping from X, to V parametrized by With 24 := (I,, y4) E Zi, we use {Z2, hg), shorthand for Ey,, he (x: )). to represent the cost of predicting he (x;) when the ground-trute truth label is y4. For example, if (xx) 0 0'x, and y R, square loss ((2,, ha) = E(y,, ho(z;) (0 x2 y;)² can be considered. In FL., all clients collaborate with server to find global model 0 such that the following sum weighted risk is minimized:
m
min PER
Ez [e(Zi,ho)]
where Ez,~P. [e(Z, he]] is client i's expected risk and A, 2 0 represents the relative "weight" of client i and Σ; | A = 1. Therefore, A == [AI..... Am]] belongs to simplex A = and A 1m 1]. Define by Pa = Di=1 AiPi the mixed clients' distribution over m domains Z == (Z,...., J. We denote by ~ PA a random data point Z generated by Px. which means that the domain of client i is chosen with probability P(Z Za) X first, then
E Rm
a data point 2; E 2, is selected with probability P(Z. = 2,), Z4 P.
While the underlying distributions P, are unknown, clients have access finite observations 2; [n:] ~ Pa. We abuse the notation [ms]] to denote the set of client i's both observable data points and their indexes. Let Pn:: := 1 >x.[(n]] 6z, be the empirical distribution of P.. where &z, is the Dirac point mass 24. In general, use the notation for quantities that dependent on the training data. Detine by PA == Ein >yPr, the mixed empirical distribution of >i= n, training data from m clients. The empirical risk minimization (ERM) problem of ① is as follows:
m
min DERd
E E ((((,h)))
Σ X.E E.E,,.P.. [e(Z,, ho)]
Σ
E(2,, ho) (2)
n
74
where d = nz/n is often chosen in ERM of the standard FL. McMahan et al. 2017)
3.2 WASSERSTEIN ROBUST RISK IN FEDERATED LEARNING
Models resulting from have been shown be vulnerable adversarial attacks and to lack of robustness to distribution shifts. We consider a robust variant of the ERM framework, involving the worst-case risk with respect to the p- Wasserstein distance between two probability measures. Given a set 2, define d 2 x 2 - J0, xc) to be the cost of "transportation" between its two points Suppose P and Q are two distributions on Z. Let ()(P.Q be the set of probability measures π on Z x Z whose marginals are P' and Q, called their couplings. In other words, (A, Z) P(A) and T(Z, A) = ((A), VA C Z. The p- Wasserstein distance between P and Q is defined as
inf WP(P,Q) TETT(P.)) (E(z,z))~z (d' (Z, Z'))) 1/P
This distance represents the minimum cost of transporting one distribution to another, where the cost of moving a unit point mass is determined by the ground metric on the space of uncertainty realizations. In this work, we mainly work with p 2. Let (P. == (Qp:PPP.Q) (). pp
The function d must satisfy non-nigativity, lower semi-continuity and d(2,2) Vz 2.
3

---

Under review as a conference paper at ICLR 2022
denote the Wasserstein ball centered at P (i.e., nomimal distribution) and having radius p 2 0. We modify 2 into the following Wasserstein robust risk minimization in FL problem:
WAFL:
min (UPQQEB(P..)) Ezzaq [e(Z', ho)l } BER
(4)
There are several merits to this framework. First, the ambiguity set Bp p) contains all (contin- uous or discrete) distributions Q that can be converted from the (discrete) nominal distribution Ps bounded transportation cost p. Second, Wasserstein distances can be approximated from the samples. Based on the non-asymptotic convergence results of Fournier & Guillin (2015), we can specify suitable value for p to probabilistically bound W (P.,Q) by the distance between their empirical distributions (P, Q) (e.g., for multi-source domain adaptation).
In any robust optimization problem, the ambiguity set is key ingredient defining the level of robustness. We will compare WAFL in with other approaches in terms of their ambiguity set.
Agnostic FL. Using this approach, existing tech- niques (Mohri et al. worst-case loss
2019; Deng et al. 2020b) minimize the
max E E, [e(Z, ho)l,
P
XEA
P
hence its distributional ambiguity set is Qa == {PA : X A}. While Agnostic FL's ambiguity set is the static convex hull of {Pr. Jeeim) WAFL's ambiguity set B(Pxx p) can be adjusted by controlling the robustness level p and by position- ing the ball center using A, which is useful for domain adapta- tion. Therefore, B Px. p) can cover QA by using appropriate values for and À (see Fig. D.
P1
Px
P2
Figure 1: Example of four FL clients with data distributions Pı... Pa.. The shaded area (Agnostic FL's ambiguity set) is covered by the blue ball with ra- dius p and centered PA (Wasserstein ambiguity sct). For domain adaptation with Q as a target domain, the nom- inal distribution (multi-source domain) shifted to Px such that W2 Q) is minimal.
Adversarial robust FL. Using this approach, Reisizadeh jet al.] (2020) combines a general affine covariate shift in stan- dard adversarial robust training with FL. Most existing tech- niques using this approach (Goodfellow et al. [015 Kurakin al. 2017; Carlini & Wagner, [017]] Madry et al.] 2019] Tramer et al. (2020) define an adversarial perturbation u at a data point Z, and minimize the worst-case loss over all perturbations: maNuEu E Ez~A, [e(Z +
u, he)], where the ambiguity set is U == {u e IR++1; ||u|| S {e}- In Appendix A we show that the Wasserstein ambiguity set also contains the perturbation points induced by the solution to this adversarial robust training problem.
3.3 WAFL: ALGORITHM DESIGN AND CONVERGENCE ANALYSIS
The original form of WAFL in (4) is not friendly for distributed algorithm design. Fortunately, the Wasserstein robust risk (or B( PAP p) -worst-case risk) has its dual formulation as follows (Gao &
Kleywegt 2016 Sinha et al.) 2020)
sup E Egrn [( ho)] = inf (rp² + E EzZA, [0, (Z,0))
(5)
QEB(Px.n)
720
where (2,, 55 0) == suPeez [E(Ç. ho) - (c, and ||r
+ kly s 12,x
0. The crux of using the dual is that the inner supremum problem (finding Φo,)) is easily solvable when its objective is well-conditioned: if l is L-smooth and d is 1-strongly convex, setting > L ensures that Ç + E(C, h4) rdf( 2) is strongly concave, and using gradient ascent for the inner supremum problem (for finding Φb,) ensures linear convergence. For each let 2 be the solution to supçez [e(S, he) - >d2 =)]. Then, Vog, 0) = V hha (Sinha et al. 2020 Lemma I),
Rather than find the optimal y in (5). we consider y as control parameter, and instead solve the following client- dccempdlccompoable problem, which is amenable to distributed algorithm design.
min 0ER
E Ez~i, 10, (2,0)
X,E z, Pr, [0,
(6)
4

---

Under review as a conference paper at ICLR 2022
Algorithm 1: Local SGD for WAFL
1:
for t = 0, 1, T do
// Global rounds
2:
Sample subset of clients St C for client E St in parallel do
[m]
3:
4:
Set local parameters: 0f 0,5.0) 0²
// Communication
5:
for k= 0, 1, ..... K - 1 do
// Local rounds
6:
Sample mini-batch D. from client i's dataset
ofc.k+1)= 0l..k) PoI Σ
Vo0(2,, k))
7:
4ED;
8: Send ((c.k) to server
// Communication canion
9: Server: update ((++1) = Eiss, ).O((.K)
Eves, A
This motivates the development of Algorithm for solving (6) The structure of WAFL similar to FedAvg with 7 communication rounds. We further note three key components of Algorithm common to FL methods. First, client sampling (line 2 refers to the partial participation of clients in each global round. Second, each client performs K local steps (line 5 before sending its local model to the server. Finally, stochastic approximation of a client's gradient using a mini-batch (lines 6 and 7) is necessary when the data size is large. The main difference between WAFL and FedAvg is that WAFL aims to minimize the risk with respect to the surrogate loss Φy, rather than l.
We show that the convergence of WAFL can be similarly characterized that of FedAvg, the de facto FL. algorithm based on local SGD updates (McMahan et al. 2017). In FedAvg optimization, seek to establish the convergence when using the original loss function f. On the other hand, in WAFL the convergence is with respect to the surrogate loss o.,, through which the local and global risks are defined by (.( == Ez,-p, [Ø,(2 (Z, 0)] and F())= == Ez~pA (Z, 0)] respectively.
We first make the following assumptions, common analyses of Wasserstein-robust optimiza- tion (Sinha et al. 2020). Unless stated otherwise, all norms are the Euclidean norm.
Assumption 1. The function Viz E Z.
Z + IR is continuous, and d(,, z0) is 1-strongly convex,
Assumption 2. The loss function Z x H + R is Lipschitz continuous as follows
(a)]l((z, he) - 0(z', (('',h)) he)ll & L1 - z/11. (b)]l((, ha) É(2, heg )l S Løll0 - 0'|).
Assumption 3. The loss function Ł 2 H + I is Lipschitz smooth as follows
Vol(z, ho) - Vol(z, ho < T. ||0 0'11, [V ho) v=((=', ho)1 S T1zz||z - 2'11,
Voll(2 ho) - Vol(z' he)ll < Lazl|z - 2/11, [V ha) - V (z, ho S L 0 - 8' ||.
Given Assumption 3 it has been shown that the mapping + (- is L-smooth with L = Lee + LocLae >-bus ? Las (Sinha et al. 2020 (more detail in Appendix Lemma 2 In addition, we make the following assumptions common to FL analysis (Wang et al., 2021).
Assumption 4. The unbiased stochastic approximation of (F), denoted by go. (0) = (od,(2i, 0), z, ~ P,, has x² o2-uniformly bounded variance, i.e., E I1ge. (0) (F)( S o2.
Assumption 5. The difference between the local gradient V (.)) and the global grudient (() is Q-uniformly bounded, i.e., max, V. FF < Ω.
Assuming complete participation of clients in every round (S, = m, Vl), using standard techniques in Wang et al. (202I). we have:
Theorem 1 (WAFL's convergence for convex loss function). Let Assumptions 5 hold and the mapping O + E(2, ho) be convex. Denote by git,k) the "shadow" sequence, defined as (it.k)
Em )1((1 and by O* the optimal solution to minger F(O). If the client learning rate satisfies
D
D
D
n & min
3L' 2VTT
43 403 K
5

---

Under review as a conference paper at ICLR 2022
then have
F(B(LA)) - F(8)) S So
LD2 KT +
@DA LiOiDE
VKT + K!T + 7$ (
E
(7)
k=0
where 6²
D |1000) 1 and A = Dim1 p
We provide a proof Theorem ! in Appendix C
4 ROBUST GENERALIZATION BOUNDS
We show the generalization and robustness properties of WAFL's output by bounding its excess risk. Denote the loss class by F := f o H := {a + {(z, h), h H). where we use f (resp. fa) E J to represent a generic loss (resp. a loss function parametrized by 0).
Definition 1. Denote the expected risk and surrogate of Wasserstein robust risk of an arbitrary f. respectively, as follows
L(Px, f) = Ez-ra [<( and
= (2,(),)) + +pp²
Then their excess risks are defined respectively as follows inf
PeF (') and == inf (Px.
PeF
If a distribution Q is the ambiguity set B(Px, p). we can bound its excess risk f(Q )) as follows. (a) holds and y 2 La/p. For all f F and for all Q B(Px.P), 87 PP. f) g(p. 7) &QQ,)) - PP,, f) ,(p,)),
Lemma 1. Let Assumption
where g(p,)) = 2L + h - --- pp?,, and 7- == arg min 220 &3 (Px.f).
We provide proof of Lemma] | in Appendix
Remark 1. Lemma shows that the lower and upper bounds for f)) f) can be analyzed using fP). and a two-component error term ((p,)) The first component, 2 says that when increased to allow for larger Wasserstein distance between the nominal Px and any worst-case distribution Q the difference between the excess risks &(Q,f) f) and (P). increases, and this error is amplified at most by the Lipschitz constant I., of the mapping = + ((2..). The second component, 1y>-"1p²> addresses the sub-optimality error of chosen value of y, which is amplified when y is drifted away from the optimal y" Note that 47' (Px, f) is the same B(Px, worstt- case risk thanks to the strong duality (5) which is obtained with > 0.
Denote by 0= E θ an E-minimizer the surrogate ERM, i.c., Ez~P, (2, fz.) infage E, Ez-P, [ø(Z. fo)] + where c. Rd is a parameter class, we obtain the following result. Theorem (Robust Generalization Bounds). Let Assumptions 2 and hold, 7 2 max and [[(2, h)] < Me. We have the following result for all Q B ((P.P))
771
48e(0)
2og (2m/8)
E(Q, S
Σ A
+ 2M
+ t + (((g))
Vni
with probability at least 1 8- where C(@) La Jo Vlog (e, e and N (e, lle.e)
denotes the c-covering number of θ w.n.f metric || ele as the norm on θ
We provide a proof of Thcorem in Appendix We sketch the proof as follows: first bound E7 (P,, using standard excess risk decomposition and uniform convergence with Rademacher complexity. Then leverage the upper-bound of Lemma | to bound E(Q, f) YQ ((Px.P). based on the bound of (P.. fz.). The result shows that using WAFL to minimize the surrogate of Wasserstein robust empirical risk also controls the robustness and generalization. As example

---

Under review as a conference paper at ICLR 2022
H ((,,)) } with -= = R |101|2 C}. The diameter of θ is supp,rees | 6'|l 2C, thus (ouss (9,11-11,,c) = (1 2C/e"" and e(@) - 3C V d/2 (Lee Raginsky 2018).
Generally, the radius of Wasserstein ball p can be considered hyperparameter that needs fine-tuning (e.g., through cross-validation which hurt the empirical average performance, but also not too small to become similar the ERM, and thus can lack robustness. From a statistical standpoint, we are interested in learning how scale p w.r.t. the sample size n4, [m]] such that the generalization of the WAFL solution OF w.r.t the true distribution Px is guaranteed, while still ensuring robustness w.r.t all distributions inside the Wasserstein ball. Using the result from Fournier & Guillin (2015). which shows that Pn.. converges in Wasserstein distance to the true P. at specific rate, we obtain the following
In principle, p should not be too large to become over-conservative,
Corollary 1. With all assumptions as in Theorem defining Pr.
Ei=1 Xipri we have
5/m
m
48E(O)
2 log( (4m/8)
E(Px. fac)
+ 2M
+ g(pn: +
Vnz
n4
min 12/
log(es/6) c2n
in log(e1/8) 02
with probability at least 1 5, where ph
1/a
logler/6) C22
if log(c1/8)
02
We provide a proof of Corollary iin Appendix
F
5 CHOOSING A: APPLICATIONS
We focus on two applications: multi-source domain adaptation and generalization to all client dis- tributions. We provide insights about choosing the client weights À for these applications.
Multi-source domain adaptation: Consider the multi-source domain distribution Px (Mansour al., 202T). Lee & Raginsky (2018) show that solving the minimax risk with the Wasserstein ambiguity set can help transfer data/ tata/knowledge from the source domain Px to different, but related, target domain Q. They bound the distance W, PP.. Q) using the triangle inequality
(PAP
(8) where PA and @ are the empirical versions of PA and Q. respectively. While Wr PA) and W,(Q.Q) Q) can be probabilistically bounded with confidence parameter d (0, according to Fournier & Guillin (015)) W,( (Px.)) can be deterministically computed using linear or convex programming (Peyré & Cuturi 2019).
Wp(Ps, S WP,,P, Px) + Wp((s. Q) + W Q),
In FL context, in order to have better bound for W2(Px, Q) similar to (8) it is straightforward to choose = arg min W2(Px- Q). To relax this problem into form solvable using existing approaches, observe that (P Q) S E A: R. Q) due to the convexity of Wasserstein
distance. We then consider the following upper-bound minxea (2(PA, Q):
min Σ 1=1 We Q) =:p". XEA
which is a linear program, considering each Wz( (P... Q) can be found by efficiently solving convex programs especially with entropic regularization and the Sinkhorn algorithm (Cuturi 2013). Corollary 2. Denote the solution to (9) by X², and assume that domain Q generates ng i.i.d. data points. With probability at least 1 6, have
(P Q) < W2(Px,, W3 (2 +)
W2(Q,
Σ m
X
b/mm
+ + 5512
The proof of this corollary is similar that of Corollary in Appendix
B
Covering all client distributions in the Wasserstein ball: Suppose we want to cover all client distributions inside a Wassertein ball so that the generalization and robustness result by WAFL in Theorem 2 applicable to all clients' distributions. We show in Appendix B that this is a problem of finding À such that the Wasserstein distance between PA and P,, Vj, is small as possible.

---

Under review as a conference paper at ICLR 2022
6 EXPERIMENTS
We first show how to change the level of worst-case perturbations by varying the robust parameter 7. To show the generalizability and robustness of WAFL, we evaluate WAFL under non-i.id. and distribution shift settings. We compare WAFL with baseline robust methods and non-robust FedAvg. Finally, in Appendix G.4 we show WAFL's capability in domain adaptation.
Experimental settings. We use two datasets in two non-i.id. settings. We distribute the first dataset MNIIST (Lecun et al.) 1998) to 100 clients and use a multinomial logistic regression model to model a convex setting. For the second dataset- CIFAR-10 (Krizhevsky) 2009) we use 20 clients and a CNN model employed in |MMMahan et al. (2177) to model non-convex setting. We set A; = n n as the client weights. In optimization, we randomly sample S. 10 clients to participate in training at each communication round. More detail can be found in Appendix G.I
Effeet of Y on the worst-case risk perturbations. We first examine the relationship between the robust parameter y and the average worst-case perturbations p. defined as 0² = E Ez~FA d (Z, (Z, Z)]. where Z is the adversarial perturbation of Z. As shown in Fig. smaller values of y correspond to larger worst-case perturbations á on both the MNIST and CIFAR-10 datasets. For the rest of the experiment, rather than control ý di- rectly, we set 7 on the opposite direction to control the level of robustness.
Average worst-case perturbation
150
125
0.75
0.50
WAFU:
0.0
WAFL: CFAR-10
15:1
1.8
307
1.0
10a
10
3/y
Figure 2: Varying ý by y.
Generalization and robustness of WAFL. We consider PA and Q as the empirical distribution of the training samples and test samples of all clients, respectively. By varying 7 we aim to train a global model robust to any empirical test distribution Q To show the generalization and robustness of WAFL, we train and evaluate it in two different scenarios. First, in the clean data scenario, the global robust model is trained with different values of y and then evaluated on given clients' test data (similar to traditional FL). Second, in the distribution shift setting, the training process is similar; however, the global robust model is evaluated when there distribution shifts at the clients' test data. To obtain these shifts, we use the common PGD attack method in Madry et al. (2019) under -oom generate e-level perturbation shows benefits in adversaries and gives large perturbations. Following Madry et al. 20119) we fix the number of gradient steps to generate adversarial examples with favd 40, 0.3. a 0.01 for MNIST and tand = 10, E = 8 8/255, a = 2/255 for CIFAR-10 10 with batch size of 64. We note that this setting is similar to the adversarial poisoning attacks and the main purpose of this setting to increase the Wasserstein distance between PA and Q, thereby verifying the robustness of WAFL.
clients' test data. We choose the 1 -norm as it
The performance of WAFL in both scenarios MNIST and CIFAR-10 10 is shown in Fig. When the clients' data is clean, the Wasserstein distance between Pa and Q is relatively small, and training WAFL with small y (large P) gives the worse performance on the test set. With a sufficiently large y, WAFL less robust and has the same generalization with FedAvg. By carefully fine-turning y in the range [0.5, 1] for MNIST and [10, 20) for CIFAR-10, WAFL shows an improvement over FedAvg. y this scenario plays the same role as a regularization parameter to handle non-i.i.d. data.
By adding distribution shifts, we increase the Wasserstein distance between P, and Q By vary- ing 7. we train global robust model with different levels of worst-case perturbation to handle the distribution shifts. Small values of 7 generate larger ambiguity sets B Px. ②) and increase the bustness of WAFL, thus increase the chance Q lies inside B(P,, P). In Fig. 3 WAFL shows better performance with smaller values. However, as in mentioned in Sec. 4 when is much larger or smaller than the level of distribution shifts (7 S 0.01 o 2 1 for MNIST and < 0.1 or y > 10 for CIFAR-10), WAFL performs inefficiently: too small 7 hurts the empirical performance, while too large makes WAFL lack robustness. For every scenario, 7 needs to be tuned correspondingly. In experiments, by carefully choosing 7. WAFL not only handles distribution shifts or com- mon data poisoning attacks but also provides better performance than FedAvg in non-i.id, data and heterogeneous settings. Specifically, we s 0.5 for MNIST and y 10 for CIFAR-10.
Comparison with other robust methods. We compare WAFL with three baselines: two common robust methods in FL called FedPGM and FedFGSM, and non-robust method Fe-
8

---

Under review as a conference paper at ICLR 2022
0.70
2.00
0.00
3.25
0.00
0.55
1.25
1.00
0.50
0.50
2.00
1.50
0.75
0.50
0.35
0.50
0.45
2.45
0.50
0.2.0.05
0.2
0.25
0.22
0.1
0.3
0.40
0.23
0.1
0.05 02
50
WAT:: db
Redfup Dasn
WATL Didhtoutoon shtte
Fealogg Edttrtution
Figure 3: Global accuracy loss of with different values 7 on MNIST and CIFAR-IO under clean data and distribution shifts (40% of clients are affected by PGD attack). The blue vertical line indicates the value of giving the same level of PGD attack(7 = 0.05 for MNIST and 7 0.5 for CIFAR-10)
MINIST
NNIST
CIARR
CIFAR
0.3
250
0.79
2.
2.25
2.00
3.-
0.50
0.45
0.40
0.75
14
2.0
Possetinn attacces Clarrs
Proport.o arccdee
Clants
Eottckedd Cierts
Arsporian Eitaceed
WATL
TecPGH
--.
Ped1GS54
---
TepPng
Figure 4: Comparison with other robust methods on MNIST and CIFAR-10 at different propartion of attacked clients (clients are affected by distribution shifts).
dAvg. FedPGM and FedFGSM are FedAvg with adversarial training using the projected gradi- ent method PGD (Madry al. 2019) and the fast-gradient method FGSM (Goodfellow et al.) 2015), respectively. In FedPGM and FedFGSM, in each local update, all clients solve d" arg m <e {E(ha(2 }) using projection onto an -nom find the worst-case per- turbation 6. While FedPGD uses lavd gradient steps to find 8*, FedFGSM uses only one gradient step. We use the same values of and from the distribution shifts setting for projection. For fair comparison, both WAFL and FedPGM have the same value land and all algorithms have the same number of local updates K. We also train WAFL with the value of y generating the same level perturbation distribution shifts in Fig. The robust accuracy of all methods decreases when the percentage of clients having distribution shifts increases. Especially, the FedAvg's performance drops dramati- cally when the percentage of attacked clients reaches 80%. For all scenarios, WAFL outperforms all the baselines. Specifically, in the case of 80% attacked clients, the improvements in accuracy of WAFL over FedPGM, FedFGSM, and FedAvg are 7%, 8%, and 33% for MNIST and 2.5%, 4.5%, and 18% for CIFAR-10, respectively.
in FedPGM and FedFGSM. We study different proportions of clients having
7 CONCLUSION
In this paper, we present WAFL, a Wasserstein distributionally robust optimization framework, to tackle the issue of statistical heterogeneity in federated learning. We first remodel the duality of the worst-case risk to an empirical surrogate risk minimization problem, then solve it using local SGD-baseed algorithm with convergence analysis. We show that WAFL is more general in terms of robustness compared to related approaches, and obtains an explicit robust generalization bound with respect to all unknown distributions in the Wasserstein ambiguity set. Through numerical experiments, we demonstrate that WAFL generalizes better than the standard FedAvg baseline in non-i.i.d. settings, and outperforms related methods with respect to robustness to distribution shifts. Moreover, WAFL reveals its capability in generalizing unseen data distributions.
9

---

Under review as a conference paper at ICLR 2022
REFERENCES
Peter Bartlett, Dylan J. Foster, and Matus Telgarsky. Spectrally margin bounds for neural networks. arXiv:1706.0844 [cs, stat], December 2017. URL http://arrii.orr////////60..
arXiv: 1706.08498.
Jose Blanchet, Yang Kang. and Karthyek Murthy. Robust Wasserstein Profile Inference and Appli- cations to Machine Learning. Journal of Applied Probability, 56(3):830-857, September 2019. ISSN 0021-9002, 1475-6072. doi: 10.101 20.1017//p:220194499 URL http://arriioorraaa////110005177 arXiv: 1610.05627.
Nicholas Carlini and David Wagner. Towards Evaluating the Robustness of Neural Networks.
arXiv:1608. 046444 [cs], March 2017. URL http://arxxiv.rr/bbb///// http://aabb/600..0066666 646404084444 arXiv: 1608.04644.
Ruidi Chen and loannis Ch Paschalidis. A Robust Learning Approach for Regression Models Based on Distributionally Robust Optimization. Journal of Machine Learning Research, 19(13):1-48,
2018. ISSN 1533-7928. URL hattp://imlr.org/aaprsss
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Proceed- ings of the 26th International Conference on Neural Information Processing Systems Volume 2, NIPS' 13, pp. 2292-2300, Red Hook, NY, USA, 2013. Curran Associates Inc.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive Personalized Feder- ated Learning. KKIv:2003.13466 [cs, stat], November 2020a. URL http://arxiv.org/ab//200..
1346 arXiv: 2303.13461.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionally Robust Feder- ated Averaging. In Advances in Neural Information Processing Systems, volume 33, pp. 15111- 15122. Curran Associates, Inc., 2020b. URL httpss/pooeeedingg..eerreeccpppp
Abbsrattht
Theo Diamandis, Yonina C. Eldar, Alireza Fallah, Farzan Farnia, and Asuman Ozdaglar. A Wasser- stein Minimax Framework for Mixed Linear Regression. arXiv:2106.07537 [cs, math, stat), June arXiv: 2106.07537.
2021. URL hRLLittt//aarxiv.og/bb//22000 07537
Canh T. Dinh, Nguyen H. Tran, and Tuan Dung Nguyen. Personalized Federated Learning with Moreau Envelopes. arXiv:2006.00 08848 [cs, stat], March 2021. URL http://arrivorr/abb//000. 08848 arXiv: 2006.08848.
Wei Du, Depeng Xu, Xintao Wu, and Hanghang Tong, Faimess-aware Agnostic Federated Learning. arXiv:2010.05057 [cs], October 2020. URL hWLLLtttp/aaaxiiorrr/////1100000 arXiv: 2010.05057.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven Distributionally Robust Optimiza- tion Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations. arXiv: /50.00510 16 [math, stat], June 2017. URL. http://arxiivorr/aa///555 500.aa/5//555515516.6 arXiv: 1505.05116.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized Federated Learning with The- oretical Guarantees: A Model-Agnostic Meta-Learning Approach. In Advances in Neural Infor- mation Processing Systems, volume 33, pp. 3557-3568. Curran Associates, Inc., 2020. URL/httss
Nicolas Fournier and Arnaud Guillin. On the rate of convergence in wasserstein distance of the empirical measure. Probability Theory and Related Fields, 162:707, 2015. doi: 10.1007/
500440-014-0583-7.
Yaroslav Ganin and Victor Lempitsky. Unsupervised Domain Adaptation by Backpropagation. In International Conference on Machine Learning, pp. 1180-1189. PMLR, June 2015. URL https:
1..mmm) html ISSN: 1938-7228.
Rui Gao and Anton J. Kleywegt. Distributionally Robust Stochastic Optimization with Wasserstein Distance. arXiv: [math], July 2016. URL http://arrvvoog//b////6400299999 arXiv: 1604.02199.
10

---

Under review as a conference paper at ICLR 2022
Rui Gao, Xi Chen, and Anton J. Kleywegt. Wasserstein Distributionally Robust Optimization and [cs, math, stat], October 2020. URL http://axxiv.orrg
Variation Regularization.
abs/1712.06050 arXiv: 1712.06050.
Jan J. Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and Hameessin Adversarial Examples. 112.:2241 [cs, stat], March 2015. URL http://aaxi..orr//b/1111..77 arXiv: 1412.6572.
Jochen Gorski, Frank Pfeuffer, and Kathrin Klamroth. Biconvex sets and optimization with biconvex functions: a survey and extensions. Mathematical Methods of Operations Research, 66(3):373- 407, December 2007. ISSN 1432-5217. doi: 10.
URL https://ooi.o
10.11077/600188-007--1----
J.J. Hull. A database for handwritten text recognition research. IEEE Transactions on Pattern Analysis and Machine Intelligence, 16(5):550-554, May 1994. ISSN 1939-3539, doi: 10.1109) 34.291440.
Jakub Konečný, H. Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated Optimiza- tion: Distributed Machine Learning for On-Device Intelligence. arXiv:1610.02527 [cs], October arXiv: 1610.02527.
2016. URL http://aaxiv.org/b////111......
Jakub Konečny, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated Learning: Strategies for Improving Communication Efficiency. arXiv:1610.05492 [cs], October 2017. http:p/p/aaxiiorr/a///////000000
arXiv: 1610.05492.
Alex Krizhevsky. Learning Multiple Layers of Features from Tiny Images. pp. 60, 2009.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Shafieezadch-Abade Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learn-
ing. a90i 08729 [cs, math, stat], August 2019. URL http://arxiv.org/ab////a http:////////////000/00//// arXiv: 1908.08729.
Alexey Kurakin, Ian Goodfellow. and Samy Bengio. Adversarial Machine Learning at Scale. arXiv: /611:/161.0136 fcs, star), February 2017. URL http://axiv..rr/abs//611111010 1.1011236 1611.01236.
arXiv:
Y. Lecun. L. Bottou, Y. Bengio, and P. Haffner. Gradient- learning applied document recog- nition. Proceedings of the IEEE, 86(11):227--222 November 1998. ISSN 1558-2256. doi: 10. 10.1109/5.726791. Conference Name: Proceedings of the IEEE.
Jaeho Lee and Maxim Raginsky. Minimax Statistical Learning with Wasserstein distances. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. 110
URL
html
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and Robust Federated Learning Through Personalization. arXiv:2012.04221 [cs, stat], June 2021. URL http://arxiv. arXiv: 2012.04221.
0gg/abs/011..02221
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang, On the Convergence of FedAvg on Non-IID Data. arXiv:1907.022 1907.02899 [cs, math, stat], June 2020. URL http://arxiv.org)
abs/ 1907 02189
arxiv: 1907.02189
Quande Liu, Cheng Chen, Jing Qin, Qi Dou, and Pheng-Am Heng. FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space. arXiv:2103.0600 [cs], March 2021.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards Deep Learning Models Resistant to Adversarial Attacks. 1106::170600088 [cs, stat], September 2019. URL http://arxiv.orr/abb//iit 706.0688) arXiv: 1706.06083.
Yishay Mansour, Mehryar Mohri, Jac Ro, and Ananda Theertha Suresh. Three Approaches for Personalization with Applications to Federated Learning. 10669 [cs, stat], July 2020. URL httth///arri..ogg/h/////...11111 arXiv: 2061.

---

Under review as a conference paper at ICLR 2022
Yishay Mansour, Mehryar Mohri, Jae Ro, Ananda Theertha Suresh, and Ke Wu. A Theory of Multiple-oooe Sorple--Sureeeeeee Adaptation with Limited Target Labeled Data. In Proceedings of The 24th Inter- national Conference on Artificial Intelligence and Statistics, pp. 2332- 2340. PMLR, March 2021. URL httpss//proceeding..mii prese 1 30/mansour221 ha.mmll ISSN: 2640-3498.
H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera Arcas. Communication-Effiiii Learning of Deep Networks from Decentralized Data. ww0ivv:602.00000 [cs], February 2017. URL http://aniiv.orr/b////60..011.1 1602.05629.
arXiv:
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic Federated Learning. arXiv:1902.00144 1414 [cs, stat), January 2019. URL http://ankv.oorgdds//900.....11 arXiv: 1902.00146.
Nicolas Papernot, Patrick McDaniel, Somesh Jha, Matt Fredrikson, Z. Berkay Celik, and Ananthram Swami. The Limitations of Deep Learning in Adversarial Settings. arXiv:1511.07528 [cs, stat], November 2015. URL http://arxiv.or//bb//511...
arXiv: 1511.05228.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances in Neural Information Processing Systems 32, Vancouver, BC, Canada, 2019.
Xingchao Peng, Zijun Huang, Yizhe Zhu, and Kate Saenko. Federated Adversarial Domain Adap- tation. arXiv: 1911 1.00054 [cs], December 2019.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data sci- ence. Foundations and Trends@ in Muchine Learning, 116-66)3555607, 2019. ISSN 1935-8237. doi: 10.1561/2200000007. U http://dxx.oi..org////o (tt.:/...151122000000000
Amirhossein Reisizadeh, Farzan Farnia, Ramtin Pedarsani, and Ali Jadbabaie. Robust Federated Learning: The Case of Affine Distribution Shifts. arXiv:2006.0 08907 [cs, math, stat], June 2020.
URL http:///arriv.og///b///000.0//9
arXiv: 2006.08907.
Filippo Santambrogio. Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling. Progress in Nonlinear Differential Equations and Their Applications. URL
Birkhäuser Basel, 2015. ISBN 978-3-319-20822
doi: 10.1007/978-3-33
Soroosh Shafieezadeh Abadeh, Peyman Mohajerin Mohajerin Esfahani, and Daniel Kuhn. Dis- tributionally Robust Logistic Regression. In Advances in Neural Information Processing Sys- tems, volume 28, Curran Associates, Inc., 2015. URL httpss:/pppprrs..ttttt///// httpss//papers.iisscccp/pppppppp https/sssscccppppe///////aaaaaa
Soroosh Shafieczadeh-AAaadee Daniel Kuhn, and Peyman Mohajerin Esfahani. Regularization via Mass Transportation. 1rXiv:1710.10016 [cs, math, stat], July 2019. URL http://aaxxi.orr/abb/
1710.10016 arXiv: 7710.10016
Shai Shalev-Shwartz and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, USA, 2014. ISBN 978-1-107-0571335.
Aman Sinha, Hongseok Namkoong, Riccardo Volpi, and John Duchi. Certifying Some Distribu- tional Robustness with Principled Adversarial Training, arXiv:1710.1057: [cs, stat], May 2020. URL http://axxiv.org/ab///11
arXiv: 1710.10571.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated Multi-Task Learning. arXiv: 10467 [cs, stat], February 2018. URL http::/arxxi.orr//bb/// 100.
arXiv: 1046.
Ananda Theertha Suresh, Felix X. Yu. Sanjiv Kumar, and H. Brendan McMahan. Distributed Mean Estimation with Limited Communication. arXiv:1611.00429 [cs], September 2017. URL http: Taaxiv.org/abs/1661..o 61.099999 arXiv: 161 1.00229.
12

---

Under review as a conference paper at ICLR 2022
Florian Tramèr, Alexey Kurakin, Nicolas Papernot, Ian Goodfellow, Dan Bonch, and Patrick Mc- Daniel. Ensemble Adversarial Training: Attacks and Defenses. vrXiv:1705.07200 [cs, stat], April 2020. URL http://araw..orr/aab////7077772
arXiv: 1705.07204.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H. Brendan McMahan, Blaise Aguera Arcas, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly. Deepesh Data, et al. A Field Guide to Federated Optimization. arXiv:2107.06917 [es]. July 2021. URL http: Warxvw.org/abs/110700601 arxiv: 2107 .06977
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated Learning with Non-IID Data. arXiv: 00582 (cs, stat], June 2018, URL http://arxxiv.or//bbb/
08582 arXiv: 1806.00582.
13

---

Under review as a conference paper at ICLR 2022
A ADVERSARIAL ROBUST FL's AMBIGUITY SET V.S. WASSERTEIN BALL
We show that using the Wasserstein ambiguity set contains the perturbation points induced by the solution to the Adversarial Robust FL approach. As we present in Sec. 3.2 existing techniques for adversarial training robust models (Goodfellow et al., 2015; Papernot et al. 2015 Kurakin et al.,
2017 Carlini& & Wagner 2017 Madry al.) 2019 Tramer al. (2020] define adversarial pertur-
bation u at data point 2, and minimize the following worst-case loss over all possible perturbations
max E Zz Pb [(Z + u, ho)]: wEl
(10)
where the ambiguity set {u E gd+ : ||u|| S e]. To compare this approach with Wasserstein- robust FL, we relate the above problem to its counterpart defined in the probability space of input as follows
max Ez 11)2
(11)
QE((c)
where Q(c) {Ō [P[IZ Z
= 11 Ps, 2 ~ @] Considering u² solution to problem (10). we see that the distribution Q' of perturbation points (i.e., Ż := (Z + u²) ~ Q')
in problem (10) belongs to the feasible set ((e) in problem (11) (If not, then P |||0|| S 6 < 1,
contradiction). Next, consider an arbitrary distribution Q E Q(c) in problem 11). with any Z ~ Q and % Ps. we have
w.p.I
2 - Z S = E Bz- -1|l Se =
inf B(z,z))= Z1| S =)
which implies that (i(PA. S v. Q(c). and thus Q(e) C (P,, We have shown that the Wasserstein ambiguity set contains the perturbation points induced by the solution to the adversarial robust training problem (T0).
B CHOOSING A: GENERALIZING TO ALL CLIENT DISTRIBUTIONS
We show that by calibrating appropriate À value, proposed algorithm will be capable of general- izing to all client distributions. Suppose we want to cover all client distributions inside a Wassertein ball so that the generalization and robustness result by WAFL in Theorem 2 is applicable to all clients' distributions. This is the problem of finding À such that the Wasserstein distance between PA and Pj, Vj, is as small as possible. Instead of directly finding the minimum Wasserstein dius that cover all client distributions, we will leverage the popular Wasserstein barycenter problem (Peyré & Cuturi 2019). Specificallyy, consider the problem
min W2 (2 P XEA
s.t.
P arg min
Σ W2 Q),
(12)
QEP
where P is the Wasserstein bary center w.r.t the solution x to this problem. Even though the solution is not straightforward, we propose solve its tractable upper-bound:
min XEAQQEP
Σ m
A:W2((n,, Q)
(13)
This is a bi-convex problem, which is convex w.r.t to À (resp. Q) when fixing Q (resp. A). Thus, use alternative minimization (Gorski et al., 2007) to find a local solution to this problem. Denoting Ā as the solution (12) and (X", P*) local solution to (13) obtain
W2 P < (P P*)
Σ m
:=1 X222(Pn, m,, p*)=:p"
(14)
Corollary 3. For all client j E [m)) with probubility at least 1 - ō, we have
W2(PA., P2) < (P PA-) (P( P*) + W2(P", P.,)) W2(Pn,, P2))
Σ
ApPri + p²
p²
+
Png
5/2
Aj
14

---

Under review as a conference paper at ICLR 2022
Proof.
The first line is by triangle inequality. The second line is by following facts: (i)
(P,,P Pa-) I P;p
S 8/2 according to (37). (ii) W2 Po.)
Wa (aa)) G. and (iii) P P[22((P)) P3) 2 Ag union bound.
5/2 Pn; 8/2 according to (3 and (iv) using
C PROOF OF THEOREM
Our proof is based on the analysis of local SGD for FL presented in Wang et al. 22021)
Fix some 24. Define Q(C, 8; 24 = E(S. ho) - yd(s, Since é Lms-Shooth and is 1-strongly convex, (p,, 2,) is (r LE= -troogly concave with respect to Ç, given that y L.z.
Lemma 2. Let 27 = arg m ((s, 0; zı). Therefore, 0,(0; 24) = ppz,, 21). Let satisfy Assumption 3
Then d,y is differentiable, and
(2
Vo, (2,, 0'1 Lll0 111,
with L = Log + Lackee when > Liz.
7-baa
The proof can be found in (Sinha et al. 2020 Lemma 1). Lemma implies that 6, is L-smooth.
Define (0) = D.l ExED. Vo (2, 0), then we have
(l))(0) [F)
1
Vo Vo (0) |
Di
EDD.
E (0)- <
(by Assumption 4
(15)
With the shadow sequence ((t.k)
= m 210(1.)) we have
((t
Σ Exoo((tt1) = Σ -
(ng: -
((c.k)
Augi (o(tk))
Lemma 3. the client learning rate satisfies 1) S 3L. then
K-1
Σ
E
F(((()))
-P())))))F(())))))))))))))))))))
f(Or))
<2n00
Σ
m
x2
122
E [10/4.5) 0cc.2)122
K
K
0=
Am
1 2nK
[18(0)
8" 9-12
E -
Proof. Since ((t,k+1) g((.h) -m >:0 (o((.))), by parallelogram law
de (B)(0(()))) glth++1) - -0r) 2n (101e.)) 00-1 [| - - -10(t+++1) -
-1104.141--00111 0--12)
(16)
Fact: F.(0) = Ez,,r (Z,, is Lipschitz smooth with = T.po + Lozkee when 7 > Lizz- 7-bes
With the assumption that 0 + E(2, ho) is convex, we have 0 + F.(O) is convex.
Since F. convex and L-smooth,
P.(0(e++1) D) S F:(O.(10)) + (VF.CO.LA)) 0,0041111 - -O(4A))
L -
I.
<F((0) + (FF(((())) 0(C,+++)) 0r)
2 (10le,tt2) 0((:2))12
g(c.++)) 8r) <F((0)+ + )
Ll g( + +L L((0((:A)) g0. (17)
15

---

Under review as a conference paper at ICLR 2022
From
16 and 117.
we have
T
F(B(Lkt1)) - F(00) Σ > 2 (
X (F((0(t.))) -,,(((+)),,0cattt))
0(aat1) - -0r) + -1I0(c+tt1) - 112 + L
>e110 - 01
(10(th) - 1-112 -110tt.t+ 0(1.8)122 [10(et+1) - -0913).
(18)
2ry
We have
- ((ck+1)- - -0r)
E
A
m
X (VEL(O!LAT) (u((!(it)), g(ehtt)) - ](
(since E (g.( -VF((((t.)) given g(r,k) 0t)
E
3
m
X (VE.(O ((((!(A)))
p1 1
$ 2"
(by Peter Paul inequality)
6t7
< 21j²²
Σ
[10(La41) - -0c.
(19)
617
Plugging 19 back to the conditional expectation of (88 and noting that 17 3L we have
E (()( - F(O0)] F(O0)]
1
E 1100..++11
0-113
-1(0(t.)) - 0-1
+
2n
22no²
Σ
m
X7
1
E
|10().0t1) - gpe.
T.
x111((( 00
3n
<2n10²
L
X )=llo((.k)
(c.k) ( 10
By convexity of F and telescoping & from 0 to K 1, we have
K-1
K-1
F(g(()))) F(())
22m02
Σ x
K
k=0
k=0
1
(IBCC.0)
0 1|
E
[((t.K) - (((1)))-0(13))
2mK
Since Ølt.0) ((t) and ((t.K) B(++1), we complete the proof.
Lemma 4 (Bounded client drift). Assuming the client learning rate satisfies 17 S 37. have
[10[([ (12 < <7(22133522 20K5?)).
where 02 := max 122)
Proof.
E ||of(att1) - -B( 1/ = -E ==|00(.0) - -g(c.)
- -nl -
(0)(0()))) -0 - (2((()1)))
=1(0[(Ab) p/2 2n (9i (B((:))))) - -VF(((((ik) ). g((c))
-of(.)))
-2n ( (o (2(( ),(((.)) .f(t
-2n (F(()(()))) k) p((.)))) of(,1)) ) + n7 - (1( (20)
16

---

Under review as a conference paper at ICLR 2022
The second term (and similarly for the third term) is bounded as follows
(F
0, 06..k))
|10
-02 glcek)
1|2
K
l1 (0fthh))
-
(by Peter Paul inequality)
6nK
2
|10((.0)
88c.k)1
112
σ²
(by 15
6nK
2
Since max, supo (VP:) V1)) - Ω (Assumption 5), the 4th-term bounded
VF
(FE(0 ). g/( olT.A))
S VP(O(( V ), ob4 ) + 202110 - -05400
1 IVF(((((A) VF(OB-h))))
20
(t.k) 1
(by smoothness and convexity)
s I IVF(O((*)) - VF(O( 112
6n |10( ,fr 2 + 6n
1
(by Young's 's inequality)
The last term is bounded as follows
119: (19 (o(t.)) 92 (=(((()))))
91 (0 (0ft.t)) VFI (1 (o(ttt)))) VF(((( ) - VF(B((.)))
+ /VVF(B(()))
(F2 )|2
() + (8.
< ( )F(( 10 + ?2 (by 15) and Assumption 5
Substituting the above four bounds back to 20 gives (note that 17 S 37.
E (1++)) o5 oft, 12 $ 1
K 110
o5 112 - ( 51) || | - -VF(S)))))) ||²
+ 6r7² KO² + 12 + 10² +
S I + +) [ -
1 + 67/² K2² + 12 + 10" 02).
Unrolling recursively, we obtain
(1 (1+1//4 1/K)* K
E ||0 [[((++++11 - 1112 S
'nn" - 12nj² Kn² + 10n + [2]]
1/K
< 12n K2 + +24r2K202 + 20m2 + 0t)
S 22 + -20K?)).
where we use the fact that (2+1/k)* K
-1 1/K
S K(e - 1) S 2K, and 22 = max 2?].
By convexity, for any i,
E [10 _ _p(
20K
Substituting the result of Lemma 4 to Lemma 3 and telescoping over 4, we obtain
D2
F(glt.k)) F((0)
<
+ 2njo A +20K0)
knKT
17

---

Under review as a conference paper at ICLR 2022
where D |10(0) 10 A = Σ X² By optimizing 7 the R.H.S, we obtain
T-1K-1
LD KT +
ADA
L
L D
F(B(L.k)) F(8*) s =o
KT
VKT
K$T
T
k=0
when
D
Då
1 = min
3L' 2VKTΛБ K
40 L
D PROOF OF LEMMA I
We first prove the following fact:
Fact 1:
(a) (7(xx.) VE
(b) inf (() < inf (T),,, PEE
Pef
For (a), we have
L(Q,5)
sup P'EB((x.M)
_(P', = inf {r'p² + +EZnP (0,
7/20
S 3p2² + Ez~PA Φ (2, = (7 f).
where the equality is due strong duality result by Gao & Kleywegt 2016). For (b), defining
Sr,, = arg min yer ('),, we have
inf =(Q.)) ( =(Q)))) ((Q. FeFF
sup P'EB(PX.P))
C(P',Sr))
(21)
{ (''p" + EzNPa [0,(Z,5r))}
= inf
(22)
7/00
Exp² + +EZ~P (e,(Z, fP.)
(23)
= inf
(P).
(24)
PEF
We next prove the second fact:
Fact
(a) ()(Px,)) S (QQ,)) + (TEF,QEB(Ps.))
(b) inf (3)P), < inf + ferr
feFF
For (a), we have:
[(P'.)] (E7(P2))) sup A(P', .n} PIEB(PA.P)
[(P'.f)
sup (PEBB(PAM))
s (((Q).)+2LL))
+ +2L.)] ( 12, p²y
min 7'20
p²n/ (22-
S <(Q0,)) + 2L. p2ly -2(1-7")) Ezap (Z, f) on. (2.
L(Q. + 2L. 12 - -7-)= + EZNr
sup { (e()h)) - -rd(s, 2)}
sup
{
<E2
(E2
= =(Q,)) + 2L. (1-7") 77) (p² -EZap
sup d2(s, 2)))
CEZ
s <L(0,)) + 21-p+1--77
18

---

Under review as a conference paper at ICLR 2022
where the first inequality is due Proposition
1 and the last inequality is because choose 7
L,/p and that fact that 7" 2 L by Lemma Lee & Raginsky (2018).
For (b), defining fq arg jen (QQ, f), we have
(2(Ps, 2 (7(2..1) PEF
inf
(25)
S S + 2L 1 - 701p²
(26)
inf f'EF
(27)
where the second line is due to Fact 2(a).
Combining all facts, we complete the proof. Specifically, adding two inequalities in Fact ((a)
and Fact 2(b), we obtain the upperbound of Lemma]] ! Similarly, adding two inequalities in Fact I(b) and Fact 2(a), we obtain the lowerbound of this lemma.
Finally, we provide the proof of the following proposition that was used in proving Fact 2(a).
Proposition 1. Let Assumption
(a) holds. For any F and for all Q B(Px. have
sup P)EB((X.P))
S () 2
Proof. Denote p arg max ((P", We have ('EB((a.P))
sup P'EB(PX.PP)
C(P', .I -)Q, sup P'EB(Pa.P)
L(P", -(Q,,)
S (Q()))) ( - L(Q,
S S(QQ,))+ +1, EZ-P- BB.,P ((Z,))/1
-Ez-q(A(Z,N)))
+ + +LWW .Q)
S + (. PA) (Px(P.Q))
(28)
- + L.2p,
where the fourth line are due Cantorovich-Rubinsteer dual representation theorem, i.e.,
Wi Q) sup h
Ezar [M(Z)] Ezzq [M()] M(-)) is 1-Lipschitz
and the fifth line is due to WPPP*,Q ()) < W2 Q) and triangle inequality.
E PROOF OF THEOREM 2
Proof. To simplify notation, we denote + = o.FF===
(2, E F} where F {fe, == f
θ mad), which represents the composition of Φb, with each of the loss function fo parametrized by 0 belonging to the parameter class θ
Defining frx E arg min fEF L3 f) and 0
E argmin E Ez-PA [6, (2. such that AE
C3(P,, fo-) inf Ezwe, (Z, 7p? decompose the excess risk as follows:
19

---

Under review as a conference paper at ICLR 2022
().) = (7) inf (Px)) JEF
= =EG(Ps f1c) SM)
= [ - -EZ(PANOO))
36 - (1)
+ (o.) - SPN
50
2 sup & cΦl
Ezwpa 10, (2, E (2, Soll +E
< 2 sup 6,E0 1=1
Xj 1 (2i. - E E2 [0, (Zi, ʃo)l +
m
sup
Ez 10y (Zz: [o)]
Ez (0g so)l 3
i=1 6y
2 2lo((2m/8)
( + 2M
+ with probability at least 1 8-8, 8, (29)
74.
where the first inequality is due to optimization error and definition of 0-. The second inequality is due to the fact that > Xaa;} S Σi= Ag|ay], Va, and A, 2 0. The third inequality because pushing the sup inside increases the value. For the last inequality, using the facts that (i) l0, (2, f) < Mt due to Me S {(z, h) S (z f) I supzez E(2, h) S Mc and (ii) the Rademacher complexity of the function class & defined by R, (8) Elsup c n, Σ TRO, (Zx: fo)] where the expecta- iid. ~ P. and i.i.d. Rademacher random variable on independent of Zx. Vk E [n.].
tion is w.r.t both Z
we have
2 llog(mm/)
sup di., E+
Ez - 1o, fo)] 12-. fo)] ( + V
(30)
with probability < d/m due to the standard symmetrization argument and McDiarmid's ' inequal- ity Shalev-Shwardd & Ben-Davie 2014, Theorem 26.5). Multiplying A; to both sides of (30
summing up the inequalities over all E [n]. and using union bound, we obtain 29
Define stochastic process (Xo.) 6,E0
X
σR fo)
Th
which is zero-mean because [X,,] = for all Φy E क. To upper-bound R, (Φ), we first show
that (Xa,) eo a sub-Gaussian process with respect to the following pseudometric
sup =e2 || o', 11 == sup (2. (x,2
(31)
For any t E R, using Hoeffding inequality with the fact that k to E [n], are i.i.d. bounded random variable with sub- Gaussian parameter 1, we have
E exp (1 (Xx, xxx))]
E
exp
ak (g) (7 fo) o (Zx, fo
k=1
T (0, (Z fo) o, (Z,, for )1 7R4
E
exp
Vni
t2
Φ o 11
< exp
2
20

---

Under review as a conference paper at ICLR 2022
Then, invoking Dudley entropy integral, we have
X S $12 l
Vn, (
sup 64
Vlog N (0, 11:1oo, e)
(32)
We will show that when 0 8(z, (()) is T Tipsiicc by Assumption 2 then + (2, f6) is also Lo Lipschih as follows.
fo) - 2r (2,
sup inf {E(C, h4) 7d((; P((', + (d(((==))
CEZ C'EZ
s sup CEZ
e(S, ho) ho) - hor
s sup CE2
((S. ho) E(S, hor)
s Loll0 0'|1,
which implies
0 - 11 S Lollo - 8'1l
Therefore, by contraction principle (Shalev-Shwarte & Ben-David 2014), we have
N ( loo S N 10.1l,,,,,/)))
(33)
Substituting 33 and
32
into
29
obtain
48E(()
2 2log(2m/0)
E7 fa.
2M
tE,
(34)
Vn,
n4
which will substituted into the upper-bound
Lemma]
1 to complete the proof.
F PROOF OF CORROLARY
We now present how we adapt the result from Fournier & Guillin Proposition 2 (Measure concentration (Fournier & Guillin 2015 Theorem 2)). Let P be a probabill-
2015) to prove Corollary !
ity distribution on a bounded set Z. Let P,, denote the empirical distribution of Zz,..., Zn itd. P. Assuming that there exist constants light-tail distribution). Then, for any > 0.
1 such that A == Ez-p [cxp([]
[1]] E2-p(cxp(112
(i.e. P is
P [WW,(Pas P) 2 S
P)
Ci exp (-conpuxx(//22))
if <
0i exp (-c2pp")
if >
where C1, C2 are constants depending on a, and d.
As consequence of this proposition, for any > 0, we have
W2 (2 P) < p 2 1 where p.. : {
log(cs/ /8)
min (2/d.1/2)
if log(er /8)
P
Can
C2
(35)
log(es/5) ngn
1/a
if log(er/8)
In Proposition 2 Fournier & Guillin (2015) show that the empirical distribution P. converges in Wasserstein distance to the true P at a specific rate. This implies that judiciously scaling the radius of Wasserstein balls according (35) provides natural confidence regions for the data-generating distribution P.
By the duality of transport cost (Santambrogio) 2015 p.261), we have
/
WF(,vv) =
sup
o + du
sup Tylm,v), Vp 2 1,
((e)+
21

---

Under review as a conference paper at ICLR 2022
which is the supremum of linear functionals T₁ P P + R defined by Ty(m.2)
((A, v), (p,, v));; therefore, (ju, v) + (p, 2)) is convex. Thus we have
(x PA) = W Σ
A Ar m. P.))
S
W, P.).
(36)
Then, we have
P[Wa(PA: (P)(PPPP))2 2 1 1 -
m6/m AiPr:
P W ( (A. PA) ス Σ
X
5/m
m
m
A22 P 14 P:) 2
Aapbl.
P ( P nu. P.) 2 Phim
P
W2 Pr.: P.) 2 P4.
1=1
8
(37)
2m
2'
i=1
where the first inequality is due to 36 the second inequality is due to the union bound, and the last inequality is due to Proposition 2 and 35
1/2
According to (28) by setting p -i >iPri complete the proof.
in Theorem and using union bound, we
G ADDITIONAL EXPERIMENTAL SETTINGS AND RESULTS
G.1 DATASETS
Table 1: Statistics of all datasets using in the experiments.
Total samples
Num labels
Samples client Mean
Dataset
client
Std
CIFAR-10 10
20
43,098
3
2154
593.8
MNIST
100
70,000
700
313.4
MNIST-M
100
70,000
2
700
322.4
We distribute all datasets to clients as follows:
MNIST: A handwritten digit dataset (Lecun al. 1998) including 70, 000 instances be- longed to 10 classes. We distribute dataset to m = clients and each client has a different local data size with only 2 of the 10 classes.
An object recognition dataset (Krizhevsky) 2009) including 60, 000 colored images belonged to 10 classes. We partition the dataset to m 20 clients and there are labels per client. Each client has different local data size.
CIFAR-10:
Three handwritten digit: MNIST, MNIST-M (Ganin & Lempitsky 2015). and USPS (Hull) 1994). We distribute
dataset, for example, MNIST to 100 source clients, and leave MNIST-M to the target client. Each source client has only 2 labels over 10 labels. As the number of data samples in USPS is relatively small amount (9,298 samples), we only use this dataset for the target client.
We standardize and randomly split all datasets with 75% and 25% for training and testing, respec- tively. In domain adaptation, the target client only has test data. The statistics of all datasets are summarized in Tab. 1
22

---

Under review as a conference paper at ICLR 2022
Distribution
Distribution
1.6
2.C
0.5
1.0
3/
52
182
150
202
750
200
22
3.0
150
200
CIPAR: Clean Data
CIFAR: Clean Data
CIFAR: Under Dietribution Shifts
CIFAR: Under Distribution Stifts
3.50
2.0
0.35
0.30
0.20
2.4
0.25
0.29
0.20
202
1.0
32
120
100
130
200
WAFL 0.505
Figure S: Convergence of WAFL.
G.2 MODELS
The details of models for each dataset is provided follows:
MNIST: We multinomial logistic regression model (MLR) with cross-entropy loss
function and an Le-regularization term.
CIFAR-10: We use a CNN model employed in McMahan et al. (2017).
Set of three handwriten digits: We use a multinomial logistic regression model (MLR) with cross-entropy loss function and an Iæ-regularizztion term.
In all settings, we randomly sample 10 clients to participate in training the global robust model in each communication round, and the number of local epochs K 2 and the number of communication rounds to 7 200. All experiments were conducted using PyTorch (Paszke et al.
2019
G.3 CONVERGENCE OF WAFL
We verify the convergence of WAFL under two cases: clean data (no attacked clients) and distri- bution shifts (where 40% of clients are attacked). In each case, we use two datasets: MNIST and CIFAR-10 and employ the same setup as in Sec, 6 Specifically, for MNIST, we distribute the dataset to 100 clients and set 7 = 0.05. For CIFAR- 10, we use 20 clients and set ? = 0.5. We use T = 200 communication iterations.
To show WAFL's 's convergence, we plot both the original loss (using the function ! and global accuracy in Fig.
G.4 DOMAIN ADAPTATION.
We show that WAFL demonstrates its superior capability in domain adaptation by collaboratively learning from a multi-soure domain Ps. and applying an unseen, but related, target domain Q. In federated domain adaptation (Peng et al. 2019 Liu et al. 2021), source domains often need access the private data of target domains find common representation to aid training. However, this violates the data privacy assumption of FL. By contrast, the construction of the model in WAFL does not involve such private data.
Our experimental set up is as follows. We consider three datasets: MNIST (mt), MNIST-M (mm) and USPS (up), which are widely used in domain adaptation. Among these datasets, consider
23

---

Under review as a conference paper at ICLR 2022
Table 2: Domain adaptation performance of WAFL and FedAvg accuracy the target dataset.
Algorithm
mf-mm
mi up
mm-
-----
Average
WAFL
40.23
66.94
62.04
79.27
62.12
FedAvg
37.71
66.08
58.60
75.80
59.55
dataset source domain (distributed to 100 clients), and another dataset target domain (assumed to be on one client). For example, if mt is used as the source and mm the domain, we mt- denote this case. Specifically, the mm dataset is distributed to 100 clients learn a global model, then the model is tested on the mi dataset. We compare WAFL against the vanilla FedAvg as baseline, and report the accuracy on the target dataset in several scenarios in Tab. 2
WAFL markedly improves from FedAvg in all cases. On average, the accuracy of WAFL is 2.5 per- centage point higher than that FedAvg. This indicates WAFL's benefits transferring knowledge from multi-sourced, private client distributions to new target distribution.
24


Reviews:
[{"review_id": "QH_oiYGIoBn", "reviewer": "Reviewer_M1m9", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "The paper is well-written and the concepts are explained in technical detail. While the ideas implemented are not surprising, appearing previously in Wasserstein Adversarial learning algorithms, the authors do a decent job of explicating the details of the method. The SGD algorithm via usage of duality properties is standard and the theoretical justification is solid. Overall, I would say this work is a reasonable contribution. I have some comments about the paper. \n\n1. Can you show simple cases of inference where Wasserstein robust optimization does better or specific signal/noise types?\n\n\n2. Suppose you have $n$ data points and you perform the WAFL method, does it give smaller number of components in terms of the maximizing argument (Q in this case). I feel that invoking Wasserstein robustness should provide a more interpretable classifier. Is that intuition justified? If so, is that confirmable through some experimental validation or otherwise?\n\n3. Please provide the computational time comparison with relevant methods.\n", "summary_of_the_paper": "The paper proposes a Wasserstein robust Distributional optimization scheme to provide robust approach for empirical risk minimization. Generalizing the concepts of Agnostic Federated Learning, the method finds applicability in domain adaptation as well as big data settings. The authors also propose and SGD algorithm to solve the problem described and provides a theoretical upper bound for the resultant estimates.", "summary_of_the_review": "Overall a solid piece of work. Results are not surprising but are reasonable. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3AeDaWjJeuL", "reviewer": "Reviewer_g9vL", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "The problem under investigation is important. However, I do not feel this paper made sufficient novel contributions to either the field of WDRO or the field of FL. The duality result (5) used by this paper was not new. The formulation (4) is the canonical form of a WDRO problem, and thus I did not see the difficulty of solving it. \n\nOne of my major concerns is in (6), where the authors did not solve \\gamma to optimality, but instead considered it as a hyper-parameter. This would break strong duality, meaning that what Algorithm 1 solved is not the original WDRO problem (4), but just a relaxed version, and we do not know the tightness of such a relaxation. Can you explain the rationale of not solving \\gamma to optimality?\n\nAlgorithm 1 does not seem to be very different from FedAvg, except that the original loss is replaced by the surrogate loss. \n\nAlso, at the end of Section 2, the authors mentioned several related works exploring Wasserstein robustness in FL. It would be good to expand a bit on what is the difference between this work and those previous works, what makes your paper novel and what are your contributions. In the experimental section, it would be good to see whether your approach improves upon those past works.\n\nThe experiments do not seem to be convincing either. FedAvg itself is not a robust algorithm, and thus I do not think comparing with it gives any insights on the robustness of the proposed approach. The improvement of WAFL starts to be significant only when the proportion of attached clients is very high (~80%), which puts the benefits of using WAFL in doubt. \n\nMinor points: how did you choose \\lambda_i? Did you set it to n_i/n?", "summary_of_the_paper": "This paper presents a Wasserstein distributionally robust optimization (WDRO) framework for federated learning to hedge against statistical heterogeneity. The authors utilized the duality result in previous WDRO works to make their formulation tractable, and proposed a distributed algorithm to solve a relaxed WDRO problem. They conducted several numerical experiments on MNIST and CIFAR-10 to compare their method with FedAvg (non-robust algorithm), FedPGM and FedFGSM (robust algorithms), and demonstrated the advantage of their approach when the proportion of attacked clients is over 80%.", "summary_of_the_review": "Overall I do not think this paper made sufficient novel contributions to the field of WDRO and FL. Several theoretical results are already well-known by the WDRO community, and the proposed algorithm is not very different from FedAvg. The numerical experiments also do not look very convincing. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-ab3aJkBWH", "reviewer": "Reviewer_1hSQ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "**Major commments:**\n\n- The theoretical part of WAFL is systematic and well-written, with analysis from both optimization and statistics aspects.\n\n- My biggest concern lies in the motivation/marginal contribution of WAFL against existing agnostic FL framework such as (Mohri et al., 2019); (Deng et al., 2020b). On page 4, the paper compares WAFL with counterparts in terms of uncertainty set, and conclude that Wasserstein ball is able to cover the uncertainty set of agnostic FL. However, a larger uncertainty ball may not be an advantage, it will also leads to over-conservative decisions. Therefore, a crucial question for the authors is: for what types of application can WAFL outperform existing method? If so, is there any theoretical or empirical evidence? \n\n- The paper didn't report a systematic way to choose $\\gamma$ in the numerical experiment part, the choices of $\\gamma = 0.5$ for MNIST and $\\gamma = 10$ for CIFAR-10 are not well-supported. \n\n- For the numerical part, it would be better to include (Mohri et al., 2019) or (Deng et al., 2020b) as benchmark as well. This would be helpful to clarify my first point.\n\n**Minor Comments:**\n\n- It is claimed on page 3 that ``In FL, some existing works have explored the Wasserstein distance to enhance robustness (Rei- sizadeh et al., 2020; Diamandis et al., 2021; Du et al., 2020; Deng et al., 2020b).`` Though Du et al., 2020 and Deng et al., 2020b are examples of distributional robust FL, I can not find their connection to Wasserstein distance. Did I miss something here?", "summary_of_the_paper": "The paper proposes a Wasserstein-based distributionally robust federated learning (WAFL) framework to address the statistical heterogeneity problem, which is also applicable to distributional shift setting and domain adaptation. Inspired by (Sinha et al., 2020), the WAFL problem is reformulated into the dual form, and the Lagrangian ($\\gamma$) is chosen as a fixed value reflecting the algorithm's conservatives. The paper develops a local SGD scheme to solve the problem and provides a theoretical convergence guarantee under a set of standard assumptions. In addition, the paper also provides a generalization bound for the proposed algorithm by resorting to standard techniques such as covering number, entropy integral, and empirical convergence rate of Wasserstein distances. Finally, the results of several numerical experiments are reported to justify the performance of the algorithm.", "summary_of_the_review": "In general, the paper is well-written. It provides a comprehensive study of the proposed framework. However, in my humble view, the motivation of the paper is not well-supported, and the comparison against existing methods is insufficient. Therefore, I rate this paper as a marginal paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The main motivation of this work is to introduce robustness in federated learning, through a Wasserstein uncertainty set. The end result, however, leaves a mixed feeling: As the reviewers pointed out, the authors, perhaps for computational convenience, forgo strong duality and treat the important variable gamma as a hyperparameter, which renders large part of the work follow immediately from existing work: essentially, we simply use a different loss function in FedAvg. While there may be advantages to choose one loss over another in any specific application, this itself is not a significant contribution. The comparison against existing FL algorithms is also a bit weak: Despite of the reviewer's request, the authors did not compare to other robust FL algorithms (e.g., AFL), thus it is not clear what is the real advantage of the proposed algorithm. As a result, we believe the current draft is not ready for publication.
Decision: Reject