Prototype-Guided Federated Learning with Performance-Based Aggregation and Improved Model Training for Imbalanced Fault Diagnosis

Individual factories often lack access to diverse fault data because faults are rare events, posing a significant challenge for effective diagnosis. Although cross-factory collaboration could help, data privacy and competitive concerns hinder direct sharing. Previous studies have explored federated learning (FL) as a decentralized solution that enables collaborative model training while preserving data locality. However, applying FL to industrial fault diagnosis remains difficult due to non-independent and identically distributed (non-IID) data and the severe class imbalance caused by infrequent fault events. To address these challenges, an improved federated framework is proposed. The local model consists of three blocks: a compact convolutional encoder, a channel-wise attention module, and a performance improvement block. This architecture enhances feature quality under class imbalance. A hybrid aggregation strategy combines sample size and training loss to adaptively weigh client updates based on their performance. This design enhances generalization and inter-client knowledge alignment under decentralized constraints. Experimental results on a real-world rolling bearing dataset demonstrate that the proposed method significantly outperforms traditional FL baselines (e.g., FedAvg), achieving up to 91% classification accuracy. These results highlight their potential for reliable deployment in industrial predictive maintenance systems under decentralized and data-limited conditions.