FjORD: Fair and Accurate Federated Learning under heterogeneous targets with Ordered Dropout

1 Introduction . Over the past few years , advances in deep learning have revolutionised the way we interact with everyday devices . Much of this success relies on the availability of large-scale training infrastructures and the collection of vast amounts of training data . However , users and providers are becoming increasingly aware of the privacy implications of this ever-increasing data collection , leading to the creation of various privacy-preserving initiatives by service providers [ 3 ] and government regulators [ 10 ] . Federated Learning ( FL ) [ 46 ] is a relatively new subfield of machine learning ( ML ) that allows the training of models without the data leaving the users ’ devices ; instead , FL allows users to collaboratively train a model by moving the computation to them . At each round , participating devices download the latest model and compute an updated model using their local data . These locally trained models are then sent from the participating devices back to a central server where updates are ∗Indicates equal contribution . †Work while intern at Samsung AI Center . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . aggregated for next round ’ s global model . Until now , a lot of research effort has been invested with the sole goal of maximising the accuracy of the global model [ 46 , 42 , 39 , 31 , 63 ] , while complementary mechanisms have been proposed to ensure privacy and robustness [ 6 , 14 , 47 , 48 , 27 , 4 ] . A key challenge of deploying FL in the wild is the vast heterogeneity of devices [ 38 ] , ranging from low-end IoT to flagship mobile devices . Despite this fact , the widely accepted norm in FL is that the local models have to share the same architecture as the global model . Under this assumption , developers typically opt to either drop low-tier devices from training , hence introducing training bias due to unseen data [ 30 ] , or limit the global model ’ s size to accommodate the slowest clients , leading to degraded accuracy due to the restricted model capacity [ 8 ] . In addition to these limitations , variability in sample sizes , computation load and data transmission speeds further contribute to a very unbalanced training environment . Finally , the resulting model might not be as efficient as models specifically tailored to the capabilities of each device tier to meet the minimum processingperformance requirements [ 34 ] . In this work , we introduce FjORD ( Fig . 1 ) , a novel adaptive training framework that enables heterogeneous devices to participate in FL by dynamically adapting model size – and thus computation , memory and data exchange sizes – to the available client resources . To this end , we introduce Ordered Dropout ( OD ) , a mechanism for run-time ordered ( importance-based ) pruning , which enables us to extract and train submodels in a nested manner . As such , OD enables all devices to participate in the FL process independently of their capabilities by training a submodel of the original DNN , while still contributing knowledge to the global model . Alongside OD , we propose a self-distillation method from the maximal supported submodel on a device to enhance the feature extraction of smaller submodels . Finally , our framework has the additional benefit of producing models that can be dynamically scaled during inference , based on the hardware and load constraints of the device . Our evaluation shows that FjORD enables significant accuracy benefits over the baselines across diverse datasets and networks , while allowing for the extraction of submodels of varying FLOPs and sizes without the need for retraining . 2 Motivation . Despite the progress on the accuracy front , the unique deployment challenges of FL still set a limit to the attainable performance . FL is typically deployed on either siloed setups , such as among hospitals , or on mobile devices in the wild [ 7 ] . In this work , we focus on the latter setting . Hence , while cloud-based distributed training uses powerful high-end clients [ 19 ] , in FL these are commonly substituted by resource-constrained and heterogeneous embedded devices . In this respect , FL deployment is currently hindered by the vast heterogeneity of client hardware [ 66 , 28 , 7 ] . On the one hand , different mobile hardware leads to significantly varying processing speed [ 1 ] , in turn leading to longer waits upon aggregation of updates ( i.e . stragglers ) . At the same time , devices of mid and low tiers might not even be able to support larger models , e.g . the model does not fit in memory or processing is slow , and , thus , are either excluded or dropped upon timeouts from the training process , together with their unique data . More interestingly , the resource allocation to participating devices may also reflect on demographic and socio-economic information of owners , that makes the exclusion of such clients unfair [ 30 ] in terms of participation . Analogous to the device load and heterogeneity , a similar trend can be traced in the downstream ( model ) and upstream ( updates ) network communication in FL , which can be an additional substantial bottleneck for the training procedure [ 55 ] . 3 Ordered Dropout . In this paper , we firstly introduce the tools that act as enablers for heterogeneous federated training . Concretely , we have devised a mechanism of importance-based pruning for the easy extraction of subnetworks from the original , specially trained model , each with a different computational and memory footprint . We name this technique Ordered Dropout ( OD ) , as it orders knowledge representation in nested submodels of the original network . More specifically , our technique starts by sampling a value ( denoted by p ) from a distribution of candidate values . Each of these values corresponds to a specific submodel , which in turn gets translated to a specific computational and memory footprint ( see Table 1b ) . Such sampled values and associations are depicted in Fig . 2 . Contrary to conventional dropout ( RD ) , our technique drops adjacent components of the model instead of random neurons , which translates to computational benefits3 in today ’ s linear algebra libraries and higher accuracy as shown later . 3.1 Ordered Dropout Mechanics . The proposed OD method is parametrised with respect to : i ) the value of the dropout rate p ∈ ( 0 , 1 ] per layer , ii ) the set of candidate values P , such that p ∈ P and iii ) the sampling method of p over the set of candidate values , such that p ∼ DP , where DP is the distribution over P . A primary hyperparameter of OD is the dropout rate p which defines how much of each layer is to be included , with the rest of the units dropped in a structured and ordered manner . The value of p is selected by sampling from the dropout distributionDP which is represented by a set of discrete values P = { s1 , s2 , . . . , s|P| } such that 0 < s1 < . . . < s|P| ≤ 1 and probabilities P ( p = si ) > 0 , ∀i ∈ [ |P| ] such that ∑|P| i=1 P ( p = si ) = 1 . For instance , a uniform distribution over P is denoted by p ∼ UP ( i.e . D = U ) . In our experiments we use uniform distribution over the set P = { i/k } ki=1 , which we refer to as Uk ( or uniform-k ) . The discrete nature of the distribution stems from the innately discrete number of neurons or filters to be selected . The selection of set P is discussed in the next subsection . The dropout rate p can be constant across all layers or configured individually per layer l , leading to pl ∼ DlP . As such an approach opens the search space dramatically , we refer the reader to NAS techniques [ 69 ] and continue with the same p value across network layers for simplicity , without hurting the generality of our approach . Given a p value , a pruned p-subnetwork can be directly obtained as follows . For each4 layer l with width5 Kl , the submodel for a given p has all neurons/filters with index { 0 , 1 , . . . , dp ·Kle − 1 } included and { dp ·Kle , . . . , Kl − 1 } pruned . Moreover , the unnecessary connections between pruned neurons/filters are also removed6 . We denote a pruned p-subnetwork Fp with its weights wp , where F and w are the original network and weights , respectively . Importantly , contrary to existing pruning techniques [ 18 , 35 , 49 ] , a p-subnetwork from OD can be directly obtained post-training without the need to fine-tune , thus eliminating the requirement to access any labelled data . 3.2 Training OD Formulation . We propose two ways to train an OD-enabled network : i ) plain OD and ii ) knowledge distillation OD training ( OD w/ KD ) . In the first approach , in each step we first sample p ∼ DP ; then we perform the forward and backward pass using the p-reduced network Fp ; finally we update the submodel ’ s weights using the selected optimiser . Since sampling a p-reduced network provides us significant computational savings on average , we can exploit this reduction to further boost accuracy . Therefore , in the second approach we exploit the nested structure of OD , i.e . p1 < p2 =⇒ Fp1 ⊂ Fp2 and allow for the bigger capacity supermodel to teach the sampled p-reduced network at each iteration 3OD , through its nested pruning scheme that requires neither additional data structures for bookkeeping nor complex and costly data layout transformations , can capitalise directly over the existing and highly optimised dense matrix multiplication libraries . 4Note that p affects the number of output channels/neurons and thus the number of input channels/neurons of the next layer . Furthermore , OD is not applied on the input and last layer to maintain the same dimensionality . 5i.e . neurons for fully-connected layers ( linear and recurrent ) and filters for convolutional layers . RNN cells can be seen as a set of linear feedforward layers with activation and composition functions . 6For BatchNorm , we maintain a separate set of statistics for every dropout rate p. This has only a marginal effect on the number of parameters and can be used in a privacy-preserving manner [ 41 ] . via knowledge distillation ( teacher pmax > p , pmax = maxP ) . In particular , in each iteration , the loss function consists of two components as follows : Ld ( SMp , SMpmax , ylabel ) = ( 1− α ) CE ( max ( SMp ) , ylabel ) + αKL ( SMp , SMpmax , T ) where SMp is the softmax output of the sampled p-submodel , ylabel is the ground-truth label , CE is the cross-entropy function , KL is the KL divergence , T is the distillation temperature [ 21 ] and α is the relative weight of the two components . We observed in our experiments always backpropagating also the teacher network further boosts performance . Furthermore , the best performing values for distillation were α = T = 1 , thus smaller models exactly mimic the teacher output . This means that new knowledge propagates in submodels by proxy , i.e . by backpropagating on the teacher , leading to the following loss function : Ld ( SMp , SMpmax , ylabel ) = KL ( SMp , SMpmax , T ) + CE ( max ( SMpmax ) , ylabel )

Summary: The paper considers system heterogeneity in federated learning where different devices may have distinct processing capabilities or network bandwidth. Prior work addressed this scenario by not considering lower-end devices or scaling down models to fit all participants. In contrast, this paper proposes a novel (ordered) dropout technique to dynamically prune a shared model without sacrificing accuracy. Unlike random dropout, ordered dropout induces hierarchical pruning, allowing for efficient training via knowledge distillation. The paper empirically validates the proposed method on various federated learning benchmarks, demonstrating superior performance over the state-of-the-art.