Gearbox compound fault diagnosis using CEEMDAN feature extraction and a dual-attention multi-scale BiLSTM model

2.1. Basic principle of CEEMDAN

Complete Ensemble Empirical Mode Decomposition with Adaptive Noise (CEEMDAN) is an improved signal decomposition algorithm based on Ensemble Empirical Mode Decomposition (EEMD) and Complementary Ensemble Empirical Mode Decomposition (CEEMD). It addresses the shortcomings of traditional decomposition methods, such as mode mixing and error accumulation, by introducing independent additive Gaussian white noise in each decomposition round, dynamically offsetting the calculation errors of each Intrinsic Mode Function (IMF), and significantly improving mode orthogonality and decomposition stability [21]. The core idea is to decompose a non-stationary original signal into a series of physically meaningful Intrinsic Mode Functions (IMFs) and residual signals. The mathematical expression is:

where $x\left(t\right)$ is the input signal; $IM{F}_k\left(t\right)$ is the $k$-th Intrinsic Mode Function; $k$ is the decomposition order (10 in this paper); and $r_k\left(t\right)$ is the residual signal.

2.2. Adaptive PCA

Principal component analysis (PCA) is a classical linear dimensionality reduction method that maps high-dimensional features to low-dimensional space through orthogonal transformation to maximize data variance in the new space, thereby retaining key information and eliminating redundancy [22]. In order to adapt to the high-dimensional characteristics of compound fault characteristics, this paper adopts an adaptive PCA strategy, and the specific steps are as follows:

(1) Missing value imputation:

Since outliers or missing values may be generated during signal decomposition and statistical feature calculation, in order to ensure that the covariance matrix can be constructed correctly, the mean imputation method is first used to handle the missing values generated during feature extraction to ensure data integrity.

where ${\stackrel{-}{x}}_j$ is the mean of the $j$-th feature.

(2) Standardization processing:

Different features have different dimensions. If the covariance matrix is calculated directly, large-scale features will dominate the principal component direction. Therefore, Z-score standardization is performed on the padded features to eliminate the impact of dimensional differences on model training:

where ${\stackrel{-}{x}}_j$ is the feature mean; ${\sigma }_j$ is the feature standard deviation.

After standardization, the mean of each feature is 0 and the variance is 1, causing the covariance matrix to degenerate into a correlation coefficient matrix, thus ensuring that each feature contributes fairly to the principal component direction.

(3) Construction and eigenvalue decomposition of the covariance matrix:

Construct the covariance matrix for the standardized matrix $Z$:

Then, eigenvalue decomposition is performed:

Where $\mathrm{\Lambda }=\text{diag}({\lambda }_1,{\lambda }_2,\dots ,{\lambda }_m)$, $U=[{\mu }_1,{\mu }_2,\dots ,{\mu }_m]$.

Sort by eigenvalue from largest to smallest:

(4) Adaptive Principal Component Selection.

The cumulative variance contribution rate is defined as:

Traditional PCA requires manual setting of the number of principal components $k$, while this paper adopts an adaptive strategy: setting the cumulative contribution rate threshold to 98 % and automatically selecting the minimum $k$ value that satisfies $R_k$:

This strategy effectively removes low-contribution noise components while preserving the dominant frequency band energy, modulation structure, and interaction features in complex faults to the greatest extent possible. This reduces the input dimensionality and parameter size of subsequent deep networks, decreases the risk of overfitting, and improves training stability.

The final dimensionality-reduced data is as follows:

where $W_k=[{\mu }_1,\dots ,{\mu }_k]$.

2.3. Convolutional neural networks

Convolutional Neural Networks (CNNs) are deep learning models inspired by biological visual systems. They are based on multi-layer convolution and pooling operations, possessing parameter sharing, sparse connectivity, and strong inductive bias capabilities. They exhibit excellent performance in feature extraction and generalization, while having a more streamlined parameter set. This study employs a multi-scale one-dimensional convolutional structure, using three different sized convolutional kernels: a 7×1 kernel to capture low-frequency modulation and global structural information, a 5×1 kernel to enhance the expression of meshing frequencies and their harmonic features, and a 3×1 kernel to focus on characterizing high-frequency impacts and transient pulses. This progressive design of convolutional kernels, from large to small, enables the model to extract features layer by layer from macroscopic modulation structures to microscopic impact details, effectively alleviating the multi-band coupling problem in complex faults and improving feature decoupling and discrimination capabilities.

