Optimal Metric-Aware Class Rebalancing in Federated Learning

Abstract

Class imbalance is a significant challenge in many practical classification tasks, particularly in federated learning (FL), where both global data imbalance and local data heterogeneity across clients worsen the problem. Addressing the class imbalance problem effectively while adhering to privacy constraints remains a formidable challenge. A common solution is to rebalance (or reweight) the samples of different classes; however, the rebalancing ratio largely depends on empirical results. In this study, we propose an analytical class rebalancing method to compute the optimal rebalancing ratios for imbalanced datasets. We first theoretically derive the relationship between evaluation metrics--such as macro-precision, macro-recall, and macro-F1--and the rebalancing ratio. Based on these findings, we devise an efficient algorithm to determine the optimal rebalancing ratio that maximizes the corresponding metrics. Our method is parameter-free and doesn't increase the complexity of existing neural models. We demonstrate that our method achieves the optimum ratio for maximizing the concerned metrics while maintaining low computational complexity, scaling linearly with both the number of clients and samples. Experimental results on different datasets validate the effectiveness of our algorithms. To the best of our knowledge, we are the first to provide an analytical and parameter-free solution to the problem.