[Paper Review] Personalized Federated Learning for Multi-task Fault Diagnosis of Rotating Machinery

Sheng Guo, Zengxiang Li+3 more|arXiv (Cornell University)|Nov 17, 2022

E-commerce and Technology Innovations4 citations

TL;DR

This paper proposes a personalized federated learning framework for multi-task fault diagnosis in rotating machinery, using federated clustering to group similar machines and adaptive hierarchical aggregation to handle data heterogeneity. The method significantly improves diagnosis accuracy—especially for machines with scarce fault samples—by enabling privacy-preserving, collaborative training across factories while maintaining model personalization.

ABSTRACT

Intelligent fault diagnosis is essential to safe operation of machinery. However, due to scarce fault samples and data heterogeneity in field machinery, deep learning based diagnosis methods are prone to over-fitting with poor generalization ability. To solve the problem, this paper proposes a personalized federated learning framework, enabling multi-task fault diagnosis method across multiple factories in a privacypreserving manner. Firstly, rotating machines from different factories with similar vibration feature data are categorized into machine groups using a federated clustering method. Then, a multi-task deep learning model based on convolutional neural network is constructed to diagnose the multiple faults of machinery with heterogeneous information fusion. Finally, a personalized federated learning framework is proposed to solve data heterogeneity across different machines using adaptive hierarchical aggregation strategy. The case study on collected data from real machines verifies the effectiveness of the proposed framework. The result shows that the diagnosis accuracy could be improved significantly using the proposed personalized federated learning, especially for those machines with scarce fault samples.

Motivation & Objective

• To address the challenge of poor generalization in deep learning-based fault diagnosis due to scarce fault samples and data heterogeneity across industrial machines.
• To enable collaborative model training across multiple factories without sharing raw vibration data, preserving data privacy.
• To improve diagnosis accuracy for machines with limited fault samples through personalized federated learning.
• To develop a multi-task deep learning model that fuses time- and frequency-domain features from vibration signals for simultaneous fault detection.
• To design an adaptive hierarchical aggregation strategy that differentiates between shared feature extraction and personalized classification components.

Proposed method

• Federated clustering is used to group machines from different factories based on similarity in vibration signal features, without exposing raw data.
• A multi-task deep learning model is constructed with a shared convolutional neural network for common feature extraction and a personalized multi-task classification head for each machine group.
• The personalized federated learning framework employs adaptive hierarchical aggregation: global aggregation of the shared feature extractor across all machines and local aggregation of personalized classification heads within each machine group.
• The framework uses a novel adaptive aggregation strategy that dynamically weights model updates based on local data distribution and model performance.
• Vibration signals are processed to extract time- and frequency-domain features, which are used as input to the deep learning model.
• The method avoids centralizing data, ensuring privacy while enabling knowledge transfer across factories with heterogeneous data distributions.

Experimental results

Research questions

• RQ1 Can federated clustering effectively group similar rotating machines across factories without sharing raw vibration data?
• RQ2 How does personalized federated learning improve fault diagnosis accuracy compared to vanilla federated learning and single-machine models?
• RQ3 To what extent does the proposed framework mitigate performance degradation in machines with scarce fault samples?
• RQ4 Does adaptive hierarchical aggregation outperform standard aggregation strategies in handling data heterogeneity across machines?
• RQ5 Can multi-task learning with fused time- and frequency-domain features enhance fault diagnosis performance in a federated setting?

Key findings

• The personalized federated learning framework achieved an average F1 score of 0.855, significantly outperforming the single-machine model (0.818) and vanilla federated learning (0.831).
• For machines with fault rates below 5%, the improvement in F1 score from single-machine models to personalized FL was highest, reaching up to 0.746 compared to 0.543.
• The framework maintained high diagnosis accuracy even when fault samples were limited, demonstrating robustness in low-data regimes.
• Federated clustering reduced data heterogeneity across machines, as evidenced by higher F1 scores on fault types 1 and 3, where the method achieved 0.851 and 0.813, respectively.
• The personalized FL method outperformed both vanilla FL and clustering FL in 90% of the tested machines and fault types, especially in cases with low fault rates.
• When fault samples were abundant (e.g., >50% fault rate), the improvement from federated learning was marginal, confirming its greatest benefit lies in low-data scenarios.