A lightweight intelligent fault diagnosis method based on fusion of hierarchical token pooling and vision transformer characteristics

Abstract To improve the accuracy of bearing fault diagnosis under small-sample conditions where fault data is difficult to acquire, this paper proposes a lightweight diagnosis model named Hierarchical Token Pooling-based Vision Transformer (HiT-PiT). This network applies a hierarchical vision transformer to bearing fault diagnosis. Addressing the multi-scale distribution characteristics of bearing fault time-frequency features, the model designs a phased hierarchical token compression and feature aggregation strategy. It adopts a progressive hierarchical token compression method to achieve layer-by-layer aggregation of fault features from local information to global semantics, while preserving local fault details. First, the continuous wavelet transform is employed to convert time-series data into 2D time-frequency images for data preprocessing, followed by the proposed network for feature extraction and fault analysis. Experimental results on the Case Western Reserve University and Huazhong University of Science and Technology bearing datasets demonstrate that the model can reliably identify different types of bearing faults and exhibits good generalization ability. The diagnostic accuracies on the two datasets reach 99.5% and 96%, respectively, showing superior performance in accuracy while providing optimization insights for rotating machinery fault diagnosis.