PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yifei Ding, Minping Jia, Qiuhua Miao, Yudong Cao

arXiv:2104.09079v3cs.AIcs.LGeess.SP

TL;DR

The paper addresses limitations of classical CNN and RNN structures and the absence of Transformer architectures in fault diagnosis. It proposes an SWT-based, end-to-end time–frequency Transformer framework with a TFR-focused tokenizer and Transformer encoder, and reports superior performance against benchmark and state-of-the-art methods on experimental bearing datasets.

  • Problem

    Classical convolutional and recurrent structures have computational-efficiency and feature-representation drawbacks, while Transformer had not been applied to fault diagnosis.

  • Method

    The framework applies SWT to vibration signals, then uses a TFR-focused tokenizer and Transformer-block encoder to classify bearing fault modes end to end.

  • Results

    The proposed TFT achieves the best prediction accuracy among four compared models, with maximum accuracy reaching 100% and the highest average accuracy.

  • Takeaways & Limitations

    Case studies report higher diagnosis accuracy, faster training, noise adaptability, and applicability across multiple working speeds than benchmark and state-of-the-art methods.

  • Takeaways & Limitations

    Future work includes applying Transformer architecture to prognostics and improving the tokenizer’s local receptive field with convolution.

Abstract

from arXiv · show

The scope of data-driven fault diagnosis models is greatly extended through deep learning (DL). However, the classical convolution and recurrent structure have their defects in computational efficiency and feature representation, while the latest Transformer architecture based on attention mechanism has not yet been applied in this field. To solve these problems, we propose a novel time-frequency Transformer (TFT) model inspired by the massive success of vanilla Transformer in sequence processing. Specially, we design a fresh tokenizer and encoder module to extract effective abstractions from the time-frequency representation (TFR) of vibration signals. On this basis, a new end-to-end fault diagnosis framework based on time-frequency Transformer is presented in this paper. Through the case studies on bearing experimental datasets, we construct the optimal Transformer structure and verify its fault diagnosis performance. The superiority of the proposed method is demonstrated in comparison with the benchmark models and other state-of-the-art methods.

1. Introduction

Rolling-bearing fault diagnosis is important for industrial machinery, but existing approaches face limitations in prior-knowledge requirements, feature extraction, and computational structure. The paper addresses these gaps by introducing a time–frequency Transformer and an end-to-end diagnosis framework based on time–frequency representations.

  • Rolling bearings are vulnerable components in widely used rotating machinery, making fault diagnosis important for facility safety and reliability.
  • Model-based diagnosis often requires substantial prior knowledge, while data-driven methods learn correlations from sensor data.
  • Existing fault-diagnosis studies use attention mainly as an auxiliary module within CNN or RNN backbones rather than as the complete architecture.
  • Although Transformer relies only on attention and can model sequence correlations, it had not been used in fault diagnosis or related fields.
  • The proposed TFT uses SWT-derived time–frequency representations, a TFR-focused tokenizer, and Transformer-block encoder to build an end-to-end bearing fault-diagnosis framework.
  • Case studies on bearing experimental datasets compare TFT with benchmark and state-of-the-art methods to assess its diagnosis performance.

2. Preliminaries

