A rolling bearing fault classification method based on feature optimization and transformer-SVM

2.1. VMD principles

VMD is a non-recursive signal decomposition algorithm that decomposes a signal into several Intrinsic Mode Functions (IMF) based on a preset decomposition level. The resulting IMFs can be expressed as:

where $u_k\left(t\right)$ represents the $k$-th decomposed component, $A_k\left(t\right)$ denotes the instantaneous amplitude, and ${\phi }_k\left(t\right)$ is the signal phase.

The objective of VMD is to ensure that the sum of the decomposed IMF components equals the original signal while minimizing the total bandwidth of all modal components. Consequently, the constrained variational model is formulated as follows:

where ${\omega }_k$ is the center frequency of each IMF component, $x\left(t\right)$ is the original signal, $\delta \left(t\right)$ represents the Dirac delta function, and $*$ denotes the convolution operation.

To solve this constrained variational problem, the Lagrangian multiplier and quadratic penalty factor are introduced, transforming the formula into:

The above formulation is solved through iterative updates of the values of $\left\{u_k\right\}$,$\mathrm{ }\left\{{\omega }_k\right\}$, $\lambda$ until the optimal solution is obtained. The update expressions are as follows:

where $\widehat{x}\left(\omega \right)$, $\widehat{u}\left(\omega \right)$ and $\widehat{\lambda }\left(\omega \right)$ are the Fourier transform results of $x\left(t\right)$, $u\left(t\right)$ and $\lambda \left(t\right)$ respectively, and $n$ is the iteration count. The iteration terminates when the following condition is met:

where $\epsilon$ is the convergence threshold.

Although VMD possesses excellent adaptive characteristics and can effectively handle nonlinear signals, its decomposition performance is influenced by the number of modal components and the penalty factor. Manual setting of these parameters often fails to achieve the desired objectives, thus necessitating suitable algorithms for optimizing VMD parameters.

2.2. OCSSA

The Sparrow Search Algorithm (SSA) was proposed in 2020 [28], inspired by sparrows’ foraging and anti-predation behaviors. This algorithm exhibits strong global exploration and local exploitation capabilities, demonstrating high stability and fast convergence, with a simple structure and few parameters. However, SSA is still prone to falling into local optima when dealing with complex functions and suffers from over-reliance on the previous generation's position update mechanism. Therefore, OCSSA introduces improvements to SSA primarily through three aspects.

First, for the initial population, logistic chaotic mapping is used to enhance the diversity of the initialized population, as shown in the following equation:

where $y_m\in \left(\mathrm{0,1}\right]$, $\mu \in \left(\mathrm{0,4}\right)$ represents the control parameter of the mapping evolution. As the value of $\mu$ increases, the mapping range expands and the distribution becomes more uniform.

Furthermore, the Osprey Optimization Algorithm (OOA) [29] is introduced to address SSA’s dependency on previous iteration positions. By simulating the movement of ospreys toward prey, the position update mechanism for explorers in SSA is enhanced, thereby improving global exploration capability. The first-phase global search strategy is as follows:

where $SF$ is the fish selected by the osprey, $r$ is a random number between [0,1], and $I$ is either 1 or 2.

Finally, the Cauchy distribution is employed to replace the position update for followers in SSA. The Cauchy distribution generates greater perturbations compared to the Gaussian distribution, which expands the search scope of SSA and helps avoid convergence to local optima by introducing perturbations throughout the update phase. The position update formula is as follows:

where $P\left(t\right)$ represents the sparrow’s position, and $Cauchy\left(\mathrm{0,1}\right)$ is a random number generated from the standard Cauchy distribution.

The VMD parameters are optimized using the OCSSA algorithm, with the detailed procedure shown in Table 1.

2.3. Transformer principles

Transformer is a novel deep neural network primarily based on the self-attention mechanism. During data processing, positional information is first embedded into the data. The data with positional encoding is then fed into the Transformer encoder, which leverages multi-head self-attention layers and multi-layer perceptrons to establish mapping relationships between inputs and outputs.

The multi-head attention mechanism enables parallel feature learning by implementing $N$ independent attention heads in parallel. Each attention head first projects the input into query, key, and value vectors. It then calculates attention weights through dot-product operations, ultimately outputting a weighted sum of the values. The computational process is as follows:

