Class Information-Guided Personalized Federated Learning for Fault Diagnosis Under Label Distribution Skew

Federated learning (FL), dedicated to ensuring interclient data privacy and leveraging the private data among clients to collectively train global models, has seen widespread research in gearbox fault diagnosis in recent years. However, in gearbox fault diagnosis, various clients typically exhibit imbalanced distributions of labels, resulting in data label distribution skew across clients, which poses a challenge to achieving convergence of the global model. Moreover, it is difficult for the global model to adapt to the personalized fault diagnosis requirements of different gearboxes. To this end, this study proposes a class information-guided personalized FL (CIGPFL) for gearbox fault diagnosis. To provide personalized diagnostic requirements for different gearbox clients, a generic layer robust feature optimization strategy and a class prototype-guided personalized layer strategy are designed in this framework. In addition, within this framework, local class information orthogonal constraint (CIOC) loss is used to mitigate the negative impact of label distribution skew. Extensive experiments are conducted on the drivetrain diagnostics simulator (DDS) dataset, and the results demonstrate the effectiveness and stability of our approach.