Harnessing Heterogeneity: Improving Convergence Through Partial Variance Control in Federated Learning

Pranab Sahoo, Ashutosh Tripathi, Sriparna Saha, Samrat Mondal

Abstract: Federated Learning (FL) has emerged as a promising paradigm for collaborative model training without sharing local data. However, a significant challenge in FL arises from the heterogeneous data distributions across participating clients. This heterogeneity leads to highly variable gradient norms in the model's final layers, resulting in poor generalization, slower convergence, and reduced robustness of the global model. To address these issues, we propose a novel technique that incorporates a gradient penalty term into partial variance control. Our method enables diverse representation learning from heterogeneous client data in the initial layers while modifying standard SGD in the final layers. This approach reduces the variance in the classification layers, aligns the gradients, and mitigates the effects of data heterogeneity. Through theoretical analysis, we establish convergence rate bounds for the proposed algorithm, demonstrating its potential for competitive convergence compared to current FL methods in highly heterogeneous data settings. Empirical evaluations on five benchmark datasets validate our approach, showing enhanced performance and faster convergence over state-of-the-art baselines across various levels of data heterogeneity.

Submission Type: Long submission (more than 12 pages of main content)

Previous TMLR Submission Url: https://openreview.net/forum?id=5ZWFANhDzT&referrer=%5Bthe%20profile%20of%20Pranab%20Sahoo%5D(%2Fprofile%3Fid%3D~Pranab_Sahoo1)

Changes Since Last Submission:

As per the conversation with the EIC, we are resubmitting the revised paper. We thank all the reviwers for the feedback.
We have updated the manuscript with the suggested change by the reviewer 4sKJ. All the minor concerned raised by reviwers are addressed. As per the request, we have added all the results with error bars (Kindly refer to section figure 8, 10, 11, 12, 13, 14 and 15) and provied the p-value for statistical significants and given the justifications of p-values (kindly refer to Table 7 and 8).

Regarding the concerns of reviewer Eyon about 'Theorem 1's convergence rate and assumptions', we have updated the manuscript. Kindly refere to section 4 in the manuscript and the full updated proof can be found in section 6 of the Appendix.

All the minor concerns are addressed and updated in the manuscript. Thank you.

Video: https://youtu.be/DpJSqlkHF5I?si=3PE3_aRkSi3oNAmc

Code: https://github.com/Ashutoshtripathi1234/FedPGVC/tree/main

Assigned Action Editor: ~Arya_Mazumdar1

Submission Number: 7394