Transformers replace recurrent and convolutional sequence processing with stacked blocks centered on multi-head self-attention and position-wise feed-forward layers. Their encoder-only form is suitable for classification, while residual connectors and layer normalization support stable training.

  • Transformer architecture: A vanilla Transformer stacks blocks containing multi-head self-attention, position-wise feed-forward layers, layer normalization, and residual connectors.Embeddings receive positional encoding before passing through the attention and feed-forward sublayers.
  • Multi-head self-attention: Scaled dot-product attention weights values using softmax-normalized query–key dot products divided by √d_k.Different learned projections enable the multiple attention heads.
  • Multi-head self-attention: Multi-head self-attention projects input embeddings into multiple query, key, and value representations, computes parallel attention outputs, and concatenates them.The projected outputs are combined through a further linear projection.
  • Feed-forward layers: The position-wise feed-forward layer applies the same two-layer transformation independently at each sequence position.Its hidden layer uses ReLU activation.
  • Training and stacking: Residual connectors and layer normalization are used around attention and feed-forward modules to reduce training difficulty and support faster, more stable convergence.The stacked Transformer blocks use the same structure but do not share parameters.
  • Transformer modes: Transformer usage modes include encoder-only for classification, decoder-only for language modeling, and encoder–decoder for sequence-to-sequence tasks.The proposed method adopts an encoder-only structure rather than the vanilla encoder–decoder arrangement.
  • Transformer motivation: Transformer addresses recurrent models’ limited training parallelization and difficulty modeling long-range dependencies.Recurrent processing is inherently sequential and cannot directly connect non-adjacent tokens.

3. Time–frequency Transformer

The proposed time–frequency Transformer is organized as a tokenizer, encoder, and classifier for processing time–frequency representations. Its architecture is presented as the paper’s TFT network.

  • Architecture: The proposed TFT network mainly consists of a tokenizer, an encoder, and a classifier.The overall architecture is shown in Fig. 2.

3.1. Tokenizer

The tokenizer converts two-dimensional time–frequency representations into position-aware token embeddings for the Transformer encoder. It segments along time to preserve instantaneous spectral information rather than using image-like grid patches.

  • Tokenization pipeline: The tokenizer flattens, segments, linearly maps, and position-encodes two-dimensional TFR data.This adapts dictionary-based tokenization for TFR inputs.
  • Patch construction: Given x ∈ R^{N_t×N_f×C}, the tokenizer reshapes it to x′ ∈ R^{N_t×(N_f·C)} and cuts patches along the time direction.C represents stacked multi-sensor channels.
  • Token embedding: Each temporal patch is linearly projected into token embeddings with dimension d_model.The learnable mapping W_t performs this projection.
  • Temporal segmentation: Temporal TFR patches represent instantaneous spectra over time, matching the Transformer’s strength in processing temporal sequences.The authors distinguish this from ViT-style grid cutting, which they argue would not retain instantaneous spectrum estimates at each time.
  • Class token: A trainable class token is prepended to the embedded token sequence, and its encoder output serves as the hidden representation of the TFR.The class token is randomly initialized and trainable.
  • Position encoding: The tokenizer adds learnable position encoding because the Transformer has no recurrence or convolution to represent token order.Both 1D and 2D learnable position encodings are considered.

3.2. Encoder

The TFT encoder extracts category-related information from token embeddings through stacked Transformer blocks. It uses multi-head self-attention, feed-forward layers, residual connections, layer normalization, and GeLU activation.

  • Encoder structure: The encoder uses N Transformer blocks to mine category-related information from the embedded sequence.Each block processes the tokenizer output through self-attention and feed-forward components.
  • Activation function: GeLU replaces ReLU in the feed-forward layers to improve network convergence.GeLU uses a Gaussian-distribution-based stochastic gating formulation, reflecting the tendency of normalized neuron inputs to follow a normal distribution.
  • Activation function: GeLU is continuously differentiable and more nonlinear than ReLU around x = 0.The two activation functions are compared in Fig. 3.
  • Encoder structure: Each Transformer block alternates multi-head self-attention with position-wise feed-forward layers using residual connections and layer normalization.The feed-forward module has d_model-dimensional input and output with a d_ff-dimensional hidden layer.
  • Class token: The class token is treated as the next token to be predicted, and its final hidden representation is passed to the classifier.This autoregressive formulation is intended to help the Transformer establish relationships among tokens for classification.

3.3. Classifier

The classifier converts the encoder’s category-informed hidden features into class probabilities. It uses a two-layer feed-forward MLP followed by softmax.

  • Classifier design: The classifier maps hidden features containing category information to one-hot class-label representations.Its output probabilities support prediction of the input sample’s category.
  • Classifier design: The classifier consists of two feed-forward MLP layers and a softmax activation.The hidden layer uses the encoder’s d_ff dimension to reduce the number of hyperparameters.
  • Prediction: The softmax output provides the probability of each category for final sample classification.The number of output categories is denoted by N_cla.

