Year: 2021

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 subnetwork 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× reduction in server-to-client communication and 18× 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× reduction in upload communication.

heading: 1 INTRODUCTION
text: Federated Learning (FL) (Li et al., 2021; Bonawitz et al., 2019; Kairouz et al., 2019; Li et al., 2019) ensures data privacy by decoupling the training on the dataset in local clients and model aggregation in the global server. The iterative process of the local updates optimized on each client and aggregating the updates by the global server assures the convergence of the global model and the excellent generalization on test datasets.
(A). Despite the benefits gained by the paradigm of local-train and global-aggregation of FL, the server-clients communication burden and limited local computing resources suspend the deployment of the State-of-The-Art (SOTA) large neural networks on a wide scale. Recently the lossy model compression techniques, such as gradient sparsification (Wangni et al., 2017; Zhou et al., 2021) and quantization (Courbariaux et al., 2016; Hubara et al., 2017), have been studied to relieve the communication issues. Practical as they announced, but the limitations are also apparent. The gradient sparsification technique reduces the communication burden by selectively updating local gradients. Hence, they only work for client-to-server communication and have no compression benefits in the downstream direction. Meanwhile, the quantization technique compresses the communication and computation by changing the floating operations into low precision operations, and it works in both directions. However, its compression ability is at most 32× by reducing 32-bits floating operations to the 1-bit operations, not to mention the performance degradation.
Absorbing the benefits from the former lossy compression techniques, FedDropout (Caldas et al., 2018) is another trial to reduce the downstream direction communication. It builds upon the basic idea of dropout (Srivastava et al., 2014) by randomly dropping some neurons of the global model at each round, resulting in a smaller sub-network with smaller weight matrices before broadcasting it to local clients. The updated sub-networks by local clients will map back to the global network.
Distillation of the ‘knowledge’ from an extensive neural network (Hinton et al., 2015) is another branch to solve local devices’ communication and computation burdens. The main idea is to train a small model to keep the output similar to the large one on a sizeable prepared dataset. It has significant performance on the recent developments (Sanh et al., 2019) (Jiao et al., 2019). However, it requires a well pre-trained large model and a vast dataset for training the smaller network, which is not always satisfactory in the context of FL.
(B). Besides considering communication and computation budget, broadcasting a raw neural network may also cause a severe generalization problem. While it is almost common sense in the deep learning community that training a well-designed deep neural network performs better than a shallow one, the statement only establishes when the training dataset is sufficient to avoid overfitting (Pitt & Myung, 2002; Hinton et al., 2012). In FL settings, each local client holds a dataset in a small ratio of the whole, which is universal when the number of clients is enormous. Hence, human labor’s well-designed deep neural network works well on the whole dataset, nevertheless easily overfit the local datasets. The overfitting will impede the number of local iterations, which is a crucial hyper-parameter in FL settings (McMahan et al., 2017), resulting in a deteriorated model performance. Kairouz et al. (2019) has a similar statement in their Network Architecture Search (NAS) section that a well-designed network architecture in centralized settings may not work well in federated learning settings.
This paper tries to solve the earlier issues by decoupling the network architectures on local clients and the global server. In detail, the global server maintains a large neural network suitable in a centralized setting, and keeps the architecture untouched throughout the entire process. During each communication round, the global server morphs a new sub-network from the maintained architecture and broadcasts it to selected clients for local optimization. The maintained network optimizes its weights at the morphing process by learning from the average aggregated network updated by local clients. Meanwhile, a newly morphed sub-network is constrained by keeping the knowledge similar to the average aggregated one while minimizing the number of its parameters. By morphing the shared neural network into a smaller size among all communication rounds, we (1) decrease the communication in both upload and download directions; (2) relieve the computation overload of local clients; (3) reduce the generalization error caused by overfitting on local datasets with a large network.

heading: 2 RELATED WORKS
text: Neural Architecture Search Our work is closely related to the NAS (Zoph & Le, 2016), whose purpose is automatically designing neural network architecture such as the number of layers and the number of neurons or filters in each layer. The search strategies include random search, evolutionary methods, Bayesian optimization, and gradient-based methods (Elsken et al., 2019). Since NAS involves a vast search space, recent works focus on improving the speed and efficiency by attentive sampling (Wang et al., 2021), untrained scheme (Mellor et al., 2021), and block-wisely search with knowledge distillation (Li et al., 2020). Besides, the network architecture search under the scenario of federated learning scenario was also explored in (Zhu et al., 2021), aiming to reduce both computation and communication.
Model Compression To reduce the complexity of deep neural networks, model compression was first proposed in (Buciluǎ et al., 2006), followed by tremendous attention in both academia and industry. One of the most straightforward method to reduce model size is parameter pruning and sharing by removing redundant parameters that are not critical to model performance (Han et al., 2015; Blakeney et al., 2020). Similarly, the informative parameters can also be measured and selected by low-rank factorization (Sainath et al., 2013; Denton et al., 2014). To simultaneously reduce the computation and storage of a deep model, approaches based on transferred or compact convolutional filters were further proposed by designing special structural kernels (Cohen & Welling, 2016; Wu et al., 2016) and achieved benefits in domains with human prior. Besides, some other works are focusing on transferring the learned knowledge of a large teacher network to a small and lightweight student network, which yields the concept of knowledge distillation (Hinton et al., 2015; Mirzadeh et al., 2020; Chen et al., 2021; Gao et al., 2021), which is suitable for small- or medium-size data sets (Cheng et al., 2018).

heading: 3 PROBLEM DEFINITION

