Paper: Making Batch Normalization Great in Federated Deep Learning (ICLR-2024)

Reviewer comment (Theory):
This paper lacks theoretical analysis for why this simple fix would work for BN in FL. It only proposes intuitive explanation.

Author response:
Although our study is primarily empirical, we respectfully do not think it’s just intuition. Rather, each step that leads to our proposed solution is supported by evidence from a carefully designed specific study. 

First, we note that table 1 reveals a crucial observation, the use of local batch statistics during local training degrades performance, as GN can recover centralized (-0% in accuracy) while BN degrades (-45%) when communicating at every step. This suggests that the use of local batch statistics in local training is what causes degradation. In other words, as long as local batch statistics are not used, the degradation can be rescued. This motivates us to design a local-statistics-free strategy in training. 

In sec. 4.1, we further examine this motivation. We examine both the BN statistics variance within local models and the gap of BN statistics between the local model and the global model. We found that while local statistics converge (fig 2b), the gap remains (fig 2a). The former suggests that if we replace the local statistics with global statistics after it has become static, then the benefits of BN will not be hindered (sec. 4.1 last paragraph); the latter essentially confirms that the local gradient issue will persist given that local gradient is highly dependent on local statistics in eq. 6 in sec. 3.5. Thus, it becomes clear that if we could choose a good moment to “fix” the BN statistics, we could not only enjoy BN benefits but also eliminate the degradation caused during local training.

We can view our two-stage training separately where the first stage serves as a pre-training stage to get a good global statistics estimate and the second stage serves as BN-free training to rescue the degradation. With this regard the effectiveness of this simple “fix” is clearly understandable and the fact that FixBN perfectly recovers centralized learning when E=1 perfectly echoes the conclusion drawn from table 1.

Additionally, we want to kindly share that developing theories is not the main goal of our paper. Our main contributions are to diagnose why and when BN fails in non-IID FL and provide a simple solution to reclaim the performance, as mentioned in the “Contributions” paragraph on page 2, sec. 1.

Reviewer comment (Novelty):
Some references are missing. For example, [1, 2] study why BN fails in FL. [3] compares BN with GN and LN in with extensive experiments.

Author response:
We will make sure all 3 references are cited in the final version. For quick clarification, in addition to the external co-variate shift proposed by [A], we also identify local training as a source of failure of BN in FL. We also want to note that Wang et al., 2023 is a concurrent work and has already been addressed in the footnote on page 5 and compared with it in table 4 (FedTAN), in which our proposed method requires 50 times fewer communication rounds to achieve similar competitive performance. We address [B] in response to Q1 below.

[A] Du, Z., Sun, J., Li, A., Chen, P.-Y., Zhang, J., Li, H., Chen, Y., et al. (2022). Rethinking normalization methods in federated learning. arXiv preprint arXiv:2210.03277. 

[B] Guojun Zhang, Mahdi Beitollahi, Alex Bie, Xi Chen, Normalization Is All You Need: Understanding Layer-Normalized Federated Learning under Extreme Label Shift, arXiv:2308.09565.

Reviewer comment (Experiments):
There is no comparison with FedBN (Li et al.), which also adapts BN to federated learning. Compared to all the previous works, I feel the contribution of this paper is limited.

Author response:
Comparison with FedBN
We agree with the reviewer that FedBN (Li et al.) is a well-known and important work in adapting BN to FL. We intentionally did not include it in our work because it is not a fair comparison. Specifically, we’d kindly point out that FedBN is proposed specifically for Personalized Federated Learning (PFL), which addresses a different learning objective than the Generic Federated Learning (GFL) setting that we consider (see eq. 3). We also explained both settings more explicitly in the “Federated Learning” paragraph in related work and made it clear that our focus is on GFL (last sentence in this paragraph). Nonetheless, we have conducted experiments with FedBN before and are happy to share them here. Following the same setting of table 4, we report the following results (in % accuracy): 

|    FL Scheme          |   IID       | non-IID
| :----:  |    :----:   |          :----: |
| FedBN  |   91.32   | 19.46 
|  Ours (FixBN)  |    91.35  |   87.71

We note that FedBN allows each client to accumulate the local BN statistics without average during the FL process. In the above experiment, we simply average the clients’ BN statistics at the end of the FL process to construct a global model. Although the results show ours is 68.25% higher than FedBN, we still do not claim that our approach “outperforms” FedBN as, again, FedBN was proposed for a different purpose: personalized FL. 


**Contribution of our work**


