Addressing Data Heterogeneity In Federated Learning With Adaptive Normalization-Free Feature Recalibration

Authors: Vasilis Siomos, Sergio Naval Marimont, Jonathan Passerat-Palmbach, Giacomo Tarroni
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Federated Learning, Deep Learning, Convolutional Neural Network, Image Classification

Abstract:
Federated learning is a decentralized collaborative training paradigm preserving stakeholders’ data ownership while improving performance and generalization. However, statistical heterogeneity among client datasets degrades system performance. To address this issue, we propose **Adaptive Normalization-free Feature Recalibration (ANFR)**, an architecture-level approach that combines weight standardization and channel attention. Weight standardization normalizes the weights of layers, making it less prone to mismatched client statistics and inconsistent averaging, ensuring robustness under heterogeneity. Channel attention produces learnable scaling factors for feature maps, suppressing inconsistencies across clients due to heterogeneity. We demonstrate that combining these techniques boosts model performance beyond their individual contributions, by improving class selectivity and channel attention weight distribution. ANFR works with any aggregation method, supports both global and personalized FL, and adds minimal overhead. When training with differential privacy, ANFR achieves an appealing balance between privacy and utility, enabling strong privacy guarantees without sacrificing performance. By integrating weight standardization and channel attention in the backbone model, ANFR offers a novel and versatile approach to the challenge of statistical heterogeneity. Extensive experiments show ANFR consistently outperforms established baselines across various aggregation methods, datasets, and heterogeneity conditions. Code is provided at https://anonymous.4open.science/r/anfr_iclr_updated/.

Decision: reject
Reviewer scores: 6, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1729764529291

thread_type:
official_review

<thread_object>
creator: reviewer_qmh4

<content>
summary:
This paper presents ANFR, a method designed to address the challenge of statistical heterogeneity in federated learning environments, with a particular focus on the limitations of batch normalization across heterogeneous client datasets. The authors propose a combination of weight standardization and channel attention mechanisms to normalize layer weights and adaptively scale feature maps, aiming to emphasize shared, informative features across clients. The approach is evaluated across various domains, including medical and common computer vision tasks.

strengths:
1. Simple and straightforward approach that can be easily integrated into existing FL frameworks

2. Comprehensive empirical evaluation across multiple datasets

3. Clear experimental setup and ablation studies

weaknesses:
1. The proposed method combines two existing techniques—weight standardization and channel attention—without substantial modifications specific to the federated learning context. The theoretical justification primarily draws from existing literature, lacking novel analysis within the federated learning framework.

2. In Section 3, the connection between feature consistency in different layers and the proposed weight assignment mechanism (A_C_R, A_C_NR) is unclear, undermining the theoretical foundation.

3. The analysis of Figure 2, which aims to demonstrate the method's advantages regarding class selectivity, requires revision. The current presentation shows minimal difference between ANFR and NF-ResNet, contradicting the authors' explanation of significant improvements.

4. (edited and moved the Questions section)

5. (deleted)

questions:
1. Improved organizational structure in Section 3, with clearer problem definition and methodology subsections, would enhance readability. Including detailed pseudo-code would also aid reproducibility.

2. The experiments in Section 4.4 seem disconnected from the core claims. It would be advisable to include content related to this experiment in Section 3.

3. In Table 1, the performance result of FedChest + FedProx + BN-ResNet is identical to that of FedChest + FedProx + ANFR. The authors state that "we employ a large batch size of 128 to maximize the probability that all classes are represented in each batch, following best practices. This choice incidentally biases the setting in favor of BN models since larger batches reduce inconsistent averaging of mini-batch statistics and BN parameters" (lines 427–429). To gain a more detailed understanding, could the authors provide additional results using smaller or larger batch sizes? Additionally, is there further analysis to explain why this result is observed only under the FedProx setting?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730070030643

thread_type:
official_review

<thread_object>
creator: reviewer_qrtp

<content>
summary:
Statistical heterogeneity among multiple client datasets in federated learning can diminish the system's effectiveness. To address this, this paper introduces Adaptive Normalization-free Feature Recalibration (ANFR), an architecture-level solution combining weight standardization and channel attention. Weight standardization normalizes layer weights, making the system more resilient to client data inconsistencies and irregular averaging. Channel attention produces learnable scaling factors for feature maps, helping suppress those that vary significantly between clients. 