The basic structure alternately stacks the convolutional layer, the pooled layer and the fully connected layer: the convolutional layer slides in the local receptive field through the convolutional kernel, extracts the input features and introduces nonlinearity through the activation function; The pooling layer downsamples the feature map, which retains the salient features while reducing the dimension and enhancing the robustness of the model. The fully connected layer maps the extracted high-order features to the probabilities of each category of the sample to complete the classification task.

2.4. Bidirectional long short-term memory network (BiLSTM)

Long short-term memory networks (LSTMs) effectively solve the gradient disappearance problem of traditional recurrent neural networks (RNNs) through the gating mechanism of input gates, forgetting gates, and output gates, but they can only use historical time series information and cannot capture future data dependencies. Bidirectional long short-term memory network (BiLSTM) realizes the synchronous extraction of bidirectional information from sequence data by constructing forward and reverse LSTM layers in parallel, and its core formula is as follows:

where $X_t$ is the input of the time $t$; $w_x$ and $w_h$ are the weight matrix input to the hidden state and the weight matrix from the previous moment to the current hidden state, respectively. $b_n$ and $b_y$ are the currently hidden bias terms and the output bias terms, respectively.

Fig. 1Structure diagram of 1D CNN

2.5. Design of dual attention mechanism

To achieve precise focusing of composite fault features and improve the model’s ability to identify complex coupled faults, this paper designs a dual-attention mechanism. This structure consists of two branches: General Multi-head Robust Attention and Compound Fault-specific Attention, achieving feature weighting and selection from two levels: general feature modeling and fault-specific feature enhancement.

Let the input feature sequence be:

where $n$ is the feature length and $d$ is the feature dimension.

(1) General Multi-head Robust Attention.

To improve the model's ability to represent fault features at different scales, this paper introduces a multi-head attention mechanism. Multiple attention heads can learn feature weights from different subspaces, thereby enhancing the model's robustness to complex signal features.

The weights of the $i$-th attention head are calculated as follows:

where: $W_{1,i}$, $W_{2,i}$ are the weight matrices of the $i$-th attention head; $b_{1,i}$, $b_{2,i}$ are bias terms; and $\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{h}(\cdot )$ is a non-linear activation function.

The weighted features obtained for each attention head are:

where$\odot$ represents the element-wise weighting operation.

To reduce the bias influence that a single attention head may bring, this paper adopts a mean fusion strategy to fuse the four attention heads:

where $H=$4 represents the number of attention heads.

This module effectively learns the global dependencies of common fault features, improving the model’s stable recognition capability for multiple types of fault features.

(2) Composite Fault-Specific Attention.

Since composite faults often contain coupled features of multiple single faults, relying solely on general attention is insufficient for precise differentiation. Therefore, this paper designs a composite fault-specific attention branch to enhance features for different fault types.

First, the $G_x$ fault attention branch and the G5 fault attention branch are constructed respectively:

The corresponding weighted features are:

where $a_{gx}$ represents the $G_x$ fault feature weight; $a_{g5}$ represents the G5 fault feature weight.

To further characterize the coupling relationship between complex faults, this paper constructs an interactive attention branch:

where $[F_{gx},F_{g5}]$ represent feature concatenation operations.

The final complex fault feature is represented as:

where:

is used to control the contribution ratio of the three types of attention weights.

(3) Dual Attention Fusion Output.

The final model fuses general attention features with fault-specific attention features:

This structure can simultaneously capture general fault pattern features, specific fault features, and the coupling relationship of complex faults, thereby significantly improving the model’s ability to identify complex faults.

2.6. Improve the mixed loss function

To enhance the model's ability to distinguish between composite faults, this paper designs an improved hybrid loss function, consisting of three parts: multi-label binary classification cross-entropy loss, composite fault contrast loss, and attention entropy regularization term.

Different composite categories (such as G1+G5 and G2+G5) are prone to feature aliasing when sharing G5 components. To strengthen the discriminative margin between different composite categories, a boundary constraint based on class response differences is introduced.

Suppose there are two types of composite faults A and B, with key discriminative dimensions $c_A$ and $c_B$, respectively. Define the average logits of the two classes of samples in their corresponding dimensions:

where $S_A$, $S_B$ represent the sample sets of the two composite fault classes, respectively.

The inter-class response difference is defined as:

To ensure a minimum separable margin $\gamma$ (set to 0.6 in this paper) between different composite classes, a hinge-like contrast loss is constructed:

A penalty is applied when the inter-class difference is less than the margin threshold; otherwise, the loss is 0.

The summation is applied pairwise for all composite classes:

where$\mathcal{P}$ represents the set of composite fault class pairs.

Theoretically, this loss is equivalent to constructing a margin-based separation constraint in the output space, ensuring that different composite classes maintain a minimum margin in the discriminative subspace, thereby enhancing inter-class separability.

Let the output weights of the general multi-head robust attention be:

satisfying:

its information entropy is defined as:

where $\epsilon$ is the numerically stable term.

Averaging the batch of samples:

The smaller the entropy, the more concentrated the attention distribution, meaning the model is more focused on key feature regions.

Let the weights of the composite fault-specific attention be:

Similarly, define its entropy as:

Since the specific attention focuses more on the interaction region between $G_x$ and G5, its entropy constraint helps to suppress weight dispersion and improve feature decoupling ability.

Complete hybrid loss function:

where $L_{BCE}$ is the cross-entropy loss of multi-label binary classification, which is used for basic fault classification tasks; The $L_{contrast}$ is the comparative loss of compound faults, with a weight $\alpha$ of 0.15, which improves the discrimination by maximizing the feature differences of different compound faults. $L_{attn1}$ and $L_{attn2}$ are the entropy regular terms of general attention and compound fault attention, respectively, with a weight $\beta$ of 0.02, which enhances the feature focusing ability by minimizing attention entropy.

3.1. Experimental data and preprocessing

To verify the effectiveness of the proposed method, the Beijing Jiaotong University BJTU-RAO bogie transmission system fault simulation test bench dataset was used for experimental validation [23]. This dataset was obtained through a fault simulation experiment on the subway train bogie transmission system and contains multi-sensor data for 51 health states. In this experiment, gearbox fault data were selected, covering six basic states: G0 (normal state), G1 (tooth root crack), G2 (tooth surface wear), G3 (tooth missing), G4 (tooth breakage), and G5 (bearing inner ring fault), as well as four compound fault states: G1 G5, G2 G5, G3 G5, and G4 G5. The data acquisition parameters were: sampling frequency 64 kHz, motor speed 20 Hz, and load 0 kN.

In order to balance feature integrity and computational efficiency, the sliding window technology is used to segment the original vibration signal: the window length is set to 2048 points, covering multiple gear meshing cycles; The step size is set to 1024 points, with an overlap rate of 50 %. For small sample problems, 3x data augmentation is performed on each raw sample, with specific strategies including:

1) Amplitude scaling: randomly scaling the signal amplitude to 0.85-1.15 times the original amplitude;

2) Gaussian noise addition: Add Gaussian white noise with a standard deviation of 1 % of the standard deviation of the original signal to enhance the model's anti-interference ability.

The hierarchical sampling strategy is used to ensure the consistent distribution of various fault samples, and the dataset is divided into training set, verification set and test set according to a ratio of 7:1:2.

Fig. 2Metro train bogie drive system fault simulation test bench

3.2. Feature extraction process

In order to comprehensively capture the multi-dimensional features of compound faults, a four-level feature extraction system is constructed, and the specific steps are as follows:

1) CEEMDAN Decomposition and Effective IMF Selection: This paper determines the optimal number of decomposition layers through sensitivity experiments. A comparison of 6-12 layer decompositions reveals that a 10-layer decomposition yields the best overall performance. Therefore, this paper performs a 10-layer CEEMDAN decomposition on the preprocessed vibration signal, obtaining 10 IMF components. Based on the matching criteria of energy proportion, correlation coefficient, and fault frequency, two effective IMF components containing key fault characteristic information are selected from the 10 IMFs to eliminate noise and spurious components, thereby improving feature effectiveness.

Fig. 3Results of the layer number sensitivity experiment

2) Basic feature extraction: For the two selected effective IMF components, 14-dimensional time-domain features (mean, standard deviation, root mean square, skewness, kurtosis, peak factor, etc.) and 14-dimensional frequency-domain features (spectral energy, spectral centroid, spectral entropy, the first three peak frequencies and amplitude, etc.) are extracted for each component, resulting in a total of 56 basic features;

