[Paper Review] TrustMAE: A Noise-Resilient Defect Classification Framework using Memory-Augmented Auto-Encoders with Trust Regions

Daniel Stanley Tan, Yi-Chun Chen+2 more|arXiv (Cornell University)|Dec 29, 2020

Industrial Vision Systems and Defect Detection41 references4 citations

TL;DR

TrustMAE is a noise-resilient, unsupervised defect classification framework that uses a memory-augmented auto-encoder with sparse memory addressing and trust-region-based memory updates to reconstruct defect-free images and detect defects via perceptual distance. It achieves robust performance even with up to 40% noisy (defective) images in training, outperforming state-of-the-art baselines in both clean and noisy settings.

ABSTRACT

In this paper, we propose a framework called TrustMAE to address the problem of product defect classification. Instead of relying on defective images that are difficult to collect and laborious to label, our framework can accept datasets with unlabeled images. Moreover, unlike most anomaly detection methods, our approach is robust against noises, or defective images, in the training dataset. Our framework uses a memory-augmented auto-encoder with a sparse memory addressing scheme to avoid over-generalizing the auto-encoder, and a novel trust-region memory updating scheme to keep the noises away from the memory slots. The result is a framework that can reconstruct defect-free images and identify the defective regions using a perceptual distance network. When compared against various state-of-the-art baselines, our approach performs competitively under noise-free MVTec datasets. More importantly, it remains effective at a noise level up to 40% while significantly outperforming other baselines.

Motivation & Objective

• To address the challenge of defective image data contamination in training datasets for product defect classification.
• To enable effective defect detection without requiring fully labeled defect annotations.
• To improve model robustness against noisy training data, including defective images mistakenly included in the training set.
• To develop a framework that can utilize both normal images and normal patches from defective images for improved data efficiency.
• To maintain high detection accuracy on textured and complex surfaces where traditional reconstruction methods fail.

Proposed method

• Employs a memory-augmented auto-encoder that reconstructs input images using stored normal features from a memory bank.
• Uses a sparse memory addressing scheme to limit updates to only a few critical memory slots, reducing over-generalization.
• Introduces a trust-region memory update mechanism that classifies and excludes defective latent features from being stored in memory.
• Applies a perceptual distance network to compute patch-wise differences between input and reconstructed images using deep features.
• Utilizes a spatial perceptual distance metric based on deep feature embeddings to detect defects in high-texture regions.
• Trains the model end-to-end using reconstruction loss and perceptual loss, with noise resilience enforced via trust-region constraints.

Experimental results

Research questions

• RQ1 Can a defect detection framework maintain high performance when trained on datasets containing up to 40% defective images?
• RQ2 How does sparse memory addressing improve generalization and reduce overfitting to defects in the training data?
• RQ3 To what extent does the trust-region memory update mechanism prevent contamination of memory slots by defective features?
• RQ4 How effective is perceptual distance in detecting defects on textured surfaces compared to pixel-wise reconstruction errors?
• RQ5 What is the optimal configuration of trust thresholds and down-sampling layers for robust defect detection across diverse product classes?

Key findings

• TrustMAE achieves competitive performance on the MVTec dataset under noise-free conditions, outperforming multiple state-of-the-art baselines.
• The framework maintains strong performance with up to 40% defective images in the training set, significantly outperforming other methods under high noise levels.
• The ablation study confirms that sparse memory addressing reduces over-generalization and limits defect feature propagation to memory slots.
• The trust-region update mechanism improves robustness, with optimal performance observed in the range δ₂ = [15, 40], achieving ~88% AUC at 30% noise.
• Fewer down-sampling layers improve performance on fine-detail classes (e.g., grid, screw, pill, capsule), as they preserve local feature variations.
• Removing the trust threshold δ₂ leads to model instability, as feature vectors are pushed to infinity, confirming its necessity for training stability.