Remb: regularized embedding memory book to extend metric learning in fault diagnosis

Abstract

Few-Shot Learning (FSL) has gained significant attention in fault diagnosis due to its ability to classify faults with limited labeled data. Metric learning-based approaches, such as prototypical networks, have demonstrated effectiveness in this domain by computing class prototypes from support examples. However, these methods often rely on high-dimensional embeddings, which can lead to computational inefficiencies and overfitting, particularly in low-data regimes. This work introduces the Regularized Embedding Memory Book (REMB), a novel module designed to improve prototype estimation and enhance classification accuracy in FSL for fault diagnosis. The REMB module constructs a memory book of embeddings for each class and refines prototype computation by selecting the most representative embeddings. Additionally, we incorporate a set of regularization techniques to improve generalization and robustness, particularly in noisy environments. We evaluate REMB on four benchmark datasets commonly used in fault diagnosis: CWRU, JNU, PU, and an additional dataset with varying training sample sizes. The experimental results demonstrate that REMB improves classification performance, particularly in low-data scenarios, and enhances model calibration. Furthermore, we analyze the impact of noise adaptation and the contribution of different regularization terms through an ablation study. Our findings suggest that the proposed approach provides a robust and efficient solution for few-shot fault diagnosis, outperforming conventional metric learning methods.