Prototype Based Personalized Federated Learning for Planetary Gearbox Fault Diagnosis

Federated Learning (FL) has been widely investigated in machine fault diagnosis as a decentralized machine-learning paradigm that addresses the fault data island problem in a privacy-preserving scheme. The existing FL methods for machine fault diagnosis are mostly based on the model parameter aggregation. However, due to the different operating conditions of machines with planetary gearbox., the collected fault data of different machines are under different domain distributions so that different machines require their customized models for planetary gearbox fault diagnosis. Those FL methods based on model parameter aggregation are not applicable for the model heterogeneous scenarios and the machines in the FL cannot fully personalize the local models for their own fault diagnosis. To address the above issues, this paper proposes a prototype based personalized federated learning (pFedProto) method for planetary gearbox fault diagnosis, where the class prototypes are utilized for aggregation instead of the model parameters to improve the tolerance to model heterogeneity. Moreover, the L2 norm regularization is added on the representations of local data to align the local prototypes of different clients, further reducing the domain distribution discrepancy of local data and improving the quality of aggregated prototypes. Our proposed pFedProto not only facilitates the personalization of local model for each client on its own diagnosis task, but also improves the classification accuracy of each local model. The comprehensive experimental studies have verified that our proposed pFedProto outperforms the other FL methods for planetary gearbox fault diagnosis in heterogeneous FL.