H-FL: A Hierarchical Communication-Efficient and Privacy-Protected Architecture for Federated Learning

1 Introduction

Federated learning (FL) is a promising distributed paradigm for training a shared model while keeping all the training data localized [Yang et al., 2019; Kairouz et al., 2019; Koneˇcn`y et al., 2016]. However, FL always involves expensive communication and privacy concerns in order to maintain a great model performance [Li et al., 2020; Zhang et al., 2021]. Therefore, how to find a great balance among model performance, communication overhead and privacy requirements is a long-term, challenging goal.

From a methodological standpoint, DGC [Lin et al., 2017] and FetchSGD [Rothchild et al., 2020] have given a good trade-off between communication overhead and model performance by compressing the gradients and giving some corrections. NbAFL [Wei et al., 2020] and DP-FedAVG

∗Contact Author

[McMahan et al., 2017b] provide strong privacy guarantees via differential privacy without undue sacrifice on model performance. SplitNN [Vepakomma et al., 2018] can achieve higher model performance in contrast to the aforementioned methods while protecting sensitive raw data. All these works try to make some trade-offs from different perspectives. However, when treating model performance, communication overhead and privacy requirements as a whole perspective, it will introduce a completely new contradiction: the contradiction between communication overhead and privacy requirements while maintaining model performance in a certain level. Since when utilizing some privacy protection methods such as differential privacy and secure multiparty computing to provide privacy guarantees, it will inevitably introduce additional communication overhead directly or slow down the convergence rate, leading to requiring extra communication rounds for FL algorithms to converge. Therefore, we cannot just do simple combinations from different perspectives.

In this paper, we develop a hierarchical federated learning architecture (H-FL) as shown in Figure 1. To counter-weigh the degradation of model performance due to statistic heterogeneity of the training data, H-FL introduces mediators to reconstruct the local distributions. We cluster the clients according to the KL divergence between local distributions of each client and a uniform distribution, as well as the information entropy of the local distributions, and then reallocate them to different mediators. When participating in federated tasks, H-FL selects mediators rather than clients and each mediator rearranges its clients to perform the training tasks in order to alleviate the statistic heterogeneity. In addition, we design a compression-correction mechanism to reduce the communication overhead without sacrificing the model performance, which significantly compresses the extracted features of the clients uploaded to mediators and corrects the corresponding gradients download from the mediators. To further provide privacy guarantees for clients, we introduce differential privacy when each client conducts its local training.

Our contributions can be summarized as following:

• To the best of our knowledge, H-FL is the first attempt to treat model performance, communication overhead and privacy requirements as a whole perspective to find a great balance among them.

• We devise a runtime distribution reconstruction strategy to alleviate the statistic heterogeneity of the training data while not compromising user privacy. Moreover, we design a compression-correction mechanism to reduce the communication overhead without sacrificing the model performance.
• Extensive experiments on different datasets show that our H-FL architecture achieves state-of-the-art performance on federated image recognition tasks.

2 Related Research

Federated learning is a collaborative distributed learning paradigm which removes the necessity to pool the raw data out from local clients. Specifically, FedAVG algorithm proposed in [McMahan et al., 2017a] aims to reduce the communication overhead while maintaining a good performance of the model on non-IID (Independent and identically distributed) training data, which is used as our baseline in Section 4. Furthermore, concurrent works such as [Lin et al., 2018; Sattler et al., 2019] have focused on further reducing communication overhead in FL via gradient sparsification, and propose solutions to counter-weigh the reduction in accuracy due to the statistic heterogeneity of the training data. Concretely, DGC [Lin et al., 2018] employs momentum correction and local gradient clipping on top of the gradient sparsification to ensure no loss of accuracy. In addition, DGC also uses momentum factor masking and warmup training to overcome the staleness problem caused by reduced communication. STC [Sattler et al., 2019] propose a sparse ternary compression (STC) framework to reduce the communication overhead in FL, which enables ternarization and optimal Golomb encoding of the weight updates and also behaves robust to non-IID training data. We conduct a comprehensive analysis and comparison with the aforementioned methods in Section 4 to illustrate the effectiveness of our HFL framework. in Section 4 to illustrate the effectiveness of our H-FL framework.

3 Our Approach

In this section, we propose a hierarchical FL architecture as shown in Figure 1 to find a great balance among model performance, communication overhead and privacy requirements.

3.1 AdversaryModel

We first assume that all the components (FL Server, Mediators, Clients) in H-FL have following abilities: 1) they are honest-but-curious, which means that they will honestly follow the designed protocol but are curious about the others' local data; 2) they have arbitrary auxiliary information to help infer a specific client's private information during the process of collaboratively building a shared model; 3) they do not collude with each other, which means that they will not provide any additional information to clients during the training.

3.2 Initialization

FL server first splits the complete model into two components: shallow model and deep model, then distributes the former one to the Aggregation Mediator (AM) and the latter one to the other Mediators (M). AMdistributes the shallow model to all the clients. At the same time, FL server initializes the global hyper-parameters such as learning rate η, sampling probability of each client P, sampling probability of each example S, global compression ratio C (C < 0.5), iterations of deep training in mediators I, ℓ2-norm of the clipped gradients L and noise level σ. Specifically, when sampling locally in practice, we randomly permute the local data and partition them into mini-batches of the appropriate sizes for efficiency.

3.3 Runtime Distribution Reconstruction

In FL settings, as the training data resident in the individual clients is collected by the clients themselves on the basis of their local environments, the distribution of the local datasets will considerably differ with each other. Considering this characteristic, we redefine the optimization objective function of federated learning training on non-IID datasets as follows:

where w(c) , p(c) , p, DKL are the weights of client c, the local distribution of client c, the distribution of potential global training data, KL divergence between p and p(c) , respectively. When the latter term is approximate to 0, it will degrade to an optimization problem under IID. In general FL settings, p(c) s are a series of different fixed distributions such that the latter term is a fixed value and the optimization objective will be consistent.

Whereas we consider p(c) s as variable distributions rather than fixed distributions in H-FL, so we can change local distributions arbitrarily. An intuitive way is to gather the clients' local data and form a series of different new distributions, each of which is approximate to the potential global distribution p, enabling the latter term in Formula (1) to be 0. However, sharing local data raises serious privacy risks and causes high communication overhead. Therefore, we introduce the runtime distribution reconstruction strategy to mitigate differences among local distributions while meeting the privacy requirements.

Specifically, a uniform distribution p(r) is initialized and broadcast among the clients. Each client calculates the information entropy H(c) of its local distribution p(c) and KL divergence DKL(p(r) ‖p(c) ) between p(r) and p(c) . Furthermore, K-means algorithm is utilized to cluster the clients according to the binary group (H(c) ,DKL(p(r) ‖p(c) )). Then H-FL randomly selects clients from each cluster, marks them as a group, and assigns the group to one of mediators. The allocation pattern loops until all the clients are assigned to the corresponding mediator.

Algorithm 1 Runtime distribution reconstruction

Input: U, M

Parameter: W (c) ,P, S, C

Output: B(m)

• 1: Randomly initialize a distribution p(r)
• 2: for each c ∈ U in parallel do
• 3: Compute H(c) , DKL(p(r) ‖p(c) )
• 4: end for
• 5: Cluster according to (H(c) ,DKL(p(r) ‖p(c) ))
• 6: for each m ∈Mdo
• 7: Randomly select clients fromeach cluster according to the same ratio 1/ |M| and assign them to m
• 8: B(m) ← ∅
• 9: end for
• 10: Mt ← (Randomly sampling mediators in M)
• 11: for each m ∈Mt do
• 12: U t ← (Randomly sampling clients in U with P)
• 13: for each c ∈ U t do
• 14: Randomly sampling a mini-batch X(c) with S
• 15: O(c) ←W (c)X(c)
• 16: k ← ⌊∣∣O(c) ∣∣ ∗ C⌋
• 17: B(m) ←B(m) ∪LF(O(c) )
• 18: end for
• 19: end for
• 20: return B(m)

When performing local training, each client utilizes the shallow model to extract features, which will be compressed by the lossy compressor (introduced in subsection 3.4) and sent to the corresponding mediator. After that, each mediator concatenates the received features through a connector (as shown in Figure 1) to obtain synthetic features. This procedure can be considered as sampling from a virtual reconstructed distribution p(m) and then conducting forward propagation using the shallow model (see Algorithm 1). Intuitively, p(m) will be more approximate to the potential global distribution p than p(c) . The optimization objective function will be changed to the following form:

Assuming that there exists enough clients, p(m) s will infinitely approximate the potential global distribution p and the latter term will be 0, which is translated to the optimization problem under IID. When finishing the distribution reconstruction, each mediator leverages the synthetic features to train the deep model and then sends back the gradients of the synthesized features to the clients to assist training the shallow model. In this way, H-FL alleviates the statistic heterogeneity of the training data while not compromising user privacy.

3.4 Compression-CorrectionMechanism

To reduce the communication overhead, each participating client compresses the extracted features through the lossy compressor in Figure 1 by:

where O is feature matrix extracted by the shallow model, UO, ΣO and V T O are the results of singular value decomposition (SVD) respectively, UO[:, : k], ΣO[: k] and VO[: k] represent the first k columns of UO, ΣO and VO respectively. In this way, the feature matrix can be changed to a low-rank matrix that can be expressed by as the product of two relatively small matrices, thus reducing the communication overhead.

For the sake of clarification, let us introduce some new representations:

When updating W (c) , we should compute dW (c) as follows according to the chain rule:

However, according to formula (3), we cannot compute ∂B/∂W (c) directly since there is no direct differentiable mapping fromW (c) to B. For convenience, ∂O/∂W (c) can be used instead of ∂B/∂W (c) , which may still work but it leads to a reduction in model accuracy.

Therefore, we design a bias corrector on clients to correct the gradients of lossy features, which could build the mapping from O to B so that we can better approximate ∂B/∂W (c) and counter-weigh the reduction. According to the feature of SVD, we can get:

where UO here is just the same thing as the UO in formula (3), Dk represents a diagonal matrix where its first k elements on the diagonal are 1 and the rest are 0. Therefore, ∂B/∂W (c) can be rewritten as:

Thus, the bias corrector can be considered as consisting of many fully connected layers stacked on top of each other, and the parameters depend on the SVD results of the features extracted from the shallow model. In other words, the parameters of the bias corrector will be updated during the procedure of forward propagation. We also compare the results for the presence or absence of the bias corrector through appropriate experiments in Section 4.

After we obtain rectified dW (c) , we conduct gradient clipping so that the ℓ2 norm of dW (c) is limited to L and then add noise for it in order to protect privacy:

where g(c) is dW (c) itself, n(c) is the size of the sampled mini-batch in client c, N is the Gaussian distribution with mean 0 and standard deviation σLI/√n.

In summary, the workflow of H-FL mainly includes runtime distribution reconstruction, training and aggregation, the pseudo-code of which is given as Algorithm 2.

Algorithm 2 The workflow for H-FL

Input: U, AM, M

Parameter: W (m)

Output: W (d)

t , W (c) t , P, S, C, I, L, σ

t+1 , W (s) t+1

Mediators:

• 1: B(m) ←Run-time data augmentation
• 2: for each m ∈M\AMin parallel do
• 3: for each epoch i from 1 to I do

• 5: end for

• 7: for each c ∈ m do
• 8: dB(c) ← dB(m) [: n(c) ]
• 9: dB(m) ← dB(m) [n(c) :]
• 10: end for
• 11: end for

Clients:

• 1: for each c ∈ U t in parallel do
• 2: B(c) ←U (c) O D(c) k U (c)T O O(c)
• 3: dW (c) t ← dB(c)∂B(c) /∂W (c) t
• 4: dW (c) t ← dW (c) t +N (0, σ2L2 I/n(c) )
• 5: W (c) t ←W (c) t - ηdW (c) t
• 6: end for

FL Server:

AM:

Theorem 1. Formula (8) satisfies differential privacy in distribute environment and the privacy loss can be tracked via moments accountant.

Proof. We can consider the first term of formula (8) as follows approximately:

where g is the gradient of backward propagation, x(c) i is the i-th example of client c and n(c) is the size of sampled mini-batch of client c. In addition, we can also consider the latter term of formula (8) as follows according to central limit theorem:

Therefore, we can rewrite formula (8) as formula (11), which satisfies example-level differential privacy for each client according to Theorem 2 [Abadi et al., 2016]. In addition, since L and σ are the same for all clients, the privacy loss accumulated via moment accountant for each client in the distributed environment is the same. It also satisfies differential privacy in the distributed environment according to differential privacy parallel principle.

4 Experimental Results

We evaluate H-FL on different datasets and compare the performance to FedAVG [McMahan et al., 2017a], STC [Sattler et al., 2019] and DGC [Lin et al., 2018] in non-IID environments. Specifically, we have trained a modified version of LeNet5 [LeCun et al., 1998] network on FMNIST [Xiao et al., 2017] and a modified VGG16 [Simonyan and Zisserman, 2014] network network on CIFAR10 [Krizhevsky et al., 2009] respectively. In addition, the first two CNN blocks of VGG16 and the first one CNN block of modified LeNet5 are set to the shallow part in practice. All the batch-normalization layers are removed in the shallow model. The experiment settings are listed in Table 2.

4.1 Behavior Of The Model Performance

Figure 2(a) and Figure 2(b) show the top-1 accuracy of LeNet5 on FMNIST after 200 communication rounds and the accuracy of VGG16 on CIFAR10 after 2000 communication rounds respectively using H-FL and the aforementioned three methods. The magenta dotted line refers to an accuracy of 80%. The experiment results show that H-FL outperforms the other methods both on convergence rate and final accuracy. The results are quite reasonable since H-FL reconstructs a series of virtual distributions p(m) s, each of which is more closer to the potential global distribution and the optimization problem under non-IID is almost turn into that under IID. Thus, H-FL have the better capability to handle the heterogeneous dataset. Specifically, we take the average of the last 10 rounds of the accuracy as the final accuracy after 200 rounds on FMNIST and 2000 rounds on CIFAR10 respectively. HFL achieves an accuracy of 88.16% on FMNIST, whereas FedAVG, DGC and STC only achieve 82.28%, 82.00%, and 82.12% respectively. Moreover H-FL achieves an accuracy of 87.28% on CIFAR10, whereas FedAVG, DGC and STC only achieve 73.83%, 81.25% and 81.24%.

• (b) Communication Overhead On FMNIST (c) Communication Overhead On CIFAR10

Figure 3: Effectiveness of the Bias Corrector and Communication Overhead

4.2 Influence Of Different Parameters For H-FL

From Figure 2(c), Figure 2(d) and Figure 2(e), we can observe that as P, S and C increase, the model performance and the convergence behavior are getting better. The phenomenon is quite reasonable because: 1) As aforementioned, the procedure of reconstructing distributions in H-FL is closely related to the training samples of each client. The larger P and S are, the more the training samples are, and the closer the reconstructed distribution is to the potential global distribution, thus reducing the impact of non-IID and obtaining a relatively ideal effect; 2) The number of training samples of each client has a great impact on noise injecting. As S increase, the number of training samples of each client becomes larger, so that the n(c) in Formula 8 will be larger and the injected noise is correspondingly small; 3) As C increases, the lossy compression becomes less and less effective and the behavior will get better. In addition, as the noise level σ increases in Figure 2(f), the oscillation amplitude of the accuracy curve becomes larger, the convergence speed becomes slower and the model performance becomes worse.

4.3 Effectiveness Of The Bias Corrector

Figure 3(a) shows the top-1 accuracy of VGG16 on CIFAR10 with and without the bias corrector. The black dotted line refers to an accuracy of 85%. As we can see, bias corrector has significant influence on the convergence behavior and the final accuracy of H-FL. When there exists the bias corrector, the accuracy of the global model converges to 85% around 1000 rounds. Whereas when we remove the bias corrector, the accuracy gradually approximates to 85% until 2000 rounds. Additionally, we take the average of the last 10 rounds of the accuracy as the final accuracy after 2000 rounds, and the bias corrector can achieve an improvement of 2.47 percentage points. The result of the experiment is in line with our expectation since the bias corrector gives a relatively precise approximation of dW (c) when it can't be calculated directly, and once we remove the bias corrector, it will obtain a biased dW (c) , leading to the degradation in model performance and other metrics (convergence behavior).

4.4 Communication Overhead

Finally, we compare the different methods with respect to the communication overhead which are required to achieve a certain target accuracy. As we can see in the Figure 2(a) and Figure 2(b), the convergence behavior is much better than other methods, which considerably reduces the communication rounds. Notice that FedAVG does not converge on CIFAR10, thus we do not show that in Figure 3(c). To compare the communication overhead, we set a window of size 10, which is utilized to calculate an average of 10 rounds. The communication overhead accumulates while moving forward the window until the average accuracy is no less than the target accuracy (80% in our experiments). Figure 3(b) and Figure 3(c) show the communication overhead required to achieve the target accuracy for the different methods on FMNIST and CIFAR10 respectively.

5 Conclusion

In this paper, we present a Hierarchical Federated Learning architecture (H-FL) to find a great balance among model performance, communication overhead and privacy requirements. Firstly, we devise a runtime distribution reconstruction strategy to counter-weigh the degradation due to nonIID. Then we design a compression-correction mechanism to reduce the communication overhead without sacrificing the model performance. The experimental results have proved that H-FL achieves the state-of-the-art performance under different federated learning tasks.

Acknowledgements

We would like to thank anonymous reviewers for their helpful comments. This work was supported by the the National Key Research and Development Program of China (No. 2018AAA0100500), NSFC Grant No. 61872285, 62072367, 61772413, the Key Research and Development Program of Shaanxi Province (2020GY-033), Key Science and Technology Project of Henan Province (201300210400), Fundamental Research Funds for the Central Universities (xzy012020112).