Please allow us to clarify. BN has been a long-standing issue in FL, and the widely accepted explanation in the community is by Hsieh et al., 2020 which attributes the cause to the mismatch between estimated global statistics and true statistics. Hsieh et al., 2020 also propose to replace BN with GN in all cases, a practice that many subsequent works have adopted. Although widely accepted, the mismatch issue was never explained by Hsieh et al., 2020. Moreover, our results show that even vanilla BN has better performance than GN in most cases (fig. 5), thus we believe bluntly replacing BN with GN in all cases is unfaithful. 

Given these motivations, our goal is to scrutinize the previous belief, diagnose why and when BN fails in non-IID FL, and provide a simple and clean solution to reclaim the performance, as mentioned in the “Contributions” paragraph on page 2, sec. 1. It is our deliberate intention to keep our solution as simple as possible.

To the best of our knowledge, our work is the first of the kind that challenges the widely adopted results of Hsieh et al., 2020. We offer a more detailed analysis and evidence-based reasons of exactly why BN degrades in FL (sec 3) followed by a straightforward solution (sec 4) that could eliminate this issue in all cases (local SGD momentum: sec. 5 & sec. 6.2; different tasks: sec. 6.1; different communication frequency: sec. 7, etc.) 

To summarize, our study addresses a long-standing issue in FL. Its timeliness, simplicity, and strong performance could bring immediate benefits for the practitioners in the FL community and the industry to adopt and build upon. In addition, it also offers a good starting point for deeper theoretical study to follow.

Reviewer comment (Theory):
The conclusion in this work seems not to agree with [3], which says with Yogi and LN, FedAvg can approximate centralized training. What is the explanation for this?

Author response:
We thank the reviewer for raising the comparison here and are happy to provide an explanation. First of all, we respectfully think our work does NOT disagree with this work. In fact, we agree with this work that LN (we refer to it as GN later as it’s a more generic representation of batch-free normalization family and LN is basically GN when #of group = 1) can help bridge FL and centralized training (aka. approximate centralized training) when communicating each mini-batch step. As we can see, in fig. 4A, the leftmost blue dot represents FedAvg + GN at 1 mini-batch per communication, and we can see while centralized training is at 91.5% in accuracy, GN achieves 87.5%, which is already a good approximation. However, our work’s main focus is on BN and investigating and correcting BN’s failure in FL. In essence, they analyze why LN/GN can approximate centralized training while our intent is to ask, can we keep BN but also achieve the same goal? From our analysis and experiment it is clear that with correct treatment BN + FedAvg (FixBN) can indeed approximate centralized learning just as well if not better than GN + FedAvg. This conclusion means BN may just as well if not better approximate centralized learning in FL than GN but it is by no means a contradiction to the claim from [B].

Reviewer comment (Experiments):
Under covariate shift, is FixBN better than FedBN? And what happens under extreme label shift?

Author response:
As mentioned in response to W3, we note again that FedBN is proposed for Personalized Federated Learning (PFL), which has a different learning objective than Generic Federated Learning (GFL) ours (eq. 3). We respectfully believe their results cannot be faithfully and fairly compared. 

We note that the results obtained in response to W3 are already conducted under the “shards” non-IID setting, where each client possesses some distinct classes. In this case, specifically, we have 5 clients each taking 2 classes for 10 total classes in the Cifar10 dataset. Under this condition, as we can see above, FixBN is 68.25% higher than FedBN in accuracy. However, we emphasize again here that we do not claim either one is “better” than the other per se due to the different objectives. Additionally, for the same reason, we believe comparing ours with it under covariate shift would not be meaningful either.

Reviewer comment (Theory):
The article falls short in providing a theoretical basis for the effectiveness of the two-stage BN training strategy in mitigating the degradation issue. A more in-depth examination of why the first stage's sensitivity to FL aggregation does not impede convergence and final performance would enhance the article's comprehensibility and credibility.

Author response:
**The article falls short in providing a theoretical basis for the effectiveness of the two-stage BN training strategy in mitigating the degradation issue.**

We respectfully believe our motivation and analysis for using our two-stage FixBN already plots a rather clear picture of its effectiveness in mitigating degradation issues. First, we note that table 1 reveals a crucial observation, the use of local batch statistics during local training degrades performance, as GN can recover centralized (-0% in accuracy) while BN degrades (-45%) when communicating at every step. This suggests that the use of local batch statistics in local training is what causes degradation. In other words, as long as local batch statistics are not used, the degradation can be rescued. This motivates us to design a local-statistics-free strategy in training. 

In sec. 4.1, we further examine this motivation. We examine both the BN statistics variance within local models and the gap of BN statistics between the local model and global model and found that while local statistics converge (fig. 2B), the gap remains (fig. 2A). The earlier suggests that if we replace the local statistics with global statistics after it has become static, then the benefits of BN will not be hindered (sec. 4.1 last paragraph); the latter essentially confirms that local gradient will be continuously affected given our analysis that BN statistics drives local gradient in eq. 6 in sec. 3.5. Thus, if we could choose a good moment to “fix” the BN statistics, we could not only enjoy BN benefits but also eliminate the degradation caused during local training.

Therefore, we can view our two-stage training separately where the first stage serves as a pre-training stage to get a good global statistics estimate and the second stage serves as BN-free training to rescue the degradation. With this understanding, our two-stage training strategy method could mitigate the degradation issue, as analyzed above, and its performance also perfectly echoes the conclusion drawn from table 1.

**Why the first stage's sensitivity to FL aggregation does not impede convergence and final performance**

As mentioned in the response above, sec. 4.1 underlines that the use of local batch statistics during local training leading to gradient bias (sec 3.5) is the main issue, and once these local batch statistics are no longer used, convergence will no longer be hindered (table 1, GN vs BN in FL setting). While the first stage’s sensitivity to FL aggregation does impede convergence in non-IID scenarios, it does not affect the performance after local batch statistics are fixed in the second stage. We’d also point out that the duration of the first stage does have some effect on the final performance, as fig. 2C shows the first stage duration needs to be chosen appropriately to fully release the full potential of FixBN.

Reviewer comment (Novelty):
This identified reason that “data heterogeneity would lead to unstable gradient estimation in federated learning” is to be expected and is also not new to the community.

Author response:
We do not claim the identified issue of "data heterogeneity leading to unstable gradient estimation in federated learning" as a novel contribution as it is known in the community. However, our main points are 1) that BN will amplify the gradient deviation (sec. 3.5) and 2) to challenge the widely accepted narrative that the degradation in BN performance is solely due to the mismatch between global and local statistics (i.e., FedAvg aggregation results in an inaccurate estimate of global statistics) raised by Hsieh et al., 2020 with our comprehensive diagnosis. As we demonstrate in sec. 3.5, rectifying the global statistics on a combined dataset yields a marginal (< 1%) improvement. We instead attribute the cause to local training.