heading: 3.1 PRELIMINARY
text: In this work, we consider the following federated learning optimization problem:
min w
{ F (w) ,
K∑ k=1 pkFk(w)
} , (1)
where K is the number of clients, and pk is the weight of the k-th client such that pk ≥ 0 and∑K k=1 pk = 1. Suppose the k-th client holds the nk training data: xk,1, xk,2, · · · , xk,nk . The local objective Fk(·) is defined by
Fk(w) , 1
nk nk∑ j=1 ` (w;xk,j) , (2)
where `(·; ·) is a user-specified loss function. Specifically we consider a C class classification problem defined over a compact space X and a label space Y = [C], where [C] = {1, . . . , C}. The data point {x, y} is a random sample over X ×Y . A function fθ : X → S maps x to the probability simplex S, where S = { z | ∑C c=1 zc = 1, zc ≥ 0,∀c ∈ [C] } , and z is a C-dimensional column vector. fθ is parameterized over the hypothesis class w, i.e., the weight of the neural network θ. We define the loss function `(w, x) with the widely used cross-entropy loss as `(w, x) = Ey [log fθ(w, x)] .

heading: 3.2 ALGORITHM DESCRIPTION
text: Definition 1. Morphing Set: Given a neural network θ with its parameter weight w, its Morphing Set (Θ, W) is defined as a set that contains every pair of sub-network and the corresponding weight of the original network θ.
We define the concept of Morphing Set for the convenience of explanation. Before the description of the proposed algorithm, we define θo and wo as the server maintained neural network and its parameter weights, and denote (Θo, Wo) as the Morphing Set of θo.
Here, we describe one round (say the t-th) of our proposed algorithm. First, the central server broadcasts the latest model, (θt,wt), to the selected clients (say the K). Second, every client (say the k-th) begins with wkt,e=0 = wt, and then performsE ≥ 1 local updates with a randomly selected batch samples ξkt,e as follows,
wkt,e+1 ←− wkt,e − ηt,e∇Fk ( wkt,e, ξ k t,e ) (3)
Third, unlike Federated Average (FedAvg) (McMahan et al., 2017) or other conventional algorithms letting w̃t+1 = ∑ k∈K wkt,E |K| be the model weights for next round, our method, in order to compress the communication, morphs a small neural network from θo by minimizing the performance divergence of the morphed model wt+1 with w̃t+1 as follows,
(θt+1,wt+1) = arg min (θ,w)∈(Θo,Wo)
Lt(w; w̃t+1)
= arg min (θ,w)∈(Θo,Wo) ∑ x∈Xv JS (fθ(w, x)||fθt(w̃t+1, x)) , (4)
where w̃t+1 = ∑ k∈K wkt,E |K| follows the FedAvg algorithm to be the average aggregated value of local clients updated parameter weights. (θ,w) ∈ (Θo,Wo) denotes the newly morphed network and its weights, which is a sub-network of the server maintained one. fθ is the corresponding mapping function defined on the morphed neural network θ. The Jensen-Shannon (JS) divergence is a well-known metric for symmetrically measuring the distance of two probabilities. We apply the JS divergence as the loss function to minimize the output performance of the morphed sub-network
and average aggregated network on the validation compact space. Here, we use the finite-sum on the validation dataset Xv (without label information) to evaluate the performance divergence. Finally, the newly morphed sub-network with the updated weights (θt+1,wt+1) are broadcasted to selected clients for next round local optimization.
The optimization objective is to maintain as much knowledge as possible from the aggregated model in previous round, while keeping the morphed network as small as possible. Therefore, we consider the network architecture as a regularizer and reformat equation (4) as follows,
(θt+1,wt+1) = arg min (θ,w)∈(Θo,Wo)
{ Lt(w; w̃t+1) + λLc(θ) } , (5)
where Lc(θ) is a regularizer to measure the constraints on the number of neural network parameters (PARAMs) or the computation floating-point operations per second (FLOPs) or any other constraints as required. λ is an adjustable non-trained variable to balance the two losses.

heading: 4 OPTIMIZATION SOLUTION AND RELAXATION

heading: 4.1 SOFT MASKS
text: The optimization process of the equation (5) is incompatible with the sophisticated deep learning frameworks such as Pytorch (Paszke et al., 2019) or TensorFlow (Abadi et al., 2016) for the reasons that 1) the network architecture is not a well-defined parameter for optimization; 2) the categorical loss of the regularizer is incompatible with the neural network optimizer on differentiable functions.
In order to solve the issues mentioned above, we apply the soft-mask trick (Jang et al., 2016; Chaudhuri et al., 2020), a popular method to relax the categorical variables into continuous. Specifically, we append a soft-mask layer after each layer in the original neural network, which requires to be constrained. The soft-mask layer is a parameter vector whose dimension equals the number of constrained values of its previous layer. For instance, the dimension of the soft-mask layer would be the
Algorithm 1 Communication efficient federated morphing neural network. 1: for t in {0, · · · , T − 1} do 2: Broadcast the masked results to selected clients; 3: Local Update wkt on the selected clients according to (3); 4: Distill knowledge from the aggregated model w̃t+1 to (wot+1, vt+1) according to (7); 5: Morph (θo, wot+1) with vt+1 to (θt+1, wt+1) for next step iteration; 6: end for 7: return (θo, woT );
number of the output channels if the previous layer is a conv2d layer; Or it would be the number of neurons if the previous layer is a linear layer. The activation outputs of the previous layer passing through the soft-mask layer have the same effect of masking directly on the previous layer.
For each variable in the soft-mask layer (parameter vector), we apply the differentiable S-shape-like functions f : R → (0, 1) to relax the categorical representations into differentiable variables. The relaxation function is as follows,
m = Sigmoid (( log( π
1− π ) + l
) /τ ) . (6)
As the temperature τ → 0, m approaches {0, 1} with probability {1 − π, π}, respectively. l ∼ logistic is randomized to increases the expressiveness of the masked channel or neuron. Specifically, it leaks information with a possibility even when the channel or neuron is decided to be masked by π. The optimizer operates on v = log( π1−π ) instead of π, as m’s learnable parameters for numerical stability. We naively choose the standard logistic function, aka sigmoid function, for our differentiable relaxation. The comparisons among different kinds of S-shape-like functions is not the main object, and we leave it for the future research.
With the help of the soft-mask layer, the calculation of the architecture constraints is relaxed to the L1 norm of the soft-mask layer’s parameters with a non-parametric multiplication factor. Taking the conv2d as an example, its number of output channels can be approximated as ∑CHo i=1 mi, with CHo being the number of output channels. Meanwhile, the output channel i would be pruned if vi < 0 (π < 0.5), resulting in a regularized neural network.
Hence, with vt+1 being the neural architecture soft mask parameters at round t+ 1, we re-organize the objective function (5) as follows:
(wot+1,vt+1) = arg min (wo,v) {Lt(wo,v; w̃t+1) + λLc(v)} (θt+1,wt+1) = (θ o,wot+1) ◦ vt+1,
(7)
where the operation ◦ means the pruning operation on the server maintained network architecture θo and the corresponding model weights wot+1 at the round t+ 1 according to the mask parameters vt+1. The whole process of the proposed algorithm can be found from Algorithm 1, and the process of FedMorph is presented in Figure 1.

