Contrastive-weighted self-supervised model for long-tailed data classification with vision transformer augmented

Abstract

In practical mechanical fault diagnosis, it’s difficult to obtain fault data and the acquisition between normal and fault data is in great imbalance, usually presenting a long-tail distribution state. When training the long-tailed data directly, the imbalanced label may cause label bias, leading to better performance on dominant classes but poorer generalization on tail classes. To address this problem, we adopt two-stage training and propose contrastive-weighted self-supervised model (CSM) with augmented vision transformer which merges the imbalanced learning strategies during the pre-training. In the pretext task, we abandon the label information and adopt contrastive learning by constructing positive and negative sample pairs of each sample. Training it approach to positive sample and away from negatives via augmented vision transformer. The imbalanced strategies are implemented by adaptively weighting to the similarity loss with effective number of samples to learn a better original representation of long-tailed data. In the downstream task, long-tailed data with label are used to fine-tune the pre-trained encoder, which can effectively achieve the final classification task. The experiments under two datasets demonstrate that the pre-training stage can effectively learn a good initialized encoder and can be used in the downstream tasks for better long-tailed data classification.