where $W^o$ is the weight matrix of the multi-head attention mechanism, $H_i$ can be represented as:

where $Q$ is the query matrix, $K$ is the key matrix, $V$ is the value matrix, $d_k$ is the dimension of key vectors.

Meanwhile, layer normalization is added after each layer. The input image is transmitted to the MLP after multi-head attention processing, where the fully connected layer completes the mapping, and finally achieves bearing fault classification through Softmax function operation to output category probabilities.

Transformer achieves efficient parallel computation through its self-attention mechanism and possesses powerful long-range dependency modeling capabilities, enabling effective feature extraction from non-stationary signals. However, the Softmax classification function it employs often fails to deliver satisfactory results when handling nonlinear problems and frequently leads to overfitting issues.

2.4. Support vector machine principles

Support Vector Machine is grounded in statistical learning theory and demonstrates high accuracy in handling small-sample classification tasks. By constructing kernel functions to map the input space into a higher-dimensional space, it creates an optimal classification hyperplane for category separation. For given data $x$ and labels $y$, the classification hyperplane can be defined as $\left(w\cdot x\right)+b=0$ linear function, and the optimal hyperplane can be constructed by solving the following optimization problem:

By converting the aforementioned optimization model into a Lagrange function for solution:

where $\alpha \ge 0$ represents the Lagrange multiplier. Its dual problem forms a typical convex quadratic optimization model, from which the SVM model parameters ${\omega }^{*}$ and $b^{*}$ can be derived. In the feature space of SVM, the distance $D$ between the original feature vector and the optimal hyperplane can be expressed as:

Maximizing this distance can enhance the robustness of the SVM.

2.5. Proposed methodology

This paper proposes a model based on OCSSA-optimized VMD for feature extraction combined with Transformer-SVM, and applies it to faulty bearing diagnosis. To address the dependency of VMD decomposition on the mode number k and penalty factor $\alpha$, the model first utilizes an optimization algorithm to adaptively select the optimal VMD parameters. The selected optimal IMFs are then used to compute a set of indicators – including mean, variance, peak value, kurtosis, RMS, crest factor, impulse factor, margin factor, and waveform factor – to construct feature vectors as input for the subsequent neural network. The global modeling capability of Transformer is then leveraged to further extract deep-level features to aid in final fault classification. Finally, the features extracted by Transformer are fed into a SVM, which utilizes the principle of structural risk minimization to achieve accurate fault pattern classification. The specific workflow is shown in Fig. 1.

Fig. 1Overall flowchart

3.1. Algorithm testing

This paper introduces and employs the OCSSA algorithm to optimize VMD parameters. In this subsection, its performance superiority is validated through comparisons with five other common algorithms: Aquila Optimizer (AO) [30], Moth Firefly Optimization (MFO) [31], Honey Badger Algorithm (HBA) [32], Marine Predator Algorithm (MPA), and Sparrow Search Algorithm (SSA) [33].

To comprehensively evaluate the performance of different optimization algorithms, six test functions listed in Table 2 are employed to assess their performance under various optimization scenarios. To ensure a fair comparison, all algorithms are uniformly configured with a population size of 30, a maximum of 1000 iterations, and a problem dimension of 30. Three-dimensional diagrams of the different test functions are shown in Fig. 2.

The solution accuracy and convergence speed of the algorithms are evaluated using the aforementioned test functions. Fig. 3 illustrates the comparison of convergence accuracy and speed among the six algorithms. The letters at the bottom of the figure correspond to the different test functions mentioned above, while the curves in different colors represent the convergence processes of the respective algorithms on each test function. The vertical axis denotes the function value, and the horizontal axis represents the number of iterations.

As observed from the figure, OCSSA demonstrates high convergence accuracy and fast convergence speed across all test functions. Its convergence curve descends rapidly and stabilizes within a relatively short time, outperforming comparison algorithms such as SSA, AO, MFO, HBA, and MPA. In test function (e), OCSSA is the first to escape local optima and continues searching toward the global optimum. In test function (f), other algorithms tend to stagnate in local optima with slow subsequent improvement in accuracy, while SSA completely fails to escape local optima. In contrast, OCSSA converges at an extremely fast rate to a significantly higher accuracy level than the other algorithms. Comprehensive analysis confirms that OCSSA achieves the fastest convergence speed and the highest convergence accuracy across various types of nonlinear optimization problems, demonstrating excellent global optimization performance.

