A novel time–frequency Transformer based on self–attention mechanism and its application in fault diagnosis of rolling bearings

The multi-head self–attention mechanism is a key defining characteristic of Transformer models, whose mechanism behind can be viewed as learning an alignment, that is, each token in the sequence attempts to gather information from others [36]. Generally, through linear transformations on a group of input embeddings $X$ with dimension $d_{\rm model}$, we get queries $Q_{s}=XW_{s}^{q}$ and keys $K_{s}=XW_{s}^{k}$ with dimension $d_{k}$, and values $V_{s}=XW_{s}^{v}$ with dimension $d_{v}$. A single-head scaled dot-product attention calculates the dot products of all queries and keys, divides each by scaling factor $\sqrt{d_{k}}$, and apply a softmax function to obtain the weights on the values.

However, instead of mapping with only one version of linear transformations, it is more beneficial to project the input tokens to different queries, keys and values $h$ times through different learned linear transformations. Thus, the so-called multi-head self–attention is introduced. Through the parallel self–attention calculation on queries, keys, and values of each projected version, $h$ different outputs ${\rm head}_{i}$ are obtained. These ${\rm head}_{i}$ are concatenated and once again projected, resulting in the multi-head self–attention

where $W_{i}^{q}\in\mathbb{R}^{d_{\rm model}\times d_{k}}$, $W_{i}^{k}\in\mathbb{R}^{d_{\rm model}\times d_{k}}$ and $W_{i}^{v}\in\mathbb{R}^{d_{\rm model}\times d_{v}}$ denote $i$-th version of linear projection on embeddings $X$ to obtain different queries, keys and values, respectively. $W^{O}\in\mathbb{R}^{h\cdot d_{v}\times d_{\rm model}}$ denotes the linear projection on concatenated multi-head. Note that $d_{k}=d_{v}=d_{\rm model}/h$ in Transformer.

The output of multi-head self–attention module is then passed through a two-layer feed-forward network, whose hidden layer is activated by ReLU. This feed-forward layer operates on each position independently hence the term position-wise.

where $W_{1}\in\mathbb{R}^{d_{\rm model}\times d_{ff}}$, $W_{2}\in\mathbb{R}^{d_{ff}\times d_{\rm model}}$, $b_{1}\in\mathbb{R}^{d_{ff}}$, $b_{2}\in\mathbb{R}^{d_{\rm model}}$ denotes the weights and bias of two layers, respectively.

Transformer composes of several stacked Transformer blocks. A Transformer block usually contains a multi-head self–attention module and a position-wise feed-forward module, both of which use a residual connector and layer normalization to get the combined output of the module.

where $X_{A}$ and $X_{FF}$ are the output of the multi-head self–attention module and the position-wise feed-forward module, respectively. Note that the stacked multiple Transformer blocks take the same structure, while do not share the same parameters.

It is important to note the differences in the mode of usage of the Transformer block. Transformers generally can be divided into three categories, named: 1) encoder-only (e.g., for classification), 2) decoder-only (e.g., for language modeling), and 3) encoder-decoder (e.g., for machine translation). The vanilla Transformer proposed by Vaswani uses an encoder-decoder structure (as shown in Fig. 1) for machine translation, which is a Seq2Seq problem. The decoder of vanilla Transformer uses Transformer blocks that are different from the aforementioned those for the encoder. The proposed method in this paper, however, adopts an encoder-only structure, so details of decoders are not introduced.

