FedMorph: Communication Efficient Federated Learning via Morphing Neural Network

Abstract: The two fundamental bottlenecks in Federated Learning (FL) are communication and computation on heterogeneous edge networks, restricting both model capacity and user participation. 
To address these issues, we present FedMorph, an approach to automatically morph the global neural network to a sub-network to reduce both the communication and local computation overloads. 
FedMorph distills a fresh sub-network from the original one at the beginning of each communication round while keeps its `knowledge' as similar as the aggregated model from local clients in a federated average (FedAvg) like way.
The network morphing process considers the constraints, e.g., model size or computation flops, as an extra regularizer to the objective function.
To make the objective function solvable, we relax the model with the concept of soft-mask.
We empirically show that FedMorph, without any other tricks, reduces communication and computation overloads and increases the generalization accuracy.
E.g., it provides an $85\times$ reduction in server-to-client communication and $18\times$ reduction in local device computation on the MNIST dataset with ResNet8 as the training network.
With benchmark compression approaches, e.g., TopK sparsification, FedMorph collectively provides an $847\times$ reduction in upload communication.

Reviews and author responses:

 We sincerely thank the objective suggestions from the reviewer. Here we need to make some clarifications for the misunderstandings of the reviewer.
After seriously double-checking our codes, we believe the results of ResNet18 on Cifar10 on NonIID partition differ from the reviewers' might due to the different hyper-parameter settings. We select the hyper-parameters which perform best for the baseline methods on the IID partition and continue to choose them for the experimental consistency.
The reasons we choose to use the test dataset for distillation are based on: 1). To make a fair comparison with the baselines, we are unable to split a validation dataset from the training dataset. Besides, our experiments are built on supervised learning, and the pixels of the test dataset is used to minimize the network overloads, which are not used for local training.
2). The NonIID partition experiments simulate that when the distribution of the local training dataset evolves and differs from the server collected. The experimental results prove our method works in such conditions.
We hope the explanations will release the concerns from the reviewer.

 We sincerely thank the objective suggestions from the reviewer. Here we need to make some clarifications for the misunderstandings of the reviewer.
The reasons we choose to use the test dataset for distillation are based on: 1). To make a fair comparison with the baselines, we are unable to split a validation dataset from the training dataset. Besides, our experiments are built on supervised learning, and the pixels of the test dataset is used to minimize the network overloads, which are not used for local training.
2). The NonIID partition experiments simulate that when the distribution of the local training dataset evolves and differs from the server collected. The experimental results prove our method works in such conditions.
We hope the explanations will release the concerns from the reviewer.

 We sincerely thank the objective suggestions from the reviewer. Here we need to make some clarifications for the misunderstandings of the reviewer.
It's always hard to collect private data from users, such as their tendency to do something (the labels of the dataset), which is why FL is a hot topic for dealing with such situations.
However, there are always some 'cheap' datasets (without labels), which is open-sourced and prepared for anyone to use.
Our method requires the server to collect a dataset similar to the global dataset of all users but does not require the labels, which is economical and applicable.
The theory bound of our method depends on the similarity of the server-owned dataset and the clients owned. 
The experiments on the Non-IID partition illustrate the effectiveness of our method when the server-owned dataset differs from the clients owned.

This paper proposes distilling the global model's knowledge in FL into a smaller sub-network. This way, communication and computation only have to deal with the small sub-networks instead of the full model.  Strengths: 
- The paper reports two to three orders of magnitudes savings in communication bandwidth using the proposed method. 

Weaknesses:
- The core idea could be presented more straightforwardly.
- It is unrealistic to assume that there exists a representative validation dataset for the global model in FL. How would one collect that data? Without a representative dataset for the global model, the knowledge distillation for the global model becomes infeasible, as the paper also points out in Theorem 2. The r.h.s. of the bound on the proposed method's error in Theorem 2 depends on the distance between the distillation and the local training datasets, which can be arbitrarily large in practice. If I have understood the paper correctly, then I think the second point mentioned in the weaknesses of this paper can be a critical issue. For that reason, I currently cannot recommend accepting this paper.  

The paper proposes FedMorph to address the communication and computation heterogeneity problem in federated learning. At every round, small morphed sub networks are send to clients for local training. The server updates the global model by distilling from the aggregated morphed subnets. The idea of training smaller networks to workaround heterogeneity is not novel, though the authors proposed a formulation that optimizes the subnetwork together with a distillation loss when updating server model parameter. Specifically, the proposed method shares a lot of similarities with federated dropout, i.e. subnetwork is obtained by dropping columns from model parameters, with the only difference being that the pruning units are are optimized through Eq 7.

Optimizing the distillation loss relies on using a validation dataset on the server, and the quality of distillation relies heavily on whether the distribution the validation dataset is close to that of the decentralized training set. This seems to be a rather strong requirement in federated learning where the data is hard to obtain and the distribution may evolve over time. Therefore it makes me question whether the distillation is a realistic proposal in practice.

Also related to the validation set used for distillation, the authors use the test dataset as the distillation dataset (Page 14, Appendix B). This seems to me a major experimentation flaw that the pixels from the test set are leaked into the training algorithm. This likely explains why in Figure 4 the reported generalization gaps of the proposed method are much smaller than those of the alternative methods. This flaw renders the experimentation results and conclusions not trustworthy. The use of the test set in the experiments as the distillation dataset is a major concern.

The nature of FL workloads poses three evident system and machine learning challenges: communication overheads to broadcast models to clients and send locally-updated models back to the server; compute overheads, since clients might be in the form of constrained devices or battery powered; and overfitting, due to the asymmetric data distribution in the clients in terms of labels and number of training examples. The work here proposed (FedMorph) addresses these three challenges by morphing/extracting sub-models from the global model and dispatch these to the clients to perform local training. Then, the morphed sub-networks get aggregated into the global model via distillation. A FL setup using FedMorph demonstrates a much reduced overffiting.  ### Strong and Weak points 

- (strong): Clearly stated the three of the challenges to push FL forward: (1) training larger models often requires broadcasting larger files across the network; (2) training larger models is computationally expensive and might overwhelm more constrained devices; (3) the capacity of the model trained on the clients should be a function of the clients' data complexity (e.g. amount of training examples, distribution over labels).

- (strong): FedMorph addresses the three points above.

- (strong): Compact but complete formulation of a generic FL setup in Section 3.1

- (strong): Results from figure 4 evidence the contribution of FedMorph towards easing the overfitting problem. This is one of the most interesting results of this work. 


- (major weak): Distillation has been studied in several works in the context of FL. I can point the authors to this one [1] which aligns with the goals of FedMorph. Authors should include in the Related Work and compare against this work (maybe others too) in terms of: (1) communication costs savings, (2) easing the overfit problem, (3) reducing the compute and memory footprints of performing local training.

- (weak): Part of the goal of FedMorph is to make lightweight training stages on the clients. This broadly aligns with other works attempting to do "efficient training". I would recommend Authors to include this in the related work. This new paragraph could include works doing sparse training [2], low-precision training [3], pruning at initialization [4,5] to directly learn a smaller model.

- (weak): The works in the second paragraph of the introduction talking about quantization do not make training more efficient. Those works learn a quantized model that can be use to accelerate inference (it results in model compression too) but the training process happens generally at full precision and the quantization is only "simulated". This is often referred to as Quantization Aware Training, which is more computationally expensive that standard training as it involves more intermediate operations. 

- (weak): Are ResNet-18 baselines in Figure 5 correctly trained?. ResNet-18 for CIFAR-10 using 100 clients and sampling 10 per round in a non-IID using DLA partitioning with alpha<1.0 setting is known to result in >80% accuracy for FedAvg. However, Figure 5 (even in sparse=100 setting) shows <75%. I believe the the performance gap with FedMorph/FedMorphProx would be significantly reduced if hyper-parameter are set correctly. I can point the Authors to this work [6] that builds on top of FedDropout and wich uses a similar FL setting and reach ~84% acc. I have also run the experiment myself (excluding all the FjORD related components) and I consistently achieve a similar value.

- (weak): I was surprised not to find results/curves showing the quality of the distillation process, which is by far of of the main challenges this work has to deal with to make FedMorph work. Or did I miss something?

- (weak): Measuring Test accuracy after each round is not a good practice. It seems this was done for the majority of the results reported. A validation set should be constructed from the training set (before or after partitioning into the N clients). ### Recommendation

For the IID setting, FedMorph shows ~10x compression and ~5x FLOPs reduction for just a 0.94% accuracy drop (last row in Table 1). This is an encouraging result. However, as I stated in my point above, the performance of FedMorph on the non-IID scenario is not proven. Other works with the exactly same FedAvg baselines report results identical (if not better) than FedMorph for ResNet-18 CIFAR-10 (i.e. the most complex setting evaluated in this work).

My recommendation is for the authors to compile a list of newer experiments that better showcase the unique issues FedMorph addresses.

### Supporting my Recommendation

Please see points above.

### Minor points

- Include comparison vs. (Zhu et al. 2021)
- The first paragraph of the intro states that FL "ensures data privacy". I think "ensures" should be replaced with "enables". This is because FL is not that private out of the box, several components need to be added and, careful engineering practices need to be inplace, as the works cited on that line suggest.

- Figure 3 shows per-layer compression ratio. I believe that the numbers aren't correct since in general a "compression ratio" <1 means that there is no compression. Probably the Authors forgot to do 1.0/"ratio"?

- Section 4 feels quite long given the amount of contributions it adds to the paper. I would advice authors to compress it and add results from the appendix. 

### References

- [1] https://arxiv.org/abs/2007.14513
- [2] https://arxiv.org/abs/2001.01969
- [3] https://papers.nips.cc/paper/2020/hash/13b919438259814cd5be8cb45877d577-Abstract.html
- [4] https://arxiv.org/abs/2002.07376
- [5] https://arxiv.org/abs/2007.00389
- [6] https://arxiv.org/abs/2102.13451

Meta-review: The paper proposes FedMorph to address the communication and computation heterogeneity problem in federated learning. The proposed FedMorph extracts sub-models from the global model and dispatch these to the clients to perform local training. Then, the morphed sub-networks get aggregated into the global model via distillation. 

The paper reports two to three orders of magnitudes savings in communication bandwidth using the proposed method. However, as agreed by all reviewers, the paper has some critical problems as listed below that prevent it being accepted at this point. 

1. The idea of training smaller networks to workaround heterogeneity is not novel, though the authors proposed a formulation that optimizes the subnetwork together with a distillation loss when updating server model parameters.  Authors should include in the Related Work and compare against other distillation-related FL work in terms of: (1) communication costs savings, (2) easing the overfit problem, (3) reducing the compute and memory footprints of performing local training.

2. Optimizing the distillation loss relies on using a validation dataset on the server, and the quality of distillation relies heavily on whether the distribution the validation dataset is close to that of the decentralized training set. This seems to be a rather strong requirement in federated learning where the data is hard to obtain and the distribution may evolve over time. Therefore it makes me question whether the distillation is a realistic proposal in practice.

3. The test dataset is used as the distillation dataset which is a major experimentation flaw that the pixels from the test set are leaked into the training algorithm. 

4. It may be unrealistic to assume that there exists a representative validation dataset for the global model in FL. The proposed method's error (in Theorem 2) depends on the distance between the distillation and the local training datasets, which can be arbitrarily large in practice.

Decision: Reject