Tackling Non-IID for Federated Learning with Components Alignment

Federated Learning (FL) is a privacy-preserving framework used to perform machine learning tasks with distributed data. One of the key challenges is heterogeneous data distributions among clients, which results in client-drift, leading to the oscillatory and low-accuracy global model. Although lots of work has been proposed to mitigate client-drift, we find there are drawbacks associated with the two common methods: feature alignment and classifier tuning. For the former, the great bias in classifiers still holds in local models and degrades global model performance. For the latter, it’s hard to obtain suitable global features to introduce external knowledge to locals. To address the above drawbacks, in this paper, we propose a privacy-preserving and effective method, named FCA, to tackle client-drift issues in Non-IID federated learning via aligning models’ components. Specifically, FCA enhances similarity among the local models’ components, i.e. feature extractors and classifiers, by utilizing the estimated global feature representations. Experimental results demonstrate that FCA achieves better performance with fewer rounds. Compared with vanilla, our method achieves from 0.4% to 7.5% performance improvement on three popular datasets with four different Non-IID scenarios.