On Bridging Generic and Personalized Federated Learning for Image Classification

“ Should we prioritize the learned model ’ s generic performance ( for future use at the server ) or its personalized performance ( for each client ) ? ” These two , seemingly competing goals have divided the community to focus on one or the other , yet in this paper we show that it is possible to approach both at the same time . Concretely , we propose a novel federated learning framework that explicitly decouples a model ’ s dual duties with two prediction tasks . On the one hand , we introduce a family of losses that are robust to non-identical class distributions , enabling clients to train a generic predictor with a consistent objective across them . On the other hand , we formulate the personalized predictor as a lightweight adaptive module that is learned to minimize each client ’ s empirical risk on top of the generic predictor . With this two-loss , two-predictor framework which we name Federated Robust Decoupling ( FED-ROD ) , the learned model can simultaneously achieve state-ofthe-art generic and personalized performance , essentially bridging the two tasks . 1 INTRODUCTION . Large-scale data are the driving forces for modern machine learning but come with the risk of data privacy . In applications like health care , data are required to be kept separate to enforce ownership and protection , hindering the collective wisdom ( of data ) for training strong models . Federated learning ( FL ) , which aims to train a model with multiple data sources ( i.e. , clients ) while keeping their data decentralized , has emerged as a popular paradigm to resolve these concerns ( Kairouz et al. , 2019 ) . The standard setup of FL seeks to train a single “ global ” model that can perform well on generic data distributions ( Kairouz et al. , 2019 ) , e.g. , the union of clients ’ data . As clients ’ data are kept separate , mainstream algorithms like FEDAVG ( McMahan et al. , 2017 ) take a multi-round approach shown in Figure 1 . Within each round , the server first broadcasts the “ global ” model to the clients , who then independently update the model locally using their own ( often limited ) data . The server then aggregates the “ local ” models back into the “ global ” model and proceeds to the next round . This pipeline is shown promising if clients ’ data are IID ( i.e. , with similar data and label distributions ) ( Stich , 2019 ; Zhou & Cong , 2017 ) , which is , however , hard to meet in reality and thus results in a drastic performance drop ( Li et al. , 2020b ; Zhao et al. , 2018 ) . Instead of sticking to a single “ global ” model that features the generic performance , another setup of FL seeks to construct a “ personalized ” model for each client to acknowledge the heterogeneity among clients ( Dinh et al. , 2020 ; Hanzely et al. , 2020 ; Smith et al. , 2017 ) . This latter setup ( usually called personalized FL ) is shown to outperform the former ( which we name generic FL ) regarding the test accuracy of each client alone . So far , these two seemingly contrasting FL setups are developed independently . In this paper , we however found that they can be approached simultaneously by generic FL algorithms like FEDAVG . Concretely , algorithms designed for generic FL ( G-FL ) often discard the local models { wm } after training ( see Figure 1 ) . As a result , when they are evaluated in a personalized setting ( P-FL ) , it is the global model w̄ being tested ( Arivazhagan et al. , 2019 ; Dinh et al. , 2020 ; Fallah et al. , 2020 ; Li et al. , 2021a ; Liang et al. , 2020 ; Smith et al. , 2017 ; Zhang et al. , 2021 ) . Here , we found that if we instead keep { wm } and evaluate them in P-FL , they outperform nearly all the existing P-FL algorithms . In other words , personalized models seem to come for free from the local training step of generic FL . At first glance , this may not be totally surprising : local training in G-FL algorithms is driven by the client ’ s empirical risk , which is what a personalized model strives to optimize1 . What really surprises us is that even without an explicit regularization term imposed by most P-FL algorithms ( Dinh et al. , 2020 ; Smith et al. , 2017 ) , the local models of G-FL algorithms can achieve better generalization performance . We conduct a detailed analysis and argue that global aggregation — taking average over model weights — indeed acts like a regularizer for local models . Moreover , applying advanced G-FL algorithms ( Acar et al. , 2021 ) to improve the G-FL accuracy seems to not hurt the “ local ” models ’ P-FL accuracy . Building upon these observations , we dig deeper into generic FL . Specifically for classification , the nonIID clients can result from non-identical class distributions or non-identical class-conditional data distributions . One way to mitigate their influences is to make the local training objectives more aligned among clients . While this can hardly be achieved for the latter case without knowing clients ’ data , we can do so for the former case by setting a consistent goal among clients — the learned local models should classify every class well , even if clients ’ data have different class distributions . We realize this by viewing each client ’ s local training as an independent class-imbalanced problem ( Cui et al. , 2019 ; He & Garcia , 2009 ) and applying objective functions dedicated to it ( Cao et al. , 2019 ; Ren et al. , 2020 ) . As will be shown in section 5 , these class-balanced objectives lead to much consistent local training among clients , making the resulting “ global ” model more robust to non-IID conditions . The use of class-balanced objectives , nevertheless , degrades the local models ’ P-FL performance . This is because the local models are no longer learned to optimize clients ’ empirical risks . To address this issue , we propose a unifying framework for G-FL and P-FL which explicitly decouples a local model ’ s dual duties : serving as the personalized model and the ingredient of the global model . Concretely , we follow the FEDAVG pipeline and train the local model with the class-balanced loss , but on top of the feature extractor , we introduce a lightweight personalized predictor and train it with client ’ s empirical risk ( see Figure 2 ) . With this two-loss , twopredictor framework which we name Federated Ro- bust Decoupling ( FED-ROD ) , the resulting global model can be more robust to non-identical class distributions ; the personalized predictor can lead to decent P-FL accuracy due to the implicit regularization and the empirical loss . Specifically for the personalized predictor , we propose to explicitly parameterize it with clients ’ class distributions via a hypernetwork ( Ha et al. , 2017 ) . That is , we learn a shared meta-model that outputs personalized predictors for clients given their class distributions . This not only enables zero-shot model adaptation to new clients ( without their data but class distributions ) , but also provides a better initialization to fine-tune the models given new clients ’ data . We validate FED-ROD on multiple datasets under various non-IID settings . FED-ROD consistently outperforms existing generic and personalized FL algorithms in both setups . Moreover , FED-ROD is compatible with and can further improve advanced generic FL algorithms like FEDDYN ( Acar et al. , 2021 ) whenever non-identical class distributions occur . Our contributions are three-folded : • Unlike most of the previous works that focus on either generic FL or personalized FL , we propose FED-ROD to excel on both at the same time . We validate FED-ROD with extensive experiments . 1However , the literature does not evaluate G-FL algorithms in the P-FL setup by their local models . • We show that strong personalized models emerge from the local training step of generic FL algorithms , due to implicit regularization . We further show that class-balanced objectives are effective for improving the generic FL performance when clients have different class distributions . • FED-ROD enables zero-shot adaptation and much effective fine-tuning for new clients . 2 RELATED WORK ( A DETAILED VERSION IS IN APPENDIX A ) . Generic federated learning . FEDAVG ( McMahan et al. , 2017 ) is the standard algorithm , and many works are proposed to improve it , either in the global aggregation step ( Chen & Chao , 2021 ; Hsu et al. , 2019 ; Lin et al. , 2020 ; Reddi et al. , 2021 ; Wang et al. , 2020a ; Yurochkin et al. , 2019 ) or local training step ( Malinovskiy et al. , 2020 ; Wang et al. , 2020b ; Yuan & Ma , 2020 ; Zhao et al. , 2018 ) . For example , to reduce local models ’ drifts from the global model , FEDPROX ( Li et al. , 2020a ) and FEDDYN ( Acar et al. , 2021 ) employed regularization toward the global model ; SCAFFOLD ( Karimireddy et al. , 2020a ) leveraged control variates to correct local gradients . We also aim to reduce local models ’ drifts but via a different way . We apply objective functions in class-imbalanced learning ( He & Garcia , 2009 ) , which are designed to be robust to class distribution changes . The closest to ours is ( Hsu et al. , 2020 ) , which used a traditional class-imbalanced treatment named re-weighting . We show that more advanced techniques can be applied to further improve the performance , especially under extreme non-IID conditions where re-weighting is ineffective . Personalized federated learning . Many approaches for personalized FL ( Kulkarni et al. , 2020 ) are based on multi-task learning ( MTL ) ( Ruder , 2017 ; Zhang & Yang , 2017 ) . For instance , Smith et al . ( 2017 ) encouraged related clients to learn similar models ; Dinh et al . ( 2020 ) ; Hanzely et al . ( 2020 ) ; Li et al . ( 2021a ) regularized local models with a learnable global model . Our approach is inspired by MTL as well but has notable differences . First , we found that global aggregation in generic FL already serves as a strong regularizer . Second , instead of learning for each client a feature extractor ( Bui et al. , 2019 ; Liang et al. , 2020 ) or an entire model , FED-ROD shares a single feature extractor among clients , inspired by Caruana ( 1997 ) ; Zhang et al . ( 2014 ) . This reduces the total parameters to be learned and improves generalization . Compared to ( Arivazhagan et al. , 2019 ; Collins et al. , 2021 ) which also learned a shared feature extractor , FED-ROD simultaneously excels in both FL setups . Instead of designing specific algorithms for personalized FL , Cheng et al . ( 2021 ) ; Wang et al . ( 2019 ) ; Yu et al . ( 2020 ) showed that performing post-processing ( e.g. , fine-tuning ) to a generic FL model ( e.g. , w̄ in FEDAVG ) leads to promising personalized accuracy . We further showed that , the local models { wm } learned in FEDAVG and other generic FL algorithms are strong personalized models . We note that , while many personalized FL algorithms also produce a global model , it is mainly used to regularize or construct personalized models but not for evaluation in the generic setup . In contrast , we learn models to excel in both setups via a single framework without sacrificing either of them . A recent work PFEDHN ( Shamsian et al. , 2021 ) also applies hypernetworks ( Ha et al. , 2017 ) but in a very different way from FED-ROD . PFEDHN learns a hypernetwork at the server to aggregate clients ’ updates and produce entire models for them for the next round . In contrast , we learn the hypernetwork locally to construct the personalized predictors , not the entire models , for fast adaptation to clients .

Summary: FL trains models both on the server and clients. Such model training becomes problematic when the data distributions on each of the clients diverge from each other. The paper answers the question of prioritizing the generic model at the server or the personalized model on each client. The paper proposes to decouple a model’s dual duties with two prediction tasks.