Learning to Detect Multi-class Anomalies with Just One Normal Image Prompt

AI解读

好的，我将尽力对这篇论文进行详细的分析，侧重于方法和技术细节，并用通俗易懂的语言进行解释。

1. 核心方法与创新点

概述：

这篇论文提出了一种名为 OneNIP (One Normal Image Prompt) 的新型无监督异常检测方法。它主要基于重建网络，利用一个正常的图像作为提示 (prompt) 来指导异常检测。简单来说，就是让模型不仅学习如何重建正常图像，还学习如何利用一个“参照物”（正常的提示图像）来更好地识别和修复异常部分。此外，论文还提出了一个监督微调器 (refiner) 来进一步提高异常分割的精度。

主要创新点：

* One Normal Image Prompt (OneNIP)：
* 创新性： 这是该论文最核心的创新。不同于以往的重建模型只关注自身的上下文信息，OneNIP引入了一个正常的图像作为全局提示。
* 技术亮点： 模型学习如何利用这个正常的提示图像来重建或修复待检测图像。通过交叉注意力机制，模型可以有效地将目标图像的特征与提示图像的特征进行交互，从而更好地捕捉全局结构信息，尤其是对于那些与周围环境“无缝”融合的伪装性异常 (camouflaged anomalies)。
* 区别： 传统的重建方法容易过度拟合，导致正常和异常区域都能被完美重建，OneNIP 通过引入提示图像，增加了重建的难度，使模型更敏感于异常。

* Unsupervised Restoration Stream (无监督修复流)：
* 创新性： 为了进一步增强正常图像提示的指导作用，论文构造了一些伪异常样本，然后让模型学习如何将这些伪异常样本“修复”回正常状态。
* 技术亮点： 这个修复过程鼓励模型更多地依赖于正常提示图像，而不是仅仅依赖自身的上下文信息。
* 区别： 其他方法可能只关注正常样本的重建，而忽略了如何利用模型来“拒绝”或“纠正”异常样本。

* Supervised Refiner (监督微调器)：
* 创新性： 为了提高像素级别的异常分割精度，论文提出了一个监督微调器，利用真实的正常样本和合成的异常样本来回归重建误差。
* 技术亮点： 该微调器可以将低分辨率的重建误差映射到高分辨率的异常分割图，从而实现更精确的异常定位。
* 区别： 许多重建方法在低分辨率的潜在空间中进行重建，导致分割精度不高。这个微调器有效地解决了这个问题。

* Unified Anomaly Detection (统一异常检测)：
* 创新性： 这项工作旨在建立一个单一的模型来检测各种类别的所有异常，而不是为每个类别单独训练模型。这种统一的方法更具挑战性，但更实用。
* 技术亮点： 该模型在三个工业异常检测基准测试中表现出色，证明了其对各种类别的异常的泛化能力。
* 区别： 大多数现有的方法都专注于为每个类别构建单独的模型，这导致存储和计算成本都很高。

2. 算法细节与流程

算法流程：

1. 特征提取：
* 对于输入图像（正常的或待检测的）以及提示图像，使用预训练的 EfficientNet-b4 模型提取多层级的特征图。

2. Unsupervised Reconstruction/Restoration (无监督重建/修复)：
* 重建流： 正常图像的特征经过一个自注意力编码器 (self-attention encoder) 来捕捉上下文信息。编码后的特征和提示图像的特征进入一个双向交叉注意力解码器 (bidirectional cross-attention decoder)，进行特征融合和重建。目标是重建原始的正常图像特征。
* 修复流： 伪异常图像的特征也经过类似的自注意力编码器和双向交叉注意力解码器。目标是将伪异常特征修复回对应的正常特征。
* 双向交叉注意力解码器： 提示图像的特征和目标图像的特征相互影响，相互更新。
* 损失函数： 使用均方误差 (MSE) 来衡量重建/修复的质量。

