用于异常检测的Transformer - InTra《Inpainting Transformer for Anomaly Detection》

原文地址

https://arxiv.org/pdf/2104.13897v1.pdf

论文阅读方法

三遍论文法

初识

GAN，AE这类基于重构的异常检测方法，缺点在于其对于异常样本的重构也非常好，这会导致检测错误。而目前出现的一些方法将生成重构的问题转换为inpainting问题来进行异常检测，inpainting就是将图像对某些区域进行覆盖，然后进行恢复，也可以视为一种自监督方法。

解决inpainting这类问题，从更大的区域捕获长距离语义信息有助于覆盖区域的重建。但CNN由于感受野的限制，其不善于捕捉长距离信息。因此，作者受到最近大火的视觉Transformer的启发，因此采用Transformer架构解决这个问题。如下图(a)所示，训练时，图像被切成大小相等的块，利用一个大区域内的其他图像块来进行inpainting。图(b)展示了重建的效果，以及根据像素级误差得到的异常得分图。

并且作者只依据MVTec AD数据集本身的少量样本进行训练，也达到了state-of-the-art的效果。

相知

2. Related Work

将当前的异常检测/分割方法主要分为了两类，一是基于重构的方法，类似AE、GAN、VAE等方法；其次是基于嵌入(Embedding)的方法，主要依据在ImageNet上预训练的CNN提取判别性特征进行比较。

随后也介绍了inpainting和transformer的一些相关方法。

3. Inpainting Transformer for Anomaly Detection

使用Transformer执行inpainting任务进行训练。测试时，同样以inpainting的方式进行重建，比较输入图像与重建图像之间的差别，得到检测结果。

3.1 Embedding Patches and Positions

如上图(a)所示，本文的方法是选择一个长度为L的正方形区域（而非ViT中的整副图像）进行inpainting，过程中有两种位置编码方式，一种是局部编码，如下图左所示，另一种是全局编码，如下图右所示。

为什么需要这两种编码模式，直觉上来说，纹理类图像(图左)不需要考虑图像块在全局的位置信息，而另一些类别则很重要(图右)。

和ViT中的设置类似，位置嵌入信息为D维，将图像块也映射到D维后，将两者进行相加即可。需要注意的是，有一个图像块$P(t,u)$被覆盖了。本文将其视为ViT中的分类头(class token)：

最后得到$L\times L$个维度为D的序列，准备送入后续的Transformer。

3.2 Multihead Feature Self-Attention

原始的MSA模块q与k都是通过一个映射维持在D维，但作者任务由于训练图像的图像块之间非常相似，这导致计算出来注意力权重几乎为恒等权重。因此作者对Transformer中的多头注意力模块做了略微修改，在计算q与k时，利用MLP进行一个非线性降维（文中设置为D/2），文中称之为MFSA (multihead feature self-attention)。

MLP隐层维度为2D D - 2D - D/2

加快模型收敛以及提升了精度，但这也增加了参数量

3.3 Network Architecture

最后整体的网络架构如下所示，图左为Transformer的一个模块，每个模块的输入输出均为$L^2\times D$。对最后一层block的输出进行平均(D)，再进行映射作为inpainting的结果($K^2\ast C$)。

也可以使用最后一层的第一个输出直接进行线性映射，这与ViT类似。

4. Training

训练时随机选择一个大小为L的窗口，然后在窗口内选择一个图像块进行覆盖，然后将该窗口内的图像块一起送入Transformer中执行inpainting任务。

损失函数用的就是像素级的L2 loss，同时还使用了SSIM与GMS两种loss。

5. Inference and Anomaly Detection

首先根据重建图像与原始图像之间的差异计算pixel-level的异常得分图，然后选取最大的一个作为image-level的检测得分。

与训练时一致，测试图像被分成NxM块，对于(t,u)位置上的图像块采用如下的公式选择其周围大小为L的窗口，(r,s)为窗口的左上角坐标。

使得图像块尽可能的在窗口中心

最后对所有NxM个图像块进行inpainting，即可得到整副图像重构后的结果。值得注意的是，作者在计算异常得分图的时候，不是使用L2距离，而是采用了GMS-based的方法，在{1/2，1/4}尺度下计算梯度幅度相似度，然后进行均值滤波和高斯滤波。

分别得到两种尺度下的异常图m1,m2，将其恢复到原始图像尺寸，计算两者的像素级均值作为异常得分图。

进一步将异常得分图减去训练集中的异常均值得分图（训练集均为正常样本）后求平方，T为训练集。

最后在得分图中选择最大的一个像素得分作为最后图像级的得分。

4. Experiments

具体结果见原文，这里介绍一点实现细节

实验细节：

从训练图像中随机选取10%作为验证集，用于控制重建结果（最多20张）。这个epoch内，每张图像会随机选取600个窗口用于训练，训练时使用随机旋转与翻转作为增广手段。

图像块大小K=16，窗口大小为L=7，对于MVTec中不同类别的数据也选择了不同的图像尺寸，{256x256，320x320，512x512}。Transformer中的维度D设置512。所有的resize操作均采用双线性插值。

优化器采用adam，但是transformer的训练时间比较长，有些时候需要超过500个epoch。当验证集loss超过50个epoch都没有明显下降，则停止训练，并选择最优的模型进行评估。

回首

整篇文章其实最大的创新点就是将Transformer引入到异常检测领域，执行inpainting任务。针对不同情况设计了local和global两种位置嵌入方式。虽然也引入了U-Net架构（但消融实验证明其并不是对所有类别都有效），以及修改MSA为MFSA，但都是小修小改，没有大的提升。

总得来说，inpainting还是基于重构的方法。而对于某些噪声区域本身就是难以重构的，比如字符区域。此外，本方法仍是patch-based，这导致一张测试样本需要执行多次模型推理才能得到最后的异常得分图。并且能观察到明显的边界效应，如下图所示：

代码

目前该论文还未开源代码