FaultFormer: Pretraining Transformers for Adaptable Bearing Fault Classification

arXiv:2312.02380 · cs.LG, eess.SP · Submitted 2023-12-04 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "FaultFormer: Pretraining Transformers for Adaptable Bearing Fault Classification".

Jane: The gist The framework introduces a novel self-supervised pretraining and fine-tuning framework based on transformer models for adaptable bearing fault classification, which allows models to improve performance on scarce,

Tom: First, who's behind it and why it matters.

Title and authors: Jane: Let’s start with the title and who put this paper out there. It’s "FaultFormer: Pretraining Transformers for Adaptable Bearing Fault Classification," and it's by Anthony Zhou and Amir Barati Farimani from Carnegie Mellon University.

Tom: Yeah, so you see the title immediately tells you what they’re doing: using transformers for bearing fault classification, but with a focus on pretraining them to be adaptable. It’s not just about training a model once; it’s about setting up a system that can handle different scenarios.

Lu: The authors are from strong backgrounds in machine learning and engineering, which makes sense because this paper dives deep into the architecture itself, looking at things like tokenization and how the transformer encoder processes the data.

Meng: It sounds like they’re trying to build a model foundation that isn't tied too tightly to one specific dataset from day one. That flexibility is what I need when I’m shipping models to production systems that have constantly changing failure modes.

Lalam: They are essentially proposing a self-supervised method where the model learns rich representations of vibration data from abundant, unlabeled samples before it ever sees a single labeled fault classification example.

Tom: Exactly, so the implication here is moving away from needing huge, perfectly curated datasets just to start training something useful for predictive maintenance.

The paper's summary: Jane: So, what’s the actual mechanism they use? Basically, they build a transformer encoder that takes raw signal data, optionally augmented and then tokenized into embeddings before it gets fed into an MLP.

Tom: Right, and then they prepended a class token to that sequence of embeddings and processed it with rotary positional embeddings to give the model some context about what it’s looking at.

Lu: The core idea of their pretraining is masked reconstruction, where they apply various masking strategies—zero masking, random masking—to a percentage of tokens in the sequence and then try to reconstruct the original signal.

Meng: That’s clever because it forces the model to learn really expressive latent representations of what a healthy or faulty vibration signature actually looks like, rather than just memorizing patterns from the training set.

Lalam: They are also testing different tokenization strategies—constant, CNN, and Fourier—and different data augmentations like adding Gaussian noise or cutting out portions of the signal to see which one helps them learn better representations.

Tom: So they’re not just sticking to one pipeline; they are systematically investigating how changing the way we represent the signal before it hits the transformer encoder impacts its ability to classify faults later on.

Jane: It sounds like their summary boils down to: pretrain on abundant unlabeled data using masked reconstruction, try different tokenization and augmentation tricks, and then fine-tune that pretrained model for a specific fault class or even a completely new dataset.

The paper's improvements: Tom: Now let’s talk about what they claim actually works better. They show that masking tokens aggressively is helpful because it lets the transformer learn more expressive latent representations, and signal data is easier to reconstruct than language tokens.

Lu: And they validated that these pretrained models perform better when fine-tuning on small amounts of unseen data compared to benchmark models that were trained from scratch on those same tiny datasets.

Meng: If a model can reach high accuracy with fewer training samples because it has this pretraining, then the computational cost and the time spent labeling new data for specific applications drops dramatically.

Lalam: They demonstrated that pretrained transformer models can achieve higher accuracies more quickly and with less computation when adapting them to new bearing faults compared to other models fully trained on those new tasks or data.

Jane: So the main improvement they highlight is that this self-supervised pretraining significantly boosts performance when we are working in low-data regimes, meaning you need fewer samples for good results.

Tom: And they also showed that these pretrained transformers can generalize to entirely new datasets in a few-shot manner, achieving greater than ninety percent accuracy after just two epochs of fine-tuning on a new dataset <ref:2312.02380#pg2>.

Conclusion: Jane: So to wrap up, the main point of FaultFormer is that this framework for pretraining and fine-tuning transformer models allows them to outperform models in low-data regimes because they are learning meaningful context from separate, unlabeled data.

Tom: It’s about getting a model foundation that is flexible enough to be quickly adapted to new bearing faults or entirely new datasets without needing massive amounts of labeled mechanical data upfront.