heading: 4.2 INTRA-LAYER CONSTRAINTS
text: We relaxed the problem by soft-masks previously, and the constraints of two connected layers are solved in this subsection. The foot-stone of recent SOTA neural architectures, e.g., ResNet, VGG for Computing Vision (CV) tasks or Transformer, BERT for Natural Language Processing (NLP) tasks, is the stacking of shallow layers or blocks in depth to increase the express-ability. Pruning the previous layer will also influence the next layer. Consider the example of a two-layers MLP f(x) = σ(x ∗W1) ∗W2, with W1 and W2 denotes the weight matrix of linear mappings, and σ denotes the non-linear activation functions. Assume m1 and m2 be the learned masking layers of W1 and W2, taking effect on the columns, whose sizes equal the number of columns of W1 and W2, respectively. In order to be coherent with matrix multiplication, the pruned columns of W1 also affect W2, which will prune the corresponding rows of W2. Moreover, we keep the final output layer unmasked to keep the coherence of the input and output of the morphed model with the server maintained one.

heading: 4.3 UNDERSTANDING FEDMORPH
text: Theorem 1. The learning of server maintained network weights wo with the knowledge distillation has the same effect of learning with random dropout layer and a regularizer loss.
Proof. The learning process of randomly dropout (Srivastava et al., 2014) is as follows, 1) training the network weights with dropout layers; 2) evaluate on the test dataset without activating the dropout layers. By setting τ → 0, the weights of soft-mask layers would generate either 1 or 0 (approximated) and mimic the random dropouts characters on the neuron level. The knowledge transfer process of the first equation in (7) equals to the optimization of wo with randomly dropout layers stacked after every learnable layer. The expectation on (v + l) determines the dropout ratio. The convergence of wot+1 at each knowledge distillation process is guaranteed with sufficient researches on deep learning with dropouts.
With Theorem 1, we know that the server maintained network will have a good performance on the validation dataset as long as the aggregated network performs well on it. Now, we want to study the generalization performance of the morphed network. In other words, how the knowledge on the training datasets of local clients is perceived by the morphed network via learning from the aggregated network on the validation dataset.
Theorem 2. Let a sample US of size m drawn from DS for knowledge distillation at round t, and δ < 1 being its upper bound of the empirical error, such that 1mL t x∈US (wt+1; w̃t+1) ≤ δ. Then the expectation evaluation error of wt+1 on DT is bounded, with the probability 1− δ:
`DT (wt+1) ≤`DT (w̃t+1) + δ
( 1 + √ −m log δ
2
) + 1
2 dH∆H (DS ,DT ) ,
where dH∆H (DS ,DT ) measures the distance of the knowledge distillation dataset distribution DS and local training dataset distribution DT .
With Theorem 2, we know that the error upper bound of the morphed neural network on the local training dataset distribution mainly consists of 1) the expectation error of the averaged aggregated model on the local dataset; 2) the empirical error δ of knowledge distillation on the validation dataset; 3) terms dependent on the divergence between local training dataset distribution and evaluation dataset distribution. If the local training dataset and validation dataset follows the same distribution, we have dH∆H (DS ,DT ) = 0, then the knowledge perceived by the newly morphed network is only determined by the empirical risk of knowledge distillation.

heading: 5 EXPERIMENTS

