Lightweight fault diagnosis of rolling bearings based on improved linear attention Transformer

The Transformer-based rolling bearing fault diagnosis algorithms have a quadratic increase in computational complexity with the input time window, leading to a decrease in the real-time performance of the model inference. To address this problem, a lightweight Transformer fault diagnosis model based on improved linear attention is proposed. Firstly, the improved linear attention is proposed to reduce the quadratic computational complexity, which uses the strategy of changing the computation order of the dot product. Secondly, by substituting the suggested feature bias mapping function for the Softmax global mapping function, the enhanced linear attention feature recovery block is suggested, which reduces the computational burden of utilizing the global acceptance field. At the same time, the bias function has an efficient feature focusing mechanism, which demonstrates significant anti-noise interference properties by strengthening the connection between similar features and weakening the coupling between dissimilar features. Then, the feature diversity recovery block is used to approximate the performance of the original self-attention after global activation and to recover the modeling ability for long-term temporal dependencies. Experiments are conducted on three mechanical failure datasets from Xi'an Jiaotong University and the University of Ottawa. Compared with seven typical models, namely CLFormer, ConvFormer-NSE, MCSwin-T, MobileNet, MobileNet-V2, ResNet18 and MK-ResCNN, the results show that the proposed model outperforms the above models in terms of accuracy and real time performance, and has good robustness in heavy noise environments at the same time. To create a comprehensible connection between the suggested approach and the prediction outcomes, visualize the feature bias mapping function's weight information. Finally, the effectiveness of the proposed modules (feature bias mapping function, feature diversity recovery module) is verified by ablation experiments.