Extended Invariant Risk Minimization for Machine Fault Diagnosis With Label Noise and Data Shift

Abstract

Incorrect labels as well as the discrepancy between training and test domain data distributions can significantly affect the effectiveness of supervised data-driven models in machine fault diagnosis applications. Such a challenge can be characterized as the noisy label-domain generalization (NL-DG) problem. In this article, the extended invariant risk minimization (EIRM) is developed, which incorporates flat minima seeking to address the NL-DG challenge. The ability of handling NL-DG is realized by shifting the gradient penalty base from the dummy classifier to the entire model. EIRM is shown to be closely related to locating a flat minimum, which is crucial for label noise (LN) robustness and model generalization. Explorations on function smoothness and algorithm convergence are offered to understand EIRM from the theoretical aspect. An efficient implementation of EIRM is also developed to construct the fault diagnosis model. The EIRM-based fault diagnosis method is compared with strong benchmarks on multiple NL-DG tasks using actuator and gearbox fault datasets. Results indicate that the EIRM-based method on average is more effective than the benchmarks. The code is available at https://github.com/mozhenling/doge-eirm.