FedAMA: Neural Collapse Inspired Federated Adaptive Margin Adjustment Algorithm for Imbalanced Data

Statistical heterogeneity poses a formidable challenge to federated learning (FL), resulting in inconsistent training objectives and serious biases in local feature representations among FL clients. Recent research on neural collapse has identified an optimal simplex equiangular tight frame (ETF) for the structure of mean class feature vectors and classifier vectors under a balanced data distribution. However, the amount of data pertaining to each class in every FL client, as well as the distribution of existing classes across all FL clients, are both imbalanced. Therefore, we first confirm the occurrence of minority collapse in imbalanced FL scenarios through experimentation. Then, we propose the FedAMA algorithm to mitigate its adverse impact by enforcing adaptive pushing forces to minority classes in the global training stage and readjusting the structure of the mean feature and classifier vectors in the local fine-tuning stage. Finally, we theoretically analyze the global convergence of FedAMA and establish its upper bound. Extensive experiments have also been carried out to demonstrate that FedAMA outperforms existing algorithms in terms of global and personalized model performance, particularly in highly heterogeneous settings.