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 aims to enable training of higher accuracy models by training a high-capacity model at the server and lower capacity models at the clients. Knowledge is shared via distillation on a weighted average of logits provided by the clients, rather than by averaging model parameters or model differences. The approach assumes the existence of a labeled public dataset. Experiments illustrate the promise of this approach.