Lu: The authors show that the ability of these pretrained models to reach high accuracies more quickly and with less computation suggests they can be rapidly adapted to new fault classes through learning meaningful context from other faults.

Meng: From a production perspective, this means we can deploy a single, well-trained foundation model and then quickly tune it for specific manufacturing needs without having to retrain the whole thing every time we get a new fault type.

Lalam: They’ve shown that pretrained transformer models can generalize to new datasets in a few-shot manner after pretraining on one dataset, like CWRU, and then fine-tuning on another, like Paderborn.

Jane: That means the implication is that we are moving toward a system where we don't need exhaustive labeling for every single application; we can leverage prior knowledge from vast amounts of unlabeled data to make these AI models highly adaptable.

ANTHONY ZHOU, AMIR BARATI FARIMANI

Department of Mechanical Engineering, Carnegie Mellon University

cs.LG, eess.SP

Submitted: 2023-12-04

Updated: 2024-05-29

Comments: 10 pages, 6 figures

Journal ref: in IEEE Access, vol. 12, pp. 70719-70728, 2024

DOI: 10.1109/ACCESS.2024.3399670

Code: https://github.com/lucidrains/xtransformers

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 78/100

The gist: The gist The framework introduces a novel self-supervised pretraining and fine-tuning framework based on transformer models for adaptable bearing fault classification, which allows models to improve

Key concepts

Self-supervised Pretraining: This method involves training a transformer model on unlabeled vibration signals by masking parts of the signal (tokens) and forcing the model to reconstruct them. This teaches the model to learn meaningful features from raw data without needing human labels for every sample, making it effective even with very little labeled data.
Transformer Encoder: The core architecture uses a transformer encoder with multi-headed self-attention layers. This structure allows the model to weigh the importance of different parts of the input vibration signal relative to each other, capturing complex sequential dependencies within the time-series data.
Signal Tokenization Strategies: Since raw signals are long, they must be converted into tokens for transformers. The paper explores three methods: constant (reshaping), CNN (reducing sequence length via convolution), and Fourier (extracting key frequency modes). Choosing the right tokenizer is crucial for effectively feeding sequential vibration data into the transformer.
Few-Shot Generalization: This refers to the model's ability to perform well on new tasks or datasets after only seeing a very small amount of new data. The pretraining allows models to quickly adapt and achieve high accuracy on novel bearing faults or entirely new machinery types with minimal additional training.

Terminology

Transformer — 4× in this paper · explained in 2 episodes
Adaptation — 2× in this paper · explained in 1 episode
Fine-Tuning — 2× in this paper · explained in 5 episodes
Transformer Model — 2× in this paper · explained in 2 episodes
Transformer Models — 2× in this paper · explained in 2 episodes
Data Augmentation — 1× in this paper · explained in 8 episodes
Generalizability — 1× in this paper · explained in 6 episodes
Predictive Maintenance — 1× in this paper · explained in 2 episodes
Tokenization — 1× in this paper · explained in 2 episodes

Summary

The gist The framework introduces a novel self-supervised pretraining and fine-tuning framework based on transformer models for adaptable bearing fault classification, which allows models to improve performance on scarce, unseen training samples and generalize to new tasks or datasets quickly.

How it works

The proposed architecture is based on a transformer encoder that can be pretrained or fine-tuned, where raw signal data is optionally augmented then tokenized before being embedded with an MLP
A class token is prepended to the sequence of embeddings and processed with rotary positional embeddings
For pretraining, a mask is applied after tokenization where a percentage of tokens are either zero masked, randomly masked, or left unchanged
The model utilizes different strategies for training transformers for vibration data, including different tokenization strategies, data augmentations, and using unlabled data through masked pretraining

Key Components and Strategies

**)& Data Augmentation: To meet the data needs of the transformer architecture and prevent overfitting, random data augmentations are proposed These include adding Gaussian noise, cutting out portions of the signal, cropping and resizing the signal, and shifting the signal across time The parameters for each augmentation are uniformly sampled from a range of values to generate a unique sample after each augmentation **

