H ETERO FL: C OMPUTATION AND C OMMUNICATION E FFICIENT F EDERATED L EARNING FOR H ETEROGENEOUS C LIENTS
Anonymous authors
Paper under double-blind review
A BSTRACT
Federated Learning (FL) is a method of training machine learning models on private data distributed over a large number of possibly heterogeneous clients such as mobile phones and IoT devices.
In this work, we propose a new federated learning framework named HeteroFL to address heterogeneous clients equipped with very different computation and communication capabilities.
Our solution can enable the training of heterogeneous local models with varying computation complexities and still produce a single global inference model.
For the ﬁrst time, our method challenges the underlying assumption of existing work that local models have to share the same architecture as the global model.
We demonstrate several strategies to enhance FL training and conduct extensive empirical evaluations, including ﬁve computation complexity levels of three model architecture on three datasets.
We show that adaptively distributing subnetworks according to clients’ capabilities is both computation and communication efﬁcient.
I NTRODUCTION
Mobile devices and the Internet of Things (IoT) devices are becoming the primary computing resource for billions of users worldwide (Lim et al., 2020).
These devices generate a signiﬁcant amount of data that can be used to improve numerous existing applications (Hard et al., 2018).
From the privacy and economic point of view, due to these devices’ growing computational capabilities, it becomes increasingly attractive to store data and train models locally.
Federated learning (FL) (Kone ˇ cn ` y et al., 2016; McMahan et al., 2017) is a distributed machine learning framework that enables a number of clients to produce a global inference model without sharing local data by aggregating locally trained model parameters.
A widely accepted assumption is that local models have to share the same architecture as the global model (Li et al., 2020b) to produce a single global inference model.
With this underlying assumption, we have to limit the global model complexity for the most indigent client to train its data.
In practice, the computation and communication capabilities of each client may vary signiﬁcantly and even dynamically.
It is crucial to address heterogeneous clients equipped with very different computation and communication capabilities.
In this work, we propose a new federated learning framework called HeteroFL to train heterogeneous local models with varying computation complexities and still produce a single global inference model.
This model heterogeneity differs signiﬁcantly from the classical distributed machine learning framework where local data are trained with the same model architecture (Li et al., 2020b;
Ben-Nun & Hoeﬂer, 2019).
It is natural to adaptively distribute subnetworks according to clients’ capabilities.
However, to stably aggregate heterogeneous local models to a single global model under various heterogeneous settings is not apparent.
Addressing these issues is thus a key component of our work.
Our main contributions of this work are three-fold.
• We identify the possibility of model heterogeneity and propose an easy-to-implement framework HeteroFL that can train heterogeneous local models and aggregate them stably and effectively into a single global inference model.
Our approach outperforms state-ofthe-art results without introducing additional computation overhead.
• Our proposed solution addresses various heterogeneous settings where different proportions of clients have distinct capabilities.
Our results demonstrate that even when the model
heterogeneity changes dynamically, the learning result from our framework is still stable and effective.
• We introduce several strategies for improving FL training and demonstrate that our method is robust against the non-IID statistical heterogeneity.
Also, the proposed method can reduce the number of communication rounds needed to obtain state-of-the-art results.
Experimental studies have been performed to evaluate the proposed approach.
R ELATED W ORK
Federated Learning aims to train massively distributed models at a large scale (Bonawitz et al., 2019).
FedAvg proposed by McMahan et al.
(2017) is currently the most widely adopted FL baseline, which reduces communication cost by allowing clients to train multiple iterations locally.
Major challenges involved in FL include communication efﬁciency, system heterogeneity, statistical heterogeneity, and privacy (Li et al., 2020b).
To reduce communication costs in FL, some studies propose to use data compression techniques such as quantization and sketching (Koneˇcn`y et al., 2016;
Alistarh et al., 2017; Ivkin et al., 2019).
To tackle system heterogeneity, techniques of asynchronous communication and active sampling of clients have been developed (Bonawitz et al., 2019;
Nishio & Yonetani, 2019).
Statistical heterogeneity is the major battleground for current FL research.
A research trend is to adapt the global model to accommodate personalized local models for non-IID data (Liang et al., 2020), e.g., by integrating FL with other frameworks such as assisted learning (Xian et al., 2020), meta-learning (Jiang et al., 2019; Khodak et al., 2019), multi-task learning (Smith et al., 2017), transfer learning (Wang et al., 2019;
Mansour et al., 2020), knowledge distillation (Li & Wang, 2019) and lottery ticket hypothesis (Li et al., 2020a).
Nevertheless, these personalization methods often introduce additional computation and communication overhead that may not be necessary.
Another major concern of FL is data privacy (Lyu et al., 2020), as model gradient updates can reveal sensitive information (Melis et al., 2019) and even local training data (Zhu et al., 2019;
Zhao et al., 2020).
To our best knowledge, what we present is the ﬁrst work that allows local models to have different architectures from the global model.
Heterogeneous local models can allow local clients to adaptively contribute to the training of global models.
System heterogeneity and communication efﬁciency can be well addressed by our approach, where local clients can optimize low computation complexity models and therefore communicate a small number of model parameters.
To address statistical heterogeneity, we propose a “Masking Trick” and demonstrate that personalization is unnecessary for non-IID data in classiﬁcation problems.
We also propose a modiﬁcation of Batch Normalization (BN) (Ioffe & Szegedy, 2015) as privacy concern of running estimates hinders the usage of advanced deep learning models.
H ETEROGENEOUS F EDERATED L EARNING H ETEROGENEOUS M ODELS
Federated Learning aims to train a global inference model from locally distributed data { X 1 , . . .
X m } across m clients.
The local models are parameterized by model parameters { W 1 , . . .
, W m } .
The server will receive local model parameters and aggregate them into a global model W g through model averaging.
This process iterates multiple communication rounds and can be formulated as W tg
= 1 m (cid:80) mi =1 W ti at iteration t .
At the next iteration, W tg is transmitted to a subset of local clients and update their local models as W t +1 i =
W tg .
In this work, we focus on the relaxation of the assumption that local models need to share the same architecture as the global model.
Since our primary motivation is to reduce the computation and communication complexity of local clients, we consider local models to have similar architecture but can shrink their complexity within the same model class.
To simplify global aggregation and local update, it is tempting to propose local model parameters to be a subset of global model parameters W t +1 i ⊆
W tg .
However, this raises several new challenges like the optimal way to select subsets of global model parameters, compatibility of the-state-of-art model architecture, and mini- [Figure] mum modiﬁcation from the existing FL framework.
We develop Heterogeneous Federated Learning (HeteroFL) to address these issues in the context of deep learning models.
A variety of works show that we can modulate the size of deep neural networks by varying the width and depth of networks (Zagoruyko & Komodakis, 2016; Tan & Le, 2019).
Because we aim to reduce the computation complexity of local models, we choose to vary the width of hidden channels.
In this way, we can signiﬁcantly reduce the number of local model parameters, while the local and global model architectures are also within the same model class, which stabilizes global model aggregation.
We demonstrate our method of selecting subsets of global model parameters W l for a single hidden layer parameterized by W g ∈ R d g × k g in Fig.
1, where d g and k g are the output and input channel size of this layer.
It is possible to have multiple computation complexity levels
W pl ⊂
W p − 1 l · · · ⊂
W 1 l as illustrated in Fig.
Let r be the hidden channel shrinkage ratio such that d pl = r p − 1 d g and k pl = r p − 1 k g .
It follows that the size of local model parameters | W pl | = r 2( p − 1) | W g | and the model shrinkage ratio R = | W pl | | W g | = r 2( p − 1) .
With this construction, we can adaptively allocate subsets of global model parameters according to the corresponding capabilities of local clients.
Suppose that number of clients in each computation complexity level is { m 1 , . . .
, m p } .
Speciﬁcally, we perform global aggregation in the following way.
[Equation]
For notational convenience, we have dropped the iteration index t .
We denote the W pi as a matrix/tensor.
The W tg
[: d m , : k m ] denotes the upper left submatrix with a size of d m × k m .
Also,
W p − 1 ,t +1 g \ W p,t +1 g denotes the set of elements included in W p − 1 ,t +1 g but excluded in W p,t +1 g .
We exemplify the above equations using Fig.
The ﬁrst part of Equation (1) shows that the smallest part of model parameters (blue, p = 3 ) is aggregated from all the local clients that contain it.
In the second part of Equation (1), the set difference between part p − 1 (orange) and p (blue) of model parameters is aggregated from local clients with computation complexity level smaller than p − 1 .
In Equation (2), the red part of model parameters can be similarly aggregated from m − m 2: p = mclients.
In Equation (3), the global model parameters W tg is constructed from the union of all disjoint sets of the partition.
In summary, each parameter will be averaged from those clients whose allocated parameter matrix contains that parameter.
Thus, a model of an intermediate complexity will have parameters fully averaged with all the other larger models but partially with smaller models (according to the corresponding upper left submatrix).
Several works show that wide neural networks can drop a tremendous number of parameters per layer and still produce acceptable results (Han et al., 2015; Frankle & Carbin, 2018).
The intuition is thus to perform global aggregation across all local models, at least on one subnetwork.
To stabilize global model aggregation, we also allocate a ﬁxed subnetwork for every computation complexity level.
Our proposed inclusive subsets of global model parameters also guarantee that smaller local models will aggregate with more local models.
Thus, small local models can beneﬁt more from global aggregation by performing less global aggregation for part of larger local model parameters.
We empirically found that this approach produces better results than uniformly sampled subnetworks for each client or computation complexity level.
4. S TATIC B ATCH N ORMALIZATION
After global model parameters are distributed to active local clients, we can optimize local model parameters with private data.
It is well-known that the latest deep learning models usually adopt Batch
Normalization (BN) to facilitate and stabilize optimization.
However, classical FedAvg and most recent works avoid BN.
A major concern of BN is that it requires running estimates of representations at every hidden layer.
Uploading these statistics to the server will cause higher communication costs and privacy issues.
We believe that private data can be readily reconstructed with these running estimates.
However, BN facilitates optimization and achieve the same performance with less number of communication rounds.
However, the transmission of running statistics as used in the vanilla BN will likely to cause privacy issues.
Therefore, we propose static Batch Normaliztion (sBN) for optimizing privacy constrained deep neural networks.
During the training phase, sBN does not track running estimates and simply normalize batch data.
We only calculate statistics of hidden representations from local data after the model converges and is ready for inference.
This method is suitable for the FL framework as local models do not need to upload running estimates during training.
Local models only upload their statistics for once after optimization is completed.
Thus, this method greatly reduces the risk of leaking private data.
We also empirically found this trick significantly outperforms other forms of normalization methods including the InstanceNorm (Ulyanov et al., 2016), GroupNorm (Wu &
He, 2018) , and LayerNorm (Ba et al., 2016) as shown in Table 4 and Table 5.
4. S CALER
There still exists another cornerstone of our HeteroFL framework.
Because we need to optimize local models for multiple epochs, local model parameters at different computation complexity levels will digress to various scales.
This known phenomenon was initially discussed by the dropout (Srivastava et al., 2014).
To directly use the full model during the inference phase, inverted dropout with dropout rate q scales representations with 11 − q during the training phase.
In practice, dropout is usually attached after the activation layer as the selection of subnetworks is performed with masking.
Our method directly selects subnetworks from the subsets of global model parameters.
Therefore, we append a Scaler module right after the parametric layer and before the sBN and activation layers.
The Scaler module scales representations by [Equation] 1 during the training phase.
After the global aggregation, the global model can be directly used for inference without scaling.
To further illustrate this point, we include a comprehensive ablation study in Tables 4 and 5.
A typical linear hidden layer used in our HeteroFL framework can be formulated as
[Equation]
where y is the output, φ ( · ) denotes a non-linear activation layer, e.g ReLU(), and W pm , b pm are the weight and bias for local model m at computation complexity level p .
With all the practical methods mentioned above, we propose the complete pseudo-code for our HeteroFL framework in Algorithm 1.
The local capabilities information L m is an abstraction of the computation and communication capabilities of a local client m .
Once this information is communicated to the server, the server
System executes:
Initialize
do
M t ← max( C · M
S t ← random set of M t clients for m ∈ t parallel do
Determine computation complexity level p based on L m r m ← r ( p − 1) , d m ← r m d g , k m ← r m k g
W tm ← W tg
[: d m , : k m ]
W t +1 m ← ClientUpdate ( m r m , W tm ) end for p do
W p − 1 ,t +1 g \ W p,t +1 g ← 1 M t − M p : P,t (cid:80) M t − M p : P,t i =
W p − 1 ,t +1 i \ W p,t
+1 i end
W t +1 g ← (cid:83)
Pp =1 W p − 1 ,t +1 g \ W p,t +1 g
Update L 1: K , η (Optional) end
Query representation Optional)
ClientUpdate ( m, r m , W m ) :
B m ← split local data X m into batches of size B for local epoch from 1 to E do for batch b m ∈ B m do
W m ← W m − η ∇ (cid:96) ( W m , r m ; b m ) end end
Return W m to server can know the model complexity that should be allocated to the client. We can also optionally update learning rates to facilitate optimization and local capabilities information if changing dynamically.  E XPERIMENTAL R ESULTS We trained over 600 individual models for exploring and demonstrating the effectiveness of our method.
We experimented with MNIST and CIFAR10 image classiﬁcation tasks and the WikiText language modeling task (LeCun et al., 1998; Krizhevsky et al., 2009; Merity et al., 2016;
Devlin et al., 2018).
Our experiments are performed with three different models including a CNN for MNIST, a preactivated ResNet (PreResNet18)
(He et al., 2016) for CIFAR10 and a Transformer (Vaswani et al., 2017) for WikiText2.
We replace BN in CNN and PreResNet18 with our proposed sBN, and attach the
Scaler module after each convolution layer.
To study federated optimization, we adopt data partition the same as in (McMahan et al., 2017; Liang et al., 2020).
We have 100 clients, and the fraction C of active clients per communication round is 0 . 1 throughout our experiments.
For IID data partition, we uniformly assign the same number of data examples for each client.
For non-IID data partition, we assume that the label distribution is skewed, where clients will only have examples at most from two classes.
We note that there exist other kinds of non-IID data partition, e.g., the feature distribution skew where clients may hold different features, and the quantity distribution skew where clients may hold unbalanced labeled dataset.
We conduct a masked language modeling task with a 15% masking rate and uniformly assign balanced data examples for each client.
It needs to point out that each client will roughly have 3000 different words in their local dataset, while the total vocabulary size is 33278.
The details regarding hyperparameters and model architecture can be found in Tableof the Appendix.
To study the effectiveness of our proposed HeteroFL framework, we construct ﬁve different computation complexity levels { a, b, c, d, e } with the hidden channel shrinkage ratio r = 0 .
We have tried various shrinkage ratios, and we found that it is most illustrative to use the discrete complexity levels 0.5, 0.25, 0.125, and 0.0625 (relative to the most complex model). For example, model ‘a’ has all the model parameters, while models ‘b’ to ‘e’ have the effective shrinkage ratios 0.5, 0.25, 0.125, and 0.0625. We note that the complexity of ‘e’ is close to a logistic regression model.
Our experiments indicated that the ratio can be arbitrary between (0 , 1] and dynamically change.
In practice, using a dictionary of discrete complexity levels are convenient for coordination purposes.
Each local client is assigned an initial computation complexity level.
We annotate Fix for experiments with a ﬁxed assignment of computation complexity levels, and Dynamic for local clients uniformly sampling computation complexity levels at each communication round.
We perform distinct experiments for Fix and Dynamic assignments.
All the ﬁgures are based on the Fix scenario, where we considered models of different sizes, and each client is allocated with a ﬁxed size. All the tables are based on the Dynamic scenario, where we randomly vary the allocation of clients’ model complexity, and the ratio of the number of weak learners is ﬁxed to 50%.
The x-axis of ﬁgures represents the average model parameters.
When 10% clients use the model ’a’ and 90% use the model ’e’, the average number of model parameters is 0 .
1 × (size of model ‘a’) + 0 .
9 × (size of model ‘e’).
We interpolate this partition from 10% to 100% with step size 10% to demonstrate the effect of proportionality of clients with various computation complexity levels.
To demonstrate the effect of dynamically varying computation and communication capabilities, we uniformly sample from various combinations of computation complexity levels.
For example, model ’a-b-c-d-e’ means that we uniformly sample from all possible available levels for every active client at each communication round.
We show the number of model parameters, FLOPs, and Space (MB) to indicate the computation and communication requirements of our methods.
For example, since we uniformly sample levels, model a − e calculates computational metrics by averaging those of model a and e .
The ratio is calculated between the number of parameters of a given model with respect to its 100% global model.
We compare our results to other baseline methods like Standalone, FedAvg, and LG-FedAvg gathered from (Liang et al., 2020).
Standalone means there is no communication between clients and the server.
In our experimental studies, we considered more complex models compared with the existing work.
In particular, the baseline models in LG-FedAvg used MLP (on MNIST) and CNN (on CIFAR10).
In terms of the number of parameters, our models ‘a-e’ (on MNIST) and ‘b-e’ (on CIFAR10) are comparable with those baselines. In terms of the FLOPs, our model ‘d-e’ (on MNIST and CIFAR10) can be compared with those baselines. The single-letter models ‘a’, ‘b’, ‘c’, ‘d’, ‘e’ are our implementations of the FedAvg equipped with the sBN and Masking CrossEntropy.
The takeaway of Table 2 is that a weak learner that can only train a small model ‘e’ (on CIFAR10)(77.09%) can boost its performance to ‘c-e’ (86.88%), ‘b-e’ (89.10%), or ‘a-e’ (90.29%), which are close to the scenario where all the learners are strong, namely c(87.55%), b(89.82%), or a(91.99%). In particular, in ‘c-e’, ‘b-e’, or ‘a-e’, half of the clients are trained with larger models ‘c’, ‘b’, or ‘a’, while the other half are trained with the model ‘e’. Only the aggregated global models ‘c’, ‘b’, or ‘a’ are used during the testing stage. Although weak clients train smaller models ‘e’, they will test with the largest models ‘c’, ‘b’, or ‘a’ to gain better performance.
Full results including other possible combinations can be found in appendix in Table 7-9.
Finally, our method is robust to dynamically varying model complexities.
It is worth noting that our method does not incur any additional computation overhead and can be readily adapted to existing applications.
We also perform experiments for non-IID data partition and provide a simple trick to achieve comparable results.
As mentioned earlier, most state-of-the-art results of non-IID datasets suggest the personalization of local models to achieve better local results (Smith et al., 2017; Liang et al., 2020;
Li et al., 2020a).
Here, the Local results assume that the training data distribution and test data
[Figure]
distribution for each local client are the same, and assign zero probability for those classes that are not presented to a client during training. The Global results were calculated from the global model applied to the test data directly. The Local results were cumulatively averaged from the performance of each data example on each local client.
Zhao et al.
(2018) showed that the failure of non-IID FL is related to the weight divergence among local model parameters trained locally for many iterations.
The weight divergence mostly occurs in the last classiﬁcation layer of networks.
Thus, instead of a full Cross-Entropy Loss for all classes, we are motivated to train each local model only with their corresponding classes.
In this way, each local model will train a subtask given locally available label information.
Speciﬁcally, we mask out the output of the model before passing it Cross-Entropy Loss, which we named as Masked Cross-Entropy Loss .
We experimented with several different ways of masking, we ﬁnd replacing the last layer outputs that are not associated with local labels with zero achieves both stable and comparable local and global results.
When aggregating local model parameters, we do not aggregate the untrained parameters in the last classiﬁcation layers.
Either the server can infer this information implicitly, or the local clients can report which classes they have to the server explicitly.
We provide a comprehensive ablation study in Tables 5.
The results show that Masked Cross-Entropy Loss significantly improve local performance and moderately global performance.
Since our primary focus is to address model heterogeneity, we leave the analysis of this trick to future work.
We show the results of Fix experiments in appendix in
Fig.
Dynamic non-IID results are also included in Table 1-3.
The results show that our method performs comparably to those with personalized local models.
Our method is readily adaptable, free of computation overhead, and only rely on the single global model for testing local and global results.
It allows local clients to switch to another subtask simply by changing its mask without querying the server for others’ personalized models.
We show the learning curves of 50% Fix and Dynamic assignments in appendix in Fig.
The learning curves show that the optimization of HeteroFL for the IID dataset is stable and efﬁcient. [Table] [Table] Full results can be found in Table
[Table]
Our method achieves better results with a fewer number of communication rounds, e
[Table] g., 800 for Heteroﬂ and 1800 for LG-FedAvg (Liang et al., 2020). We empirically discover gradient clipping stabilizes the optimization of HeteroFL as it prevents small models from gradient explosion.
We can therefore adopt a universal learning rate for heterogeneous local models.
It is also perceivable that aggregation of model parameters trained with non-IID data makes the optimization less stable.
Results of Dynamic show that global aggregation of dynamically varying computation complexities is stable.
[Table]
Full can be found in
[Table]
C ONCLUSIONS AND F UTURE W ORK
We propose Heterogeneous Federated Learning (HeteroFL), which shows the possibility of coordinatively training local models much smaller than a global model to produce a single global inference model.
Our experiments show that FL can be made more practical by introducing HeteroFL and sBN and Mased Cross-Entropy Loss, as HeteroFL fully exploits local clients’ capabilities and achieves better results with a fewer number of communication rounds.
We demonstrate our results with various model architectures, including CNN, PreResNet18, and Transformer, and show that our method is robust to statistical heterogeneity and dynamically varying local capabilities.
A future direction is to distinct model classes as well as model heterogeneity.
Also, the proposed methods may be emulated to address heterogeneous few-shot learning, multi-modal learning, and multi-task