To obtain the tokens sequence, TFR is segmented into several patches along the time direction. Specifically, given a TFR ${\rm x}\in\mathbb{R}^{N_{t}\times N_{f}\times C}$, where $N_{t}$ and $N_{f}$ are the length in time and frequency direction, respectively. $C$ is the number of channels, which generally refers to the stacked layers of multi-sensor signals. We first reshape $\rm x$ into sequence ${\rm x}^{\prime}\in\mathbb{R}^{N_{t}\times(N_{f}\cdot C)}$ to flatten multiple channels, and cut it along the time direction to get a patch sequence $\left[{\rm x}_{p}^{1},{\rm x}_{p}^{2},\dots,{\rm x}_{p}^{N_{t}}\right]$ with length $N_{t}$, where ${\rm x}_{p}^{i}\in\mathbb{R}^{N_{f}\cdot C}$, $i=1,\dots,N_{t}$ . Then, to obtain the sequence of token embeddings with dimension $d_{\rm model}$, a learnable linear transformation is used to obtain the projected patches sequence ${\rm x}_{t}$. This process is expressed as:

where $W_{t}\in\mathbb{R}^{N_{f}\times d_{\rm model}}$ denotes the learnable linear mapping. Such processing is based on the view that TFR is formed by splicing the instantaneous spectrum of the signal over a period of time. This is different from gridding segmentation on images by ViT [28], because we believe that grid cutting TFR will not retain instantaneous spectrum estimation at a certain time in each patch. The segmented TFR patches can be regarded as a sequence of instantaneous spectrum in a period of time, and processing such temporal sequence is the strength of Transformer-based structure.

Similar to the class token in BERT [39], a randomly initialized trainable embedding ${\rm x}_{t}^{0}={\rm x}_{\rm class}\in\mathbb{R}^{d_{\rm model}}$ is added to the beginning of the embedded token sequence. Thus, an embedding sequence ${\rm x}_{t}=\left[{\rm x}_{\rm class};{\rm x}_{p}^{1}W_{t},{\rm x}_{p}^{2}W_{t},\dots,{\rm x}_{p}^{N_{t}}W_{t}\right]$ with length $N_{t}+1$ is obtained. Note that output ${\rm z}_{N}^{0}$ of the class token after processed by the subsequent Transformer encoder will serve as the hidden representation of TFR.

Since the Transformer contains no recurrence or convolutional operations, to make full use of the sequence order, we should inject some information about the relative or absolute position of the tokens into the embedding sequence. Vanilla Transformer employs a kind of sinusoid position encoding based on corpus dictionary and token location [26], which is not suitable for the problem we are trying to solve. We propose to use a learnable position encoding $E_{\rm pos}\in\mathbb{R}^{(N_{t}+1)\times d_{\rm model}}$ to extract position information more flexibly through the learning process. The position encodings and token embeddings are added to get the input embeddings

Two types of learnable position encodings, 1D and 2D, are considered. 1) The 1D position encoding $E_{\rm 1d\ pos}\in\mathbb{R}^{(N_{t}+1)\times d_{\rm model}}$ is broadcasted from a vector $e_{\rm 1d}\in\mathbb{R}^{N_{t}+1}$, i.e., $E_{\rm 1d\ pos}={\rm broadcast}(e_{\rm 1d})$. 1D position encoding can only encode the relative and absolute position information among $N_{t}+1$ tokens since the vector elements in the same token share the same encoding. 2) 2D position encoding $E_{\rm 2d\ pos}\in\mathbb{R}^{(N_{t}+1)\times d_{\rm model}}$ is a learning matrix with dimensions $(N_{t}+1)\times d_{\rm model}$, so it can simultaneously encode the location information between and within the tokens. Subsequent case studies will analyze the performance of these two encodings.

To improve the convergence of the network, we adopt Gaussian error linear units (GeLU) activation [47] in feed-forward layers instead of ReLU activation used in vanilla Transformer. GeLU is defined as the product of input $x$ and mask $m\sim{\rm Bernouli(\Phi(x))}$, where $\Phi(x)=P(X\leq x)$, $X\sim\mathcal{N}(0,1)$ is the cumulative distribution function of the standard normal distribution. This distribution is chosen since neuron inputs tend to follow a normal distribution, especially with layer normalization. In this setting, inputs have a higher probability of being “dropped” as $x$ decreases, so the transformation applied to $x$ is stochastic yet depends upon the input.

We can estimate GeLU as

