本文介绍 2023 年一篇异常检测中基于扩散模型实现异常检测的工作 —— DiffusionAD。
项目 | 内容 | 备注 |
---|---|---|
方法名称 | DiffusionAD | |
论文题目 | DiffusionAD: Norm-guided One-step Denoising Diffusion for Anomaly Detection | |
论文连接 | | |
开源代码 | | |
发表时间 | 2023.08.29 | |
方法类别 | 深度学习 -> 基于重构 -> 扩散模型 | |
Detection AU-ROC | 99.7% | |
Segmentation AU-ROC | 98.7% | |
Segmentation AU-PRO | 95.7% | |
FPS | 23.5 | |
核心思想 | 用扩散模型作为数据重构子模块,结合一个分割网络用于异常检测。 | |
提出的基于生成模型的框架由两个组件组成:重建子网络和分割子网络
通过引导扩散模型实现重建子网络,该模型将重建过程重新表述为噪声到正常图像。
输入图像 x_0 ,向其中逐步加入 t 次高斯噪声得到 x_t ,该加噪的图像会逐渐失去原始图像的信息,变为一个随机的高斯噪声。
网络先向 x_0 输入两种程度的噪声得到 x_{t_s} (弱)和 x_{t_b} (强)。
文章没有像传统扩散模型逐步加噪恢复图像,而是通过一步加噪直接恢复原始图像的样貌,这样可以节约大量运算量。
将 x_{t_s} 和 x_{t_b} 输入扩散模型,预测可以将两组噪声图像恢复正常的添加噪声 \epsilon_{\theta}\left(x_{t_s}\right) 和 \epsilon_{\theta}\left(x_{t_b}\right)
文章发现如果直接将 x_{t_s} 和 x_{t_b} 恢复为正常图像,由于 \hat{x_{0_s}} 带有更多原始图像的特征,由他恢复出的图像细节和原始图像更接近,但是扩散模型一步难以恢复较严重的异常,而 \hat{x_{0_b}} 基本上是正常图像的样貌。
也就是说 x_{t_s} 细节较好, x_{t_b} 的大局观更好,为了结合二者的优点,文章使用了带引导的扩散模型恢复 \hat{x_{0_s}} 。
具体来说就是将 \hat{x_{0_b}} 作为预测 \epsilon_{\theta}\left(x_{t_s}\right) 的输入加进去,这样预测出的 \hat{x_{0_b}} 质量高且具有大局观(\hat{x_{0_g}} )。
分割子网络为 U-Net 结构,包含编码器、解码器和层级连接,网络接受 x_0 和 \hat{x_{0_g}} 作为输入,直接预测异常分值。
网络学习到的不一致性减少了正常区域与其重建之间轻微的像素差异引起的误报,并突出显示了异常的区域。
扩散模型学习正常样本 (y=0 ) 的重建过程:
$$ \mathcal{L}_{noise}=\frac{(1-y)(|\epsilon_{t_s}-\epsilon_\theta(x_{t_s})|^2+|\epsilon_{t_b}-\epsilon_\theta(x_{t_b})|^2)}{2}. $$
分割网络直接预测异常得分 Mask:
$$ \mathcal{L}_{mask}=\text{Smooth}_{\mathcal{L}1}\left(M,\hat{M}\right)+\gamma\mathcal{L}_{focal}\left(M,\hat{M}\right) $$
其中 M 是 GT 异常 mask, \hat{M} 是分割网络的输出,平滑 L1 损失和 Focal Loss 同时应用,以减少对异常值的过度敏感并准确分割困难异常样本。
将扩散模型与分割网络一起训练:
$$ \mathcal{L}_{total}=\mathcal{L}_{noise}+\mathcal{L}_{mask}. $$
上述过程中需要大量带异常的数据才可以训练起来,论文中使用了人工制造异常的方式实现。
通过随机创造一些 mask 将其他数据集的数据以一定透明度叠加到正常图像上,得到异常数据: