[论文解读] HaloAE: An HaloNet based Local Transformer Auto-Encoder for Anomaly Detection and Localization

Émilie Mathian, Hengming Liu+4 more|arXiv (Cornell University)|Aug 6, 2022

Anomaly Detection Techniques and Applications被引用 6

一句话总结

HaloAE 提出了一种新颖的局部 Transformer 自编码器，采用 HaloNet 的分块自注意力机制，实现无监督异常检测与定位的联合优化。通过结合局部自注意力、卷积特征提取与自监督学习，其在 MVTec 基准测试中达到最先进性能，图像级平均 ROC-AUC 达 91.4%，像素级平均 ROC-AUC 达 91.2%。

ABSTRACT

Unsupervised anomaly detection and localization is a crucial task as it is impossible to collect and label all possible anomalies. Many studies have emphasized the importance of integrating local and global information to achieve accurate segmentation of anomalies. To this end, there has been a growing interest in Transformer, which allows modeling long-range content interactions. However, global interactions through self attention are generally too expensive for most image scales. In this study, we introduce HaloAE, the first auto-encoder based on a local 2D version of Transformer with HaloNet. With HaloAE, we have created a hybrid model that combines convolution and local 2D block-wise self-attention layers and jointly performs anomaly detection and segmentation through a single model. We achieved competitive results on the MVTec dataset, suggesting that vision models incorporating Transformer could benefit from a local computation of the self-attention operation, and pave the way for other applications.

研究动机与目标

• 解决工业与医学影像中因异常标注稀缺而带来的无监督异常检测与定位挑战。
• 克服标准自编码器在重建异常区域时泛化能力过强、导致检测性能下降的局限性。
• 通过结合卷积特征提取与分块自注意力机制，整合局部与全局上下文信息，提升重建保真度与异常定位能力。
• 构建统一模型，实现无需基于图像块推理的异常检测与像素级分割。
• 通过基于 Cut&Paste 的人工缺陷增强策略，实施自监督学习，提升模型泛化性与鲁棒性。

提出的方法

• 提出 HaloAE，一种混合架构，结合预训练的 CNN 特征提取器与基于 HaloNet 的局部 2D 自注意力模块，以建模块内与模块间的依赖关系。
• 采用受先前工作启发的多尺度特征图生成策略，以捕捉分层且富含语义的表征。
• 通过 Cut&Paste 增强方法引入自监督代理任务，模拟缺陷并正则化自编码器。
• 设计具有自适应加权的多任务损失函数，结合分类损失（在增强图像上）、特征图重建损失与图像重建损失。
• 应用局部自注意力机制，计算图像块内及块间的注意力，相比全局 ViT 风格注意力，显著降低计算成本。
• 采用端到端训练策略，结合 L2 与 SSIM 损失进行重建，最终输出同时用于异常检测与基于像素重建误差的分割。

实验结果

研究问题

• RQ1 基于 HaloNet 的局部自注意力机制是否相比标准卷积自编码器，能提升异常检测与定位性能？
• RQ2 通过人工缺陷增强实现的自监督学习是否能增强自编码器在异常检测中的泛化性与鲁棒性？
• RQ3 单一统一模型是否能有效实现无需基于图像块推理的图像级异常检测与像素级定位？
• RQ4 在多个目标（分类、特征图、图像重建）之间采用自适应损失加权，对模型性能有何影响？
• RQ5 通过局部注意力建模块内相关性，在多大程度上提升了重建质量与异常定位精度？

主要发现

• HaloAE 在 MVTec 数据集上实现 91.4% 的图像级平均 ROC-AUC 与 91.2% 的像素级平均 ROC-AUC，表现具有竞争力。
• 该模型优于标准卷积自编码器（如 AE-L2：70.0% 图像级 AUC），甚至优于基线 HaloNet 自编码器（75.6% 图像级 AUC），证明了架构融合的优势。
• 在相同架构中，将 HaloNet 模块替换为卷积自编码器，图像级 AUC 提升 6.4 个百分点（从 83.1 提升至 89.5），证实局部自注意力的优势。
• 自适应损失加权策略（如不确定性加权）在文本类别上实现 97.2% 的图像级 AUC，表明多任务学习的有效性。
• 通过 Cut&Paste 实现的自监督学习提升了模型泛化能力，体现在检测与分割任务得分均更高。
• 消融实验证实，包含所有组件（特征提取器、局部注意力、SSL 与多损失）的完整 HaloAE 模型在所有指标上均取得最佳性能。