GLAD:利用全局和局部自适应扩散模型实现更好的无监督异常检测重建
论文标题:GLAD: Towards Better Reconstruction with Global and Local Adaptive Diffusion Models for Unsupervised Anomaly Detection
论文地址:https://arxiv.org/abs/2406.07487
研究背景
工业异常检测在计算机视觉的研究和工业应用中起到了非常重要的作用,视觉异常检测的主要目标是识别图像中的异常区域的位置。现有的异常检测方法中基于重建的方法取得了很好的效果,这种方法通过在正常样本上训练学习正常样本的特征,将异常样本重建为正常样本后与原始异常样本进行比较,从而确定异常的位置。
最近,扩散模型以其强大的生成能力在无监督异常检测任务上取得了非常好的效果。由于只使用正常图像进行模型训练,扩散模型倾向于重建添加一定噪声的测试图像的正常区域。但是这种方法对所有的潜在异常都一样,会导致两个主要的问题:1.从全局视角,重建有着不同异常的图像的困难是不同的,修复一块确实的元素比一个划痕困难,因此需要更多的降噪步骤。2.从局部视角,即使是同一张图像重建异常和正常区域的难度是不同的。理论上,扩散模型每一步预测一个满足高斯分布的噪声,然而,由于异常与潜在正常对应点之间的差异,异常区域的预测噪声不可避免地会偏离标准高斯分布。
为了解决以上问题,本文提出了一种用于无监督异常检测的全局和局部自适应扩散模型GLAD,引入了极高的灵活性并在实现无异常重建的同时尽可能多的保存正常信息。本文提出不对所有样本使用相同的设置,而是通过评估图像内容与从扩散模型提取的先验之间的差异来预测每个样本的特定去噪步骤。为此,本文提出在训练中引入合成异常样本,以促使扩散模型突破标准高斯分布的限制,并在推理过程中采用空间自适应特征融合方案。
通过在三个常用的异常检测数据集(MVTec-AD、MPDD 和 VisA)和一个印刷电路板数据集(PCB-Bank)进行了大量的实验,证明了本文所提方法的有效性。
相关工作
异常检测
主流的无监督异常检测方法可以被分为两类:基于重建的方法、基于嵌入的方法。
-
基于重建的方法:
基于重建的方法基于一种模型在正常样本上训练只能重建正常区域但不能重建异常区域的假设,通过对比重建前后的图像进行异常检测。早期的方法尝试使用VAE重建样本。OCR-GAN将图像解耦为不同发的频率组成部分,并将重建过程建模为并行全频图像恢复的组合。一些其他的工作探索了更广泛的自监督学习异常重建方法。UniAD使用transformers重建带有 masked self-attention 的样本特征,并提出了一种特征抖动策略来解决 shortcut 问题。DRAEM 通过生成异常图像训练 UNet 架构将异常图像映射到正常图像,并使用合成的异常图像和重建图像训练鉴别器以区分异常区域。DiffAD将测试图像的潜在表征扩散为噪声条件嵌入,这有助于在保留测试样本正常信息的同时生成高质量的重建图像。DDAD使用基于分数的函数在去噪过程中重新整合测试样本的信息。然而,这些方法在去噪过程中增加了固定的噪声步骤,不适用于各种异常类型。
然而,这种假设可能并不总是成立,有时训练的模型的泛化性能比较好,能很好地重建异常区域,从而导致导致错误地检测。
-
基于嵌入的方法:
基于嵌入的方法通过提取图像特征来评估异常区域,大多数方法采用在ImageNet数据集上预训练的网络进行特征提取。基于知识蒸馏的方法 首先用正常样本训练学生网络。然后将预训练的教师网络的特征与学生网络的特征进行比较,以检测和定位异常。Reverse distillation利用教师和学生的不同架构来保持异常的区别。PaDiM为正常样本的块特征建立多元高斯分布,并使用马哈拉诺比斯距离作为异常分数。PatchCore使用一个存储库来保存正常图像的特征,并将其与测试图像的特征图进行比较,以区分正常特征和异常特征之间的差异。
然而,工业图像通常与ImageNet有不同的分布。直接使用这些有偏差的特征可能会导致匹配问题。此外,统计算法总是存在高计算复杂性或高内存消耗的问题。
扩散模型
受非平衡热力学原理的启发,扩散模型 (DM)被提出,并用于图像生成等多种下游任务。去噪扩散隐式模型 (DDIM) 将 D扩散模型的逆过程视为非马尔可夫过程,极大提高了推理速度。潜在扩散模型 (LDM) 使用预训练的 的VAE 在潜在空间进行训练和推理,进一步降低了模型需要资源和时间。此外,Text inversion和 Dreambooth 学习给定参考集中主体的外观,并在不同情境中合成它们的新再现。这些方法遵循预测高斯噪声的训练方法,并从高斯噪声开始去噪,因此这些方法不能实现自适应去噪从而将异常图像重建为正常图像。
前置知识
Diffusion Process
在传统的扩散模型中,一个满足高斯分布的随机噪声 ϵ 被添加到样本 𝑥 中,在 t 步后得到的结果可以用以下公式表示:
Intermediate Result Visualization
根据公式 (1) 可以得到第 t 步的中间结果的无噪声版本,即:
Generation Process
生成阶段的每个步骤都可以用以下公式来表达:
注意,扩散模型预测的变量以∧标记,例如, 是通过直接添加随机噪声获得的(扩散过程),而 是通过从更大的步骤去噪获得的(生成过程)。
重建误差的公式分析
现有的基于DM的UAD方法的常见方式是对测试样本添加特定的噪声再执行去噪步骤。定义有异常的测试样本为,将其潜在正常样本表示为,则该过程可描述为
,理想上,由于异常的检测与定位是基于和的不同,需要,其中_τ_是手动设置的阈值,用于区分正常样本和异常样本。由于DM是预训练的并在生成过程中是参数固定的。定义和的不同为,即,根据公式 (1) 可以得到:
定义从步骤t的生成过程为,通过结合公式(1)和公式(4),误差可以表示为:
其中
表示近似成比例,这是基于
足够平滑的合理假设。接下来,本文将努力减少公式 (5) 中的误差,以获得更好的重建质量。
本文方法
-
自适应降噪步骤(ADS)
现有的基于扩散模型的方法中,噪声总是满足标准高斯分布,因此忽略和 之间的区别,公式(5)中仅剩下考虑到,对于更大的,希望更小,由于与负相关,更小的意味着更大的,以上分析为每个样本设置适当的去噪步骤的动机提供了公式化的解释。
基于以上的步骤,一个直观的方式是通过评估的数值大小决定正确的步骤。然而,是一个不可以获得的概念。在生成过程中,体现在中,可以用于比较。
如图2左部分所示,从输入样本开始,通过添加大量的噪声获得,接下来的每一步比较和 的特征,由于从生成,它遵循现有的基于DM的方法并倾向于为正常样本。对比之下, 直接向中添加噪声获得,异常在一定程度得到保留。考虑到和 都包含噪声,并且噪声可能会因去噪过程而相互偏离,本文将它们通过公式(2)转换为无噪声版本,即和他们之间的不同通过异常得分计算。如图2所示,如果不同小于阈值,继续去噪以进行进一步的比较;如果差异在第 t 步出现(即 > δ),可以以为起点进行 t + n 步去噪,其中 n 是一个较小的数字以保留一些冗余。这样,正常区域的细节得到最大程度的保留,并且可以重建异常区域。
2. 空间自适应特征融合(SAFF)
对整个图像设置一样的降噪步骤(t+n)不合适,基于以上分析,本文只需要为异常区域设置更大的降噪步骤。本文认为,对于正常区域可以减少降噪步骤为t个,同时为异常区域保存较大的降噪步骤。如图2所示,本文已经从T执行了降噪步骤,这些步骤在这个过程中可以重复使用。具体来说,可以得出一个掩码 m,它表示像素成为异常部分的可能性,然后可以将这两个特征与掩码结合起来,即:
注意,公式(6)中仍然存在噪声偏差的问题,本文按照ADP中的策略添加无噪声的版本并按照公式(1)添加相同的噪声 ϵ ,以获得一致的噪声,即:
-
面向异常的训练范式(ATP)
根据ADS和SAFF,本文根据异常检测任务的特性调整了重建过程。然而,公式(5)中还存在一个没有解决的不兼容问题。具体来说,更好的重建质量意味着公式(5)的值更低,即:
在公式(8)中,
是手动设置的数值,模型唯一需要估计的值是
,通过重写公式(8)可以得到:
在DM的设置下, 是一个满足标准高斯分布的随机噪声。一旦是非零的(即存在异常),可以得到结论: 偏离了标准高斯分布。
对于一个只使用正常样本训练的DM,所有步骤的噪声满足标准高斯分布。也就是说,预测超越了ϵ的范围。因此本文提出在训练中引入异常,这允许扩散模型突破标准高斯分布的限制并符合公式(9)。在公式(9)的基础上,学习目标可以表示为:
可以得到公式(10)甚至比有监督的异常检测对数据集有更高的要求,因为需要相应的正常样本,这是困难甚至不可行的。为了解决这个问题,本文遵循MemSeg使用正常样本合成异常样本,使得模型训练以无监督的方式执行。
公式(10)明显具备更好的通用性。对于正常区域,公式可以退化为原始的扩散损失;对于异常区域,该公式迫使模型预测非高斯分布的噪声从而得到相应的正常区域。
-
异常评分和地图构建
遵循AprilGAN,本文将重建的图像作为参考图像,与测试图像作对比从而获得异常图。预训练模型DINO用于提取测试图像的多层特征和重建图像的多层特征。层的异常图基于第层的特征的余弦相似度计算,即:
其中 (i, j) 为坐标,〈x, y〉 计算 x 与 y 的所有元素的余弦相似度,min(·) 返回最小值。最后将各层的异常图相加,得到异常图,即:
整个图像的异常得分是M的前K个最大值的平均值。
实验
接下来,首先介绍实验设置及其实施细节。然后,将本文提出的方法与最先进的 (SOTA) 无监督异常检测方法进行比较,以证明本文提出方法的优越性。此外,为了评估本文所提出方法的有效性,本文对所提出的方法的各个组件进行了消融实验。
数据集
MVTec-AD
该数据集模拟了真实的工业生产场景,填补了无监督异常检测的空白。它包含 5 种纹理和 10 种物体,共计 5,354 张来自不同领域的高分辨率图像。训练集包含 3,629 张仅有无异常样本的图像。测试集包含 1,725 张图像,包括正常和异常样本。提供了像素级标注用于异常定位评估。
MPDD
包含 6 类金属部件,训练集有888个正常样本,测试集有458个正常或异常样本。由于不同光强度下多个物体相对于相机的空间方向、位置和距离各不相同,并且背景不均匀,该数据集更具挑战性。
VisA
该数据集包含共计 10,821 张高分辨率图像,包括 9,621 张正常图像和 1,200 张带有 78 种异常的图像。VisA 数据集由 12 个子集组成,每个子集对应一个独特的物体。12 个物体可分为三种不同的类型:复杂结构、多实例和单实例。
PCB-Bank
是整合的印刷电路板数据集,包含7个不同的类别。训练集有4214个正常样本,测试集中有2253个正常或异常样本。数据集的样本有不同的清晰度、分辨率和拍摄角度。异常类型主要包括划痕、结构缺陷、弯曲等。
实现细节
本文使用预训练的潜在扩散模型 (LDM) 并对 UNet 进行微调以适应数据进行重建。采用ViT-B/8架构的 DINO 作为特征提取模型。
由于本文使用的数据集与预训练数据集之间的差异较大,本文仅针对 VisA 和 PCB-Bank 数据集对 VAE 和 DINO 进行了微调。为了与预训练的 VAE 保持一致,图像的大小被裁剪为 512×512 的分辨率。在表 5 中展示了 256×256 分辨率的结果。第 3、6、9 和 12 层的特征用于构建异常图。UNet 经过 4 000 次迭代训练,batch-size为 32,学习率设置为 5×10−6。
对于 MVTec-AD 数据集,T 通常设置为 750,对于 MPDD 数据集,T 设置为 500,对于 VisA 和 PCB-Bank 数据集,T 设置为 450。为避免遗漏异常,对于 MVTec-AD 数据集和 MPDD 数据集,将最小步长设置为 350,对于 VisA 和 PCB-Bank 数据集,T 设置为 200。使用 σ = 6 的高斯滤波器平滑异常定位得分,并将 M 的前 250 个最大值的平均值视为异常得分。在第 3.3 节中,本文使用第 12 层的特征和前 10 个最大值来计算差异。
对比实验
在MVTec-AD数据集上的对比实验,对于图像级异常检测任务,本文的方法在 15 个类别中的 9 个类别中实现了最高的 I-AUROC。对于像素级异常定位任务,本文的方法在 P-AUROC 方面优于所有类型方法中的 SOTA。
在MPDD数据集上的实验,在所有类型的方法中,本文的方法在 I-AUROC 和 P-AUROC 方面均优于 SOTA
在VisA数据集上的实验,本文的方法在异常检测和定位任务上也取得了 SOTA 性能。
在PCB-Bank数据集上的实验,本文的方法在异常检测和定位任务上也取得了 SOTA 性能。
表5展示了 SOTA 方法与 GLAD 之间的全面比较。I-AUROC、I-AP 和 I-F1-max 显示在每种方法的第一行原始数据中,用于评估异常检测性能,而 P-AUROC、P-AP、P-F1-max 和 PRO 显示在第二行原始数据中,用于评估异常定位性能。总体而言,在四个数据集的平均值上,本文的方法在所有指标上都优于现有方法。特别是,由于可以将大规模异常区域重建为正常区域,本文的方法在异常定位结果方面具有很大的优势。
图 3 展示了 MVTec-AD 和 MPDD 数据集上的重建和定位结果。对于重建,其他方法通常无法将大规模异常重建为正常区域。相反,本文的方法可以得到很好的重建结果。在 (f)-(h) 列定位结果中,其他方法忽略了一些异常区域,由于异常区域重建失败,导致定位不准确。如 (i) 列所示,由于本文的方法可以保证无异常重建,可以准确定位异常。这也解释了本文的方法与表 5 中其他基于重建的方法相比,定位的量化结果更优。
消融实验
本文的方法主要包括三个组件,即自适应去噪步骤 (ADS)、空间自适应特征融合 (SAFF) 和面向异常的训练范式 (ATP)。接下来将进行消融研究以验证本文提出的组件的有效性。
1.自适应去噪步骤 (ADS)
基于本文提出的 ADS,DM 可以选择适当的噪声步骤并从相应的噪声样本中进行去噪,从而确保无异常重建并尽可能多地保留有关正常区域的信息。在表 6 中,ADS 带来的改进显而易见,实验结果证明了所提方法的优越性。
2.空间自适应特征融合 (SAFF)
在SAFF中,本文将预测样本和测试样本的特征与掩码 m 融合,该掩码 m 是在适当的去噪步骤下使用异常图生成的。表 6 显示了 SAFF 的效果,SAFF 可以进一步去除残留异常信息并提高检测和定位性能。
图4中展示了不同类型的异常和ADS选择的适当步骤。对于不明显的异常,如示例(a)、(c)和(e),ADS倾向于选择足以进行重建的小步骤。对于大规模异常,如示例(b)、(d)和(f),ADS从步长噪声较大的样本中去噪,以确保重建无异常图像。此外,ATP可以进一步减少去噪步骤。图4展示了ADS以及ADS和ATP组合的重建。与不使用ATP相比,该模型可以在推理过程中更好地去除异常。因此,用于比较的清晰图像在每个去噪步骤中包含的异常信息较少,并且在较小的步骤中检测到差异。这允许保留更多的正常信息以产生更准确的重建,同时确保完全去除异常。图 4 中用红色圆圈标记了一些细节上的差异。使用 ATP 重建的 ADS 与正常区域的输入更相似。性能可以进一步提高,如表 6 所示。
3.面向异常的训练范式(ATP)
由于本文的ATP在训练中添加了合成的异常样本,训练数据和测试数据之间的数据分布方差缩小了。因此,DM 也可以将异常区域映射为正常区域。表 6 显示了基线和所提出的训练范式的性能比较,即 LDM(基线)和 ATP(提出的面向异常的训练范式)。本文还在图 5 中提供了直观的定性比较。那里没有使用自适应去噪步骤。为基线和 ATP 添加相同的噪声步骤。在(b)列和(e)列中,基线的结果仍然包含异常区域。这表明基线的重建能力有限。在使用 ATP 训练之后,如图 5 中的(c)列和(f)列所示,该模型可以很好地将异常区域映射到正常区域。
影响、限制、未来工作
本文研究了无监督异常检测问题并在四个数据集上实现了SOTA的性能。由于不需要使用真实场景中的异常样本,本文的研究有可能在真实的场景中得到有效利用。尽管本文提出的方法的性能较好,本文的方法在推理中引入了评估比较,这会导致额外的时间成本。另外,为了确定ADS的去噪步骤,本文进行了逐步的比较,但是降噪的步骤并未减少。未来,作者计划设计一个轻量级的评估比较方法,并尝试使用有限的降噪步骤来预测降噪步骤,这有可能提高方法的效率。
结论
本文提出了一种全局和局部的自适应扩散模型GLAD,用于提高无监督异常检测的重建质量。作为一个全局的配置,为每个样本设置自适应降噪步骤以适应不同的异常。另外,考虑到重建异常区域与正常区域不同,一个自适应特征融合方案被提出用于移除异常区域的残留异常,并提出了一个面向异常的训练范式提高扩散模型重建异常为正常区域的能力。本文在 MVTec-AD、MPDD、VisA 和 PCB-Bank 数据集上进行了大量的实验。定量和定性结果证明本文方法的优越性和有效性。