EdgeFD: An Edge-Friendly Drift-Aware Fault Diagnosis System for Industrial IoT

Mechanical equipment often exhibits drift in the underlying data distribution of their operational conditions. Factors like varying revolutions, measured in rpm (revolutions per minute), or different loads, measured in Hp (horsepower), can instigate such drifts.

As mechanical systems function, the FD model is iteratively fine-tuned over a series of tasks, represented as $\{\mathcal{T}^{(1)},\mathcal{T}^{(2)},\cdots,\mathcal{T}^{(T)}\}$, where ${\mathcal{T}^{(t)}}$ is a labeled dataset of the $t^{th}$ task, and $\mathcal{T}^{(t)}=\{\bm{x}_{i},\bm{y}_{i}\}_{i=1}^{N_{t}}$, which consists of $N_{t}$ pairs of instances $\bm{x}_{i}$ and their corresponding labels $\bm{y}_{i}$. Assuming the most realistic situation, we consider the case where the task sequence is a task stream with an unknown arriving order, such that the model can access $\mathcal{T}^{(t)}$ only at the training period of this task $\mathcal{T}^{(t)}$, which will become inaccessible afterward. Given $\mathcal{T}^{(t)}$ and the model learned so far, the learning objective at $t$ is as follows:

where $\mathcal{L}(\cdot;\cdot)$ is the loss function that quantifies the difference between predictions made using the current model parameters $\bm{\theta}^{(t)}$ and the ground truth data associated with task $\mathcal{T}^{(t)}$. The role of $\bm{\theta}^{(t-1)}$ is to provide an initial point for the optimization process, enabling the model to adapt quickly and effectively to the new task $\mathcal{T}^{(t)}$.

To solve the optimization problem described above, various methods have been developed that leverage the knowledge contained in $\bm{\theta}^{(t-1)}$ to enhance the learning process for the new task $\mathcal{T}^{(t)}$. Particularly, in the context of diagnosis model training, the techniques based on TL and ML have been mostly explored.

Transfer Learning: The goal of TL is to speed up the learning of new tasks by sharing knowledge between different tasks. The optimization goal of TL can be expressed as minimizing empirical risk, for example using a cross-entropy loss function

where $f(\cdot;\cdot)$ is the corresponding classifier function. The model parameter for the $t$-th task is the combination of the backbone network $\bm{u}^{(t)}$ and the classifier $\bm{v}^{(t)}$: $\bm{\theta}^{(t)}:=(\bm{u}^{(t)}\circ\bm{v}^{(t)})$. The TL process freezes the backbone network $\bm{u}^{(t)}=\bm{u}^{(t-1)}$ and updates only the parameters of the classification head $\bm{v}^{(t)}$. $\ell(\bm{y}_{i},f(\bm{x}_{i};\bm{\theta}))$ is the loss function, for measuring the difference of the predicted output $f(\bm{x}_{i};\bm{\theta})$ against true labels $\bm{y}_{i}$.

Meta-Learning: The goal of ML is to learn the updating rules from multiple tasks, optimizing the learning process for subsequent tasks. A popular algorithm in this domain is Model-Agnostic Meta-Learning (MAML), where the optimization objective is

where $\bm{\theta}^{\prime}=\bm{\theta}-\alpha\nabla_{\bm{\theta}}\ell(f_{\bm{\theta}}(D_{t}^{s}),D_{t}^{q})$, and $f_{\bm{\theta}}(D_{t}^{s})$ denotes the model prediction on task $\mathcal{T}^{(t)}$. In (2), $\bm{\theta}$ represents the global model parameters, $D_{t}^{s}$ is the support set for the task $\mathcal{T}^{(t)}$ (used to fine-tune the model), $D_{t}^{q}$ is the query set (used to evaluate the performance), $\alpha$ is the learning rate.

While achieving certain successes, it is noteworthy that the existing TL-based FD methods often suffer from catastrophic forgetting (as shown in Fig. 2), leading to continuous model adaptation and increased computational burden on edge devices. Meanwhile, ML-based approaches struggle to handle sequential tasks and involve expensive second-order derivative computations. Moreover, these studies assume the learner’s awareness of task distribution changes, allowing the computation of a new posterior approximation. However, this assumption diverges from real-world IIoT scenarios encompassing gradual, sudden, and repeated drifts.

Therefore, data drift detection is necessary and important before model adaptation, and the detection of data drift can serve as a valuable cue for model adaptation. For example, the drift intensity helps to reduce the search space and ease the learning process [14]. Specifically, by dynamically adjusting the consolidation weights of previous parameters based on the drift detection results, we can effectively address the challenges posed by significant data drifts.