3.4. Training of TFT

TFT is trained by minimizing cross-entropy risk with gradient-based optimization. Adam, dropout, and label smoothing are used to improve convergence and reduce overfitting.

  • Optimization: TFT training minimizes empirical risk using stochastic gradient descent and error back-propagation.The training set consists of samples and their labels, and the model parameters are trainable.
  • Loss function: The cross-entropy loss function is used because TFT addresses a classification problem.The loss compares expected and estimated outputs over the training samples.
  • Optimization: Adam is used to accelerate training-error convergence through an adaptive gradient strategy.The optimizer is intended to help the training trajectory cross non-smooth regions of the loss landscape.
  • Regularization: Dropout is introduced into each network sub-module during training to learn more robust parameters and reduce overfitting.Dropout cuts off connections from some neurons during training.
  • Regularization: Label smoothing uses softened one-hot targets to add noise and suppress overfitting.It reduces the weight assigned to the real sample-label category in the loss calculation.
  • Training procedure: The detailed TFT optimization procedure is specified in Algorithm 1.The procedure includes constructing input sequences, obtaining the class-token representation, and updating weights and biases through loss-driven descent.

4. Fault diagnosis framework based on TFT

The proposed fault-diagnosis framework converts rolling-bearing vibration signals into time-frequency representations and feeds them to a trained TFT model. The model then identifies test-sample faults and its performance is evaluated.

  • Framework workflow: The framework collects vibration signals from rolling bearings as the diagnostic input.This is the first step of the proposed TFT-based fault-diagnosis process.
  • Framework workflow: Synchrosqueezed wavelet transforms convert vibration signals into time-frequency representations for labeled training samples.These TFRs provide the representation processed by TFT.
  • Framework workflow: The TFT hyperparameters and model structure are determined before the model is fully trained.The established model is then applied to identify test samples.
  • Framework workflow: The framework outputs fault-diagnosis results and evaluates the proposed method’s performance.The overall process is illustrated in Fig. 4.

5. Case studies and analysis

The first case study evaluates TFT on the ABLT-1A Bearing Dataset 6308, selecting its structure and hyperparameters before comparing diagnosis accuracy, stability, model size, training time, and feature quality.

  • Experimental setup: The experiment uses ABLT-1A data, with vibration signals from HRB6308 sampled at 12,800 Hz across seven bearing conditions.The dataset includes normal and multiple fault states; SWT is used to construct time-frequency representations.
  • Data preprocessing: SWT produces 1280 × 2560 time-frequency representations, which are resized to 224 × 224 to reduce network scale and computational expense.The resized representations are used as network inputs.
  • Model selection: Cross-validation selects embedding and hidden dimensions, attention heads, and Transformer blocks while assessing effects on model size and generalization.Too-small dimensions limit parameterization, whereas excessive dimensions can over-parameterize the network and degrade generalization.
  • Diagnosis results: After about 40 epochs, TFT training and validation accuracy and loss stabilize, indicating convergence during training.The authors attribute the early training-loss pattern to dropout limiting model capacity.
  • Diagnosis results: TFT achieves the highest average accuracy among TFT, RNN, DNN, and CNN, with maximum prediction accuracy reaching 100%.TFT also has lower accuracy variance than RNN, indicating more stable prediction results.
  • Diagnosis results: TFT combines temporal-information modeling with parallel computation, yielding higher accuracy than DNN and CNN and faster training than recurrent models.The comparison attributes RNN’s long training time to non-parallel computation.
  • Diagnosis results: TFT reaches 100% accuracy in the best confusion-matrix result and 99.88% in the worst, with errors concentrated in weak-fault classes.The worst case still identifies the three non-weak fault states and the normal state.
  • Feature analysis: Only TFT’s t-SNE hidden features accurately separate all faults, while its attention weights become fault-specific through successive Transformer layers.The first Transformer block applies nearly identical attention distributions across fault samples.