if greater feedforward speed is worth the cost of exactness. GeLU and ReLU are plotted in Fig. 3. It can be seen that GeLU activation is continuously differentiable and has more obvious nonlinearity than the non-differentiable ReLU activation at $x=0$.

The embedding sequence ${\rm z}_{0}=\left[{\rm z}_{0}^{0};{\rm z}_{0}^{1};\dots;{\rm z}_{0}^{N_{t}}\right]$ obtained by tokenizer goes through multiple Transformer blocks to extract the connection among the tokens using self–attention mechanism. An encoder contains $N$ Transformer blocks can be represented as

where $A_{h}(\cdot)$ denotes the multi-head self–attention function with $h$ heads in Eq. (2). $FF(\cdot)$ denotes the position-wise feed-forward module with $d_{\rm model}$ dimensions input and output, and $d_{ff}$ dimensions hidden layer. That is, the multi-head self–attention modules and the feed-forward modules are alternately connected. The extensive use of residual connectors and layer normalization reduces the training difficulty of deep network, which is conducive to faster and more stable convergence.

In TFT, the class token used to output classification features is regarded as the next token to be predicted, that is, the output of token characterizing category is regarded as an auto-regressive prediction problem. This kind of setting is born out of the idea that obtaining the information needed by the words to be predicted from the front words when Transformer is first used in NLP tasks. When applied to classification problems, this solution also enables Transformer to better establish the relationship between the input sequence and the class token to be predicted. Thus, the hidden features ${\rm z}_{N}^{0}$ of the class token output through each layer of the encoder can contain sufficient information in the input sequence. In the last layer of the encoder, the hidden features ${\rm z}_{N}^{0}$ will be served as the output of the encoder for subsequent classification.

The rolling bearing HRB6308 is selected as the experimental bearing in this experiment. The faulty bearing is installed in the first channel of the sensor, and the other three normal bearings are installed in the remaining channels of the sensor. The vibration signal of the fault or normal rolling bearing is collected by the monoaxial vibration acceleration sensor with a sampling frequency of 12,800Hz. The experimental dataset is detailed as follows. Under zero-load conditions, seven types of fault conditions such as normal (6308N), inner ring fault (6308IRF), inner ring weak fault (6308IRWF), outer ring fault (6308ORF), outer ring weak fault (6308ORWF), inner and outer ring compound fault (6308IORF), and inner and outer ring compound weak fault (6308IORWF) are simulated. Accordingly, vibration data were collected for each failure type of the bearings, which were running at 1,050 rpm. For each failure mode interception, 2000 samples of the length 1024 were obtained, totaling $2000\times 7=14000$ samples. In addition, 60% of all datasets are used as training dataset, 20% as validation dataset for model selection and cross validation, and 20% as test dataset for final test. In each training and test, the datasets are randomly divided to ensure the comprehensive evaluation of the model performance.

While the raw vibration signal contains sufficient health condition information of bearings, it is not clear to be used for fault diagnosis directly [24]. Considering the non-stationarity of vibration signals, synchrosqueezed wavelet transforms (SWT) is used to process the raw data to obtain TFRs. The resolution of TFRs obtained by SWT is $1280\times 2560$, which is too large to input a network. This will significantly increase scale of the network and computational expense. Similar to Ref. [53], the bicubic interpolation algorithm [54] is used to resize the resolution to $224\times 224$, which is the common input size. Finally, input images with a shape of $224\times 224\times 1$ are obtained. Vibration signals and corresponding TFRs of bearings with different fault modes are drawn in Fig. 6. These TFRs will be the input of the proposed fault diagnosis model.

In this section, we will determine the structure and hyperparameter selection of the model based on the cross-validation on ABLT-1A Bearing Dataset 6308, and evaluate the influence of these hyperparameter selections on the model size and generalization performance. To increase the feasibility of the research work, all parameters and diagnostic results are cross-validated for multiple average verification and analysis. The hyperparameters we need to determine include: (A) embedding dimension $d_{\rm model}$ and hidden dimension $d_{ff}$, (B) number of attention heads $h$, (C) number of Transformer blocks $N$, (D) dropout rate $r_{dp}$, and (E) selection of position encoding. Each model was trained 10 times for cross-validation. Specifically, the average of the results of cross-validation is employed. The test results of each group are shown in Table 1, where base denotes the final selected model for subsequent research, average acc denotes mean testing accuracy, and params denotes the total number of trainable parameters in the model. The empty item in the params column indicates that this hyperparameter does not affect the total number of trainable parameters. It can be seen from the table that the selection of these parameters has a certain impact on the model scale and performance. In particular, different dimensions and encoder layers will greatly affect the scale of the model.

In addition, for the embedding dimension $d_{\rm model}$, hidden layer dimension $d_{ff}$ and the number of attention heads $h$ that have a great impact on performance, the statistical box chart of test results is shown in Fig. 7. Too small embedding dimension and hidden layer dimension cannot provide the network with enough parameterization ability, so the network performance is poor. In contrast, too large dimension will lead to over-parameterization of the network. At this time, our samples will not be enough to train the over-parameterized network, which will lead to the degradation of the generalization performance. This is also called over-fitting. Besides, we notice that unreasonable selection of hyperparameters will not only reduce the average accuracy, but also make the network performance more unstable, that is, the error distribution is more dispersed. The number of attention heads also shows a similar pattern. Finally, the optimal network structure and hyperparameter selection of TFT are shown in Table 2.

For a fair comparison, the parameter settings of the three benchmark models are also standardized. The detailed model structure and hyperparameter settings are shown in Table 3.

Based on the optimal network structure and hyperparameter setting, the TFT model is trained on ABLT-1A Bearing Dataset 6308. As can be seen from Fig. 8, after about 40 epochs, accuracy and loss values of the training and validation datasets have been very stable, and the model has started to converge, indicating the strong convergence ability of the TFT. Note that in the early stage of training, the training loss is larger than the validation loss because the use of dropout limits the model capacity in training. With the process of training, dropout will drive the network to learn more robust features. Finally, the training loss and validation loss of the network are basically stable at the same value, indicating the good generalization ability of the network. The regularization technique we used fully guarantees the robust generalization of the network.

Then, the trained model is used to classify the testing dataset to evaluate performance. Training and test have been repeated 10 times on TFT and three benchmark models. We draw the results in Table 4. Besides, the total number of trainable parameters and the average training time are also calculated to comprehensively compare the performance of the models. Comparing the test performance of the 4 models, the proposed TFT achieves the best prediction accuracy. Its maximum prediction accuracy can reach 100%, and the average accuracy is also the highest. The accuracy variance of TFT is smaller, indicating that the prediction result is more stable. The performance of RNN is second only to TFT, whose maximum accuracy is 100% and average accuracy is lower than that of TFT. Meanwhile, the variance of RNN is larger, so the result is not as stable and reliable as TFT. The prediction accuracy of DNN and CNN is significantly lower than that of TFT and RNN, which is obviously related to the characterization ability of the model itself. DNN and CNN are just simple multiplication or convolution operation on inputs, lacking the grasp of temporal information.

Furthermore, we compare the scale and training time of these models. DNN contains the most trainable parameters because of its internal fully connected structure. This will lead to the model over parameterized, thus make it easier to overfit when the number of samples is limited. Interestingly, the parameters of CNN are slightly less than that of DNN, but the training time is much longer than that of DNN. This is obviously due to the fact that although convolutional structure reduces the number of trainable parameters through weight sharing, calculation amount corresponding to each parameter increases significantly. The number of trainable parameters in RNN is much less than that in CNN and DNN, but the training time of RNN is the longest for its non-parallel computation.

Generally speaking, the accuracy of TFT and RNN which can grasp the temporal information is higher. But for RNN, this comes at the cost of a very long training time. The results show that the proposed TFT method based on the Transformer structure can process the time series information well with higher prediction accuracy. Moreover, TFT can realize parallel computing, so its training is much faster than models with recurrent structure.

