[论文解读] Puzzle-AE: Novelty Detection in Images through Solving Puzzles

Mohammadreza Salehi, Ainaz Eftekhar+3 more|arXiv (Cornell University)|Aug 29, 2020

Anomaly Detection Techniques and Applications参考文献 51被引用 14

一句话总结

Puzzle-AE 提出了一种新颖的异常检测框架，通过将 U-Net 训练为解决图像拼图任务（作为自监督的预训练任务），提升了正常数据上的重建质量与泛化能力，同时防止过拟合。该方法实现了最先进（SOTA）的性能，具备稳定、数据高效训练的优势，且无需启发式早停策略，在 CIFAR-10、MNIST 以及 MVTecAD 和医学图像等真实世界数据集上均优于先前方法。

ABSTRACT

Autoencoder, as an essential part of many anomaly detection methods, is lacking flexibility on normal data in complex datasets. U-Net is proved to be effective for this purpose but overfits on the training data if trained by just using reconstruction error similar to other AE-based frameworks. Puzzle-solving, as a pretext task of self-supervised learning (SSL) methods, has earlier proved its ability in learning semantically meaningful features. We show that training U-Nets based on this task is an effective remedy that prevents overfitting and facilitates learning beyond pixel-level features. Shortcut solutions, however, are a big challenge in SSL tasks, including jigsaw puzzles. We propose adversarial robust training as an effective automatic shortcut removal. We achieve competitive or superior results compared to the State of the Art (SOTA) anomaly detection methods on various toy and real-world datasets. Unlike many competitors, the proposed framework is stable, fast, data-efficient, and does not require unprincipled early stopping.

研究动机与目标

• 解决自编码器（AE）在复杂数据集（如 CIFAR-10 和真实世界图像）上存在的过拟合与重建质量差的问题。
• 通过将拼图求解作为预训练任务，利用自监督学习（SSL）提升基于 AE 模型的泛化能力与语义特征学习能力。
• 消除自监督学习任务（如拼图）中的捷径学习问题，避免模型利用低层次统计特征而非有意义的语义特征。
• 开发一种稳定、数据高效且可复现的异常检测框架，无需非原则性的超参数调优或早停策略。
• 在实际评估标准（如 FPR = 99.5% 时的高 TPR）下，于玩具数据集（如 CIFAR-10、MNIST）和真实世界数据集（如 MVTecAD、医学影像）上展示优越性能。

提出的方法

• 训练一个 U-Net 作为生成器，用于求解 4x4 图像拼图任务，其中输入为分割成块并打乱的图像，输出为重建的原始图像。
• 在类似 GAN 的训练设置中引入判别器，以提升重建质量，并促使模型学习更具语义意义的特征。
• 应用鲁棒的对抗性训练（如 PGD 和 FGSM），自动消除模型可能利用的捷径解法（如基于边缘或强度的启发式规则）。
• 通过结合色彩化与拼图求解的多任务预训练策略，进一步增强特征抽象能力与模型鲁棒性。
• 采用自监督训练目标，使模型在无需任何人工标注的情况下，学习预测正确的图像块排列。
• 使用重建损失作为异常分数：异常输入的损失更高，表示新颖性；而正常输入则能被准确重建。

实验结果

研究问题

• RQ1 将拼图求解作为自监督预训练任务，是否能提升基于自编码器的异常检测模型的重建质量与泛化能力？
• RQ2 对抗性训练在消除拼图求解 SSL 任务中的捷径解法方面有多有效，尤其是在复杂数据集上？
• RQ3 将拼图求解与额外的 SSL 任务（如色彩化）结合，是否能进一步提升模型在真实世界异常检测基准上的鲁棒性与性能？
• RQ4 所提出的框架是否能在无需启发式早停或大量超参数调优的情况下，实现最先进性能？
• RQ5 在严格的实际评估标准下（如在医学与工业缺陷检测数据集中 FPR = 99.5% 时的高 TPR），模型表现如何？

主要发现

• 在 11 个数据集的基准测试中，Puzzle-AE 的平均 AUC 为 72.47，当仅使用 1/12 的数据训练时，显著优于 LSA（88.89）和 DSVDD（92.61）在 MNIST 上的表现，平均 AUC 达 94.90。
• 在 MNIST 上，Puzzle-AE 的平均 AUC 为 94.90，优于 LSA 的 88.89 和 DSVDD 的 92.61，证明其在数据效率与性能上的优越性。
• 在 FPR = 99.5% 时，Puzzle-AE 的平均 TPR 为 0.2932，优于 LSA 的 0.4261，表明其在高特异性条件下的更强检测能力。
• 该方法对灰度输入表现出鲁棒性，在 CIFAR-10 的“car”类别上性能下降极小，表明对颜色不敏感，泛化能力更强。
• 在存在强捷径的数据库（如 MNIST）上，加入对抗性训练（PGD/FGSM）使性能提升约 1%，显著增强了模型鲁棒性。
• 该框架在 MVTecAD 和医学影像数据集上实现了最先进性能，甚至优于 GT（一种先前的一类 SSL 方法），并完全消除了对早停的需求。