Reviewer comment (Novelty):
The proposed method is too straightforward and lacks technical depth. To be precise, adopts the original BN in the early stage of training and then fixes the BN layers. The authors are recommended to submit this paper to some workshops instead of the main conference.

Author response:
We respectfully disagree a straightforward solution is a weakness. Our method is simple but solves the issue directly after an in-depth analysis of the cause of why BN fails in sec. 3 and the motivations for fixing BN statistics in sec. 4.1. As machine learning researchers, we believe in Occam's razor: ``More things should not be used than are necessary’’, even in the deep learning era. We respectfully believe that if a simple method built upon solid empirical observations, motivations, and insights could be effective, then we really should not seek an even more complicated method even though it will make our work look fancier (and more novel superficially).

We emphasize our contribution to its timeliness, simplicity, and effectiveness, being also easily adaptable for practitioners in the FL community. For example, we’d like to acknowledge a concurrent work by Wang et al., 2023, which also addresses the BN issue, but compared to it, our proposed method requires 50 times fewer communication rounds to achieve competitive performance, as shown in table 4. Besides simplicity, FixBN is also robust and effective across various conditions, including different normalization techniques (table 5), degrees of non-IID (fig. 3), client participation levels (fig. H), and numbers of local steps (fig. H). The simplicity and effectiveness of our approach, along with its ease of integration into existing FL systems, make it a timely and valuable contribution to the community, particularly for practitioners.

For full disclosure, we received almost the same reviews in our previous submissions to other conferences. We truly appreciate the reviewer’s feedback and we have tried our best to incorporate them into this submission. We hope the reviewer can kindly reconsider the updated paper and the significance of our method to the FL practitioners during the evaluation.

Reviewer comment (Novelty):
Therefore, the contribution of this paper is limited.

Author response:
We respectfully disagree a straightforward solution is a weakness. Our method is simple but solves the issue directly after an in-depth analysis of the cause of why BN fails in sec. 3 and the motivations for fixing BN statistics in sec. 4.1. As machine learning researchers, we believe in Occam's razor: ``More things should not be used than are necessary’’, even in the deep learning era. We respectfully believe that if a simple method built upon solid empirical observations, motivations, and insights could be effective, then we really should not seek an even more complicated method even though it will make our work look fancier (and more novel superficially).