heading: 5.1 SETUP
text: Datasets and Networks We evaluate the performance of our method on multiple tasks on different networks architectures, specifically on the dataset of CIFAR-10 (Krizhevsky & Hinton, 2010) with the network of ResNet18 (He et al., 2016) and VGG16 (Simonyan & Zisserman, 2014). All networks with the corresponding dataset are initialized randomly. We refer the readers to find the experimental results and details of other datasets and network architectures in the Appendix.
Baselines. We focus on testing our methods against already well-established FL benchmarks. In particular, we restrict our baseline algorithm to the use of FedAvg, the same algorithm we use to calculate w̃t+1 in equation 4. To fairly evaluate the compression ratio of our proposed method, we combine the tricks of sparsification for client-to-server communication, precisely the Top-K sparsification trick proposed in (Lin et al., 2017), with the baseline algorithm. We also compare our method with FedDropout (Caldas et al., 2018), a strategy to compress communications in the serverto-client direction. Besides, we consider a new benchmark method, ‘FedPruned,’ which applies the FedAvg algorithm with a sub-network randomly pruned from the global network at initialization. The FedAvg, FedDropout, and FedPruned are also evaluated under the Top-K sparsification trick without losing fairness.
Hyper-Parameters. The dataset is partitioned within 100 clients at the initialization procedure in an IID fashion, and we randomly sample ten candidates for local optimization during each communication round. The optimized model gradients would be sparsified by the TopK method at a ratio of {10%, 25%} before updating to the global server. For local training at each client, we use the SGD optimizer with momentum to be 0.9 and weight decay to be 5e− 4. We set the static learning rates of 0.05 for CIFAR-10. Each selected client trains for one epoch per round using a batch size of 100. We set the knowledge distillation learning rates as 0.005 and the batch size as 50 while keeping other hyper-parameters following the local training procedure. We use a static temperature τ of 1e − 3 for all experiments during the whole process. We set the regularizer as the morphed model parameters and evaluate our method on model PARAMs and FLOPs, and set the regularizer ratios λ of CIFAR-10 for ResNet18 and VGG16 as 1.2e − 6, and 1e − 6, respectively. The soft-mask layer parameter v is initialized randomly following the logistic distribution with its mean value being the preset model compression ratio. In detail, we set the mean value of v to be 0.3 for both networks. For FedDropout and FedPrune, we set the same initialized compression ratios of 0.6. Moreover, we set the random seed to be ‘12345’ and built all experiments on a single GPU V100 to make the results reproducible.
Training Tricks. The knowledge distillation process that happens on the server-side at each communication round is computation-heavy. Therefore, to overcome the computation overloads, we 1) morphs the network when the aggregated model weights have changed significantly and 2) initialize the mapped part of the server maintained network with the average aggregated weights before the knowledge distillation procedure. In practice, we distill knowledge from the aggregated model every ten communication rounds and set the optimization epochs as 2 for every dataset.

heading: 5.2 TEST ACCURACY COMPARISON
text: We present the accuracy evaluation results as in Figure 2. In order to make a clear tendency, we apply the moving average on the result data points, saying to replace the point with the unweighted mean of its nearest 10 points. Besides, we report the maximal accuracy and the average compressed PARAMs and FLOPs in Table 1 for further clarification. The three main takeaways from these experiments are: 1) our proposed method produce competitive results (accuracy) comparing to the benchmarks; 2) FedMorph can further compress PARAMs and FLOPs without degrading the test accuracy; 3)FedMorph is compatible with the gradient sparsification methods, hence can further reduce the communication burden in FL settings. With the compressed PARAMs and FLOPs from Table 1, in other words, we achieves 10.0884 ≈ 11× and 1 0.0884 ∗ 10 ≈ 113× reduction of communication for CIFAR-10 with VGG16 in both server-to-clients and clients-to-server sides respectively. Meanwhile, the computation FLOPs achieves 10.2300 ≈ 4× reduction for it. Similar analysis on CIFAR-10 dataset with ResNet18 network, without degrading the test accuracy comparing with Fe-
dAvg, we achieves around 10×, 102×, and 5× reduction of server-to-clients side communication, clients-to-server side communication, and local clients computation FLOPs, respectively.

heading: 5.3 LAYER COMPRESSION RATIO
text: Another observation from Table 1 is that FedMorph has unbalanced PARAMs and FLOPs compression ratios. Therefore, we plot the compressed ratio of each masked layer as presented in Figure 3. The labels on the x-axis denote the masked layers from shallow to deep. The red bar, blue bar, and blue line denote the initialized compression ratio (soft-mask layer parameters), optimized compression ratio (soft-mask layer parameters when test accuracy reaches the maximum), and the trendline of the blue bar with moving average, respectively. It is clear to observe that FedMorph tends to compress the deep layers while leaving the shallow layers uncompressed for both network architecture, even though the mask layer parameters are initialized uniform-randomly. The reason is that the architectures of ResNets and VGGs are like pyramids, which have heavy parameters at deep layers. The unbalanced compression ratio of each layer proves the effectiveness of the PARAMs regularizer in our objective function. Besides, by comparing the test accuracy with FedPruned, a method with randomly pruned neurons at initialization, we know that the PARAMs regularizer can not only reduce the model size but also aim at not decreasing its accuracy.

heading: 5.4 GENERALIZATION GAP
text: It is also necessary to know why FedMorph generalizes better than other baselines. We use the gap between the training accuracy and test accuracy to measure the generalization performance of the experimented algorithms, as presented in Figure 4. We calculate the training accuracy by the mean of local clients’ accuracy, who has been selected for local updates between two consecutive times of the evaluation on the test dataset. For both datasets, Fedmorph has a significantly smaller generalization gap than that of FedAvg and FedDropout. It is worthwhile noticing that comparing to FedMorph, the training accuracy of FedAvg quickly reaches near to 1, and the training accuracy of
FedDropout is less than its test accuracy. FedAvg and FedDropout perform worse than FedMorph because FedAvg quickly gets over-fitted on the local dataset, which impedes its generalization well on the test dataset, while FedDropout somehow under-fits on the local training dataset. FedMorph balances well on the training and the evaluation parts, away from either over-fitting or under-fitting, making it generalizes the best among its competitors.

