Year: 2021

Abstract:
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 first 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 five 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 efficient.

heading: 1 INTRODUCTION
text: 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 significant 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) (Konečnỳ 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 significantly 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 significantly from the classical distributed machine learning framework where local data are trained with the same model architecture (Li et al., 2020b; Ben-Nun & Hoefler, 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 balanced 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.

heading: 2 RELATED WORK
text: 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 efficiency, 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 (Konečnỳ et al., 2016; Alistarh et al., 2017; Ivkin et al., 2019), and some propose to adopt split learning (Thapa et al., 2020). 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), metalearning (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 first 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 efficiency 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” for balanced non-IID data partition in classification problems. We also propose a modification of Batch Normalization (BN) (Ioffe & Szegedy, 2015) as privacy concern of running estimates hinders the usage of advanced deep learning models.

heading: 3 HETEROGENEOUS FEDERATED LEARNING

heading: 3.1 HETEROGENEOUS MODELS
text: Federated Learning aims to train a global inference model from locally distributed data {X1, . . . , Xm} across m clients. The local models are parameterized by model parameters {W1, . . . ,Wm}. The server will receive local model parameters and aggregate them into a global model Wg through model averaging. This process iterates multiple communication rounds and can be formulated as W tg = 1 m ∑m i=1W t i at iteration t. At the next iteration, W t g is transmitted to a subset of local clients and update their local models as W t+1i =W t g .
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+1i ⊆ 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 minimum modification 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 significantly 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 Wl for a single hidden layer parameterized by Wg ∈ Rdg×kg in Fig. 1, where dg and kg 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 1l as illustrated in Fig. 1. Let r be the hidden channel shrinkage ratio such that d p l = r
p−1dg and kpl = r
p−1kg . It follows that the size of local model parameters |W pl | = r2(p−1)|Wg| and the model shrinkage ratio R = |W p l |
|Wg| = 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 {m1, . . . ,mp}. Specifically, we perform global aggregation in the following way.
W pl = 1
m m∑ i=1 W pi , W p−1 l \W p l =
1
m−mp m−mp∑ i=1 W p−1i \W p i , . . . (1)
W 1l \W 2l = 1
m−m2:p m−m2:p∑ i=1 W 1i \W 2i (2)
Wg =W 1 l =W p l ∪ (W p−1 l \W p l ) ∪ · · · ∪ (W 1 l \W 2l ) (3)
For notational convenience, we have dropped the iteration index t. We denote the W pi as a matrix/tensor. The W tg [: dm, : km] denotes the upper left submatrix with a size of dm × km. Also, W p−1,t+1g \W p,t+1g denotes the set of elements included in W p−1,t+1g but excluded in W p,t+1g . We exemplify the above equations using Fig. 1. The first 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 −m2:p = m1 clients. 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 fixed 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 benefit 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.

heading: 3.2 STATIC BATCH NORMALIZATION
text: 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 Andreux et al. (2020) proposes to track running statistics locally.
We highlight an adaptation of BN named as static Batch Normaliztion (sBN) for optimizing privacy constrained heterogeneous models. During the training phase, sBN does not track running estimates and simply normalize batch data. We do not track the local running statistics as the size of local models may also vary dynamically. This method is suitable for HeteroFL as every communication round is independent. After the training process finishes, the server sequentially query local clients and cumulatively update global BN statistics. There exist privacy concerns about calculating global statistics cumulatively and we hope to address those issues in the future work. 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.

heading: 3.3 SCALER
text: 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 1rp−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
y = φ(sBN(Scaler(XmW pm + b p m))) (4)
where y is the output, φ(·) denotes a non-linear activation layer, e.g ReLU(), and W pm, bpm 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 Lm is an abstraction of the computation and communica-
Algorithm 1: HeteroFL: Heterogeneous Federated Learning Input: Data Xi distributed on M local clients, the fraction C of active clients per
communication round, the number of local epochs E, the local minibatch size B, the learning rate η, the global model parameterized by Wg , the channel shrinkage ratio r, and the number of computation complexity levels P .
System executes: Initialize W 0g and local capabilities information L1:K for each communication round t = 0, 1, 2, . . . do
Mt ← max(C ·M, 1) St ← random set of Mt clients for each client m ∈ St in parallel do
Determine computation complexity level p based on Lm rm ← r(p−1), dm ← rmdg , km ← rmkg W tm ←W tg [: dm, : km] W t+1m ← ClientUpdate(m, rm,W tm)
end for each computation complexity level p do
W p−1,t+1g \W p,t+1g ← 1Mt−Mp:P,t ∑Mt−Mp:P,t i=1 W p−1,t+1 i \W p,t+1 i
end W t+1g ← ⋃P p=1W p−1,t+1 g \W p,t+1g
Update L1:K , η (Optional) end Query representation statistics from local clients (Optional)
ClientUpdate (m, rm,Wm): Bm ← Split local data Xm into batches of size B for each local epoch e from 1 to E do
for batch bm ∈ Bm do Wm ←Wm − η∇`(Wm, rm; bm)
end end Return Wm to server
tion capabilities of a local client m. Once this information is communicated to the server, the 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.

heading: 4 EXPERIMENTAL RESULTS
text: We trained over 600 individual models for exploring and demonstrating the effectiveness of our method. We experimented with MNIST and CIFAR10 image classification tasks and the WikiText2 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 balanced non-IID data partition, we assume that the label distribution is skewed, where clients will only have examples at most from two classes and the number of examples per class is balanced. We note that there exist other kinds of non-IID data partition, e.g., the unbalanced non-IID data partition where clients may hold unbalanced labeled dataset and the feature distribution skew where clients may hold different features. 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 Table 6 of the Appendix.
To study the effectiveness of our proposed HeteroFL framework, we construct five different computation complexity levels {a, b, c, d, e} with the hidden channel shrinkage ratio r = 0.5. 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 fixed 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 figures are based on the Fix scenario, where we considered models of different sizes, and each client is allocated with a fixed 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 fixed to 50%.
The x-axis of figures 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 balanced non-IID data partition and provide a simple trick to achieve comparable results. As mentioned earlier, most state-of-the-art results ofbalanced 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 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 classification 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 sub-task given locally available label information. Specifically, 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 find 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 classification 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 of balanced non-IID data partition task. 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. 5-8. 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. 9-11. The learning curves show that the optimization of HeteroFL for the IID dataset is stable and efficient. Our method achieves better results with a fewer number of communication rounds, e.g., 800 for Heterofl 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.

heading: 5 CONCLUSIONS AND FUTURE WORK
text: 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 learning.

heading: ACKNOWLEDGMENTS
text: This work was supported by the Office of Naval Research (ONR) under grant number N00014-181-2244, and the Army Research Office (ARO) under grant number W911NF-20-1-0222.

heading: A APPENDIX
text: The appendix contains supplementary experimental results. In Table 5 we show the ablation study of Non-IID experiments. Compared to results of IID experiments shown in Table 4, Table .5 shows the ablation study Maksed CrossEntropy which is shown beneficial for balanced Non-IID data partition. In Table 6, we show the hyperparameters adopted in our experiments. In Figure 3 and 4, we show the Fix complexity assignments of MNIST and WikiText with iid data partition experiments. From Figure 6 to 8, we show the Fix complexity assignments for all balanced Non-IID data partition experiments. Figure 9 to 11, we show the learning curve of experiments with Dyanmic complexity assignments. The complete results for all experiments with Dyanmic complexity assignments can be found in Table 7, 8, and 9.

Review summaries:
This work presents a novel FL algorithm named HeteroFL (the name might sounds weird to some peoples) and 3 different simple methods to improve FL in heterogeneous conditions (i.e. both in term of clients and data partitioning). These tricks are: 1. A revised batchnormalisation; 2. a pre-activity scaling; 3. a masked loss (i.e only consider local classes)  to help with non-IID datasets. All these modifications have been tested on 3 different datasets and 2 different tasks. From the results, we can see that the proposed approach works better. Although, it is not clear from where the benefit comes. 

This paper proposes a new federated learning framework called HeteroFL, which supports the training of different sizes of local models in heterogeneous clients. Clients with higher computation capability can train larger models while clients with less computation capability train smaller models, and all these model architectures belong to the same model class. This approach dramatically benefits clients with limited computation capability and fully exploits their computation power.