**)& Signal Tokenization: The need for a tokenizer is apparent when considering the length of input vibration signals, leading to three tokenization strategies These strategies are constant, CNN, and Fourier The constant tokenizer simply reshapes the input signal by adding a channel dimension in addition to the time dimension The CNN tokenizer is more sophisticated and adds a channel dimension as well as reduces the sequence length through convolving the input signal Lastly, the Fourier tokenizer extracts the top Fourier modes of each token The real and imaginary amplitude as well as frequency of these modes are stacked in the channel dimension **

**)& Transformer Encoder: After producing tokens, a transformer encoder is used to process sequential vibration data The architecture consists of multiple stacked layers, with each layer consisting of a multi-headed self attention mechanism followed by a fully connected layer Within each layer, the self-attention mechanism assigns an attention score to the current token based on relevance to tokens in the input sequence The output of the multi-head attention block is then passed through a feed-forward network consisting of two linear layers and a nonlinearity σ **

Experimental Validation and Findings

**)& Baseline Results: A baseline is first established by training transformer models end-to-end on a standard 80/20 split of the CWRU dataset High-performing models make use of frequent data augmentations, and transformer models can perform just as well as benchmark models, and outperform other models when analyzing data in the Fourier domain Simple MLP architectures can reach accuracies of 99%, a model that is not conventionally considered in literature Transformer models struggle without a tokenization strategy because the attention mechanism can only resolve relationships between and not within tokens **

**)& Masked Pretraining: Aggressively masking tokens is helpful due to learning more expressive latent representations, as well as signal data being easier to reconstruct than language tokens Qualitative results are presented in Figure 4, and show that the transformer encoder is able to learn to nearly perfectly reconstruct masked signals **

**)& Data Scarcity: Self-supervised pretraining increases performance when fine-tuning on small amounts of unseen data Pretrained models can reach optimal performance with fewer samples than benchmark models, and transformers without pretraining do not seem to perform well possibly since transformer models are known to need large quantities of data due to a lack of inductive bias **

**)& Task Adaptation: Pretrained transformer models can achieve higher accuracies more quickly The benefit of pretraining is in its ability to reach high accuracies more quickly and with less computation, suggesting that pretrained models can be quickly adapted to new bearing faults through learning meaningful context from other faults **

**)& Dataset Adaptation: Pretrained models can generalize to new datasets in a few-shot manner After two epochs, pretrained transformer models can achieve greater than 90% accuracy, which is 5 times faster than the CNN benchmark **

Conclusion

In this study a framework for pretraining and fine-tuning transformer models for bearing fault classification is introduced Pretrained models are able to outperform models in low-data regimes due to learning meaningful context from separate, unlabeled data

Improvements for AI systems

The system can be pre-trained on unlabeled data from different bearings, faults, and machinery and then quickly deployed to new, data-scarce applications to suit specific manufacturing needs. This allows a single model to perform across a variety of scenarios without extensive labeled data.
The framework can achieve better final performance as well as quickly adapt to new, unseen tasks by leveraging the ability of pretrained models to reach high accuracies more quickly and with less computation. The model can also quickly extrapolate to fault classes outside of the pretraining distribution.
The system can generalize in a few-shot manner, as it is shown that pretrained transformer models can generalize to new datasets in a few-shot manner after pretraining on one dataset (CWRU) and fine-tuning on another (Paderborn).

Abstract

The growth of global consumption has motivated important applications of deep learning to smart manufacturing and machine health monitoring. In particular, analyzing vibration data offers great potential to extract meaningful insights into predictive maintenance by the detection of bearing faults. Deep learning can be a powerful method to predict these mechanical failures; however, they lack generalizability to new tasks or datasets and require expensive, labeled mechanical data. We address this by presenting a novel self-supervised pretraining and fine-tuning framework based on transformer models. In particular, we investigate different tokenization and data augmentation strategies to reach state-of-the-art accuracies using transformer models. Furthermore, we demonstrate self-supervised masked pretraining for vibration signals and its application to low-data regimes, task adaptation, and dataset adaptation. Pretraining is able to improve performance on scarce, unseen training samples, as well as when fine-tuning on fault classes outside of the pretraining distribution. Furthermore, pretrained transformers are shown to be able to generalize to a different dataset in a few-shot manner. This introduces a new paradigm where models can be pretrained on unlabeled data from different bearings, faults, and machinery and quickly deployed to new, data-scarce applications to suit specific manufacturing needs.