5.2. ABLT-1A Bearing Dataset 6205

The second case study tests TFT on multichannel HRB6205 signals collected at four rotational speeds. TFT benefits from multichannel input and maintains high accuracy when samples from different speeds are mixed.

  • Experimental setup: HRB6205 signals are collected with a triaxial acceleration sensor at 12,000 Hz across 1200, 1500, 1750, and 2000 rpm.The three sensor channels are converted into a three-channel digital signal.
  • Data preprocessing: SWT processes each channel separately, after which the resulting time-frequency representations are stacked into multichannel inputs for TFT.This preprocessing is intended to use the dataset’s multichannel information.
  • Benchmark test: TFT accuracy is higher with multichannel data, while the larger input does not significantly increase TFT training time because of parallel computation.CNN also directly accepts multichannel data, whereas larger inputs limit DNN and RNN generalization through increased network size.
  • Multiple-speed test: Mixing and scrambling samples across speeds creates a dataset of 24000 samples for the multiple-speed test.The model is trained repeatedly under the optimal TFT structure and hyperparameter setting.
  • Multiple-speed test: 99.87% average and 99.92% highest classification accuracy are obtained on the mixed-speed dataset.Only a few samples with the real IORF label are misclassified as BF in the best result.
  • Feature analysis: TFT hidden features distinguish health states well across different rotational speeds in the t-SNE visualization.The figure encodes health states by color and rotational speeds by marker.

5.3. Comparison with state-of-the-art

The paper compares TFT with several state-of-the-art fault-diagnosis methods on Datasets 6308 and 6205. TFT achieves the highest classification accuracy on both datasets, including against methods using time-frequency inputs.

  • Comparison scope: The comparison evaluates methods including ELM, SVM, DCNN, CNN, DBN, LSTM, and GRU-based approaches on Datasets 6308 and 6205.The methods use different data-processing procedures and diagnostic models.
  • Comparison results: Models that capture temporal information, including RNNs and their variants, provide greater performance gains than traditional ELM and SVM methods.This pattern is reported across the state-of-the-art comparison.
  • Comparison results: TFT outperforms the other deep-learning approaches even when all receive time-frequency inputs.The authors report that this agrees with the earlier case-study benchmarks.
  • Comparison results: TFT achieves the highest classification accuracy among the compared state-of-the-art solutions on both datasets.The authors present this as further evidence of TFT’s superiority for rolling-bearing fault diagnosis.

6. Conclusion

The paper presents an end-to-end rolling-bearing fault diagnosis framework that converts vibration signals into multi-channel TFRs for TFT-based feature extraction and classification. Case studies report higher accuracy, faster training, noise adaptability, and operation across multiple speeds, while future work targets prognostics and improved local receptive fields.

  • Framework: The framework processes rolling-bearing vibration signals with SWT, feeds multi-channel TFRs into TFT, and classifies fault modes.TFT extracts discriminative hidden features from the time-frequency representations.
  • Model characteristics: TFT removes recurrence and convolution, relying on multi-head self-attention and feed-forward layers to improve parallel computing ability and reduce network scale.Residual self-attention also focuses feature extraction on effective areas and establishes relationships across input sequences.
  • Evaluation: Case studies on experimental data verify the method's effectiveness and superiority over benchmark models and state-of-the-art methods.The reported advantages include higher diagnosis accuracy and faster training than classical deep-learning methods.
  • Operating conditions: The framework supports effective fault diagnosis under strong noise and multiple working conditions, including different speeds.It also makes better use of collected multi-channel signals, which the paper links to diagnostic accuracy and efficiency.
  • Future work: Future work will apply Transformer architecture to prognostics and use convolution in the tokenizer to improve the local receptive field.
Loading 2104.09079v3…