heading: 5.5 FEDMORPH VARIANT
text: This section evaluates FedMorphProx, a variant of FedMorph, by replacing the FedAvg algorithm for aggregation to FedProx (Li et al., 2018), on the NonIID partitions. FedProx limits the over-fitting of FedAvg on local datasets by constraining the local updated parameters not far away from the received model parameters at the beginning. In experimental details, we set the number of samples of each class in each client follows a Dirichlet distribution with α = 0.4 and set the proximal ratio µ in FedProx as 1e − 2 for CIFAR-10 while keeping other hyper-parameter untouched. We present the test accuracy results in Figure 5. The FedMorphProx and FedMorph algorithms perform significantly better than either FedAvg or FedProx. Besides, FedMorphProx performs as well as FedMorph, proving our method works for FedAvg and other FedAvg-like FL algorithms.

heading: 6 CONCLUSION
text: In this paper, we proposed that sharing a sub-network among the clients and the server can reduce edge devices’ communication and computation overloads and increase the generalization accuracy by avoiding over-fitting of FL on local datasets. The empirical experiments on varied datasets and network architectures with different FL settings, e.g., IID partition and Dirichlet partition, proved the superior performance of FedMorph. By analyzing the optimized morphed network, our work might also guide the design of network architectures in FL settings in the future.

heading: A DETAILS OF INTRAL-LAYER CONSTRAINTS
text: We give details of the intra-layer constraints on these modules, which cannot be sequentially stacked, e.g., the skip connection. Skip connections, successfully applied in ResNet, a sub-module connecting two unconnected layers of the backbone network, have increased the depth of neural networks and improved the benchmark results of deep learning. In mathematics, skip connection can be presented as out = out + skip(x), where x is the input of its previous connected layer, and out is the output of its last connected layer. The skip connection sub-module can have several stacked simple layers inside it or output the input directly. For the first condition, we attach the mask of its final layer with its last connected layer’s mask. For the second condition, we constrain the mask of its last connected layer to be its previous connected layer’s mask.

heading: B EXPERIMENTAL SETUP
text: B.1 SUMMARY OF DATASETS AND NETWORKS
We conduct the experiments on the MNIST (Deng, 2012), EMNIST (Cohen et al., 2017), and CIFAR-10 (Krizhevsky & Hinton, 2010) to consider the tasks on multiple visual computation tasks. For MNIST and EMNIST, we apply ResNet8 (He et al., 2016) as global network architecture. For CIFAR-10, we apply ResNet18 (He et al., 2016) and VGG16 (Simonyan & Zisserman, 2014) as the global networks. The network architectures used in our experiments are simply presented in Table 7, Table 8, and Table 9 at the end of this appendix. We abbreviate the details of datasets as follows,
• MNIST: is a handwritten digits (0-9) classification task with 60,000 data points in the training dataset and 10,000 data points in the test dataset. The input of the dataset is a flattened 784-dimensional (28 × 28) image, and the output is a class label between 0 and 9.
• EMNIST: is a set of handwritten character digits derived from the NIST Special Database 19 and converted to a 28x28 pixel image format. One unbalanced split, named ByClass, has 62 unbalanced class labels with 697,932 training samples and 116,323 test samples.
• CIFAR-10: consists of 60,000 32x32 colour images in 10 classes, with 6,000 images per class. There are 50,000 data points in the training dataset and 10,000 data points in the test dataset.
In abbreviation, the corresponding datasets and the networks are summarized in Table 2. The experimental settings can roughly evaluate the performance of FedMorph within different conditions:
• multiple tasks with the corresponding network architecture and dataset; • the same task with different network architectures, e.g., CIFAR-10 with ResNet18 and
VGG16; • the same network architecture on different tasks, e.g., MNIST and EMNIST with ResNet8; • the huge number of available clients, e.g., EMNIST with 5000 available clients.
It is always artistic to find the optimal hyper-parameters applied in deep learning. Instead of optimizing the hyper-parameters, we assign them casually but reasonably and present the detailed hyper-parameters applied in our experiments. The dataset is partitioned within 100 clients in an artificially IID fashion at the initialization procedure for MNIST and CIFAR-10, and 5000 clients
for EMNIST. The updated model would be sparsified by the TopK method at a ratio of {10%, 25%} before aggregation. For local training at each client, we use the SGD optimizer with momentum to be 0.9 and weight decay to be 5e − 4. We set the static learning rates of 0.1 for MNIST and EMNIST. For CIFAR-10, we select 0.05 as the learning rate. Ten clients are randomly selected to train for one epoch per communication round using a batch size of 100 for MNIST and CIFAR-10. For EMNIST, local clients train for five epochs per round. We set the knowledge distillation dataset as the test dataset without label information, and its training procedure follows the same SGD settings as the local training procedure. We set the distillation learning rates as one-tenth of that for local training and the batch sizes as 50, 500, and 50 for MNIST, EMNIST, and CIFAR-10, respectively. We use a static temperature τ of 1e−3 for all experiments during each round. We set the regularizer as the morphed model parameters and evaluate our method on model PARAMs and FLOPs. We set the regularizer ratios λ to be 1.6e − 6 and 2e − 5 for MNIST and EMNIST. For CIFAR-10 with ResNet18 and VGG16, λ is 1.2e − 6, and 1e − 6, respectively. The soft-mask layer parameter v is initialized randomly following the logistic distribution, with its mean value being the preset model compression ratio. Specifically, we set the mean value of v to be 0.1, 0.2, 0.3, and 0.3, respectively. For FedDropout and FedPrune, we set the same initialized compression ratios of 0.25, 0.35, 0.6, and 0.6, respectively, for the four tasks. Moreover, the maximum communication rounds are 1500, 3000, 4000, and 4000. To make the results reproducible, we set the random seed to be ‘12345’ for all experiments and build the experiments on GPU V100 within Pytorch. In conclusion, the corresponding hyper-parameters are listed in Table 3. If not specified, the hyper-parameters are kept unchanged.

heading: C ADDITIONAL EXPERIMENTAL EVALUATIONS
text: This section presents the additional experimental results missed in the main text due to the page limitation. In the following, we evaluate the performance of FedMorph on the datasets of MNIST and EMNIST in terms of test accuracy, compression ratio, and generalization gap. At last, we present the comparison results of FedMorph and its transient networks to help understand the algorithm.
C.1 TEST ACCURACY COMPARISON
The results of test accuracy versus the number of communication rounds are presented in Figure 6. FedMorph converges on the dataset of MNIST and EMNIST with ResNet8, and generalizes slightly better than the benchmarks. The results of MNIST and EMNIST on the same network architecture of ResNet8 prove that FedMroph generates good results irrelevant to datasets. The average compression ratio when reaches the maximal test accuracy is presented in Table 4. For MNIST dataset with ResNet8, we achieves 10.0118 ≈ 85× and 1 0.0118 ∗ 10 ≈ 847× reduction of server-to-client side communication and client-to-server side communication, and 10.0570 ≈ 18× reduction of local clients computation. For EMNIST dataset with ResNet8, we achieves 10.0901 ≈ 1× and 10.0901 ∗ 10 ≈ 111× reduction of server-to-client side communication and client-to-server side communication, and 10.3463 ≈ 3× reduction of local clients computation.
C.2 LAYER COMPRESSION RATIO
We present the layer compression ratio of FedMorph on MNIST and EMNIST in Figure 7. As expected in the analysis of CIFAR-10 with ResNet18 and VGG16, FedMorph compresses the deep layers of ResNet8 while keeping the shallow layers untouched to maintain its performance and the model size.
C.3 GENERALIZATION GAP
In order to prove that a small morphed network generates better generalization performance, we present the generalization gap of MNIST and EMNIST in Figure 8. FedMorph has a significantly smallest generalization gap for MNIST among the competitors. Comparing with FedMorph, FedAvg
is over-fitted as its training accuracy quickly reaches 1. While FedDropout is under-fitted as its training accuracy converges below that of FedMorph. For EMNIST, the dataset of size 697,932 is partitioned to 5000 clients, resulting in that each client has lesser than 140 samples. Furthermore, we set the number of local optimization epochs to be 5. Hence, all methods are easily over-fitted. However, comparing with FedAvg and FedDropout, FedMorph still performs better than them.
C.4 FEDMORPHPROX WITH NONIID SETTINGS
We set the number of samples of each class in each client follows a Dirichlet distribution with α = 0.4, and set the proximal ratio µ in FedProx to be 1e−2 for all tasks. For MNIST and EMNIST, we decrease the local learning rates to be 0.01 for a smooth convergence while keeping other hyper-
parameter untouched. We present the result of test accuracy versus the number of communication rounds for MNIST and EMNIST in Figure 9. In our FL settings, FedMroph and FedMorphProx have slightly better accuracy results than that of FedAvg and FedProx. Besides, FedMorphProx performs as well as FedMorph, proving our method works for FedAvg and other FedAvg-like FL algorithms. It is also worthwhile pointing that FedMorph and FedMorphProx converge smoother than FedAvg and FedProx. We present the maximal test accuracy and the compressed PARAMs and FLOPs of all datasets in Table 5 and Table 6 for clarification.
C.5 TEST ACCURACY OF THE TRANSIENT NETWORK
We study the performance of the transient networks, specifically, the morphed network wt of each communication round t and the average aggregated network w̃t of the morphed networks from local updates, and denote them as ‘Morphed’ and ‘Averaged.’ The results of the test accuracy versus the number of rounds for different datasets are presented in Figure 10. Interestingly, FedMorph has lower test accuracy than the other two transient networks but reaches the highest accuracy when converging. The explanation is that FedMorph learns knowledge from the smaller network ‘Averaged’ at an idle state, which causes its accuracy to be lower than that of ‘Averaged’ at the begging phase. By morphing a new network before each communication round, FedMorph can learn knowledge from offbeat teachers and store the knowledge inside the different parts of its larger network. When the smaller teacher network becomes converged, FedMorph combines all teachers’ knowledge, hence generating better performance than the ‘Averaged’ when converging.

heading: D ABLATION STUDY
text: This section studies the convergence issues with different hyper-parameters, e.g., the initial values of v and the regularizer ratio λ. The following experiments are built on CIFAR-10 with ResNet18.
D.1 THE CONVERGENCE ISSUES WITH DIFFERENT v
For different initial v, we present the test accuracy versus the number of communication rounds and the maximal test accuracy versus varied v in Figure 11. We keep the hyper-parameter unchanged and set λ = 1.2e − 6 for different experiments. The left-hand-side (LHS) figure presents that test accuracy converges faster when the initial v is higher. The right-hand-side (RHS) figure presents the maximal test accuracy and the average PARAMs compression ratio for different v when it reaches the maximal test accuracy. We use the size of the triangle and the color bar to denote the value of PARAMs. It appears that a more significant initial v leads to a higher compression ratio. Moreover, the maximal test accuracy increases with the initial v and decreases after reaching optimal. In this condition, the optimal initial v is around 0.3.
D.2 THE CONVERGENCE ISSUES WITH DIFFERENT λ
We present the results of test accuracy versus communication rounds of different regularizer ratios λ in Figure 12. We set the initial v to be 0.3 for all experiments and keep the hyper-parameters unchanged. The LHS figure presents that increasing the regularizer ratio to a large value will significantly deteriorate the test performance. The right-hand-side (RHS) figure presents the maximal test accuracy and the average PARAMs compression ratio when it reaches the maximal test accuracy for different regularizer ratios. As can be found from the color bar of triangles, a larger λ generates a smaller compression ratio but leads to terrible test accuracy. Moreover, an over-small λ can generate a consistent performance but increases the compression ratio. Therefore, an optimal λ exists to balance the test accuracy and compression ratio. In this condition, the optimal λ is around 1e− 6.