Inspired by the CUSUM-type [15] detection method, we propose a sliding window data drift detection method based on the Beta distribution, as illustrated in Fig. 3. In the realm of statistical modeling, the Beta distribution stands out when dealing with random variables bounded within a fixed interval, especially those representing proportions or probabilities. This distribution is naturally defined between 0 and 1, making it a prime choice for situations where we model uncertainties of confidence scores. A key strength of the Beta distribution lies in its versatility; its shape can adapt from U-shaped to bell-shaped, contingent upon its parameters. In our context, the confidence scores from model predictions fall between 0 and 1. As data drift manifests, these confidence score distributions might shift. By leveraging the adaptability of the Beta distribution, we can effectively capture and compare these shifts, making it a powerful tool for drift detection in our methodology.

Within a sliding window of length $N$, the tuple $[\hat{y}_{i},q_{i}]$ denotes the model’s prediction category and the corresponding confidence score for the input data sample $\bm{x}_{i}$. The sliding window is partitioned into two segments: the reference window $Q_{r}$, housing historical confidences, and the target window $Q_{t}$, which contains the latest confidences. It is assumed that these confidence scores in the two windows follow different Beta distributions. When data drift occurs, the CNN model’s classification accuracy declines, causing small confidence values to appear in the target window, leading to discrepancies between the two Beta distributions. To detect data drift, we compare the cumulative probability density function (PDF) values of the confidences in the target window derived from these two Beta distributions.

Step 1. Estimate Beta Distribution Parameters: Using the Maximum Likelihood Estimation (MLE) method, we estimate the parameters of the Beta distribution. Next, the PDF of the beta distribution is provided by

where,

Step 2. Calculate Dissimilarity Score: Compute the dissimilarity score for PDF values on the two Beta distributions. Sum these scores to obtain the cumulative dissimilarity score $s_{k}$.

Step 3. Detect Drift: If the cumulative dissimilarity score $s_{k}$ surpasses a predefined threshold $T_{h}$, data drift is confirmed. The prefixed threshold $T_{h}$ is a critical value used to determine the occurrence of data drift. Its value is inversely related to sensitivity $\lambda$, i.e., a smaller $\lambda$ results in a higher threshold, making the drift detection more stringent.

Step 4. Determine Drift Location: Resize the sub-windows and shift the divider, recalculating the dissimilarity scores. If the new cumulative dissimilarity score $s_{f}^{{}^{\prime}}$ is larger, it implies the new divider is closer to the drift point. This process continues until the maximum dissimilarity score $s_{f}$ is reached, pinpointing the position of the data drift occurrence.

Algorithm 1 provides a structured procedure to detect data drift using the cumulative dissimilarity score. It iteratively evaluates each data point in the sliding window against a threshold to determine the onset of drift. Given a sliding window $Q$, sensitivity to change $\lambda$, padding $\delta$, and a maximum window size $N_{max}$, the algorithm determines whether a drift has occurred by iterating through the sliding window and evaluating the cumulative dissimilarity score against a threshold.

Once the drift is identified using our approach, the next challenge is fine-tuning the model without compromising or forgetting the previously acquired knowledge. Zhang et al. [16] presented a CL method centered on data prototype replay to address this. While effective, this method comes with the downside of increased storage and computational costs. On the other hand, regularization, a tool traditionally employed to mitigate model overfitting, has garnered attention in the CL domain. Within this sphere, the chief avenues for regularization stem from parameter importance estimates [17, 18] and knowledge distillation[19].

Building upon these insights, we introduce our approach. Drawing from the Elastic Weight Consolidation (EWC) [17] methodology, we employ an approximate Bayesian CL strategy, aiming to seamlessly adapt models in continuous learning scenarios. Let $\bm{\theta}$ be the parameter vector, and consider the posterior of $\bm{\theta}$:

Here, the factorization in (5) arises due to the conditional independence assumption of task data. However, computing the exact posteriors is challenging, leading to Laplace’s approximation:

where mean $\bm{\theta}_{t}^{*}$ centered at the maximum a posteriori parameter when learning task $t$, and the precision given by the Fisher information matrix $\mathcal{F}_{t}$ evaluated at $\bm{\theta}_{t}^{*}$.This matrix approximates the Hessian of the negative log-likelihood, ensuring positive semidefiniteness.

