Year: 2021

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.

heading: 1 INTRODUCTION
text: Federated Learning (FL) (Konečný et al., 2016; McMahan et al., 2017) has emerged as a cuttingedge technique in distributed and privacy-preserving machine learning. The nature of highly nonindependent and identically distributed (non-i.i.d.) data in clients’ devices poses an important challenge to FL commonly called statistical heterogeneity. The global model trained on this data using the de facto FedAvg algorithm (McMahan et al., 2017) has been shown to generalize poorly to individual clients’ data, and further to unseen distributions on new clients as 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 et al., 2018) are client-adaptive approaches, where a personalized model is adapted to each client. From another perspective, distributionally robust FL trains a model using a worst-case objective over an ambiguity set (Mohri et al., 2019; Du et al., 2020; Reisizadeh et al., 2020; Deng et al., 2020b). This approach is client-uniform because a single global model is judiciously learned to deliver uniformly good performance not only for all training clients but also for new/unseen clients with unknown data distributions. It 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 error also bound the test error (generalization bounds)? In FL, Mohri et al. (2019) proposed agnostic FL where a model is designed to be robust against any ambiguity set as a 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 et al. (2020) use the margin-based technique developed by Bartlett et al. (2017).
In this work, we take a different approach called WAsserstein distributionally robust FL (WAFL for short). The ambiguity set in WAFL is a Wasserstein ball of all adversarial distributions in close proximity to the nominal data distribution at the center. Our main contributions are:
• We propose a distributionally robust optimization problem to address statistical heterogeneity in FL. By controlling the center and radius of the Wasserstein ball, we show that WAFL is robust to a wider range of adversarial distributions than is agnostic or adversarial FL.
• To make WAFL more amenable to distributed optimization, we transform the original problem into a minimization of the empirical surrogate risk. We propose a local SGD-based algorithm to solve this surrogate problem. With additional Lipschitz smoothness conditions, standard techniques can be 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 is applicable to the true (unknown) data distribution among all clients.
• We show WAFL’s extra flexibility in 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, we discuss how to control this radius for a given center location by finetuning 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 settings. 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.