3. Supervised Refiner (监督微调)：
* 计算原始特征和重建/修复后的特征之间的绝对误差。
* 使用一个轻量级的卷积神经网络（包含转置卷积层）将低分辨率的重建误差上采样到高分辨率。
* 使用 Dice 损失函数来训练这个微调器，使其能够准确地预测像素级别的异常分割图。

4. Inference (推理)：
* 对于待检测的图像，从一个预先构建的提示图像池中选择最相似的提示图像。
* 使用重建流来重建待检测图像的特征。
* 计算重建误差，并将其上采样到原始图像的分辨率。
* 使用训练好的微调器来预测异常分割图。
* 将重建误差图和预测的分割图进行加权融合，得到最终的异常分割图。

技术优势和创新之处：

* 全局信息融合： 通过引入正常提示图像，模型可以有效地利用全局结构信息，从而更好地检测伪装性异常。
* 动态特征交互： 双向交叉注意力解码器允许提示图像的特征和目标图像的特征相互影响，相互更新，从而实现更灵活的特征融合。
* 高精度分割： 监督微调器可以将低分辨率的重建误差映射到高分辨率的分割图，从而实现更精确的异常定位。
* 端到端学习： 整个框架可以进行端到端的训练，无需手动设计特征或规则。

3. 详细解读论文第三部分 (Methods)

论文的第三部分详细描述了OneNIP方法的三个主要组成部分：无监督重建、无监督修复和监督微调器。

3.1 Reconstruction with Normal Image Prompt