heading: E UNDERSTANDING FEDMORPH
text: In this section, we want to understand
• why learning from a small average aggregated sub-network guarantees the performance of the global model;
• how much ‘knowledge’ is transmitted from the average aggregated model in the previous round to the next round.
In the following, we first analyze the convergence of server maintained network weights by considering the morphing process similar to adding dropout layers in Theorem 1. Then we analyze the bound of the expected error of the morphed model on both the training and test dataset in Theorem 2 and Theorem 3 with theorems from Domain Adaptation (Ben-David et al., 2010).
E.1 ANALYSIS ON SERVER MAINTAINED NETWORK
Theorem 1. The learning of server maintained network weights wo with the knowledge distillation has the same effect of learning with random dropout layer and a regularizer loss.
Proof. The learning process of randomly dropout (Srivastava et al., 2014) is as follows, 1) training the network weights with dropout layers; 2) evaluate on the test dataset without activating the dropout layers. By setting τ → 0, the weights of soft-mask layers would generate either 1 or 0 (approximated) and mimic the random dropouts characters on the neuron level. The knowledge transfer process of the first equation in (7) equals to the optimization of wo with randomly dropout layers stacked after every learnable layer. The expectation on (v + l) determines the dropout ratio. The convergence of wot+1 at each knowledge distillation process is guaranteed with sufficient researches on deep learning with dropouts.
E.2 ANALYSIS ON MORPHED NETWORK
We want to study how the morphed neural network perceives the knowledge of local training datasets from the distillation on the average aggregated model with the validation dataset. Our theories and proofs referred a lot from Ben-David et al. (2010).
Before giving any theories, we make some important definitions to clarify the problem. Define a domain as a pair consisting of a distributionD on inputs X and the labeling function f : X → [0, 1], which can have a fractional (expected) value when labeling occurs non-deterministically. Initially, we construct two domains with the validation dataset for knowledge distillation and the training datasets for local optimization, denoted by 〈DS , fS〉, and 〈DT , fT 〉, naming as source domain and target domain, respectively. The labeling function fDS (x) and fDT (x) on the sample x generate the True label of x in corresponding datasets.
Define a hypothesis as a function h : X → {0, 1}. The probability according to the distribution DS that a hypothesis h disagrees with a labeling function f (which can also be a hypothesis) is defined as
`S(h, f) = Ex∼DS [|h(x)− f(x)|]
When we want to refer to the source error (sometimes called risk) of a hypothesis, we use the shorthand `S(h) = `S(h, fS). We write the empirical source error as ˜̀S(h). The similar definitions can be built on `T (h), `T (h, fT ) and ˜̀T (h).
Proposition 1. (Hoeffding’s inequality) Let x1, . . . , xm be independent bounded random variables with xi ∈ [a, b] for all i, where −∞ < a ≤ b <∞. Then, for all ζ ≥ 0 we have,
Pr (∣∣∣∣∣ 1m m∑ i=1 (xi − E [xi]) ∣∣∣∣∣ ≥ ζ ) ≤ exp ( − 2mζ 2 (b− a)2 ) Remark 1. Hoeffding’s inequality provides an upper bound on the probability that the sum of bounded independent random variables deviates from its expected value by more than a certain
amount. Furthermore, if x1, . . . , xm are sampled from the same distribution such that x ∼ D, then Hoeffding’s inequality can be simplized as follows,
Pr (∣∣∣∣∣ 1m m∑ i=1 xi − Ex∼D [x] ∣∣∣∣∣ ≥ ζ ) ≤ exp ( − 2mζ 2 (b− a)2 ) Before giving the theories, we refer Ben-David et al. (2010) to define the symmetric difference hypothesis spaceH∆Hand the bounding of any two hypothesises in different domains as follows. Definition 2. For a hypothesis spaceH, the symmetric difference hypothesis spaceH∆H is the set of hypotheses g ∈ H∆H ⇐⇒ g(x) = h(x)⊕ h′(x) for some h, h′ ∈ H where ⊕ is the XOR function. In words, every hypothesis g ∈ H∆H is the set of disagreements between two hypotheses inH.
The following simple lemma shows how we can make use of theH∆H-divergence in bounding the error of our hypothesis. Lemma 1. For any hypotheses h, h′ ∈ H, the error difference in any two domains S and T is bounded as follows,
|`S (h, h′)− `T (h, h′)| ≤ 1
2 dH∆H (DS ,DT )
Proof. By the definition ofH∆H-distance, we have dH∆H (DS ,DT ) = 2 sup
h,h′∈H |Prx∼DS [h(x) 6= h′(x)]− Prx∼DT [h(x) 6= h′(x)]|
= 2 sup h,h′∈H
|`S (h, h′)− `T (h, h′)| ≥ 2 |`S (h, h′)− `T (h, h′)| (1)
With the bounding of Lemma 1, we give the following two lemmas on domain adaptation on different hypothesises. Lemma 2. For any two hypothesises h, h′ ∈ H,
`T (h) ≤ `T (h′) + `S(h, h′) + 1
2 dH∆H (DS ,DT )
Proof. This proof relies on the triangle inequality for classification error, which implies that for any labeling functions f1, f2, and f3, we have `(f1, f2) ≤ `(f1, f3) + `(f3, f2). Therefore, we have,
`T (h) ≤ `T (h′) + `T (h, h′)
≤ `T (h′) + `S(h, h′) + 1
2 dH∆H (DS ,DT )
(2)
Lemma 3. Define λ , `S(h∗) + `T (h∗) with h∗ be the optimized hypothesis of the combined datasets, then for any hypothesis h, h̃ ∈ H,
`S(h) ≤ `T (h′) + `S(h, h′) + 1
2 dH∆H (DS ,DT ) + λ
Proof. This proof relies on the triangle inequality for classification error, which implies that for any labeling functions f1, f2, and f3, we have `(f1, f2) ≤ `(f1, f3) + `(f3, f2). Therefore, we have,
`S(h ′) ≤ `S (h∗) + `S (h′, h∗) ≤ `S (h∗) + `T (h′, h∗) + |`S (h′, h∗)− `T (h′, h∗)|
≤ `S (h∗) + `T (h′, h∗) + 1
2 dH∆H (DT ,DS)
≤ `S (h∗) + `T (h′) + `T (h∗) + 1
2 dH∆H (DT ,DS)
= `T (h ′) +
1 2 dH∆H (DS ,DT ) + λ
(3)
Combining the above inequality with the truth `S(h) ≤ `S(h′)+`S(h, h′), we obtain the lemma.
Remark 2. The Lemma 2 and Lemma 3 give the bounding of hypothesis h on both the source and target domain with another hypothesis h′ on the target domain. With two lemmas mentioned above, we can study the error of the morphed network on the local training datasets and validation dataset bounded by the average aggregated network on local datasets.
The following lemma studies the expected error bound of the knowledge distillation process given the empirical error on m data samples with Hoeffding’s inequality. Lemma 4. Let a sample US of size m drawn from DS for knowledge distillation at round t, and δ < 1 being its upper bound of the empirical error, such that 1mL t x∈US (wt+1; w̃t+1) ≤ δ. Then, with the probability 1− δ its expectation error is bounded:
E [ Lt(wt+1; w̃t+1) ] ≤ δ ( 1 + √ −m log δ
2
) .
Proof. Because of the truth that Ltx(wt+1; w̃t+1) ≥ 0 for ∀x ∈ US and the empirical loss of Ltx∈US (wt+1; w̃t+1) is bounded by δ, we obtain
Ltx(wt+1; w̃t+1) ≤ mδ, ∀x ∈ US (4) It means Ltx(wt+1; w̃t+1) is bounded by [0,mδ] for any random variable x ∈ US . Apply the Hoeffding’s Inequality here, we obtain as follows,
Pr (∣∣Lt(wt+1; w̃t+1)− E [Lt(wt+1; w̃t+1)]∣∣ ≥ ζ) ≤ 2 exp(−2ζ2m
(mδ)2
) (5)
Let δ = exp(−2ζ 2 mδ2 ), we easily obtain that ζ = δ √ −m log δ
2 . Hence, we know that at least at probability 1− δ that
E [ Lt(wt+1; w̃t+1) ] ≤ Lt(wt+1; w̃t+1) + δ
√ −m log δ
2 (6)
It is also worthy noticing that, δ √ −m log δ
2 → 0 as long as δ → 0 .
Theorem 2. Let a sample US of size m drawn from DS for knowledge distillation at round t, and δ < 1 being its upper bound of the empirical error, such that 1mL t x∈US (wt+1; w̃t+1) ≤ δ. Then the expectation evaluation error of wt+1 on DT is bounded, with the probability 1− δ:
`DT (wt+1) ≤`DT (w̃t+1) + δ
( 1 + √ −m log δ
2
) + 1
2 dH∆H (DS ,DT ) ,
where dH∆H (DS ,DT ) measures the distance of the knowledge distillation dataset distribution DS and local training dataset distribution DT .
Proof. Let `DT (wt+1), `DS (w̃t+1) be the error of the morphed model and the average aggregated model on local training datasets. Combine with Lemma 2, we have
`DT (wt+1) ≤ `DT (w̃t+1) + `DS (wt+1, w̃t+1) + 1
2 dH∆H (DS ,DT ) , (7)
where `DS (wt+1, w̃t+1) = EDS [Lt(wt+1; w̃t+1)]. Combine with Lemma 4, we obtain the theorem.
Theorem 3. Let a sample US of size m drawn from DS for knowledge distillation at round t, and δ < 1 being its upper bound of the empirical error, such that 1mL t x∈US (wt+1; w̃t+1) ≤ δ. Then the expectation evaluation error of wt+1 on DS is bounded, with the probability 1− δ:
`DS (wt+1) ≤`DT (w̃t+1) + δ
( 1 + √ −m log δ
2
) + 1
2 dH∆H (DS ,DT ) + λ ,
where dH∆H (DS ,DT ) measures the distance of the validation dataset DS and local training datasetDT , and λ = infh∈H (`DT (h) + `DS (h)) is the optimal loss on the combined dataset distribution;
Proof. Let `DT (wt+1), `DT (w̃t+1) be the error of the morphed model and the average aggregated model on local training datasets. Combine with Lemma 3, we have
`DS (wt+1) ≤ `DT (w̃t+1) + `DS (wt+1, w̃t+1) + 1
2 dH∆H (DS ,DT ) + λ, (8)
where `DS (wt+1, w̃t+1) = EDS [Lt(wt+1; w̃t+1)]. Combine with Lemma 4, we obtain the theorem.

Review summaries:
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 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.