FedDrop: Trajectory-weighted Dropout for Efficient Federated Learning

1 INTRODUCTION . In the light of the importance of personal data and the recent strict privacy regulations , e.g . the General Data Protection Regulation ( GDPR ) of the European Union ( Voigt & Von dem Bussche , 2017 ; Wolters , 2017 ; Politou et al. , 2018 ) , there is now a great amount of risk , responsibility ( Edwards et al. , 2016 ; Culnan & Williams , 2009 ) and technical challenges for securing private data centrally ( Sun et al. , 2014 ) ; it is often impractical to upload , store and use data on central servers . To this end , federated learning ( FL ) ( McMahan et al. , 2017 ; Li et al. , 2019a ) enables multiple edge compute devices to learn a global shared model collaboratively in a communication-efficient way without collecting their local training data . When compared against naïve decentralized SGD , Federated averaging ( FedAvg ) ( McMahan et al. , 2017 ) and subsequent FL algorithms ( Li et al. , 2020 ; Karimireddy et al. , 2020 ) reduce the burden of data transmission by many orders of magnitude . While these communication-efficient methods can notably alleviate the difficulties of using FL in scenarios with limited bandwidth or data-quota , they , however , entail a drastic increase in computation cost , which has rarely been addressed by previous literature on FL . It has been argued that communication is several orders of magnitude more expensive than computation on edge devices ( Li et al. , 2019a ; Huang et al. , 2013 ) . Yet existing FL methods optimize for communication so aggressively that an iPhone 12 Pro running as a FedAvg client in our test1 would spend at least 16 minutes on training , before even transmitting 180 MB of data , which only takes fraction of a second under the modern 5G infrastructure . The savings in communication thus are completely dwarfed by the expensive computation . In a FL setup , models trained by the clients may disagree on the gradient direction to update the same neurons if the user data distributions are non-IID . In Figure 1 , we demonstrate this effect with an initial round of FedAvg training . It turns out that clients sharing similar data distributions tend to produce similar update trajectories to the same channel , and different distributions observe disparate trajectories . From this example , it can be observed that the consequence of non-IID data is two-fold . First , a naïve averaging of client parameters may cause these conflicting signals to cancel out each other , resulting in a slow convergence of the global model . Second , neurons in a layer tend to learn distinct features , and yet they can not learn meaningfully if these features are absent in the client ’ s training data . Since neuron training is heavily dependent on the client ’ s local data , it presents us an 1See Appendix F.6 for the setup . opportunity : can we concentrate training effort to neurons that are correlated to the current data distribution of the client , while paying less attention to neurons that are less relevant to the client ? To leverage this , we introduce FedDrop , which introduces SyncDrop layers to structurally sparsify client models , and on the server-side an inter-client trajectory-based optimization of the dropout probabilities used by SyncDrop to speed up model training . FedDrop brings two-fold advantages : dropped neurons can simply be skipped , and thus reduce the training FLOPs required per step ; and dropout probabilities of each neuron can be tuned individually to minimize the impact on the global averaged model updates to assist convergence . Overall , as evinced by our experiments , FedDrop enables a much improved communication/computation trade-off relationship than traditional FL approaches . Our contributions in this paper are as follows : • We present SyncDrop layers , which are synchronous structural dropouts with adaptive keep probabilities , to reduce the computational cost required by FL . • With SyncDrop , we formally derive the FedDrop objective that adjusts each neuron ’ s dropout probability to improve convergence by minimizing the disparities among inter-client update trajectories . We further introduce FLOPs-based constraints to enforce sparsity per FL round , allowing the trade-off between FLOPs and communication to be tuned easily . • Empirical results reveal that the combined method , FedDrop , attains a substantially better communication/computation trade-off in comparison to other FL methods . 2 RELATED WORK . Federated learning . Distributed machine learning has a long history of progress and success ( PeteiroBarral & Guijarro-Berdiñas , 2013 ; Li et al. , 2014 ) , yet it mainly focuses on training with IID data . The Federated Learning ( FL ) paradigm and the Federated Averaging algorithm ( FedAvg ) initially introduced by McMahan et al . ( 2017 ) allow clients to train collaboratively without sharing the private data in a communication-efficient manner , To further tackle data heterogeneity , FedProx ( Li et al. , 2020 ) introduces new regularizations , and SCAFFOLD ( Karimireddy et al. , 2020 ) presents control variates to account for client drifts and reduce the inter-client variance . While being effective at reducing communication , the above methods neglected the computational costs associated with the training process . Computation vs. communication during training . There are a few precursory methods that focus on the joint optimization of computation and communication costs during training . Caldas et al . ( 2018a ) introduced federated dropout , which prunes parameters following a uniform random distribution , whereas PruneFL ( Jiang et al. , 2019 ) proposes a greedy gradient-magnitude-based unstructured pruning . In each FL round , both methods produce a shared pruned model with fine-grained sparsity for all clients . Such unstructured sparsity is difficult to utilize for training acceleration , and a shared global model can not exploit the data distribution of individual clients . Adaptive federated dropout ( Bouacida et al. , 2020 ) partially addresses the latter issue by allowing each client to select a sub-model to join training . FjORD ( Horvath et al. , 2021 ) introduces ordered dropout to tackle the problem of system heterogeneity in FL with sub-model training that dynamically changes model capabilities using uniform sampling . FedGKT ( He et al. , 2020 ) transmits shallow layer activations to offload training of subsequent layers to the server . However , the privacy implications were not well explored , and the trained models were substantially larger , which may limit their inference speed on edge devices . It is noteworthy that FedDrop differs from these approaches as it takes into account the non-IID client data distributions in a FL setting , and computes inter-client channel selection decisions that would minimize the impact on convergence . Dropout algorithms in centralized training . Dropout ( Hinton et al. , 2012 ; Srivastava et al. , 2014 ) may improve neural network training by reducing the overfitting effect and allow it to generalize better . It is done by randomly setting parts of the connections or weights to zero in a network , and scaling the remaining values accordingly . Since its advent , there have been an increasing interest in applying dropout in a structural fashion ( Huang et al. , 2015 ; Ghiasi et al. , 2018 ; Hou & Wang , 2019 ) . Nonetheless , the above methods missed the opportunity to explore the implications of a structural dropout on the FLOPs consumed by training or inference . Structural pruning . A closely related topic is structural neuron pruning/selection for accelerated inference . These work propose new ways to extract a smaller accurate network from the original ( He et al. , 2018 ; Wu et al. , 2019 ; Li et al. , 2019b ; Herrmann et al. , 2020 ) , and can even do so dynamically with an input-dependent policy ( Gao et al. , 2019 ; Hua et al. , 2019 ; Wang et al. , 2020c ) . 3 THE FEDDROP METHOD 3.1 HIGH-LEVEL OVERVIEW FedDrop complements existing FL methods by adding new client- and server-side components . During client training , convolutional layers are sparsified by interleaving them with channel-wise dropout layers , namely the SyncDrop layers . Section 3.3 discusses in depth how SyncDrop layers compute dropout decisions . Additional optimization stages during server aggregation are added to encourage collaboration between sparse clients . Figure 2 shows how FedDrop extends traditional FedAvg ( McMahan et al. , 2017 ) . After each round of client training , the server begins by identifying parameter update directions of each client from the previous round , and computes a cross-client trajectory similarity matrix for each channel neuron . This is then followed by an optimization stage that iteratively minimizes the trajectory disparities by tuning dropout keep probabilities for each channel neuron on each client for the next training round ( Section 3.6 ) . This process also takes into consideration the FLOPs budget constraints to sparsify client models ( Section 3.4 ) . Finally , following FedAvg , the server broadcasts the weighted-average model parameters to all training clients in the new round , and finally sends the optimized probabilities to the respective clients . 3.2 PRELIMINARIES AND DEFINITIONS . FedDrop complements most FL algorithms . In this paper , we focus on its use on FedAvg ( McMahan et al. , 2017 ) . We assume the training loss function of a client c ∈ C to be ` c ( θc ) , where θc comprises the parameters { θ [ l ] c : l ∈ L } of all layers in the model of client c. In each FedAvg training round r , clients begin by training on the loss function , with initial parameters θ ( r ) received from the server for this round : θ ( r+1 ) c = SGDc ( ` c , θ ( r ) , η , E ) . ( 1 ) Here SGDc indicates that client c carries out stochastic gradient descent ( SGD ) on ` c ( θ ( r ) c ) locally , and it uses a learning rate η forE epochs . The FedAvg server then aggregates client model parameters after the rth training round , by taking the weighted average of them : θ ( r+1 ) = ∑ c∈C λcθ ( r+1 ) c , ( 2 ) where λc is the weight of client c and is proportional to the size of its training set |Dc| with∑ c∈C λc = 1 . Finally , the ( r + 1 ) th training round starts by repeating the above procedure . 3.3 SYNCHRONIZED DROPOUT . To induce sparsity in a stochastic manner , and subsequently introduce training correlation across clients , we designed a threshold-based dropout layer , SyncDrop , to synchronize dropout decisions across multiple clients . Initially for all sampled clients c ∈ C ⊆ C , the server provides them with their corresponding dropout keep probabilities pc ∈ [ 0 , 1 ] N , and we assume p ∈ [ 0 , 1 ] C×N to be the concatenated probabilites from all sampled clients . Each element pnc in p denotes the keep probability of the channel neuron n ∈ N in client c ∈ C. We additionally use p [ l ] c , a slice of pc , to indicate the probabilities of all channels in the lth layer . During training , the lth layer samples t [ l ] n from a uniform distribution U ( 0 , 1 ) for all channel neurons n in the layer , where t [ l ] n is shared across clients by using the same random seed . If pnc < t [ l ] n , a channel n in client c is dropped , i.e . we set the entire channel map to zero ; otherwise , the channel values are scaled by 1/pnc during training 2 . Formally , for each client c , the lth dropout layer computes for the input x [ l ] : drop ( x [ l ] , p [ l ] c ) , d [ l ] c ◦ x [ l ] ◦ p [ l ] c ◦−1 , where d [ l ] c , 1 [ t [ l ] < p [ l ] c ] . ( 3 ) Here , 1 [ z ] is the ( element-wise ) indicator function , which is equal to 1 when the condition z is met and 0 otherwise , ◦ refers to the element-wise product , the term p [ l ] c ◦−1 represents the element-wise inverse of p [ l ] c , and finally all elements of t [ l ] are independently sampled from U ( 0 , 1 ) and shared across clients . Figure 3a provides a high-level overview of the procedure described above . From the local perspective of a client c , the dropout decision d [ l ] c is equivalent to the Bernoulli distributions B ( p [ l ] c ) . An important distinction from independent Bernoulli distributions is that the distribution of d [ l ] c is correlated across clients c ∈ C. Given a pair of clients ( i , j ) for channel n : Et∼U ( 0,1 ) [ dni d n j ] = ∫ 1 0 1 [ t < pni ] · 1 [ t < pnj ] dt = min ( pni , p n j ) . ( 4 ) This enables us to adjust the correlation of the same channel neurons between any pairs of clients , and Section 3.6 makes use of this property to minimize the model update disparities across all clients . 2This is known as “ inverted dropout ” and implemented by PyTorch ( 2021 ) and TensorFlow ( 2021 ) . Figure 7 provides ablation results on scaling for dropout to justify the design choice . We define f̂ to be the sparsified variant of the original model f , where SyncDrop layers are placed immediately after each convolutional layer with ReLU activation ( Figure 3b ) . This enables convolutional layers to be doubly accelerated by taking advantage of sparsities at both ends and skipping dropped channels in input/output feature maps . Recalling the client model loss function ` c , we additionally use ˆ̀c to denote the accelerated variant of the sparse model f̂ with probabilities pc . It is noteworthy that active channels are sampled for each local training mini-batch . In addition , the synchronization of dropout decisions is carried out with identical random seeds for all clients , and thus it eliminates the need to communicate between clients . Finally , for inference , one may choose to enable the SyncDrop layers for speed , or skip these layers entirely . In our evaluations , SyncDrop layers are disabled for improved test accuracies .

Summary: The paper proposes a new method where local workers will drop part of their model using a shared dropout probability received from the server at each communication round. The dropout probabilities are computed by solving an optimization to promote similarity across agent's update. Compared with popular baselines in federated learning, the proposed method demonstrates better performance in terms of number of FLOPs.