By applying weight standardization and channel attention, ANFR can enhance model performance by boosting class selectivity and optimizing channel attention weight distribution, delivering benefits that surpass their individual effects. To improve the privacy guarantee, ANFR with differential privacy achieves the balance between privacy and utility.

strengths:
This paper gives some interesting experimental results, such as performance comparison, pFL aggregation method comparison, and differential privacy training.

weaknesses:
- The proposed approach appears to integrate weight standardization and channel attention in a relatively straightforward manner, with few technical complexities or novel mechanisms to address the combination. Moreover, the experimental results are not very exciting (e.g., the improved performance over FedChest is less than 1%.)
- In Section 4.3, ANFR combines differential privacy to enhance the protection of the model. However, the work does not include a formal analysis or theoretical proof to rigorously substantiate the differential privacy guarantees, causing a lack of clarity on the degree of privacy achieved from the theoretical perspective.
- The experimental results look incremental on CIFAR-10 and FedChest for performance comparison. For example, the improved performance is mostly less than 1%.
- Writing: Regarding the writing style, Section 3 delves directly into technical details without offering an introductory overview. This sudden shift into specifics could benefit from a preliminary summary that provides context and a roadmap for the technical content, enhancing the logical flow of the section.

questions:
1. Why is the improved performance in pFL aggregation not impressive? For example, compared with NF-ResNet, the improvements are from 84.2 to 84.9 and 83.7 to 83.8.

2. Although the authors introduced "... train with strict sample-level privacy guarantees, employing a privacy budget of $\epsilon=1$, followed by training without privacy constraints to illustrate the privacy/utility trade-off of each model...", it remains unclear what the configurations are for $\delta,\sigma$ and other common parameters used in differentially private learning.

3. The paper lacks a formal analysis of differential privacy and an insightful explanation of "enabling strong privacy guarantees without sacrificing performance." Given my understanding, the authors should theoretically prove the improved privacy-utility trade-off.

4. The authors stated the contribution "offers a robust and flexible solution to the challenge of statistical heterogeneity". However, I can not find any theoretical analysis of robustness and privacy relevant to statistical heterogeneity.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626016456

thread_type:
official_review

<thread_object>
creator: reviewer_8bws

<content>
summary:
The paper addresses the challenge of data heterogeneity in federated learning by introducing feature recalibration through two main strategies: weight standardization (which normalizes the weights of layers rather than the activations) and channel attention (which suppresses feature maps inconsistent across clients due to heterogeneity). This approach is agnostic to aggregation methods and performs effectively in standard and personalized federated learning settings.

strengths:
- This paper focuses on an important challenge: statistical heterogeneity in federated learning. Introduces weight standardization and channel attention. 
- The experimental evaluation in the paper seems exhaustive, covering a wide range of baseline methods and datasets.

weaknesses:
**Unrealistically high results:** Table 1 illustrates very high and seemingly unrealistic results, which appear to outperform even centralized training outcomes (not directly provided in the paper but evident from the literature). 

**Lack of sufficient evidence:** 
- The primary concern is with the CIFAR-10 experiments. Achieving 97.42% accuracy is challenging even in the centralized settings, let alone in federated settings. This makes the results highly questionable. 
- The reported increase in Fed-ISIC experiment is unusually significant and better than what was presented in the original paper [1]. In the original work, the accuracy achieved using FedAvg, FedProx, and SCAFFOLD does not surpass 60%, and centralized training does not exceed 70%, where [1] uses the EfficientNet architecture with a similar fine-tuning approach. 
- The performance of the SCAFFOLD method on the ResNet-50 architecture appears unusually high. In nature, SCAFFOLD is designed to work well in convex settings but is generally expected to struggle in highly non-convex and non-smooth scenarios.

[1] Ogier du Terrail, Jean, et al. "Flamby: Datasets and benchmarks for cross-silo federated learning in realistic healthcare settings." Advances in Neural Information Processing Systems 35 (2022): 5315-5334.

questions:
See weaknesses and provide the following details: 

- Include a performance plot for the experiments conducted, showing results for both **Centralized** and **Individual** accuracies. 
- Conduct and include an experiment that demonstrates how your method performs when you train from scratch (i.e., without pre-training weights), starting with random weights, and provide those results for comparison.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730661840292