3) Fault-specific feature extraction: For composite fault modes, 7-dimensional fault-specific features are extracted from the original signal and key IMF, including 4-dimensional features specific to G1-G4 (spectral peak width coefficient, harmonic distribution entropy, impulse factor, pulse duration) and 3-dimensional features specific to G5 (fault frequency energy ratio, sideband energy ratio, impulse response spectrum peak value).

4) Coupling interaction feature extraction: extract the two-dimensional composite fault interaction features ($G_x$ and G5 frequency-energy ratio, composite fault entropy) to describe the nonlinear characteristics of multi-fault coupling;

5) Global feature extraction: 10-dimensional global features are directly extracted from the original signal, including global time-domain features (mean, root mean square, peak factor, etc., 5 dimensions) and global frequency-domain features (meshing frequency energy, fault frequency ratio, etc., 5 dimensions), to supplement the overall information that may be lost during the decomposition process.

6) Feature Fusion and Adaptive PCA Dimensionality Reduction: The extracted 56-dimensional basic features, 7-dimensional specific features, 2-dimensional interaction features, and 10-dimensional global features are fused to form an original feature set with a total dimension of 75. Finally, adaptive PCA dimensionality reduction is performed on this 75-dimensional feature set, setting the cumulative variance contribution rate threshold to 98 % to automatically determine the number of principal components, resulting in a final low-redundancy feature set for subsequent model training.

3.3. Diagnostic model construction and overall architecture

The overall architecture of the intelligent diagnostic model proposed in this paper is shown in Fig. 4. Its data processing flow is as follows: First, the original vibration signal is decomposed by CEEMDAN, subjected to multi-dimensional feature extraction, and reduced by adaptive PCA to obtain a low-dimensional feature vector. This vector is then reshaped and input into a three-layer multi-scale one-dimensional convolutional module (with kernel sizes of 7×1, 5×1, and 3×1) to extract deep fault features across multiple frequency bands. Subsequently, an SE channel attention module is introduced to adaptively recalibrate the feature channels output by the convolutional layers, strengthening key feature channels. Next, the feature sequence is input into a two-layer BiLSTM network to capture its temporal dependencies. The output of the BiLSTM is fed in parallel into two attention branches: a general multi-head robust attention branch and a composite fault-specific attention branch. The output features of the two branches are concatenated to form the final fault feature representation. Finally, this feature representation is fed into a classifier consisting of three fully connected layers to output multi-label diagnostic results.

Fig. 4Overall model architecture diagram

3.4. Model training and weight analysis

Model training uses the AdamW optimizer with a learning rate of 3e-4 and a weight decay of 5e-5. Training is conducted for 100 epochs with an early stopping strategy (patience = 15) to prevent overfitting, and mixed precision training is employed to improve computational efficiency.

To determine the optimal values of the composite fault loss weight α and the attention entropy regularization term weight $\beta$, this paper conducts grid search experiments on two sets of hyperparameters under fixed network structure and training parameters. The parameter ranges are set as follows: $\alpha \in$ [0.05, 0.25], $\beta \in$ [0.01, 0.05], where $\alpha$ controls the weight of the composite fault loss in the total loss function to enhance the model’s learning ability for composite fault samples; $\beta$ constrains the entropy value of the attention distribution to prevent excessive concentration of attention weights, thereby improving the model’s generalization ability.

The Macro-F1 scores, composite fault recognition accuracy, and optimal validation set F1 results of the model under different parameter combinations are shown in the Figs. 5-6.

Experimental results show that when $\alpha =$ 0.15 and $\beta =$ 0.02, the model achieves a Macro-F1 score of 0.9986 on the test set, while maintaining a composite fault identification accuracy of 1.0 on the validation set. Compared to other parameter combinations, this setting achieves the best balance between overall diagnostic performance and composite fault identification capability.

Fig. 5Heatmap of macro-F1 under different hyperparameter combinations of α and β

Fig. 6Heatmap of composite fault accuracy under different hyperparameter combinations of α and β

4.1. Single fault diagnosis results

To verify the model’s basic performance, experiments were first conducted on a single fault dataset. Experimental results show that the model’s training and validation F1 curves converged rapidly and synchronously without significant overfitting, and the loss curve decreased smoothly. The confusion matrix shows that the diagnostic accuracy for various single faults is at a high level. Overall, the model demonstrates excellent accuracy and robustness in gearbox single fault diagnosis tasks, achieving a Macro-F1 score of 0.9868 and a perfect match rate of 0.9860, which can meet the practical needs of intelligent diagnosis of single gearbox faults in industrial scenarios.

Fig. 7Gearbox single fault F1 curve

Fig. 8Loss curve of single gearbox failure

Fig. 9Confusion matrix for single gearbox fault

4.2. Compound fault diagnosis results

The proposed model was validated on a dataset containing four types of composite faults (G1+G5, G2+G5, G3+G5, G4+G5).

Fig. 10Gearbox compound fault F1 curve

Fig. 11Gearbox compound fault loss curve

Fig. 12Gearbox compound fault confusion matrix

From the F1 curve, it can be seen that the Macro-F1 score of the training set quickly converges to more than 0.99 within 5 Epoch, and the F1 synchronization of the verification set is stable around 0.98. In the loss curve, the training set loss continues to decrease with the Epoch iteration and approaches –0,05, and the validation set loss is stable at about 0.05, reflecting the convergence and stability of the training process. Combined with the confusion matrix analysis, the accuracy of various composite faults is at a high level, which verifies the effectiveness of the “exclusive attention mechanism of compound fault interaction features”, and the model can effectively distinguish the multi-fault superposition mode by fusing single fault and coupling features. In summary, the model shows excellent accuracy and robustness in the composite fault diagnosis task, with a Macro-F1 score of 0.9918 and a perfect matching rate of 0.9896, which can meet the intelligent diagnosis needs of gearbox multi-fault coupling state in industrial scenarios.

4.3. Ablation experiment

To verify the independent contribution of each core module in the proposed method to the composite fault diagnosis performance, this section designs ablation experiments. Using a controlled variable method, the CEEMDAN decomposition module, the dual attention branch, and the improved hybrid loss function are removed sequentially. The diagnostic performance of each variant model is compared under the same dataset and training strategy, thereby quantifying the actual role of each module.

The baseline model removes all core modules, retaining only the basic 1D-CNN and BiLSTM structures, and uses the standard BCE loss function to directly extract features from the original signal. Variant models: with ceemdan adds only the CEEMDAN multi-scale decomposition module to the baseline model. With attention adds only the dual attention branch to the baseline model. With loss uses only the improved hybrid loss function to the baseline model. The full model is the complete model proposed in this paper, containing all core modules.

The key performance indicators of each model are shown in the Table 2.

Ablation experiments show that each core module proposed in this paper contributes positively to the model performance. The complete model achieves optimal results on all metrics, fully demonstrating the effectiveness and synergistic effect of the module combination.

4.4. Comparative experiments

To verify the effectiveness of the proposed composite fault attention model, SVM, basic CNN, basic LSTM, and several advanced diagnostic methods (VMD-CNN, RSSD-CYCBD, ResNet-RFA) were selected as comparison models, and comparative experiments were conducted on the same dataset. All deep learning models used the same optimizer parameters, number of training epochs, and early stopping strategy; the SVM employed the RBF kernel function and a OneVsRest multi-label strategy. The experimental results show that the compound fault accuracy of the model in this paper is superior to all the compared deep learning models, demonstrating stronger feature extraction and fault pattern differentiation capabilities. Although the Macro-F1 value of SVM is slightly higher than that of the model in this paper, this is because SVM, as a traditional machine learning model, has stronger generalization ability in small-sample, high-dimensional data scenarios. The preprocessed experimental data is more linearly separable, and its simple structure can better balance the precision and recall of various fault types. The proposed model, being more complex in structure, may have slight overfitting for some small-sample fault categories during training, resulting in a slightly lower Macro-F1 value. However, SVM performs worse than the proposed model in composite fault accuracy and relies on manual feature engineering, making it difficult to capture the deep correlations among multiple faults in composite faults and hard to adapt to practical scenarios. The proposed model can automatically learn multi-scale features of composite faults through a deep network, making it more suitable for practical fault diagnosis needs. Therefore, the model in this paper was chosen over the SVM model.

4.5. Variable operating conditions experiment results

To verify the robustness of the model, variable operating conditions experiments were conducted to test the model performance under different motor speeds and lateral loads. The results are shown in Table 4. When the motor speed is 20 Hz and the load is 0 kN, the model performs optimally, with a Macro-F1 score of 0.9944. As the speed increases and the load grows, the model performance slightly decreases, but the Macro-F1 score remains above 0.97, indicating that the model has good adaptability under complex operating conditions.