* Revisiting UniAD: 首先回顾了 UniAD，强调其作为统一异常检测的强大基线。UniAD 依赖于带有邻域掩码注意力 (NMA) 和逐层查询解码器 (LQD) 的编码器-解码器 Transformer。
* Unidirectional decoder with static prompt: 作者首先尝试了一种更简单的解码器结构，通过用正常图像提示的编码器输出 $p_e$ 直接替换 LQD 中的查询嵌入 $q_i$。这使得提示和目标特征能够交互。该过程通过以下公式描述：
$$
\begin{aligned}
q' &= \text{softmax}\left(\frac{p_e x_e^T}{\sqrt{c}}\right) x_e, \\
x_{i+1}^d &= \text{softmax}\left(\frac{q' x_i^{dT}}{\sqrt{c}}\right) x_i^d,
\end{aligned}
$$
其中 $c$ 是 $x_e$ 的维度，而 $x_0^d$ 在 LQD 的第一个块中被初始化为 $x_e$。这个公式用multi-head self-attention实现。

* Bidirectional decoder with dynamic prompt: 作者进一步开发了一种双向解码器，以动态更新提示和目标特征。该解码器包括一对双向交叉注意力模块，用以下公式表示：
$$
\begin{aligned}
p_{i+1}^d &= \text{softmax}\left(\frac{p_i^d x_i^{dT}}{\sqrt{c}}\right) x_i^d, \\
x_{i+1}^d &= \text{softmax}\left(\frac{x_i^d p_{i+1}^{dT}}{\sqrt{c}}\right) p_{i+1}^d,
\end{aligned}
$$
其中 $p_0^d$ 和 $x_0^d$ 分别从编码器模块初始化为 $p_e$ 和 $x_e$。这种双向交互包括 prompt-to-features 和 features-to-prompt，允许更灵活的特征融合，并使提示特征能够适应目标特征的分布变化。
* 最后，使用最终的交叉注意力模块来更新双向解码器输出上的提示 tokens，并将其输出用作原始特征 ($F_n$) 的重建 ($\hat{F}_n$)。重建损失函数计算重建特征和原始特征之间的均方误差 (MSE)：
$$
L_{rec} = \frac{1}{c \times h \times w} \sum_{i=1}^{c} \sum_{j=1}^{h} \sum_{k=1}^{w} (F_{n_{i,j,k}} - \hat{F}_{n_{i,j,k}})^2.
$$

3.2 Restoration with Normal Image Prompt

* 为了增强图像提示的指导作用，论文引入了人工合成的伪异常图像 $I_a$，这些图像是通过对正常训练图像 $I_n$ 添加损坏或扰动生成的。然后，作者将之前的重建任务转换为修复问题，期望使用正常图像提示 $I_p$ 将异常特征 $F_a$ 修复为正常特征 $F_n$。

* 修复过程与重建过程类似。首先，将成对图像 ($I_a$ 和 $I_p$) 输入到预训练的骨干网络中，以提取离线特征 ($F_a$ 和 $F_p$)，然后通过自注意力编码器和双向交叉注意力解码器，依次将离线成对特征应用于自注意力编码器和双向交叉注意力解码器，最终获得修复后的特征 $\hat{F}_a$。具体来说，自注意力编码器并行处理 $F_a$ 和 $F_p$ 作为输入，并输出为 $\bar{x}_e$ 和 $p_e$。接下来，双向解码器用 $\bar{x}_e$ 和 $p_e$ 初始化，并使用等式（3）动态更新。

* 修复损失函数计算修复后的特征 ($\hat{F}_a$) 与相应的原始正常特征 ($F_n$) 之间的MSE：
$$
L_{res} = \frac{1}{c \times h \times w} \sum_{i=1}^{c} \sum_{j=1}^{h} \sum_{k=1}^{w} (F_{n_{i,j,k}} - \hat{F}_{a_{i,j,k}})^2.
$$

3.3 Supervised Refiner

* 鉴于正常训练图像 $I_n$ 及其对应的异常掩码 $M_n$（所有元素均为零），我们可以合成一个异常图像 $I_a$ 并将其异常掩码表示为 $M_a$。然后，将正常图像和合成图像 $\\{I_t\\}_{t=\\{n, a\\}}$ 输入到预训练的骨干网络中，并将其离线表示导出为 $\\{F_t\\}_{t=\\{n, a\\}}$。然后，使用建议的重建流将正常 $F_n$ 重建为 $\hat{F}_n$，并使用建议的修复流将异常 $F_a$ 恢复为 $\hat{F}_a$。

* 然后，使用原始特征和重建（恢复）特征的绝对元素减法来测量它们的差异，即：
$$
E_t = |F_t - \hat{F}_t|, \quad t = \\{n, a\\}.
$$

* 论文利用像素级别的异常掩码 $M_a$ 通过基于重建误差设计轻量级的像素级别微调器以执行异常分割来进一步优化从低到高的重建误差。微调器由几个转置卷积块组成，后跟一个 1 × 1 卷积层。论文中，每个转置卷积块将重建误差 $E_t$ 上采样 2 倍，并且它由 3 × 3 卷积、BatchNorm、ReLU 和 2 × 2 反卷积组成。在实验中，我们采用了两个转置卷积块，因此将重建误差相对于输入图像从 1/16 向上缩放到 1/4。最后，1 × 1 卷积层将上采样重建误差的通道数转换为 1，并获得估计的异常图作为 $\hat{M}_t$。

* 考虑到异常像素通常在异常检测中占少数，因此使用 Dice 损失函数（有效学习来自极其不平衡的数据），如下所示：
$$
L_{seg} = 1 - \frac{2 \cdot \sum_{i=1}^{H} \sum_{j=1}^{W} \hat{M}_{t_{i,j}} \cdot M_{t_{i,j}}}{\sum_{i=1}^{H} \sum_{j=1}^{W} (\hat{M}_{t_{i,j}})^2 + \sum_{i=1}^{H} \sum_{j=1}^{W} (M_{t_{i,j}})^2},
$$
其中 ($i, j$) 表示 $M_t$ 或 $\hat{M}_t$ 中的空间位置。

3.4 Training Loss

训练期间，在给定特定类别的图像的情况下，默认情况下，从该类别的所有训练图像中随机抽取一个正常图像作为提示。此外，论文还探索了其他提示策略，例如固定模式，这意味着在训练期间每个类别仅使用一个固定图像提示。考虑到 OneNIP 中的所有三个目标，包括无监督重建、无监督修复和监督微调器，总训练损失为
$$
L = L_{rec} + L_{res} + \lambda L_{seg},
$$
其中 $\lambda > 0$ 是平衡两种类型的损失函数 $L_{rec}+L_{res}$ 和 $L_{seg}$ 的重要性的权重。

3.5 Inference

在推理时，首先为每个类随机选择一个正常的训练图像，然后预提取离线提示特征以构建类感知提示池 $\\{P_i\\}_{i=1}^{C}$。给定测试图像 $I$ 及其特征 $F$，可以通过计算测试特征 $F$ 和提示池之间的余弦相似度来得出适当的提示，因为测试图像的类是不可知的。

Pixel-Level Anomaly Segmentation: 对于无监督重建，异常评分图计算为重建误差的 L2 范数，如
$$
S_{rec} = ||F - \hat{F}||_2 \in \mathbb{R}^{h \times w}.
$$

对于有监督的微调器，异常评分图被预测为 $\hat{M} \in \mathbb{R}^{H \times W}$。最后，将 $S_{rec}$（调整到原始分辨率）和 $\hat{M}$ 组合在一起，并将其作为最终异常分割图，即
$$
S = (1 - \alpha) \cdot S_{rec} + \alpha \cdot \hat{M},
$$
其中 $\alpha \in [0, 1]$ 是一个权重。

Image-Level Anomaly Classification: 异常分类旨在检测图像是否包含异常区域。按照之前的工作，取 $S$ 的最大值作为图像级别的异常评分。

4. 实现细节与注意事项

* 预训练模型： 论文使用了预训练的 EfficientNet-b4 模型。选择合适的预训练模型对性能至关重要。
* 特征图大小： 论文将不同层级的特征图resize到14x14，concatenate在一起，形成最终特征。合适的特征大小，对模型性能有影响。
* 编码器和解码器的层数： 平衡性能和计算成本。
* 伪异常样本生成： 使用 CutPaste 和 DRAEM 等方法来生成伪异常样本。这些方法在生成伪异常时需要一定的随机性，以避免模型过度拟合特定的伪异常模式。
* Dice 损失： 使用 Dice 损失函数来处理异常检测中样本不平衡的问题。
* 推理时的提示选择： 如何选择合适的提示图像是一个关键问题。论文使用余弦相似度来选择最相似的提示图像。
* 超参数调整： 损失函数的权重、学习率等超参数需要仔细调整。

实现难点和解决方案：

* 提示图像的选择：
* 问题： 如何选择合适的提示图像？如果提示图像与待检测图像差异太大，可能会导致性能下降。
* 解决方案： 构建一个提示图像池，并使用余弦相似度等方法来选择最相似的提示图像。

* 伪异常样本的生成：
* 问题： 如何生成具有代表性的伪异常样本？如果伪异常样本过于简单或过于复杂，可能会影响模型的训练效果。
* 解决方案： 结合多种伪异常生成方法，并进行随机化，以增加伪异常样本的多样性。

* 计算资源：
* 问题： 训练深度神经网络需要大量的计算资源。
* 解决方案： 使用 GPU 加速训练，并采用分布式训练等方法来提高训练效率。

优化建议和最佳实践：

* 尝试不同的预训练模型： 例如，可以使用 ResNet 或 ViT 等模型来替代 EfficientNet-b4。
* 探索不同的提示选择策略： 例如，可以使用聚类算法来构建提示图像池。
* 设计更有效的伪异常生成方法： 例如，可以根据具体的应用场景来设计特定的伪异常生成方法。
* 优化网络结构： 例如，可以尝试使用更先进的 Transformer 结构或卷积神经网络结构。

希望这个详细的分析能够帮助你更好地理解这篇论文！