Fig. 23D diagrams of test functions

a)

b)

c)

d)

e)

f)

Fig. 3Convergence curves of different optimization algorithms

a)

b)

c)

d)

e)

f)

3.2. Huazhong University of Science and Technology data validation

The experiment utilizes the bearing fault dataset published by Huazhong University of Science and Technology [34]. This dataset was collected using a Spectra-Quest mechanical fault simulator and includes nine condition categories comprising normal, minor faults, and severe faults under different rotational speeds. The signal acquisition setup incorporates triaxial accelerometers and a tachometer. The experimental apparatus is shown in Fig. 4.

Fig. 4Bearing test rig at Huazhong University of Science and Technology

The bearing model is ER-16K with a sampling frequency of 25.6 kHz and a single sampling duration of 10.2 seconds. The vibration signal is segmented into data slices every 0.2 seconds with an overlap of 0.05 seconds. Each condition generates 190 samples, with each sample containing 5120 data points. The bearing dataset covers 9 distinct conditions. The complete dataset comprises 1710 samples, of which 1440 are used as the training set and 270 as the test set, as detailed in Table 3.

For the partitioned dataset, following the procedure in Algorithm 1, OCSSA is applied with the objective of minimizing envelope entropy to evaluate the concentration of impact features across IMF components. A lower envelope entropy value indicates that signal energy is more concentrated within fewer impact moments, reflecting more distinct impact characteristics. Based on this, the optimal VMD parameters are obtained and used to decompose the data, from which the optimal IMF component is identified. Nine features are then extracted from this component to form the fault feature vector. The signal resulting from the optimized VMD decomposition is illustrated in Fig. 5.

Fig. 5Comparison of original and reconstructed signals

The figure displays time-domain comparisons of five signal types: normal bearing condition, severe inner race fault, severe outer race fault, severe ball fault, and severe compound fault, showing original signals (left) and reconstructed signals (right). The figure demonstrates that the optimized VMD effectively reduces noise interference. Therefore, selecting the optimal IMF components as computational signals can effectively enhance fault characteristics.

The signal features constructed after OCSSA-VMD processing are input into the established Transformer-SVM network model, with the number of attention heads set to 8, initial learning rate to 0.001, learning rate decay factor to 0.1, and Adam optimizer. The resulting fault classification results are shown in Fig. 6. Since this study primarily focuses on investigating the impact of different features and network models on classification accuracy, all model parameters are assigned fixed values without further parameter optimization.

Fig. 6 displays the classification results of nine bearing fault types under different diagnostic methods. The proposed method achieves the highest accuracy of 99.63 %, significantly outperforming other approaches and demonstrating exceptional capability in identifying subtle fault characteristics. In comparison, the methods in Fig. 6(b) and 6(c) achieve accuracies of only 94.81 % and 90 %, respectively. Both misclassify moderate ball faults as severe inner race faults, primarily due to the oversimplified linear classification layer in the Transformer architecture and the insufficient feature mining by SVM. The method in Fig. 6(d) achieves an accuracy of 90.74 % and shows errors in classifying severe compound faults and normal signals, indicating that standalone VMD decomposition has limitations in extracting critical fault components, thereby affecting subsequent feature learning and classification performance. Comprehensive comparison confirms that the proposed method exhibits significant advantages in both feature extraction and classification decision-making, effectively improving the accuracy of bearing fault diagnosis.

Fig. 6Confusion matrices of different algorithms

a) OCSSA-VMD-Transformer-SVM

b) OCSSA-VMD-Transformer

c) OCSSA-VMD-SVM

d) VMD-Transformer-SVM

Fig. 7Classification accuracy of different algorithms

Fig. 7 compares the classification performance of eight diagnostic methods, with all results presented as means and error bars based on five repeated experiments. Compared with conventional CNN and Transformer models, the proposed OCSSA-VMD-Transformer-SVM algorithm achieves the highest classification accuracy of approximately 99.1 %. From the perspective of input features, the performance of SVM, CNN, and Transformer models is significantly dependent on feature quality. When features are extracted using the VMD method optimized by OCSSA, the diagnostic accuracy is markedly higher than that achieved by the non-optimized VMD method, along with a notably narrower error range. This demonstrates that combining deep learning methods with optimized feature extraction strategies can effectively improve recognition accuracy and enhance model stability. From the perspective of network architectures, the Transformer-SVM structure consistently delivers the best performance across different feature inputs. The proposed architecture integrates the strength of Transformer in deep feature extraction with the generalization capability of SVM, thereby significantly improving the model's ability to discriminate multiple fault types under complex working conditions.

The time costs of various algorithms are summarized in Table 4, with results averaged over five repeated experiments. While SVM exhibits the lowest time cost, its accuracy is also the lowest; hence, only the time costs of different deep learning algorithms are presented here. Compared to other methods, the proposed OCSSA-VMD-Transformer-SVM model not only achieves the highest classification accuracy but also maintains a time cost comparable to that of the baseline CNN network. In summary, the proposed method optimizes the model from both the feature and architecture levels by integrating improved network architecture and optimized feature extraction, achieving higher fault classification accuracy under comparable time costs.

3.3. Case Western Reserve University data validation

The experiment utilizes the bearing fault dataset released by Case Western Reserve University [35]. The test rig consists of a motor, torque transducer, dynamometer, and control system. Single-point faults were introduced using electric discharge machining with fault diameters of 0.007 inches, 0.014 inches, and 0.021 inches. The fault types include inner race faults, outer race faults, and ball faults, totaling 9 fault conditions. Combined with the normal bearing condition, this forms a 10-class bearing fault dataset. The experimental setup is shown in Fig. 8.

Fig. 8Case Western Reserve University bearing test rig

The bearing used is SKF6205 with a sampling frequency of 12,000 Hz. The individual data sampling duration varies between 10-20 seconds. Therefore, 2024 data points are uniformly extracted to form one data sample using a sliding window of 1000 data points. Each fault category yields 120 data samples, with 100 allocated for training and 20 for testing. The dataset comprises 1200 samples in total, with 1000 designated as the training set and 200 as the test set, as detailed in Table 5.

Similarly, the classified data is processed using OCSSA-VMD, and the corresponding signal features are extracted as input vectors. The processing methods and parameters remain consistent with those described in Section 3.2. The classification results of four comparative methods are shown in Fig. 9.

Fig. 9Confusion matrices of different algorithms

a) OCSSA-VMD-Transformer-SVM

b) OCSSA-VMD-Transformer

c) OCSSA-VMD-SVM

d) VMD-Transformer-SVM

Fig. 9 displays the classification results of the ten-class bearing fault data using different methods. The proposed method achieves an accuracy of 99.5 %, with only one normal sample misclassified as a minor ball fault. In comparison, the methods in Fig. 9(b) and 9(c) achieve accuracies of only 97 % and 87 % respectively, validating the effectiveness of the Transformer-SVM fusion strategy in enhancing classification performance. The method in Fig. 9(d) achieves an accuracy of 92 %, indicating that the unoptimized VMD decomposition underperforms in fault feature extraction, thereby limiting subsequent classification performance. A comprehensive comparison of the four methods demonstrates the significant advantage of the proposed approach in classification accuracy.

Fig. 10 compares the classification performance of eight diagnostic methods on the Case Western Reserve dataset, with all results presented as the average and error bars from five repeated experiments. Compared to CNN and Transformer models under different feature inputs, the proposed OCSSA-VMD-Transformer-SVM model achieved the highest classification accuracy of approximately 99.2 %. From the perspective of feature extraction methods, when features were extracted using the VMD method optimized by OCSSA, diagnostic accuracy was significantly higher than that of the non-optimized VMD method, with a noticeably narrower error range. This indicates that optimizing the input features can more effectively highlight fault characteristics, thereby improving the diagnostic performance of the network. From the dimension of network architecture, the Transformer-SVM structure performed optimally across different feature input conditions; the diagnostic accuracy of the standalone Transformer model and CNN decreased relatively; using SVM alone for classification yielded the poorest results.

Fig. 10Classification accuracy of different algorithms

The time costs of different algorithms are presented in Table 6. Compared to other methods, the model proposed in this paper exhibits a time cost that is comparable to other approaches, while achieving the highest classification accuracy. In summary, the proposed method can achieve higher fault classification accuracy and stability than other methods under comparable time costs.