The overall optimization objective is to minimize the empirical risk while simultaneously accounting for the significance of parameters from prior tasks through a regularization term:

where $\lambda$ is a hyperparameter that controls the penalty on important parameters from previous tasks. The term $\mathcal{F}_{i}$ is the importance-weighted regularization term for the previous task’s $\mathcal{T}^{(t)}$ parameter. $\mathcal{T}^{(t)}$ represents the current task’s data. This regularization mechanism effectively preserves prior knowledge while adapting to new tasks.

More detailed mathematical derivation of Eqs. (5)-(7) and additional experimental validations will be included in the extended version of this paper, which will provide an in-depth exploration of the presented methodology.

Experimental Setup: All experiments are conducted on a system equipped with an NVIDIA RTX 3090 GPU, using PyTorch version 1.9.

Dataset: This study employs the publicly accessible CWRU Bearing Dataset [21] encompassing 10 classification tasks. These tasks correspond to different failure sites (Inner Race, Outer Race, Rolling Body) and sizes (0.007’, 0.014’, 0.021’). The dataset encompasses bearing data with failures under diverse loads (0 hp, 1 hp, 2 hp, 3 hp), reflecting variations in real-world data collection conditions. As a result, we segment the drifting task based on these distinct load conditions.

Backbone Network: Inspired by [22], we adopt the WDCNN as our core network, recognized for its distinct attributes: (1) a broad initial layer featuring convolutional kernels sized $64\times 1$; (2) stacking multiple compact $3\times 1$ convolutional kernels following the wide kernels. The wide initial convolutional kernels capture broad features, while the stacked smaller kernels delve into finer details.

Training Setting: From the total of 5000 data samples, each load condition has 1250 samples. We employ an 80-20 split strategy: 1000 samples from each load condition were used for model training, while the remaining 250 samples (per condition) serve as the test set to evaluate model performance.

Performance Metric: We use two metrics from [23] to evaluate algorithms when the number of tasks is large, i.e., Average Accuracy (AA) and Average Forgetting (AF).

To ensure the stability and reliability of our results, we conduct three independent repetitions of the experiments. The outcomes of these repeated trials showed some standard deviation, with precise values provided in parentheses.

Within our experiments, we focus on investigating four distinct heterogeneous/drifting tasks. After completing training on 4000 samples, we aggregate the experimental outcomes as presented in Table I. Encouragingly, our DAWC showcase a remarkable performance. Specifically, DAWC demonstrates a notable 4.64% increase in accuracy compared to the current SOTA baseline method (i.e., MAML). It’s worth noting that DAWC not only excelled in terms of performance but also exhibited significant advantages in terms of forgetting rate. This underscores the validation of the efficacy of our proposed weight consolidation strategy.

Through Fig. 4, we can observe the trends in AA variations throughout the learning process for the four different methods. Notably, our DAWC approach demonstrates superior performance, notably in the face of multiple instances of data drift. This serves as evidence of the considerable potential of DAWC in handling heterogeneous tasks.

In summary, the DAWC method outperformed the existing baseline methods in terms of both accuracy and forgetting rate, thereby validating the effectiveness of the proposed weight consolidation strategy in handling drifting tasks.

Effect of Drift Detection: We conduct experiments using the parameters $\lambda=0.05$, $N_{max}=1000$, and $\delta=100$. Each alteration in bearing load signifies a shift in the fundamental data distribution, indicating data drift. The instances where these shifts occur are 1000, 2000, and 3000. In Fig. 5, the vertical dashed lines indicate where a drift is detected by applying our detection method. We can notice that all drifts are detected shortly after they theoretically occur.

Effect of Weight Consolidation: We evaluate the accuracy of the four methods for each task, as depicted in Fig. 6. It is observed that the accuracy steadily improves with fewer fluctuations as the training for new tasks progresses, highlighting the robustness of our method.

The effectiveness of our approach in mitigating forgetting can be attributed to efficient weight consolidation, which helps preserve previously learned information.

Computational Costs: Fig. 7 illustrates the accuracy variations for each task throughout the learning process for the three methods. When confronted with multiple instances of data drift, both TL and DAWC demonstrate a certain reduction in computational costs for edge devices compared to the STL method. However, our approach outperforms TL in this aspect.

This superiority arises from the continuous triggering of model fine-tuning thresholds in TL, where parameters are adjusted to accommodate data drift. In contrast, our DAWC scheme rapidly adapts to various data drift scenarios without the need to trigger model fine-tuning thresholds. As a result, it significantly alleviates computational expenses, and it is hence edge-friendly.