Fig. 9 shows the confusion matrix of the best and the worst test results of TFT, with the accuracy of 100% and 99.88%, respectively. The columns represent the predict labels while the rows represent the true labels for different health states. In the best case, the accuracy is 100% for each health condition. In the worst one, the accuracy is 100% for most conditions except for IORWF and IRWF. A small number of samples with IORWF and IRWF were misjudged as IRWF and ORWF, respectively, indicating greater challenge of weak fault classification. It is also worth noting that, even in the worst case, TFT does not fail to identify the three non-weak fault states and the normal state.

To visualize the extracted features of these models, t-distributed stochastic neighbor embedding (t-SNE) [55] is used to simplify the extracted high-dimensional features of the last hidden layer extracted by the above four methods into two-dimensional vector distribution, the visualization results of the test samples are shown in Fig. 10. As indicated that only the hidden features extracted by the TFT algorithm can accurately separate all faults.

Since the proposed TFT extracts features from TFRs completely based on attention mechanism, it is necessary to explore its mechanism by attention visualization. The attention mechanism adopted by Transformer is mainly used to exert different degrees of attention on the tokens and form the relationship between different tokens. Here, we try to show the attention degree of attention mechanism on different tokens, namely attention weight. Firstly, the attention weight tensors of the first and last self–attention layers are derived. Note that the attention weight is not the output of the attention layer, but the weight ${\rm softmax}\left(Q_{s}K_{s}^{\top}/\sqrt{d_{k}}\right)$ of the input in Eq. 2. Since the calculation results of multi-head attention are concatenated in the network, we sum the weight matrix of $h$ attention heads. Furthermore, the target of attention heads is tokens, and our TFT uses time–frequency cutting patches as tokens. Therefore, the attention mechanism actually works along the time direction. The attention mechanism and the TFRs of bearing signals are drawn in Fig. 11. These figures show the normalized attention weight of the first and last attention layers on different tokens. The larger the value is, the greater the attention weight is.

It can be seen from Fig. 11 that the attention weight distribution of different fault samples in the first transformer block is almost the same. This obviously adheres to our intuition since the network cannot distinguish different fault types in the input layer, but “observe” different samples with the same strategy. With the layer-by-layer attention processing, the network will be able to attach different attention weights to different fault types. Combined with the TFRs, it can be found that the last Transformer block focuses on the tokens with larger values, that is, it pays more attention to the time when the amplitude is more obvious. Through such concentration of attention, the TFT model can effectively grasp the characteristic information from the TFRs, and observe each token with different weights. Thus, TFT can accurately extract the key features of different fault types and avoid the interference of fault independent factors.

To evaluate the anti-noise performance and robustness of the designed algorithm and fault diagnosis method, we operate an anti-noise robustness test. Specially, different degrees of signal-to-noise (SNR) is added into the original bearing signals. The same experimental setup was used to evaluate the variation in performance of each fault diagnosis methods above with different SNR Settings. By adding different degrees of noise to the vibration signal, this experiment evaluates the ability of these fault diagnosis methods to adapt to a noise-changing environment. As shown in Fig. 12, it is worth noting that the proposed method achieves the best performance. When SNR is bigger than 5, that is, under the condition of relatively small noise, the proposed TFT can provide an acceptable prediction accuracy. With the SNR decreasing, the accuracy of all these methods decreases. When the SNR is smaller than -5, the accuracy of DNN and CNN almost decreased to a meaningless value. Note that the meaningless value means the accuracy almost equal to 14.29% of guess blindly. To sum up, the proposed TFT method has relatively high diagnosis accuracy under noise interference environment, and its accuracy is higher than those of other benchmark methods. Also, the proposed TFT method is less disturbed by noise, and the accuracy decreases more slowly with the increase of noise.