thread_type:
official_review

<thread_object>
creator: reviewer_fxjs

<content>
summary:
The paper presents an approach for improving collaboration across different clients in FL when the clients are heterogeneous in terms of their data distributions. The contribution of this paper is the method ANFR that combines weight standardization and channel attention to reduce the affect of the data heterogeneity in collaboration. The weight standardization makes the local models' training independent of the batch statistics and channel attention can focus or suppress important and irrelevant features respectively. The usage of CA and its beenfit seems intuitive.

strengths:
1. The authors mention that this approach can be used for existing global and personalized FL algorithms with different aggregation mechanisms.
2. Experiments and ablations - The authors conduct extensive evaluation along with ablations demonstrating the effect of different components.

weaknesses:
1. Adding an extra attention module on channels might not be practical for FL when the devices do not have sufficient compute.

questions:
1. The impact of CA is clear with ablations, can you please provide more insights on how SWS helps the learning? Is it required for CA to work?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732827549480

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
List of Changes to the Manuscript

comment:
Following our submission of the revised manuscript, we would like to highlight the changes made as a result of discussions with the Reviewers. Below, we detail these changes and the specific concerns they aim to address:

---

1. **Code Repository and Experimental Logs**:
   - A link to an anonymous repository containing our code, selected experimental logs, and performance plots has been added to the abstract and reproducibility statement.
   - This addresses Reviewer 8BWs' doubts about the veracity of our results and findings and helps resolve Reviewer QmH4's question about reproducibility.

2. **Fed-ISIC2019 Experiments Using FLamby Hyper-Parameters**:
   - Additional experiments using the FLamby hyper-parameters have been added in Appendix B.1.
   - These experiments further address Reviewer 8BWs' concerns regarding our Fed-ISIC2019 results in comparison to those reported in the original FLamby paper.
   - Results show that ANFR outperforms the baselines with an even wider margin under these settings.

3. **Random Initialization Experiments**:
   - At the request of Reviewer 8BWs, we performed additional experiments using randomly initialized models (as opposed to ImageNet-pretrained models) on three datasets: CIFAR-10, Fed-ISIC, and FedChest.
   - These results, presented in Appendix B.2, demonstrate that ANFR continues to outperform the baselines in these settings as well.

4. **Performance Plots for Training Convergence**:  
   - Performance plots for training on CIFAR-10 and Fed-ISIC have been added to Appendix B.4.  
   - These plots demonstrate the stable performance increase and convergence of ANFR, addressing Reviewer 8BWs' request.

5. **Improved Section 3 Structure**:
   - Sub-sections have been added to Section 3 to improve readability and address formatting concerns raised by Reviewers QmH4 and QRtp.

6. **Fed-ISIC Hyper-Parameter Tuning Favoring ANFR**:
   - A new experiment on Fed-ISIC2019 has been added to Appendix B.3, where we tune hyper-parameters specifically in favor of ANFR (as opposed to the baseline BN-ResNet, as done in the main paper).
   - Results demonstrate that the reported improvements in the main paper represent a conservative floor of the gains achievable with ANFR, addressing Reviewer QRtp's comment that our results were not “too exciting.”

7. **Batch Size Ablation Study for FedChest**:
   - Following discussions with Reviewer QmH4, we removed the comment on the effect of batch size on FedChest results (lines 427–429 in the original submission).
   - This was replaced with a batch size ablation study in Appendix D.3, providing clearer insights into the role of batch size in this scenario.
   - Results show that ANFR consistently outperforms baselines, particularly in small batch size regimes.

8. **Low-Resource Computational Overhead Benchmark**:
   - A computational overhead benchmark in low-resource settings has been added in Appendix A.3 (Table 7), directly addressing Reviewer fxJS's concern regarding practicality.
   - Results show that ANFR introduces minimal overhead, at most 15%, compared to a vanilla ResNet, confirming its suitability for low-resource settings.

9. **Conciseness and Length Adjustments**:
   - Word- and sentence-level edits have been made throughout the manuscript to accommodate the increased content while maintaining compliance with length restrictions.
</content>
</thread_object>

</thread>


Metareview: