ON THE GENERALIZATION OF WASSERSTEIN ROBUST FEDERATED LEARNING

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 . 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 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 ) .

Summary: 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.