The rolling bearing HRB6205 is selected as the experimental bearing in this experiment. Different the previous case, a triaxial acceleration sensor with the sampling frequency of 12,000Hz is used to collect vibration signals of fault or normal rolling bearing. Thus, the original vibration signal is converted into three-channel digital signal by a data acquisition card. In addition, we have taken into account the different rotational speeds of 1200 rpm, 1500 rpm, 1750 rpm, 2000 rpm. The specific experimental dataset is described as follows. Under zero-load conditions, six types of fault conditions such as normal (6205N), inner ring fault (6205IRF), outer ring fault (6205ORF), ball fault (6205BF), inner and outer ring compound fault (6205IORF), and outer ring and ball compound fault (6205ORBF) are simulated respectively. We obtained 1000 samples with a length of 1024 for each health state and each speed. That is, a total of $1000\times 6\times 4=24000$ samples. To distinguish the samples under different speeds, the dataset is divided into four subsets according to the speed, and each subset contains 6000 samples. The benchmark will be conducted on each sub-dataset. 60% of each subset set is used as training dataset, 20% as validation dataset, and 20% as test dataset for final test. In each training and testing, the subsets are randomly divided to ensure the comprehensive evaluation of the model performance.

The proposed TFT model and three benchmark methods in Case 1 are used in the benchmark test. The model structure and hyperparameter setting are basically the same as the optimal setting in Case 1. Similarly, SWT is used to process the vibration signals of each sample to obtain the TFRs. Differently, it should be noted that the signals from the three channels are processed separately and then stacked into multi-channel TFRs to make full use of the multi-channel information in the dataset. After downsampling by bicubic interpolation algorithm, the input data with a shape of $224\times 224\times 3$ is obtained. The proposed TFR and CNN can directly process multi-channel input, while DNN and RNN need flattening operation in the input layer.

Based on the optimal network structure and hyperparameter setting, the TFT model is trained respectively on four subsets of Bearing Dataset 6205. In addition, the three benchmark models are fully trained for comparison. The fully trained models are applied to the diagnosis of test set samples in each subset, and the average results are shown in Table 5. Note that the test for each subset is repeated five times, and the table shows the average accuracy of all subsets.

It can be seen that the accuracy of TFT is higher when multi-channel data is used. Both TFT and CNN can directly input multi-channel data and construct feature graphs fusing multi-channel information. When DNN and RNN are dealing with larger dimension input data, the large network size obviously limits the generalization ability. Besides, the large amount of multi-channel input data does not significantly increase the training time of TFT and CNN due to the parallel computing capability. Particularly, hardly has the training time of TFT increased. Since RNN has no parallel processing ability and no learnable embedding module, its training time is greatly increased. It can be seen that compared with other benchmark models, the proposed TFT can better deal with multi-channel TFRs and obtain better diagnosis performance with lower computational cost and network size.

Bearing Dataset 6205 contains the vibration signals of bearings collected at different speeds, which enables us to test the performance of TFT diagnosis under multiple speed conditions. Different from the previous benchmark test, which divided the dataset into four subsets, we mixed and scrambled all samples with different speeds in this test. Therefore, the mixed dataset contains 24000 samples in total. Based on the optimal TFT structure and hyperparameter setting, the model training is repeated 10 times and used for inference of test samples. The final result shows that the average classification accuracy is 99.87%, and the highest accuracy is 99.92%. The confusion matrix of the best results is shown in Fig. 13. Only a few samples with the real label of IORF were misclassified into BF.

To further investigate the fault feature extraction of TFT under multiple speed conditions, t-SNE is used to visualize the hidden features as shown in Fig. 14. Different colors in the figure represent different health states, while different markers represent different speeds. It can be seen that the hidden features after highly abstracted with TFT have good distinguishability. The samples of different health states can be well distinguished. Furthermore, the samples of the same state, even if with different speeds, can concentrate well. Only in the samples with IORF and ORBF faults, the 1200r samples are slightly separated from the samples with other speeds, but this does not prevent the classifier from classifying them into one category. In addition, several IORF 1500r samples that were misclassified as BF failures can also be clearly observed.