We emphasize our contribution to its timeliness, simplicity, and effectiveness, being also easily adaptable for practitioners in the FL community. For example, we’d like to acknowledge a concurrent work by Wang et al., 2023, which also addresses the BN issue, but compared to it, our proposed method requires 50 times fewer communication rounds to achieve competitive performance, as shown in table 4. Besides simplicity, FixBN is also robust and effective across various conditions, including different normalization techniques (table 5), degrees of non-IID (fig. 3), client participation levels (fig. H), and numbers of local steps (fig. H). The simplicity and effectiveness of our approach, along with its ease of integration into existing FL systems, make it a timely and valuable contribution to the community, particularly for practitioners.

For full disclosure, we received almost the same reviews in our previous submissions to other conferences. We truly appreciate the reviewer’s feedback and we have tried our best to incorporate them into this submission. We hope the reviewer can kindly reconsider the updated paper and the significance of our method to the FL practitioners during the evaluation.

Reviewer comment (Theory):
Some recent work also notices BN can degrade the performance of FL as well, e.g, Wang, Yanmeng, Qingjiang Shi, and Tsung-Hui Chang. "Why batch normalization damage federated learning on non-iid data?." arXiv preprint arXiv:2301.02982 (2023). This work contains comprehensive theoretical analysis of issue of BN and emphasizes that there exists deviation of gradient. However, authors do not include theoretical comparison with this work. Authors also mention that biased gradient can cause the failure of performance in FL, which is similar to the above work. Theoretical analysis is required to validate the conclusion.

Author response:
We thank the reviewer for mentioning this reference. This is a concurrent work; we recognize its contribution and have already cited it in the footnote on page 5 and compared our method with it in table 4 (see the row of “FedTAN”). Although their work did theoretically indicate the gradient deviation problem similar to our sec. 3.5, we further provide a more comprehensive analysis to verify the true cause of the failure of BN in FL, including the BN training dynamics caused by the mini-batch statistics in sec. 3.3 and fig. 2, re-examining the statistics mismatch issue pointed out by Hsieh et al., 2020 in sec. 3.4, and provide more extensive experiments on the effects of local SGD momentum (sec. 5 & sec. 6.2), different tasks (sec. 6.1), communication frequency (sec. 7), etc. 

We also agree with the reviewer that more theoretical analysis could elevate our work, but we want to kindly share that developing theories is not the main goal of our paper. Our main contributions are to diagnose why and when BN fails in non-IID FL and provide a simple solution to reclaim the performance, as mentioned in the “Contributions” paragraph on page 2, sec. 1. Compared to FedTAN (Wang et al., 2023), our proposed method requires 50 times fewer communication rounds to achieve competitive performance, as shown in table 4.

Reviewer comment (Experiments):
Table 2 is an extreme case where single step of mini-batch SGD is performed. It is not surprising that such kind of high variance update can degrade the performance in distributed setting. More general setting should be considered to illustrate the issue of BN.

Author response:
We thank the reviewer for raising this issue. We would like to clarify that we did provide extensive studies on different/more general settings and the observations are consistent. Specifically, in sec. 7 (GN vs BN), we consider different networks (CNN & ResNet), different datasets (Cifar100 & Tiny-ImageNet), different communication frequencies (ranging from single mini-batch to 2500 mini-batches per communication), and different non-IID distributions (from shards to IID). 

The reason we focus on the extreme case of communicating every single step in some experiments is to highlight the negative effects of BN in FL. As elaborated in the second paragraph and table 1  in sec. 3.2, if we could communicate fast at every single step of mini-batch SGD in local training, then theoretically speaking standard FedAvg should already be able to recover centralized learning as this essentially emulates a distributed SGD scenario (training on multi-GPUs with local shuffling). As stated on page 2 of our paper, such a finding sharply contradicts what is observed on DNNs without BN.

Reviewer comment (Reproducibility):
How many epochs are performed in stage 1 and stage 2 respectively in the experiments?

Author response:
All experiments are conducted with each stage having 50% of the total epochs of the whole federated training (i.e., the “fix” is performed from halfway through training). This design choice is based on our observation that local batch statistics have likely converged at this point (see fig. 2B and loss curve fig. 2C, the curve flattens roughly halfway). We found it a safe and stable choice.