heading: 2 RELATED WORK
text: Federated Learning was introduced in response to three challenges of machine learning at scale: massive data quantities at the edge, communication-critical networks of participating devices, and privacy-preserving learning without central data storage (Konečný et al., 2016; McMahan et al., 2017). The de facto federated optimization algorithm – FedAvg (McMahan et al., 2017) – is based on local stochastic gradient descent (SGD) and averaging and is often considered a baseline in FL.
Most challenges of FL are 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-efficient methods (Konečný et 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 on any data distribution. To address this, many ideas have been introduced. Li et al. (2020) provided much theoretical analysis of FL non-i.i.d. settings. Zhao et al. (2018) proposed an FL framework which globally shares a small subset of data among clients to train the model with non-i.i.d. 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. (2020); Deng et al. (2020a); Dinh et al. (2021); Li et al. (2021).
Wasserstein Distributionally Robust Optimization (WDRO) aims to lean a robust model against adversarially manipulated data. The unknown data distribution is assumed to lie within a Wasserstein ball centered around the empirical distribution (Kuhn et 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 certain machine learning tasks (Shafieezadeh Abadeh et al., 2015; Gao & Kleywegt, 2016; Esfahani & Kuhn, 2017; Chen & Paschalidis, 2018; Sinha et al., 2020; Blanchet et al., 2019; ShafieezadehAbadeh et al., 2019; Gao et al., 2020). For instance, Shafieezadeh Abadeh et al. (2015) considered a robust logistic regression model under the assumption that the probability distributions lie in a Wasserstein ball. Chen & Paschalidis (2018); Blanchet et al. (2019); Gao et al. (2020) leveraged WDRO to recover regularization formulations in classification and regression. Gao & Kleywegt (2016) proposed a minimizer based on a tractable approximation of the local worst-case risk. Esfahani & 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 certificate based on a Lagrangian relaxation of the loss function which provably robust against adversarial input distributions within a Wasserstein ball centered around the original input distribution.
In FL, only a few works have explored the Wasserstein distance (Reisizadeh et al., 2020; Diamandis et al., 2021). Specially, Reisizadeh et al. (2020) proposed FedRobust based on adversarial robust training to enhance robustness. Regarding distributionally robust learning, Deng et al. (2020b) proposed DRFA, a communication efficient distributed algorithm. Besides, based on the agnostic FL framework suggested by Mohri et al. (2019), Du et al. (2020) introduced AgnosticFair, a two-player adversarial minimax game between the learner and the adversary, to achieve fairness.

heading: 3 WASSERSTEIN ROBUST FEDERATED LEARNING

heading: 3.1 EXPECTED RISK AND EMPIRICAL RISK MINIMIZATION IN FEDERATED LEARNING
text: In a federated setting, there are m clients, and each client i ∈ [m] := {1, . . . ,m} has its data generating distribution Pi supported on domain Zi := (Xi,Yi). Consider the parametrized hypothesis class H = { hθ | θ ∈ Rd } , where each member hθ is a mapping from Xi to Yi parametrized by θ. With zi := (xi, yi) ∈ Zi, we use `(zi, hθ), shorthand for `(yi, hθ(xi)), to represent the cost of predicting hθ(xi) when the ground-truth label is yi. For example, if hθ(xi) = θTxi and yi ∈ R, a square loss `(zi, hθ) = `(yi, hθ(xi)) = (θTxi−yi)2 can be considered. In FL, all clients collaborate with a server to find a global model θ such that the following sum weighted risk is minimized:
min θ∈Rd m∑ i=1 λiEZi∼Pi [ `(Zi, hθ) ] , (1)
where EZi∼Pi [ `(Zi, hθ) ] is client i’s expected risk and λi ≥ 0 represents the relative “weight”
of client i and ∑m i=1 λi = 1. Therefore, λ := [λ1, . . . , λm]
> belongs to a simplex ∆ :={ λ ∈ Rm : λ < 0 and λ>1m = 1 } . Define by Pλ := ∑m i=1 λiPi the mixed clients’ distribution over m domains Z := {Z1, . . . ,Zm}. We denote by Z ∼ Pλ a random data point Z generated by Pλ, which means that the domain of client i is chosen with probability P(Z = Zi) = λi first, then a data point zi ∈ Zi is selected with probability P(Zi = zi), Zi ∼ Pi. While the underlying distributions Pi are unknown, clients have access to finite observations zi ∈ [ni] ∼ Pλ. We abuse the notation [ni] to denote the set of client i’s both observable data points and their indexes. Let P̂ni := 1 ni ∑ zi∈[ni] δzi be the empirical distribution of Pi, where δzi is the Dirac
point mass at zi. In general, we use the notation ̂ for quantities that are dependent on the training data. Define by P̂λ := ∑m i=1 λiP̂ni the mixed empirical distribution of n = ∑m i=1 ni training data from m clients. The empirical risk minimization (ERM) problem of (1) is as follows:
min θ∈Rd
{ EZ∼P̂λ [ `(Z, hθ) ] = m∑ i=1 λiEZi∼P̂ni [ `(Zi, hθ) ] = m∑ i=1 ni n ( 1 ni ∑ zi∈[ni] `(zi, hθ) )} , (2)
where λi = ni/n is often chosen in ERM of the standard FL (McMahan et al., 2017).

heading: 3.2 WASSERSTEIN ROBUST RISK IN FEDERATED LEARNING
text: Models resulting from (2) have been shown to be vulnerable to 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 Z , define d : Z × Z → [0,∞) to be the cost of “transportation” between its two points.1 Suppose P and Q are two distributions on Z . Let Π(P,Q) be the set of probability measures π on Z ×Z whose marginals are P and Q, called their couplings. In other words, π(A,Z) = P (A) and π(Z, A) = Q(A),∀A ⊂ Z . The p-Wasserstein distance between P and Q is defined as
Wp(P,Q) = inf π∈Π(P,Q)
( E(Z,Z′)∼π [ dp(Z,Z ′) ])1/p . (3)
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 Bp(P, ρ) := {Q : Wp(P,Q) ≤ ρ}
1The function d must satisfy non-negativity, lower semi-continuity and d(z, z) = 0, ∀z ∈ Z .
denote the Wasserstein ball centered at P (i.e., nomimal distribution) and having radius ρ ≥ 0. We modify (2) into the following Wasserstein robust risk minimization in FL problem:
WAFL: min θ∈Rd
{ supQ∈B(P̂λ,ρ) EZ′∼Q [ `(Z ′, hθ) ]} . (4)
There are several merits to this framework. First, the ambiguity set Bp(P̂λ, ρ) contains all (continuous or discrete) distributions Q that can be converted from the (discrete) nominal distribution P̂λ at a bounded transportation cost ρ. Second, Wasserstein distances can be approximated from the samples. Based on the non-asymptotic convergence results of Fournier & Guillin (2015), we can specify a suitable value for ρ to probabilistically bound Wp(P,Q) by the distance between their empirical distributions Wp(P̂ , Q̂) (e.g., for multi-source domain adaptation).
In any robust optimization problem, the ambiguity set is a key ingredient to defining the level of robustness. We will compare WAFL in (4) with other approaches in terms of their ambiguity set.
Agnostic FL. Using this approach, existing techniques (Mohri et al., 2019; Deng et al., 2020b) minimize the worst-case loss
max λ∈∆ EZ∼P̂λ
[ `(Z, hθ) ] ,
hence its distributional ambiguity set is Q∆ := { P̂λ : λ ∈
∆ } . While Agnostic FL’s ambiguity set is the static convex
hull of { P̂ni } i∈[m], WAFL’s ambiguity set B(P̂λ, ρ) can be adjusted by controlling the robustness level ρ and by positioning the ball center using λ, which is useful for domain adaptation. Therefore, B(P̂λ, ρ) can cover Q∆ by using appropriate values for ρ and λ (see Fig. 1).
Adversarial robust FL. Using this approach, Reisizadeh et al. (2020) combines a general affine covariate shift in standard adversarial robust training with FL. Most existing techniques using this approach (Goodfellow et al., 2015; Kurakin et al., 2017; Carlini & Wagner, 2017; Madry et al., 2019; Tramèr et al., 2020) define an adversarial perturbation u at
a data point Z, and minimize the worst-case loss over all perturbations: maxu∈U EZ∼P̂λ
[
`(Z +
u, hθ) ] , where the ambiguity set is U := { u ∈ Rd+1 : ‖u‖ ≤ } . 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(P̂λ, ρ)-worst-case risk) has its dual formulation as follows (Gao & Kleywegt, 2016; Sinha et al., 2020)
sup Q∈B(P̂λ,ρ)
EZ′∼Q [ `(Z ′, hθ) ] = inf γ≥0 { γρ2 + EZ∼P̂λ [ φγ(Z, θ) ]} , (5)
where φγ(zi, θ) := supζ∈Z [ `(ζ, hθ) − γd2(ζ, zi) ] , and d2(z, z′) = ‖x − x′‖2 + κ|y − y′|2, κ > 0. The crux of using the dual is that the inner supremum problem (finding φγ) is easily solvable when its objective is well-conditioned: if ` is L-smooth and d is 1-strongly convex, setting γ > L ensures that ζ 7→ `(ζ, hθ)− γd2(ζ, z) is strongly concave, and using gradient ascent for the inner supremum problem (for finding φγ) ensures linear convergence. For each zi let z∗i be the solution to supζ∈Z [ `(ζ, hθ)− γd2(ζ, zi) ] . Then,∇θφγ(zi, θ) = ∇θ`(z∗i , hθ) (Sinha et al., 2020, Lemma 1).
Rather than find the optimal γ in (5), we consider γ as a control parameter, and instead solve the following client-decomposable problem, which is amenable to distributed algorithm design.
min θ∈Rd
{ EZ∼P̂λ [ φγ(Z, θ) ] = m∑ i=1 λiEZi∼P̂ni [ φγ(Zi, θ) ]} . (6)
Algorithm 1: Local SGD for WAFL 1: for t = 0, 1, ..., T − 1 do // Global rounds 2: Sample a subset of clients St ⊂ [ m ]
3: for client i ∈ St in parallel do 4: Set local parameters: θ(t,0)i = θ
t // Communication 5: for k = 0, 1, ...,K − 1 do // Local rounds 6: Sample a mini-batch Di from client i’s dataset 7: θ(t,k+1)i = θ (t,k) i − η |Di| ∑ zi∈Di ∇θφ(zi, θ(t,k)i ) 8: Send θ(t,K)i to server // Communication 9: Server: update θ(t+1) = ∑ i∈St λiθ (t,K) i / ∑ i∈St λi
This motivates the development of Algorithm 1 for solving (6). The structure of WAFL is similar to FedAvg with T communication rounds. We further note three key components of Algorithm 1 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 φγ , rather than `.
We show that the convergence of WAFL can be similarly characterized as that of FedAvg, the de facto FL algorithm based on local SGD updates (McMahan et al., 2017). In FedAvg optimization, we seek to establish the convergence when using the original loss function `. On the other hand, in WAFL the convergence is with respect to the surrogate loss φγ , through which the local and global risks are defined by Fi(θ) := EZi∼Pi [ φγ(Zi, θ) ] and F (θ) := EZ∼Pλ [ φγ(Z, θ) ] , respectively.
We first make the following assumptions, common to analyses of Wasserstein-robust optimization (Sinha et al., 2020). Unless stated otherwise, all norms are the Euclidean norm. Assumption 1. The function d : Z × Z → R+ is continuous, and d(·, z0) is 1-strongly convex, ∀z0 ∈ Z . Assumption 2. The loss function ` : Z ×H → R is Lipschitz continuous as follows
(a)‖`(z, hθ)− `(z′, hθ)‖ ≤ Lz‖z − z′‖, (b)‖`(z, hθ)− `(z, hθ′)‖ ≤ Lθ‖θ − θ′‖.
Assumption 3. The loss function ` : Z ×H 7→ R is Lipschitz smooth as follows
‖∇θ`(z, hθ)−∇θ`(z, hθ′)‖ ≤ Lθθ‖θ − θ′‖, ‖∇z`(z, hθ)−∇z`(z′, hθ)‖ ≤ Lzz‖z − z′‖, ‖∇θ`(z, hθ)−∇θ`(z′, hθ)‖ ≤ Lθz‖z − z′‖, ‖∇z`(z, hθ)−∇z`(z, hθ′)‖ ≤ Lzθ‖θ − θ′‖.
Given Assumption 3, it has been shown that the mapping θ 7→ φγ(·, θ) is L-smooth with L = Lθθ +
LθzLzθ γ−Lzz , γ > Lzz (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 ∇Fi(θ), denoted by gφi(θ) := ∇θφγ(zi, θ), zi ∼ Pi, has σ2-uniformly bounded variance, i.e., E [ ‖gφi(θ)−∇Fi(θ)‖2 ] ≤ σ2.
Assumption 5. The difference between the local gradient ∇Fi(θ) and the global gradient ∇F (θ) is Ω-uniformly bounded, i.e., maxi supθ‖∇Fi(θ)−∇F (θ)‖ ≤ Ω.
Assuming complete participation of clients in every round (St = m,∀t), using standard techniques in Wang et al. (2021), we have: Theorem 1 (WAFL’s convergence for convex loss function). Let Assumptions 1–5 hold and the mapping θ 7→ `(z, hθ) be convex. Denote by θ̄(t,k) the “shadow” sequence, defined as θ̄(t,k) =∑m i=1 λiθ (t,k) i and by θ ∗ the optimal solution to minθ∈Rd F (θ). If the client learning rate satisfies
η ≤ min { 1
3L ,
D
2 √ KTΛσ̂
, D
2 3
48 1 3K 2 3T 1 3L 1 3 Ω̄ 2 3
, D
2 3
40 1 3KT 1 3L 1 3 Ω̄ 2 3
} ,
then we have E [ 1
KT T−1∑ t=0 K−1∑ k=0 F (θ̄(t,k))− F (θ∗) ] ≤ O ( LD2 KT + σ̂DΛ 1 2 √ KT + L 1 3 Ω̄ 2 3D 4 3 K 1 3T 2 3 + L 1 3 Ω̄ 2 3D 4 3 T 2 3 ) , (7)
where σ̂2 := σ 2 |Di| , D := ‖θ (0) − θ∗‖ and Λ := ∑m i=1 λ 2 i .
We provide a proof of Theorem 1 in Appendix C.

heading: 4 ROBUST GENERALIZATION BOUNDS
text: We show the generalization and robustness properties of WAFL’s output by bounding its excess risk. Denote the loss class by F := ` ◦ H := {z 7→ `(z, h), h ∈ H}, where we use f (resp. fθ) ∈ F to represent a generic loss (resp. a loss function parametrized by θ). Definition 1. Denote the expected risk and surrogate of Wasserstein robust risk of an arbitrary f , respectively, as follows
L(Pλ, f) := EZ∼Pλ [ `(Z, h) ] and Lγρ(Pλ, f) := EZ∼Pλ [ φγ(Z, f) ] + γρ2.
Then their excess risks are defined respectively as follows
E(Pλ, f) := L(Pλ, f)− inf f ′∈F L(Pλ, f ′) and Eγρ(Pλ, f) := L γ ρ(Pλ, f)− inf f ′∈F Lγρ(Pλ, f ′).
If a distribution Q is in the ambiguity set B(Pλ, ρ), we can bound its excess risk E(Q, f) as follows. Lemma 1. Let Assumption 2 (a) holds and γ ≥ Lz/ρ. For all f ∈ F and for all Q ∈ B(Pλ, ρ),
Eγρ(Pλ, f)− g(ρ, γ) ≤ E(Q, f) ≤ Eγρ(Pλ, f) + g(ρ, γ),
where g(ρ, γ) := 2Lzρ+ |γ − γ∗|ρ2, and γ∗ := arg minγ′≥0 Lγ ′ ρ (Pλ, f).
We provide a proof of Lemma 1 in Appendix D. Remark 1. Lemma 1 shows that the lower and upper bounds for E(Q, f) can be analyzed using Eγρ(Pλ, f) and a two-component error term g(ρ, γ). The first component, 2Lzρ, says that when ρ is increased – to allow for larger Wasserstein distance between the nominal Pλ and any worst-case distribution Q – the difference between the excess risks E(Q, f) and Eγρ(Pλ, f) increases, and this error is amplified at most by the Lipschitz constant Lz of the mapping z 7→ `(z, ·). The second component, |γ−γ∗|ρ2, addresses the sub-optimality error of a chosen value of γ, which is amplified when γ is drifted away from the optimal γ∗. Note that Lγ ∗
ρ (Pλ, f) is the same as B(Pλ, ρ)-worstcase risk thanks to the strong duality (5), which is obtained with ρ > 0. Denote by θ̂ ∈ Θ an ε-minimizer to the surrogate ERM, i.e., EZ∼P̂λ [ φγ(Z, fθ̂ε) ] ≤ infθ∈Θ EZ∼P̂λ [ φ(Z, fθ) ] + ε, where Θ ⊂ Rd is a parameter class, we obtain the following result.
Theorem 2 (Robust Generalization Bounds). Let Assumptions 2 and 3 hold, γ ≥ max { Lzz, Lz/ρ } , and |`(z, h)| ≤M`. We have the following result for all Q ∈ B(Pλ, ρ)
E(Q, fθ̂ε) ≤ m∑ i=1 λi [ 48C(Θ) √ ni + 2M` √ 2 log(2m/δ) ni ] + ε+ g(ρ, γ)
with probability at least 1−δ, where C(Θ) := Lθ ∫∞
0
√ logN (Θ, ‖ · ‖Θ, )d andN (Θ, ‖ · ‖Θ, )
denotes the -covering number of Θ w.r.t a metric ‖ · ‖Θ as the norm on Θ.
We provide a proof of Theorem 2 in Appendix E. We sketch the proof as follows: first bound Eγρ(Pλ, fθ̂ε) using standard excess risk decomposition and uniform convergence with Rademacher complexity. Then leverage the upper-bound of Lemma 1 to bound E(Q, f),∀Q ∈ B(Pλ, ρ), based on the bound of Eγρ(Pλ, fθ̂ε). The result shows that using WAFL to minimize the surrogate of Wasserstein robust empirical risk also controls the robustness and generalization. As an example
H = { 〈θ, ·〉 , θ ∈ Θ } with Θ = { θ ∈ Rd : ‖θ‖2 ≤ C } . The diameter of Θ is supθ,θ′∈Θ‖θ − θ′‖ = 2C, thus N (Θ, ‖ · ‖2, ) = (1 + 2C/ )d, and C(Θ) ≤ 3CLθ √ d/2 (Lee & Raginsky, 2018).
Generally, the radius of Wasserstein ball ρ can be considered hyperparameter that needs fine-tuning (e.g., through cross-validation). In principle, ρ should not be too large to become over-conservative, which can hurt the empirical average performance, but also not too small to become similar to the ERM, and thus can lack robustness. From a statistical standpoint, we are interested in learning how to scale ρ w.r.t. the sample size ni, i ∈ [m], such that the generalization of the WAFL solution θ̂ε w.r.t the true distribution Pλ 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 P̂ni converges in Wasserstein distance to the true Pi at a specific rate, we obtain the following
Corollary 1. With all assumptions as in Theorem 2, defining ρn := √∑m i=1 λiρ̂ δ/m ni , we have
E(Pλ, fθ̂ε) ≤ m∑ i=1 λi [ 48C(Θ) √ ni + 2M` √ 2 log(4m/δ) ni ] + g(ρn, γ) + ε
with probability at least 1− δ, where ρ̂δn :=  ( log(c1/δ) c2n )min{2/d,1/2} if n ≥ log(c1/δ)c2 ,(
log(c1/δ) c2n )1/α if n < log(c1/δ)c2 .
We provide a proof of Corollary 1 in Appendix F.
5 CHOOSING λ: APPLICATIONS
We focus on two applications: multi-source domain adaptation and generalization to all client distributions. We provide insights about choosing the client weights λ for these applications.
Multi-source domain adaptation: Consider the multi-source domain distribution Pλ (Mansour et al., 2021). Lee & Raginsky (2018) show that solving the minimax risk with the Wasserstein ambiguity set can help transfer data/knowledge from the source domain Pλ to a different, but related, target domain Q. They bound the distance Wp(Pλ, Q) using the triangle inequality
Wp(Pλ, Q) ≤Wp(Pλ, P̂λ) +Wp(P̂λ, Q̂) +Wp(Q̂,Q), (8) where P̂λ and Q̂ are the empirical versions of Pλ and Q, respectively. While Wp(Pλ, P̂λ) and Wp(Q̂,Q) can be probabilistically bounded with a confidence parameter δ ∈ (0, 1) according to Fournier & Guillin (2015), Wp(P̂λ, Q̂) can be deterministically computed using linear or convex programming (Peyré & Cuturi, 2019). In FL context, in order to have a better bound for W2(Pλ, Q) similar to (8), it is straightforward to choose λ = arg minλ′∈∆W2(P̂λ′ , Q̂). To relax this problem into a form solvable using existing approaches, observe that W2(P̂λ, Q) ≤ ∑m i=1 λiW2(P̂ni , Q) due to the convexity of Wasserstein distance. We then consider the following upper-bound to minλ∈∆W2(P̂λ, Q̂):
min λ∈∆ ∑m i=1 λiW2(P̂ni , Q̂) =: ρ ?, (9)
which is a linear program, considering each W2(P̂ni , 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 λ?, and assume that domain Q generates nQ i.i.d. data points. With probability at least 1− δ, we have
W2(Pλ? , Q) ≤W2(Pλ? , P̂λ?) +W2(P̂λ? , Q̂) +W2(Q, Q̂) ≤ √∑m
i=1 λ?i ρ̂ δ/m ni + ρ ? + ρ̂δ/2nQ .
The proof of this corollary is similar to that of Corollary 3 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 is applicable to all clients’ distributions. We show in Appendix B that this is a problem of finding λ such that the Wasserstein distance between Pλ and Pj ,∀j, is as small as possible.

heading: 6 EXPERIMENTS
text: We first show how to change the level of worst-case perturbations by varying the robust parameter γ. To show the generalizability and robustness of WAFL, we evaluate WAFL under non-i.i.d. 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.i.d. settings. We distribute the first dataset – MNIST (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 McMahan et al. (2017) to model a non-convex setting. We set λi = ni/n as the client weights. In optimization, we randomly sample St = 10 clients to participate in training at each communication round. More detail can be found in Appendix G.1.
Effect of γ on the worst-case risk perturbations. We first examine the relationship between the robust parameter γ and the average worst-case perturbations ρ̂, defined as ρ̂2 = EZ∼P̂λ [ d2(Ẑ, Z) ] , where Ẑ is the adversarial perturbation of Z. As shown in Fig. 2, smaller values of γ correspond to larger worst-case perturbations ρ̂ on both the MNIST and CIFAR-10 datasets. For the rest of the experiment, rather than control ρ̂ directly, we set γ on the opposite direction to control the level of robustness.
Generalization and robustness of WAFL. We consider P̂λ and Q̂ as the empirical distribution of the training samples and test samples of all clients, respectively. By varying γ, 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 γ 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 are distribution shifts at the clients’ test data. To obtain these shifts, we use the common PGD attack method in Madry et al. (2019) under l∞-norm to generate an -level perturbation on clients’ test data. We choose the l∞-norm as it shows benefits in adversaries and gives large perturbations. Following Madry et al. (2019), we fix the number of gradient steps to generate adversarial examples with tavd = 40, = 0.3, α = 0.01 for MNIST and tavd = 10, = 8/255, α = 2/255 for CIFAR-10 with a batch size of 64. We note that this setting is similar to the adversarial poisoning attacks and the main purpose of this setting is to increase the Wasserstein distance between P̂λ and Q̂, thereby verifying the robustness of WAFL.
The performance of WAFL in both scenarios on MNIST and CIFAR-10 is shown in Fig. 3. When the clients’ data is clean, the Wasserstein distance between P̂λ and Q̂ is relatively small, and training WAFL with small γ (large ρ̂) gives the worse performance on the test set. With a sufficiently large γ, WAFL is less robust and has the same generalization with FedAvg. By carefully fine-turning γ in the range [0.5, 1] for MNIST and [10, 20] for CIFAR-10, WAFL shows an improvement over FedAvg. γ in 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 varying γ, we train a global robust model with different levels of worst-case perturbation to handle the distribution shifts. Small values of γ generate larger ambiguity sets B(P̂λ, ρ̂) and increase the robustness of WAFL, thus increase the chance Q̂ lies inside B(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 (γ ≤ 0.01 or γ ≥ 1 for MNIST and γ ≤ 0.1 or γ > 10 for CIFAR-10), WAFL performs inefficiently: too small γ hurts the empirical performance, while too large γ makes WAFL lack robustness. For every scenario, γ needs to be tuned correspondingly. In our experiments, by carefully choosing γ, WAFL not only handles distribution shifts or common data poisoning attacks but also provides better performance than FedAvg in non-i.i.d. data and heterogeneous settings. Specifically, we set γ = 0.5 for MNIST and γ = 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 one non-robust method Fe-
dAvg. FedPGM and FedFGSM are FedAvg with adversarial training using the projected gradient method PGD (Madry et al., 2019) and the fast-gradient method FGSM (Goodfellow et al., 2015), respectively. In FedPGM and FedFGSM, in each local update, all clients solve δ∗ = arg max‖δ‖∞≤ { `(hθ(z + δ), y) } using projection onto an l∞-norm to find the worst-case perturbation δ. While FedPGD uses tavd gradient steps to find δ∗, FedFGSM uses only one gradient step. We use the same values of and α from the distribution shifts setting for projection. For a fair comparison, both WAFL and FedPGM have the same value tavd and all algorithms have the same number of local updates K. We also train WAFL with the value of γ generating the same level perturbation of in FedPGM and FedFGSM. We study different proportions of clients having distribution shifts in Fig. 4. The robust accuracy of all methods decreases when the percentage of clients having distribution shifts increases. Especially, the FedAvg’s performance drops dramatically 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.

heading: 7 CONCLUSION
text: 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 a local SGD-based 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 to unseen data distributions.

heading: A ADVERSARIAL ROBUST FL’S AMBIGUITY SET V.S. WASSERTEIN BALL
text: 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 et al., 2019; Tramèr et al., 2020) define an adversarial perturbation u at a data point Z, and minimize the following worst-case loss over all possible perturbations
max u∈U EZ∼P̂λ
[ `(Z + u, hθ) ] , (10)
where the ambiguity set U := { u ∈ Rd+1 : ‖u‖ ≤ } . To compare this approach with Wassersteinrobust FL, we relate the above problem to its counterpart defined in the probability space of input as follows
max Q̃∈Q( ) EZ̃∼Q̃
[ `(Z̃, hθ) ] , (11)
where Q( ) := { Q̃ : P [ ‖Z̃ − Z‖ ≤ ] = 1, Z ∼ P̂λ, Z̃ ∼ Q̃ } . Considering u∗ as a solution to
problem (10), we see that the distribution Q′ of perturbation points (i.e., Z̃ := (Z + u∗) ∼ Q′) in problem (10) belongs to the feasible set Q( ) in problem (11) (If not, then P [ ‖u∗‖ ≤ ] < 1, a
contradiction). Next, consider an arbitrary distribution Q̃ ∈ Q( ) in problem (11), with any Z̃ ∼ Q̃ and Z ∼ P̂λ, we have
‖Z̃ − Z‖ w.p.1 ≤ =⇒ EZ∼P̂λ,Z̃∼Q̃
[ ‖Z̃ − Z‖ ] ≤ =⇒ inf
π∈Π(P̂λ,Q̃) E(Z,Z′)∼π
[ ‖Z̃ − Z‖ ] ≤ ,
which implies that W1(P̂λ, Q̃) ≤ ,∀Q̃ ∈ Q( ), and thus Q( ) ⊂ B1(P̂λ, ). We have shown that the Wasserstein ambiguity set contains the perturbation points induced by the solution to the adversarial robust training problem (10).
B CHOOSING λ: GENERALIZING TO ALL CLIENT DISTRIBUTIONS
We show that by calibrating appropriate λ value, our proposed algorithm will be capable of generalizing 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 Pλ and Pj ,∀j, is as small as possible. Instead of directly finding the minimum Wasserstein radius that cover all client distributions, we will leverage the popular Wasserstein barycenter problem (Peyré & Cuturi, 2019). Specifically, consider the problem
min λ∈∆
W2(P̂λ, P ♥) s.t. P♥ = arg min
Q∈P
∑m i=1 λiW2(P̂ni , Q), (12)
where P♥ is the Wasserstein bary center w.r.t the solution λ∗ to this problem. Even though the solution is not straightforward, we propose to solve its tractable upper-bound:
min λ∈∆,Q∈P ∑m i=1 λiW2(P̂ni , Q). (13)
This is a bi-convex problem, which is convex w.r.t to λ (resp. Q) when fixing Q (resp. λ). Thus, we can use alternative minimization (Gorski et al., 2007) to find a local solution to this problem. Denoting λ̃ as the solution to (12) and (λ∗, P ∗) as a local solution to (13), we obtain
W2(P̂λ̃, P ♥) ≤W2(P̂λ∗ , P ∗) ≤ ∑m i=1 λ∗iW2(P̂ni , P ∗) =: ρ∗. (14)
Corollary 3. For all client j ∈ [m], with probability at least 1− δ, we have
W2(Pλ∗ , Pj) ≤W2(Pλ∗ , P̂λ∗) +W2(P̂λ∗ , P ∗) +W2(P ∗, P̂nj ) +W2(P̂nj , Pj) ≤ √∑m
i=1 λ∗i ρ̂ δ/m ni + ρ
∗ + ρ∗
λj + ρ̂δ/2nj
Proof. The first line is by triangle inequality. The second line is by following facts: (i) P [ W2(Pλ∗ , P̂λ∗) ≥ √∑m i=1 λ ∗ i ρ̂ δ/m ni ] ≤ δ/2 according to (37), (ii) W2(P ∗, P̂nj ) = λjW2(P ∗,P̂nj )
λj ≤ ρ
∗
λj , and (iii) P [ W2(P̂nj , Pj) ≥ ρ̂ δ/2 nj ] ≤ δ/2 according to (35), and (iv) using
union bound.

heading: C PROOF OF THEOREM 1
text: Our proof is based on the analysis of local SGD for FL presented in Wang et al. (2021).
Fix some zi. Define ϕ(ζ, θ; zi) := `(ζ, hθ) − γd(ζ, zi). Since ` is Lzz-smooth and d is 1-strongly convex, ϕ(ζ, θ; zi) is (γ − Lzz)-strongly concave with respect to ζ, given that γ > Lzz . Lemma 2. Let z∗i = arg maxζ∈Z ϕ(ζ, θ; zi). Therefore, φγ(θ; zi) = ϕ(z∗i , θ; zi). Let ` satisfy Assumption 3. Then φγ is differentiable, and
‖∇φγ(zi, θ)−∇φγ(zi, θ′)‖ ≤ L‖θ − θ′‖,
with L = Lθθ + LθzLzθγ−Lzz when γ > Lzz .
The proof can be found in (Sinha et al., 2020, Lemma 1). Lemma 2 implies that φγ is L-smooth. Define gi(θ) := 1|Di| ∑ zi∈Di ∇θ φγ(zi, θ), then we have
E [ ‖gi(θ)−∇Fi(θ)‖2 ] = E [∥∥∥ 1|Di| ∑ zi∈Di ∇θ φγ(zi, θ)−∇Fi(θ) ∥∥∥2]
≤ 1 |Di| E [∥∥gφi(θ)−∇Fi(θ)∥∥2] ≤ σ2|Di| := σ̂2 (by Assumption 4.) (15)
With the shadow sequence θ̄(t,k) = ∑m i=1 λiθ (t,k) i , we have
θ̄(t,k+1) = m∑ i=1 λiθ (t,k+1) i = m∑ i=1 λi ( θ (t,k) i − ηgi(θ (t,k) i ) ) = θ̄(t,k) − η m∑ i=1 λigi(θ (t,k) i ).
Lemma 3. If the client learning rate satisfies η ≤ 13L , then
1
K K−1∑ k=0 E [ F (θ̄(t,k))− F (θ∗) ] ≤2ησ̂2 ( m∑ i=1 λ2i ) + L m∑ i=1 λi K−1∑ k=0 1 K E [ ‖θ(t,k)i − θ̄ (t,k)‖2 ]
+ 1
2ηK
( ‖θ(t) − θ∗‖2 −E [ ‖θ(t+1) − θ∗‖2 ]) .
Proof. Since θ̄(t,k+1) = θ̄(t,k) − η ∑m i=1 λigi(θ (t,k) i ), by parallelogram law
m∑ i=1 λi 〈 gi(θ (t,k) i ), θ̄ (t,k+1) − θ∗ 〉 = 1 2η ( ‖θ̄(t,k) − θ∗‖2 − ‖θ̄(t,k+1) − θ̄(t,k)‖2 − ‖θ̄(t,k+1) − θ∗‖2 ) .
(16)
Fact: Fi(θ) = EZi∼Pi [ φγ(Zi, θ) ] is Lipschitz smooth with L = Lθθ + LθzLzθγ−Lzz when γ > Lzz . With the assumption that θ 7→ `(z, hθ) is convex, we have θ 7→ Fi(θ) is convex. Since Fi is convex and L-smooth,
Fi(θ̄ (t,k+1)) ≤ Fi(θ(t,k)i ) + 〈 ∇Fi(θ(t,k)i ), θ̄ (t,k+1) − θ(t,k)i 〉 + L
2 ‖θ̄(t,k+1) − θ(t,k)i ‖ 2
≤Fi(θ∗) + 〈 ∇Fi(θ(t,k)i ), θ̄ (t,k+1) − θ∗ 〉 + L
2 ‖θ̄(t,k+1) − θ(t,k)i ‖ 2 ≤Fi(θ∗) + 〈 ∇Fi(θ(t,k)i ), θ̄ (t,k+1) − θ∗ 〉 + L‖θ̄(t,k+1) − θ̄(t,k)‖2 + L‖θ(t,k)i − θ̄ (t,k)‖2. (17)
From (16) and (17), we have F (θ̄(t,k+1))− F (θ∗) = m∑ i=1 λi ( Fi(θ̄ (t,k+1))− F (θ∗) )
≤ m∑ i=1 λi 〈 ∇Fi(θ(t,k)i )− gi(θ (t,k) i ), θ̄ (t,k+1) − θ∗ 〉 + L‖θ̄(t,k+1) − θ̄(t,k)‖2 + L m∑ i=1 λi‖θ(t,k)i − θ̄ (t,k)‖2
+ 1
2η
( ‖θ̄(t,k) − θ∗‖2 − ‖θ̄(t,k+1) − θ̄(t,k)‖2 − ‖θ̄(t,k+1) − θ∗‖2 ) . (18)
We have E [ m∑ i=1 λi 〈 ∇Fi(θ(t,k)i )− gi(θ (t,k) i ), θ̄ (t,k+1) − θ∗ 〉]
= E [ m∑ i=1 λi 〈 ∇Fi(θ(t,k)i )− gi(θ (t,k) i ), θ̄ (t,k+1) − θ̄(t,k) 〉] (since E [ gi(θ (t,k) i ) ] = ∇Fi(θ(t,k)i ) given θ̄ (t,k), θ∗)
≤ 3 2 η ·E [ ‖ m∑ i=1 λi ( ∇Fi(θ(t,k)i )− gi(θ (t,k) i ) ) ‖2 ] + 1 6η E [ ‖θ̄(t,k+1) − θ̄(t,k)‖2 ] (by Peter Paul inequality)
≤ 2ησ̂2 ( m∑ i=1 λ2i ) + 1 6η E [ ‖θ̄(t,k+1) − θ̄(t,k)‖2 ] , (19)
Plugging (19) back to the conditional expectation of (18), and noting that η ≤ 13L , we have
E [ F (θ̄(t,k+1))− F (θ∗) ] + 1
2η
( E [ ‖θ̄(t,k+1) − θ∗‖2 ] − ‖θ̄(t,k) − θ∗‖2 )
≤2ησ̂2 (
m∑ i=1 λ2i
) − ( 1
3η − L
) E [ ‖θ̄(t,k+1) − θ̄(t,k)‖2 ] + L
m∑ i=1 λi‖θ(t,k)i − θ̄ (t,k)‖2
≤2ησ̂2 ( m∑ i=1 λ2i ) + L m∑ i=1 λi‖θ(t,k)i − θ̄ (t,k)‖2
By convexity of F and telescoping k from 0 to K − 1, we have
1
K K−1∑ k=0 E [ F (θ̄(t,k))− F (θ∗) ] ≤2ησ̂2 ( m∑ i=1 λ2i ) + L m∑ i=1 λi K−1∑ k=0 1 K E [ ‖θ(t,k)i − θ̄ (t,k)‖2 ]
+ 1
2ηK
( ‖θ̄(t,0) − θ∗‖2 −E [ ‖θ̄(t,K) − θ∗‖2 ]) .
Since θ̄(t,0) = θ(t) and θ̄(t,K) = θ(t+1), we complete the proof.
Lemma 4 (Bounded client drift). Assuming the client learning rate satisfies η ≤ 13L , we have E [ ‖θ(t,k)i − θ̄ (t,k)‖2 ] ≤ η2(24K2Ω̄2 + 20KΩ̄2).
where Ω̄2 := max { σ̂2,Ω2 } .
Proof.
E [∥∥θ(t,k+1)1 − θ(t,k+1)2 ∥∥2] = E[∥∥∥θ(t,k)1 − θ(t,k)2 − η(g1(θ(t,k)1 )− g2(θ(t,k)1 ))∥∥∥2] =‖θ(t,k)1 − θ (t,k) 2 ‖2 − 2η 〈 g1(θ (t,k) 1 )−∇F1(θ (t,k) 1 ), θ (t,k) 1 − θ (t,k) 2
〉 − 2η 〈 ∇F2(θ(t,k)1 )− g2(θ (t,k) 2 ), θ (t,k) 1 − θ (t,k) 2
〉 − 2η 〈 ∇F1(θ(t,k)1 )−∇F2(θ (t,k) 2 ), θ (t,k) 1 − θ (t,k) 2 〉 + η2‖g1(θ(t,k)1 )− g2(θ (t,k) 2 )‖2. (20)
The second term (and similarly for the third term) is bounded as follows − 〈 g1(θ (t,k) 1 )−∇F1(θ (t,k) 1 ), θ (t,k) 1 − θ (t,k) 2 〉 ≤ 1
6ηK ∥∥θ(t,k)1 − θ(t,k)2 ∥∥2 + 3ηK2 ∥∥g1(θ(t,k)1 )−∇F1(θ(t,k)1 )∥∥2 (by Peter Paul inequality) = 1
6ηK ∥∥θ(t,k)1 − θ(t,k)2 ∥∥2 + 3ηK2 σ̂2 (by (15)) Since maxi supθ‖∇Fi(θ)−∇F (θ)‖ ≤ Ω (Assumption 5), the 4th-term is bounded as
− 〈 ∇F1(θ(t,k)1 )−∇F2(θ (t,k) 2 ), θ (t,k) 1 − θ (t,k) 2 〉 ≤− 〈 ∇F (θ(t,k)1 )−∇F (θ (t,k) 2 ), θ (t,k) 1 − θ (t,k) 2 〉 + 2Ω‖θ(t,k)1 − θ (t,k) 2 ‖
≤ − 1 L ‖∇F (θ(t,k)1 )−∇F (θ (t,k) 2 )‖2 + 2Ω‖θ (t,k) 1 − θ (t,k) 2 ‖ (by smoothness and convexity) ≤− 1 L ‖∇F (θ(t,k)1 )−∇F (θ (t,k) 2 )‖2 + 1 6ηK ‖θ(t,k)1 − θ (t,k) 2 ‖2 + 6ηKΩ2 (by Young’s inequality)
The last term is bounded as follows
‖g1(θ(t,k)1 )− g2(θ (t,k) 2 )‖2 ≤ 5 (∥∥g1(θ(t,k)1 )−∇F1(θ(t,k)1 )∥∥2 + ‖∇F1(θ(t,k)1 )−∇F (θ(t,k)1 )‖2 + ‖∇F (θ(t,k)1 )−∇F (θ(t,k)2 )‖2
+ ‖∇F (θ(t,k)2 )−∇F2(θ (t,k) 2 )‖2 + ∥∥g2(θ(t,k)2 )−∇F2(θ(t,k)2 )∥∥2) ≤ 5‖∇F (θ(t,k)1 )−∇F (θ (t,k) 2 )‖2 + 10(σ̂2 + Ω2) (by (15) and Assumption 5)
Substituting the above four bounds back to (20) gives (note that η ≤ 13L )
E [ ‖θ(t,k+1)1 − θ (t,k+1) 2 ‖2 ] ≤ ( 1 + 1
K
) ‖θ(t,k)1 − θ (t,k) 2 ‖2 − η ( 2 L − 5η )∥∥∇F (θ(t,k)1 )−∇F (θ(t,k)2 )∥∥2 + 6η2Kσ̂2 + 12η2KΩ2 + 10η2(σ̂2 + Ω2)
≤ ( 1 + 1
K
) ‖θ(t,k)1 − θ (t,k) 2 ‖2 + 6η2Kσ̂2 + 12η2KΩ2 + 10η2(σ̂2 + Ω2).
Unrolling recursively, we obtain E [∥∥θ(t,k+1)1 − θ(t,k+1)2 ∥∥2] ≤ (1 + 1/K)K − 11/K [6η2Kσ̂2 + 12η2KΩ2 + 10η2(σ̂2 + Ω2)]
≤ 12η2K2σ̂2 + 24η2K2Ω2 + 20η2K(σ̂2 + Ω2) ≤ η2(24K2Ω̄2 + 20KΩ̄2).
where we use the fact that ( 1+1/K )K −1
1/K ≤ K(e− 1) ≤ 2K, and Ω̄ 2 := max
{ σ̂2,Ω2 } .
By convexity, for any i, E [ ‖θ(t,k+1)i − θ̄ (t,k+1)‖2 ] ≤ η2(24K2Ω̄2 + 20KΩ̄2).
Substituting the result of Lemma 4 to Lemma 3, and telescoping over t, we obtain
E [ 1 T T−1∑ t=0 1 K K−1∑ k=0 F (θ̄(t,k))− F (θ∗) ] ≤ D 2 2ηKT + 2ησ̂2Λ + η2L(24K2Ω̄2 + 20KΩ̄2),
where D := ‖θ(0) − θ∗‖, Λ := m∑ i=1 λ2i . By optimizing η on the R.H.S, we obtain
E [ 1 KT T−1∑ t=0 K−1∑ k=0 F (θ̄(t,k))− F (θ∗) ] ≤ O ( LD2 KT + σ̂DΛ 1 2 √ KT + L 1 3 Ω̄ 2 3D 4 3 K 1 3T 2 3 + L 1 3 Ω̄ 2 3D 4 3 T 2 3 ) ,
when
η = min
{ 1
3L ,
D
2 √ KTΛσ̂
, D
2 3
48 1 3K 2 3T 1 3L 1 3 Ω̄ 2 3
, D
2 3
40 1 3KT 1 3L 1 3 Ω̄ 2 3
} .

heading: D PROOF OF LEMMA 1
text: We first prove the following fact:
Fact 1: (a) L(Q, f) ≤ Lγρ(Pλ, f), ∀f ∈ F , Q ∈ B(Pλ, ρ). (b) inf
f ′∈F L(Q, f ′) ≤ inf f ′∈F Lγρ(Pλ, f ′), ∀Q ∈ B(Pλ, ρ).
For (a), we have
L(Q, f) ≤ sup P ′∈B(Pλ,ρ) L(P ′, f) = inf γ′≥0
{ γ′ρ2 + EZ∼Pλ [ φγ(Z, f) ]} ≤ γρ2 + EZ∼Pλ [ φγ(Z, f) ] =: Lγρ(Pλ, f),
where the equality is due to strong duality result by Gao & Kleywegt (2016).
For (b), defining fPλ := arg minf ′∈F L γ ρ(Pλ, f ′), we have
inf f ′∈F L(Q, f ′) ≤ L(Q, fPλ) ≤ sup P ′∈B(Pλ,ρ) L(P ′, fPλ) (21)
= inf γ′≥0
{ γ′ρ2 + EZ∼Pλ [ φγ(Z, fPλ) ]} (22)
≤ γρ2 + EZ∼Pλ [ φγ(Z, fPλ) ] (23)
= inf f ′∈F
Lγρ(Pλ, f ′). (24)
We next prove the second fact:
Fact 2: (a) Lγρ(Pλ, f) ≤ L(Q, f) + 2Lzρ+ |γ − γ∗|ρ2, ∀f ∈ F , Q ∈ B(Pλ, ρ) (b) inf
f ′∈F Lγρ(Pλ, f ′) ≤ inf f ′∈F L(Q, f ′) + 2Lzρ+ |γ − γ∗|ρ2.
For (a), we have:
Lγρ(Pλ, f) =
{ sup
P ′∈B(Pλ,ρ) L(P ′, f)
} + { Lγρ(Pλ, f)− sup
P ′∈B(Pλ,ρ) L(P ′, f) } ≤ { L(Q, f) + 2Lzρ } + { EZ∼Pλ [φγ(Z, f)] + ρ
2γ − min γ′≥0
{ ρ2γ′ + EZ∼Pλ [φγ′(Z, f)] }} ≤ L(Q, f) + 2Lzρ+ ρ2(γ − γ∗) + EZ∼P [ φγ(Z, f)− φγ∗(Z, f)
] = L(Q, f) + 2Lzρ+ ρ 2(γ − γ∗) + EZ∼P [
sup ζ∈Z
{ `(ζ, h)− γd(ζ, Z) } − sup ζ∈Z { `(ζ, h)− γ∗d2(ζ, Z) }] = L(Q, f) + 2Lzρ+ (γ − γ∗) ( ρ2 −EZ∼P [ sup ζ∈Z d2(ζ, Z) ])
≤ L(Q, f) + 2Lzρ+ |γ − γ∗|ρ2,
where the first inequality is due to Proposition 1, and the last inequality is because we choose γ ≥ Lz/ρ and that fact that γ∗ ≤ Lz/ρ by Lemma 1 of Lee & Raginsky (2018). For (b), defining fQ := arg minf∈F L(Q, f), we have
inf f ′∈F
Lγρ(Pλ, f ′) ≤ Lγρ(Pλ, fQ) (25)
≤ L(Q, fQ) + 2Lzρ+ |γ − γ∗|ρ2 (26) = inf f ′∈F L(Q, f ′) + 2Lzρ+ |γ − γ∗|ρ2, (27)
where the second line is due to Fact 2(a).
Combining all facts, we complete the proof. Specifically, by adding two inequalities in Fact 1(a) and Fact 2(b), we obtain the upperbound of Lemma 1. Similarly, adding two inequalities in Fact 1(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 2 (a) holds. For any f ∈ F and for all Q ∈ B(Pλ, ρ), we have
sup P ′∈B(Pλ,ρ)
L(P ′, f) ≤ L(Q, f) + 2Lzρ.
Proof. Denote P ∗ := arg max P ′∈B(Pλ,ρ) L(P ′, f). We have
sup P ′∈B(Pλ,ρ) L(P ′, f) = L(Q, f) + sup P ′∈B(Pλ,ρ) L(P ′, f)− L(Q, f)
≤ L(Q, f) + |L(P ∗, f)− L(Q, f)|, ≤ L(Q, f) + Lz ∣∣EZ∼P∗[`(Z, h)/Lz]−EZ∼Q[`(Z, h)/Lz]∣∣ ≤ L(Q, f) + LzW1(P ∗, Q) ≤ L(Q, f) + Lz [ W2(P ∗, Pλ) +W2(Pλ, Q) ]
(28)
≤ L(Q, f) + Lz2ρ,
where the fourth line are due to Kantorovich-Rubinstein dual representation theorem, i.e.,
W1(P,Q) = sup h
{ EZ∼P [ h(Z) ] −EZ∼Q [ h(Z) ] : h(·) is 1-Lipschitz } and the fifth line is due to W1(P ∗, Q) ≤W2(P ∗, Q) and triangle inequality.

heading: E PROOF OF THEOREM 2
text: Proof. To simplify notation, we denote Φ := φγ◦F = {z 7→ φγ(z, f), f ∈ F}whereF = { fθ, θ ∈
Θ ⊂ Rd }
, which represents the composition of φγ with each of the loss function fθ parametrized by θ belonging to the parameter class Θ.
Defining fPλ ∈ arg minf∈F Lγρ(Pλ, f) and θ̂∗ ∈ argmin θ∈Θ EZ∼P̂λ
[ φγ(Z, fθ) ] such that
Lγρ(P̂λ, fθ∗) = inf θ∈Θ [ EZ∼P̂λ [ φγ(Z, fθ) ] + γρ2 ] , we decompose the excess risk as follows:
Eγρ(Pλ, fθ̂ε) = L γ ρ(Pλ, fθ̂ε)− inff∈F L γ ρ(Pλ, f)
= Lγρ(Pλ, fθ̂ε)− L γ ρ(Pλ, fPλ) = [ Lγρ(Pλ, fθ̂ε)− L γ ρ(P̂λ, fθ̂ε) ] + [ Lγρ(P̂λ, fθ̂ε)− L γ ρ(P̂λ, fθ̂∗) ] ︸ ︷︷ ︸
≤ε + [ Lγρ(P̂λ, fθ̂∗)− L γ ρ(P̂λ, fPλ) ] ︸ ︷︷ ︸
≤0
+ [ Lγρ(P̂λ, fPλ)− Lγρ(Pλ, fPλ) ]
≤ 2 sup φγ∈Φ ∣∣∣EZ∼Pλ [φγ(Z, fθ)]−EZ∼P̂λ [φγ(Z, fθ)]∣∣∣+ ε ≤ 2 sup
φγ∈Φ m∑ i=1 λi ∣∣∣EZi∼Pi [φγ(Zi, fθ)]−EZi∼P̂i [φγ(Zi, fθ)]∣∣∣+ ε ≤ 2
m∑ i=1 λi sup φγ∈Φ ∣∣∣EZi∼Pi [φγ(Zi, fθ)]−EZi∼P̂i [φγ(Zi, fθ)]∣∣∣+ ε ≤
m∑ i=1 λi [ 4Ri(Φ) + 2M` √ 2 log(2m/δ) ni ] + ε with probability at least 1− δ, (29)
where the first inequality is due to optimization error and definition of θ̂∗. The second inequality is due to the fact that | ∑m i=1 λiai| ≤ ∑m i=1 λi|ai|,∀ai ∈ R and λi ≥ 0. The third inequality is because pushing the sup inside increases the value. For the last inequality, using the facts that (i) |φγ(z, f)| ≤ M` due to−M` ≤ `(z, h) ≤ φγ(z, f) ≤ supz∈Z `(z, h) ≤M` and (ii) the Rademacher complexity of the function class Φ defined by Ri(Φ) = E[supφγ∈Φ 1 ni ∑ni k=1 σkφγ(Zk, fθ)] where the expectation is w.r.t bothZk i.i.d.∼ Pi and i.i.d. Rademacher random variable σk independent ofZk,∀k ∈ [ni], we have
sup φγ∈Φ
∣∣∣EZi∼Pi [φγ(Zi, fθ)]−EZi∼P̂i [φγ(Zi, fθ)]∣∣∣ ≥ 2Ri(Φ) +M` √ 2 log(2m/δ)
ni (30)
with probability ≤ δ/m due to the standard symmetrization argument and McDiarmid’s inequality (Shalev-Shwartz & Ben-David, 2014, Theorem 26.5). Multiplying λi to both sides of (30), summing up the inequalities over all i ∈ [n], and using union bound, we obtain (29).
Define a stochastic process ( Xφγ ) φγ∈Φ
Xφγ := 1 √ ni ni∑ k=1 σkφγ(Zk, fθ)
which is zero-mean because E [ Xφγ ] = 0 for all φγ ∈ Φ. To upper-bound Rn(Φ), we first show
that ( Xφγ ) φγ∈Φ
is a sub-Gaussian process with respect to the following pseudometric∥∥φγ − φ′γ∥∥∞ := sup z∈Z ∣∣∣φγ(z, fθ)− φγ(z, fθ′)∣∣∣. (31) For any t ∈ R, using Hoeffding inequality with the fact that σk, k ∈ [n], are i.i.d. bounded random variable with sub-Gaussian parameter 1, we have
E [ exp ( t ( Xφγ −Xφ′γ ))] = E [ exp ( t √ ni ni∑ k=1 σk (φγ (Zk, fθ)− φ (Zk, fθ′)) )]
= ( E [ exp ( t √ ni σ1 (φγ (Z1, fθ)− φγ (Z1, fθ′)) )])ni ≤ exp ( t2 ∥∥φγ − φ′γ∥∥2∞
2
) .
Then, invoking Dudley entropy integral, we have
√ niRi(Φ) = E sup
φγ∈Φ Xφγ ≤ 12 ∫ ∞ 0 √ logN (Φ, ‖·‖∞, )d (32)
We will show that when θ 7→ `(z, hθ) is Lθ-Lipschitz by Assumption 2, then θ 7→ φγ(z, fθ) is also Lθ-Lipschitz as follows.∣∣∣φγ(z, fθ)− φγ(z, fθ′)∣∣∣ = ∣∣∣ sup
ζ∈Z inf ζ′∈Z
{ `(ζ, hθ)− γd(ζ, z)− `(ζ ′, hθ′) + γd(ζ ′, z) }∣∣∣ ≤ ∣∣∣ sup ζ∈Z { `(ζ, hθ)− `(ζ, hθ′)
}∣∣∣ ≤ sup ζ∈Z
∣∣∣`(ζ, hθ)− `(ζ, hθ′)∣∣∣ ≤ Lθ‖θ − θ′‖,
which implies ∥∥φγ − φ′γ∥∥∞ ≤ Lθ‖θ − θ′‖. Therefore, by contraction principle (Shalev-Shwartz & Ben-David, 2014), we have
N (Φ, ‖·‖∞, ) ≤ N (Θ, ‖·‖, /Lθ) . (33) Substituting (33) and (32) into (29), we obtain
Eγρ(Pλ, fθ̂ε) ≤ m∑ i=1 λi [ 48C(Θ) √ ni + 2M` √ 2 log(2m/δ) ni ] + ε, (34)
which will be substituted into the upper-bound in Lemma 1 to complete the proof.

heading: F PROOF OF CORROLARY 1
text: We now present how we adapt the result from Fournier & Guillin (2015) to prove Corollary 1 Proposition 2 (Measure concentration (Fournier & Guillin, 2015, Theorem 2)). Let P be a probability distribution on a bounded set Z . Let P̂n denote the empirical distribution of Z1, . . . , Zn
i.i.d.∼ P. Assuming that there exist constants a > 1 such that A := EZ∼P [ exp(‖Z‖a) ] < ∞ (i.e., P is a light-tail distribution). Then, for any ρ > 0,
P [ Wp(P̂n, P ) ≥ ρ ] ≤ { c1 exp ( −c2nρmax{d/p,2} ) if ρ ≤ 1
c1 exp (−c2nρa) if ρ > 1
where c1, c2 are constants depending on a,A and d.
As a consequence of this proposition, for any δ > 0, we have
P [ W2(P̂n, P ) ≤ ρ̂δn ] ≥ 1− δ where ρ̂δn :=  ( log(c1/δ) c2n )min{2/d,1/2} if n ≥ log(c1/δ)c2 ,(
log(c1/δ) c2n )1/α if n < log(c1/δ)c2 . (35)
In Proposition 2, Fournier & Guillin (2015) show that the empirical distribution P̂n converges in Wasserstein distance to the true P at a specific rate. This implies that judiciously scaling the radius of Wasserstein balls according to (35) provides natural confidence regions for the data-generating distribution P .
By the duality of transport cost (Santambrogio, 2015, p.261), we have
W pp (µ, ν) = sup ϕ(x)+ψ(y)≤dp(x,y)
∫ ϕdµ+ ψ dν = sup
ϕ(x)+ψ(y)≤dp(x,y) Tf (µ, ν), ∀p ≥ 1,
which is the supremum of linear functionals Tf : P × P 7→ R defined by Tf (µ, ν) = 〈(µ, ν), (ϕ,ψ)〉; therefore, (µ, ν) 7→W pp (µ, ν) is convex. Thus we have
W 22 (Pλ, P̂λ) = W 2 2 ( m∑ i=1 λi(P̂ni , Pi) ) ≤ m∑ i=1 λiW 2 2 (P̂ni , Pi). (36)
Then, we have P [ W2(Pλ, P̂λ) ≥ √∑m i=1 λiρ̂ δ/m ni ] = P [ W 22 (Pλ, P̂λ) ≥ m∑ i=1 λiρ̂ δ/m ni ] ≤ P
[ m∑ i=1 λiW 2 2 (P̂ni , Pi) ≥ m∑ i=1 λiρ̂ δ/m ni ] ≤
m∑ i=1 P [ W 22 (P̂ni , Pi) ≥ ρ̂δ/mni ] =
m∑ i=1 P [ W2(P̂ni , Pi) ≥ ρ̂δ/2mni ] ≤
m∑ i=1 δ 2m = δ 2 , (37)
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).
According to (28), by setting ρ = (∑m
i=1 λiρ̂ δ/m ni
)1/2 in Theorem 2 and using union bound, we
complete the proof.

heading: G ADDITIONAL EXPERIMENTAL SETTINGS AND RESULTS
text: G.1 DATASETS
We distribute all datasets to clients as follows:
• MNIST: A handwritten digit dataset (Lecun et al., 1998) including 70, 000 instances belonged to 10 classes. We distribute dataset tom = 100 clients and each client has a different local data size with only 2 of the 10 classes.
• CIFAR-10: 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 3 labels per client. Each client has a different local data size.
• Three handwritten digit: MNIST, MNIST-M (Ganin & Lempitsky, 2015), and USPS (Hull, 1994). We distribute one 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, respectively. In domain adaptation, the target client only has test data. The statistics of all datasets are summarized in Tab. 1.
G.2 MODELS
The details of models for each dataset is provided as follows:
• MNIST: We use a multinomial logistic regression model (MLR) with a cross-entropy loss function and an L2-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 a cross-entropy loss function and an L2-regularization term.
In all settings, we randomly sample 10 clients to participate in training the global robust model in each communication round, and set the number of local epochs to K = 2 and the number of communication rounds to T = 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 distribution 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 γ = 0.05. For CIFAR-10, we use 20 clients and set γ = 0.5. We use T = 200 communication iterations.
To show WAFL’s convergence, we plot both the original loss (using the function `) and global accuracy in Fig. 5.
G.4 DOMAIN ADAPTATION.
We show that WAFL demonstrates its superior capability in domain adaptation by collaboratively learning from a multi-source domain P̂λ, and applying it to an unseen, but related, target domain Q̂. In federated domain adaptation (Peng et al., 2019; Liu et al., 2021), source domains often need access to the private data of target domains to find a 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, we consider
one dataset as a source domain (distributed to 100 clients), and another dataset as a target domain (assumed to be on one client). For example, if mt is used as the source and mm the domain, we use mt→mm to denote this case. Specifically, the mm dataset is distributed to 100 clients to learn a global model, then the model is tested on the mt dataset. We compare WAFL against the vanilla FedAvg as a 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 percentage point higher than that of FedAvg. This indicates WAFL’s benefits in transferring knowledge from multi-sourced, private client distributions to a new target distribution.

Review summaries:
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.

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%.

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.