UniAD：实现多类别异常检测的统一模型

news2024/10/7 10:25:16

来源：投稿作者：Mr.Eraser
编辑：学姐

论文标题：用于多类异常检测的统一模型

论文链接：https://arxiv.org/abs/2206.03687

论文贡献：

提出UniAD，它以一个统一框架完成了多个类别的异常检测。
提出分层的查询解码器来帮助建立多类分布的模型。
采用邻居屏蔽的注意力模块，以进一步避免从输入特征到重构输出特征的信息泄漏。
提出一种特征抖动策略，即使在有噪声输入的情况下也能敦促模型恢复正确的信息

1.引言

异常检测在制造业缺陷检测、医学图像分析和视频监控中得到了越来越广泛的应用。考虑到高度多样化的异常类型，一个常见的解决方案是对正常样本的分布进行建模，然后通过寻找离群值来识别异常样本。因此，学习正常数据的紧凑边界至关重要，如图a所示。

当涉及到多类情况时，我们希望模型能同时捕捉到所有类别的分布，这样它们就能像图b那样共享同一个边界。但是如果我们专注于某一特定类别，比如说图b中的绿色类别，那么其他类别的所有样本都应该被视为异常，无论它们本身是正常的（即蓝圆圈）还是异常的（即蓝三角形）。从这个角度来看，如何准确地对多类分布进行建模变得至关重要。

综上所述，本文摆脱了传统的为每一类训练一个模型的范式（图c ），并设法用一个统一的框架（图d）完成对各种类的异常检测。

2.相关工作分析

MLP, CNN, Transformer和提出的UniAD在MVTec-AD数据集上的比较：图a表示训练损失（蓝色）, 异常检测（绿色）和定位（红色）的测试AUROC。在MLP, CNN和Transformer的训练过程中，正常样本的重建误差不断变小，但在异常情况下的性能在达到峰值后却严重下降。这是由于模型学习了一个 “相同的捷径”，它倾向于直接复制输入作为输出，不管它是正常的还是异常的。图b“捷径”的视觉解释，异常样本可以被很好地恢复，因此变得很难从正常样本中检测出来。

相比之下，UniAD克服了上述问题，并设法将异常样本重建为正常样本。值得注意的是，所有的模型都是为特征重建而学习的，并且采用了一个单独的解码器来渲染来自特征的图像。这个解码器只用于可视化。

3.UniAD 方法概览

UniAD 的框架如下，由一个Neighbor Masked Encoder (NME)和一个Layer-wise Query Decoder (LQD)组成。LQD的每一层都采用了可学习的查询嵌入来帮助建立复杂的训练数据分布模型。Transformer中的全部注意力被邻居屏蔽的注意力所取代，以避免信息从输入到输出的泄漏。特征抖动策略鼓励模型在嘈杂的输入中恢复正确的信息。所有这三项改进都有助于模型避免学习 “相同的捷径“。

3.1Neighbor masked attention

我们怀疑原始Transformer中的full- attention对 “相同的捷径 ”有所贡献。在full- attention中，一个token被允许看到自己，所以它将很容易通过简单的复制来进行重建。

此外，考虑到特征tokens是由CNN骨干提取的，相邻的tokens一定有很多相似之处。

因此，提出在计算注意力图时对相邻token进行屏蔽，即Neighbor Masked Attention (NMA)。请注意，邻居区域是在二维空间定义的，如图4所示。

3.2 Neighbor masked encoder & Layer-wise query decoder

Neighbor masked encoder：

编码器遵循原始Transformer中的标准架构。每层由一个注意力模块和一个前馈网络（FFN）组成。然而，full-attention被提出的NMA取代，以防止信息泄露。

Layer-wise query decoder：

查询嵌入可以帮助防止重建异常，但原始Transformer中只有一个查询嵌入。因此，我们设计了一个分层查询解码器（LQD）来加强查询嵌入的使用，如图3所示。具体来说，在LQD的每一层，可学习的查询嵌入首先与编码器嵌入融合，然后与前一层的输出融合（第一层的自我融合）。特征融合是由NMA实现的。在原始Transformer之后，一个2层的FFN被应用于处理这些融合的tokens，并且利用残余连接来促进训练。LQD的最终输出作为重构的特征。

3.2 Layer-wise query decoder 选择策略

层级查询解码器（LQD）的各种设计选择，其中每层使用两个注意力模块。为了简单起见，省略了残差连接、层规范化和前馈网络。

对分层查询解码器（LQD）的设计选择进行消融研究。每个设计的概念可以在上图中找到。分别报告了异常检测/定位方面的表现。

3.3 Feature jittering

受去噪自动编码器（DAE）的启发，我们向特征tokens添加扰动，引导模型通过去噪任务学习正常样本的知识。具体来说，对于一个特征token， $f_{tok}\in\mathbb{R}^C$ ，我们从高斯分布中对干扰D进行采样：

其中，α是控制噪声程度的抖动尺度。同时，采样干扰以固定的抖动概率p添加到 $f_{tok}$ 中。

3.4 实施细节

特征提取：

我们采用在ImageNet上预训练的固定的EfficientNet-b4作为特征提取器。从阶段1到阶段4的特征被选中。这里的阶段指的是具有相同大小的特征图块的组合。然后，这些特征被调整为相同的大小，并沿通道维度串联起来，形成一个特征图，之后，这个特征图被标记为具有 $C_{org}$ 通道的H×W特征tokens。