BalanceFL: Addressing Class Imbalance in Long-Tail Federated Learning

Abstract

Federated Learning (FL) is an emerging learning paradigm that enables the collaborative learning of different nodes without ex-posing the raw data. However, a critical challenge faced by the current federated learning algorithms in real-world applications is the long-tailed data distribution, i.e., in both local and global views, the numbers of classes are often highly imbalanced. This would lead to poor model accuracy on some rare but vital classes, e.g., those related to safety in health and autonomous driving applications. In this paper, we propose BalanceFL, a federated learning frame-work that can robustly learn both common and rare classes from a long-tailed real-world dataset, addressing both the global and local data imbalance at the same time. Specifically, instead of letting nodes upload a class-drifted model trained on imbalanced private data, we design a novel local update scheme that rectifies the class imbalance, forcing the local model to behave as if it were trained on ideal uniform distributed data. To evaluate the performance of BalanceFL, we first adapt two public datasets to the long-tailed federated learning setting, and then collect a real-life IMU dataset for action recognition, which includes more than 10,000 data sam-ples and naturally exhibits the global long tail effect and the local imbalance. On all of these three datasets, BalanceFL outperforms state-of-the-art federated learning approaches by a large margin.