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 proposed a scheme for training deep models in the setup of federated learning, where updates to different nested submodels can be trained at different clients with different hardware capabilities.  (+) I appreciate solid grounding in practical motivation with potential significant practical value, and I would love to see further works in this direction.  (-) I feel execution on these ideas is not good enough, especially in terms of matching empirical baselines.