A Novel domain generalization method for fault diagnosis based on Mamba and causal representation learning

Abstract Current deep learning-based domain generalization (DG) methods for fault diagnosis primarily focus on learning the correlation between features and labels while neglecting the causality between them. Moreover, in equipment fault diagnosis, the quality of feature extraction directly determines model performance, and extracting better time-series features remains a key challenge. To address these challenge, we propose a causal deep DG model with strong temporal feature extraction capabilities for equipment fault diagnosis. First, a selective feature extraction (SFE) module is constructed, which integrates MambaV2 to build a feature extraction framework from the perspective of time-varying systems. This significantly enhances the model’s adaptability to operating condition variations. Second, a structured causal factorization module is designed. Here, domain-related factors such as rotational speed and load are defined as non-causal elements, while fault-related factors are treated as causal elements. By incorporating causal intervention mechanisms, we formulate a causal factorization loss function and optimize SFE module by minimizing this loss to extract invariant causal representations containing fault information. Finally, the causal factors already learned for fault diagnosis are fed into a fully connected layer for fault classification and minimizes the diagnostic classification loss to optimize the classifier and to improve diagnostic accuracy. We have designed three sets of experiments with varying domain gaps using the Case Western Reserve University Bearing Vibration Dataset and a subway train transmission system dataset. Through comparative and ablation experiments, we have validated the effectiveness of our approach. This method will enhance fault diagnosis accuracy in complex working environments, contributing significantly to the advancement of intelligent maintenance in industrial equipment.