FedGEMS: Federated Learning of Larger Server Models via Selective Knowledge Fusion

1 INTRODUCTION . Nowadays , powerful models with tremendous parameters trained with sufficient computation power are indispensable in Artificial Intelligence ( AI ) , such as AlphaGo ( Silver et al. , 2016 ) , Alphafold ( Senior et al. , 2020 ) and GPT-3 ( Brown et al. , 2020 ) . However , billions of resource-constrained mobile and IoT devices have become the primary data source to empower the intelligence of many applications ( Bonawitz et al. , 2019 ; Brisimi et al. , 2018 ; Li et al. , 2019a ) . Due to privacy , security , regulatory and economic considerations ( Voigt & Von dem Bussche , 2017 ; Li et al. , 2018 ) , it is increasingly difficult and undesirable to pool data together for centralized training . Therefore , federated learning approaches ( McMahan et al. , 2017 ; Smith et al. , 2017 ; Caldas et al. , 2018 ; Kairouz et al. , 2019 ; Yang et al. , 2019 ) allow all the participants to reap the benefits of shared models without sharing private data have become increasingly attractive . In a typical Federated Learning ( FL ) training process using FedAvg ( McMahan et al. , 2017 ) , each client sends its model parameters or gradients to a central server , which aggregates all clients ’ updates and sends the aggregated parameters back to the clients to update their local model . Because FL places computation burden on edge devices , its learnability is largely limited by the edge resources , on which training large models is often impossible . On the other hand , the server only directly aggregates the clients ’ models and its computation power is not fully exploited . In this work , we investigate the paradigm where the server adopts a larger model in FL . Here we use larger server model , or GEM to denote the setting where the size of the server model is larger than that of the clients ( See Table 1 ) . Making FL trainable with GEM is desirable to break through model capacity and enable collaborative knowledge fusion and accumulation at server . One feasible approach to bridge FL with GEM is through knowledge distillation ( KD ) ( Hinton et al. , 2015 ) , where clients and the server transfer knowledge through logits . For example , FedGKT ( He et al. , 2020a ) adopts a server model as a downstream sub-model and transfers knowledge directly from smaller edge models . In FedGKT the large server model essentially learns from one small teacher at a time and doesn ’ t learn consensus knowledge from multiple teachers . In FL , KD has also been applied ( Li & Wang , 2019 ; Lin et al. , 2020 ) to transfer ensemble knowledge of clients through consensus of output logits rather than parameters . These works either assume no model on the server ( Li & Wang , 2019 ) or a prototype server model of the same architecture as the client model ( Lin et al. , 2020 ) . Contributions . In this paper , we first propose a new paradigm to bridge FL with GEM , termed FedGEM , which can learn effectively and efficiently from fused knowledge by resource-constrained clients , and is also able to transfer knowledge back to clients with heterogeneous architectures . To further prevent negative and malicious knowledge transfer , we carefully design a selection and weighting criterion to enhance our knowledge transfer protocol , termed FedGEMS . We demonstrate with extensive experiments on various image classification tasks that our results significantly surpass the previous state-of-the-art baselines in both homogeneous and heterogeneous settings . Furthermore , thanks to our effective and selective protocol , our framework improves the robustness of FL on various malicious attacks and significantly reduces the overall communication . In summary , we propose a new framework to bridge FL with larger server models and simultaneously consolidate several benefits altogether , including superior performance , robustness of the whole system , and lower communication cost . 2 RELATED WORK . Federated Learning with GEM . FL is a collaborating learning framework without sharing private data among the clients . The classical method FedAvg ( McMahan et al. , 2017 ) and its recent variations ( Mohri et al. , 2019 ; Lin et al. , 2018 ; Li et al. , 2019b ) directly transfer the clients ’ parameters or gradients to the server nodes . To tackle the performance bottleneck by training resource-constrained clients in FL , there are two lines of work to bridge FL with GEM . The first line of studies adopts model compression ( Han et al. , 2015 ; He et al. , 2018 ; Yang et al. , 2018 ) , manually designed architectures ( Howard et al. , 2017 ; Zhang et al. , 2018 ; Iandola et al. , 2016 ) or even efficient neural architecture search ( Tan & Le , 2019 ; Wu et al. , 2019 ) to adapt a GEM to on-device learning . Another line is adopting knowledge distillation ( Hinton et al. , 2015 ) to transfer knowledge through output logits rather than parameters between a client model and a GEM ( He et al. , 2020a ) . However , FedGKT ( He et al. , 2020a ) ’ s focus is to transfer knowledge directly from clients to server without considering the consensus knowledge fused from clients . Therefore its performance is limited . Federated Learning with Knowledge Distillation . In fact , ensemble knowledge distillation has been shown to boost collaborative performance in FL . Specifically , FedMD ( Li & Wang , 2019 ) adopts a labeled public dataset and averaged logits to transfer knowledge . FedDF ( Lin et al. , 2020 ) proposes ensemble distillation for model fusion by aggregating both logits and models from clients . In addition , KD is used to enhance robustness in FL . Cronus ( Chang et al. , 2019 ) and DS-FL ( Itahara et al. , 2020 ) utilize a public dataset with soft labels jointly with local private dataset for local training , and combine with Cronus or entropy reduction aggregation , respectively , to defend against poisoning attacks in FL . In this work , we propose to exploit the benefit of both a larger server model and client knowledge fusion . Our work is also motivated by the recent studies in KD ( Qin et al. , 2021 ; You et al. , 2017 ; Li et al. , 2021 ; Yuan et al. , 2021 ; Wang et al. , 2021 ) , which show that student models can have larger capacities by learning from multiple teacher models . 3 METHODOLOGY . 3.1 PRELIMINARIES . We assume that there are K clients in federated learning process . The kth client has its own private labeled dataset Xk : = { ( xki , yki ) } N k i=1 that can be drawn from the same or different distribution , where xki is the ith training sample in the kth client model , y k i is its corresponding ground truth label , and Nk denotes the total number of samples . Each client also trains its own model fkc which can be of the same architecture ( homogeneous ) or different architecture ( heterogeneous ) . There is also a public dataset X0 : = { ( x0i , y0i ) } N 0 i=1 which is accessible to both server and clients . On the server side , we assume that there is a larger server model to be trained , denoted as fs . Ls and Lkc denotes the logit tensors from the server and the kth client model . 3.2 FEDGEMS FRAMEWORK . We illustrate our overall framework in Fig . 1 and summarize our training algorithm in Algorithm 1 . FedGEM . During each communication round , all client models first use private datasets to train several epochs , then transfer the predicted logits on public dataset as knowledge to the server model . The server model aggregates the clients ’ logits and then trains its server model with the guidance of fused knowledge . After training , the server model then transfers its logits back to all client models . Finally , each client model distills knowledge from received logits and continues their training on private datasets . Continuously iterating over multiple rounds , both the server and client models mutually learn knowledge from each other . Through this alternating training processing , we can obtain a large server model with accumulated knowledge and an ensemble of high-performance client models . FedGEMS : In FedGEMS , the server adopts a selection and weighting criterion to select knowledgeable clients for aggregation , which is detailed in the next section 3.3 . To illustrate the features of our framework , we compare it with the related studies of KD-based methods in federated learning in Table 1 on the following aspects : whether they use a labeled , unlabeled or no public dataset , whether client models can have heterogeneous architectures , the aggregation strategy on the server , and whether the server has a larger model . Note FedGEM can be regarded as placing a larger server model on top of the FedMD framework while keeping other settings the same . 3.3 SELECTIVE KNOWLEDGE FUSION IN SERVER MODEL . Since clients ’ knowledge may negatively impact the server model in the heterogeneous or malicious setting and vice versa , we further propose selective strategies on both server and client sides to enforce positive knowledge fusion into the large server model as shown in Fig . 2 . Algorithm 1 Illustration of the Framework of FedGEMS . T is the number of communication rounds ; X0 and Y 0 denotes the images and their corresponding labels in public dataset ; Xk and Y k denotes the private dataset in the kth client model ; fs and fkc are the server with parameter Ws and the kth client model with parameters W kc ; LGlobal indicates the global logits to save correct logits ; Ls and Lkc are the logit tensors from the server and the kth client model . 1 : ServerExecute ( ) : 2 : for each round t = 1 , 2 , ... , T do 3 : // Selective Knowldge Fusion 4 : for idx , x0 , y0 ∈ { X0 , Y 0 } do 5 : if fs ( Ws ; x0 ) == y0 then 6 : Ls ← LS1 ( x0 , y0 ) ▷ in Eq . 1 7 : LGlobal [ idx ] ← Ls [ idx ] 8 : else if idx in LGlobal then 9 : Ls ← LS2 ( x0 , y0 , LGlobal ) ▷ in Eq . 2 10 : else 11 : Lc ← ClientSelect ( idx ) 12 : Ls ← LS3 ( x0 , y0 , Lc ) ▷ in Eq . 3 , 4 , 5 13 : Ws ←Ws − ηk∇Ls 14 : Ls [ idx ] ← fs ( Ws ; x0 ) 15 : // Transfer Knowledge to Client Models 16 : ClientTrain ( Ls ) 17 : ClientTrain ( Ls ) : 18 : for each client kth in parallel do 19 : // Knowledge Distillation on Clients 20 : for x0 , y0 ∈ { X0 , Ls , Y 0 } do 21 : Lc ← LC ( x0 , y0 , Ls ) ▷ in Eq . 6 22 : W kc ←W kc − ηk∇Lc 23 : // Local Training on Clients 24 : for xk , yk ∈ { Xk , Y k } do 25 : Lc ← LCE ( xk , yk ) 26 : W kc ←W kc − ηk∇Lc 27 : ClientSelect ( idx ) : 28 : // Selective Transfer to Server 29 : for each client kth in parallel do 30 : L0c [ idx ] ← fkc ( W kc ; x0 [ idx ] ) 31 : Return Lc to server

Summary: This paper tries to address the big model on the server-side in federated learning (FL) via knowledge distillation. Selective transfer knowledge happens on both server and clients sides. The proposed model maintains a set of global, correct logits so that it can be used as memories for the server's distillation. When the server predicts wrong on the public datasets, it transfers the knowledge from the correct, weighted clients. The proposed model is evaluated on image datasets with many FL baselines.