摘要:本文提出了一个无需对目标类别进行再训练的新型少样本异常分割方法——DictAS,它通过从少量正常图像构建“字典”并利用字典查找机制识别查询图像中异常区域,从而在多个工业和医疗数据集上持续超越了现有的最先进 FSAS 方法。
作者丨qzzzxiaosheng@知乎(已授权)
来源丨https://zhuanlan.zhihu.com/p/1942894159007547647
编辑丨极市平台
导读:本文提出了一个无需对目标类别进行再训练的新型少样本异常分割方法——DictAS,它通过从少量正常图像构建“字典”并利用字典查找机制识别查询图像中异常区域,从而在多个工业和医疗数据集上持续超越了现有的最先进 FSAS 方法。
题目: DictAS: A Framework for Class-Generalizable Few-Shot Anomaly Segmentation via Dictionary
关键词: 少样本, 异常分割,类别泛化
论文链接:https://www.arxiv.org/pdf/2508.13560
代码: https://github.com/xiaozhen228/DictAS
近期的视觉-语言模型(如 CLIP)在小样本异常分割(Few-Shot Anomaly Segmentation,FSAS)任务中展现出了对未见类别的显著泛化能力,主要依赖于在已见类别上的监督提示学习或微调。然而,它们的跨类别泛化能力在很大程度上依赖于已见异常样本的先验知识。为此,中国科学院自动化研究所工业视觉团队提出提出了一种新颖的框架 DictAS,该框架能够在无需针对目标数据重新训练的情况下,仅通过少量正常参考图像作为视觉提示,实现对未见类别的视觉异常检测。DictAS 的核心思想是通过自监督学习将“字典查找(dictionary lookup)”的能力迁移到 FSAS 任务中,而不是仅仅记忆训练集中正常与异常的特征模式。
具体而言,DictAS 主要包含以下三个部分:1) 字典构建(Dictionary Construction)-- 使用正常参考图像的特征来模拟真实字典的索引和内容;2) 字典查找(Dictionary Lookup) -- 通过稀疏查找策略从字典中检索查询区域特征,当查询特征无法被检索到时,即判定为异常;3) 查询判别正则化(Query Discrimination Regularization) -- 通过使异常特征更难从字典中检索来增强异常判别能力。为此,本文进一步提出了对比查询约束(Contrastive Query Constraint)和文本对齐约束(Text Alignment Constraint)。在七个公开的工业和医疗数据集上的大量实验表明,DictAS 在性能上持续超越了现有的最先进 FSAS 方法。
Few-shot (few-normal-shot) anomaly segmentation 旨在在仅提供有限数量正常样本的情况下,识别图像中的异常区域。该任务在训练数据稀缺且像素级标注有限的场景中尤为重要,例如工业缺陷检测和医学图像分析。
图1 :(a) 两种少样本异常分割的不同设定 (b)所提出的DictAS的动机
现有的 FSAS 方法通常遵循类别依赖或类别可泛化的设置,如图 1(a) 所示。类别依赖的方法需要在每个未见类别上使用有限数量的正常样本进行微调。由于训练图像类别与测试图像类别相同,这类方法主要集中于建模正常图像的分布,或通过正则化学习正常与异常内容之间的边界,而未充分考虑跨类别存在的显著域间差异。因此,这些方法在应用于隐私敏感的医学场景或生产线频繁变化的动态工业场景时会面临重大挑战。相比之下,如图 1(a) 右侧所示,类别可泛化的方法旨在开发一个统一模型,无需在目标数据上重新训练,仅依赖少量正常样本作为视觉提示即可检测未见类别中的异常。
现有的类别泛化方法比较有限,主要有基于元学习的方法(如 MetaUAS, RegAD),基于特征残差的方法(如 InctRL, ResAD),基于视觉语言模型的方法(如WinCLIP, APRIL-GAN)。 但是,基于元学习的方法构建辅助数据集繁琐,其他方法整体采用记忆库(memory-bank)的思路来保证模型在新类别的泛化能力。 然而,即使是一个新手人类检验员,也能够仅凭少量正常样本作为参考,在未见类别中识别异常,而无需丰富的先验经验。我们将这一直觉近似为字典查找:如果查询图像中的某个区域无法从字典中检索到,则判定为异常;否则为正常。
受此启发,本文提出了一种基于 CLIP 的新型自监督框架 DictAS,用于类别可泛化的 FSAS。该框架包括三个组成部分:字典构建(Dictionary Construction)、字典查找(Dictionary Lookup)和查询判别正则化(Query Discrimination Regularization),如图 1(b) 所示。我们的设计动机是将异常分割重新表述为一个字典查找任务——判断查询特征是否存在于字典中。通过自监督训练,模型能够获得与特征无关且动态自适应的字典查找能力,从而实现类别可泛化的 FSAS。
问题描述: 类别可泛化的 FSAS 是一项具有挑战性的任务,其目标是在未见类别上实现高性能,而无需对目标数据进行微调。辅助训练集中的已见类别记为 $C^s$ ,测试集中的未见类别记为 $C^u$ ,它们必须满足 $C^s \cap C^u=\emptyset$ 。最终目标是开发一个统一模型,仅依赖与目标类别相对应的 $k$ 张正常样本作为视觉提示,就能够对新类别 $C^u$ 中的视觉异常进行分割。
整体框架&自监督范式: DictAS的整体框架如下图2所示。在少样本学习中,待检测或者定位异常的图像一般称为查询图像(Query Image), 提供的少量正常图像一般称为参考图像(Reference Image)或者支撑图像 (Support Image)。DictAS框架的优点在于不需要想其他方法一样组织数集为专门的格式进行元训练,而是任何图像都可以充当辅助的训练集,包括自然场景数据集。
自监督方法的本质在于训练过程中的监督信号来自于自身或者合成的信号。如图2所示,在给定任意一张原始图像,DictAS 通过异常合成算法来生成Query Image,通过数据随机变换(随机旋转、亮度)等生成参考图像。辅助训练过程中的监督信号均来自于图像自身或者生成的伪异常。
图2 :DictAS框架图
字典构建(Dictionary Construction)参考图像特征被转化为字典索引(Dictionary Key) $\mathbf{F}K^l$ 和值(Dictionary Value) $\mathbf{F}V^l$ ,查询图像特征则生成查询向量 $\mathbf{F}_Q^l$ :
其中 $F_q^l, F_n^l$ 分别为查询与参考图像在第 $l$ 层提取的特征。AttnBlock 是一个普通的 transformer的多头自注意力模块,设计如下:
import torch from torch import Tensor, nn from torch.nn import functional as F from torch.nn.functional import scaled_dot_product_attentionclass Attn_Block(nn.Module):def __init__(self, dim, num_heads = 4, qkv_bias = False, qk_scale = None, proj_drop = 0.1):super.__init__self.num_heads = num_heads head_dim = dim // num_headsself.scale = qk_scale or head_dim ** -0.5self.q_proj = nn.Linear(dim, dim, bias = qkv_bias)self.k_proj = nn.Linear(dim, dim, bias = qkv_bias)self.v_proj = nn.Linear(dim, dim, bias = qkv_bias)self.norm1 = nn.LayerNorm(dim)self.TwoLayerMLP = MLP(dim, 0.1)def forward(self, fea):q = k = v = self.norm1(fea)res = q.cloneB, N, C = q.shape assert k.shape == v.shape B, M, C = k.shape q = self.q_proj(q).reshape(B, N, self.num_heads, C // self.num_heads)k = self.k_proj(k).reshape(B, M, self.num_heads, C // self.num_heads)v = self.v_proj(v).reshape(B, M, self.num_heads, C // self.num_heads)# Way1 #q = q.transpose(1, 2)k = k.transpose(1, 2)v = v.transpose(1, 2)x = scaled_dot_product_attention(q, k, v, is_causal=False) # [B, num_heads, N, head_dim]x = x.transpose(1, 2).reshape(B, N, C)# Way1 #'''Way 1 and Way 2 are essentially the same, but scaled_dot_product_attention supports testing with more support samples due to its internal optimization. This helps alleviate high GPU memory usage when the dictionary is very large.'''# Way2 ## attn = torch.einsum('bnkc,bmkc->bknm', q, k) * self.scale# attn = attn.softmax(dim = -1)# x = torch.einsum('bknm,bmkc->bnkc', attn, v).reshape(B, N, C)# Way2 #x = self.TwoLayerMLP(x, res)return x字典查询(Dictionary Lookup)字典查询主要是利用生成的Dictionary Query $\mathbf{F}Q^l$ 和 Dictionary Key $\mathbf{F}K^l$ 进行匹配,然后利用生成的权重对Dictionary Value $\mathbf{F}V^l$ 进行加权组合,本质类似于cross-attention结构。这样的结构保证了最后的检索结果(Retrieved Result) $\mathbf{F}r^l$ 完全是由 k 张 正常的支撑图像特征的线性组合产生的。而类别泛化的关键便是线性组合的权重 $\hat{\mathbf{w}}$ 是根据Dictionary Query(来自查询图像)和 Dictionary Key(来自正常的支撑图像)动态适配调整的。为了表述方便,文章中使用了查询图像中的一个patch $\mathbf{x}_Q^l$ 来表述这个过程:
其中根据 $\hat{\mathbf{W}}$ 的获取方式,分为最大查询,密集查询以及稀疏查询三种方式:
整个字典查询的过程可以简单的用代码表示为(self.num_heads = 1):
def Lookup(self, patch_feature_query, patch_feature_support): B, N, C = patch_feature_query.shape B, M, C = patch_feature_support.shape # Dictionary ConstructionF_Q = self.Query_Generator(patch_feature_query).reshape(B, N, self.num_heads, C // self.num_heads)F_K = self.Key_Generator(patch_feature_support).reshape(B, M, self.num_heads, C // self.num_heads)F_V = self.Value_Generator(patch_feature_support).reshape(B, M, self.num_heads, C // self.num_heads)# Dictionary Lookupattn = torch.einsum('bnkc,bmkc->bknm', F_Q, F_K) * self.scaleattn = self.SPM(attn, adim = -1)x = torch.einsum('bknm,bmkc->bnkc', attn, F_V).reshape(B, N, C) return x问: 那么通过这种字典查询定位到未知类别表面异常的底层逻辑是什么?
答:查询结果是 来自 正常图像中特征的加权组合,如果query patch是正常的,那么自然能够通过正常图像中的特征加权组合得到相近(余弦距离更小)的结果。反之,如果query patch是异常的, 那么来自参考图像中正常的特征加权组合更难重建出query patch (余弦距离更大)。
损失函数 损失函数包括三个部分:Query Loss $\mathcal{L} q$(主损失),Contrastive Query Constraint $\mathcal{L}{C Q C}$(正则化损失),Text Alignment Constraint $\mathcal{L}{T A C}$(正则化损失)。其中主损失 $\mathcal{L} q$ 强调Query Image中的正常区域和 Retrieved Result中对应的部分余弦距离最小:正则化损失主要鼓励Retrieved Result 因该为正常的特征(因为它是由正常的参考图像线性组合得到的),主要包含两个部分:
Contrastive Query Constraint $\mathcal{L}{C Q C}$鼓励Query Image Feature和 Retrieved Result Feature中异常patch之间的平均余弦距离应该大于正常patch之间的平均余弦距离:
Text Alignment Constraint $\mathcal{L}_{T A C}$ 鼓励Retrieved Result Feature 在全局应该始终与正常的文本嵌入空间对齐:
最终的辅助训练总损失函数为:
在测试过程中,给定真实的测试图像和k张正常的样本,最终的异常图可以被计算为:
与大部分的零样本异常检测方法相同,采取在A数据集训练,在B数据集测试的设定。其中A和B中的类别应该没有重叠。文章中采用的工业数据集有:MVTecAD, VisA, MPDD, BTAD, MVTec3D,采用的医疗数据集有: RESC, BrasTS。 文章使用工业数据集 VisA 中的所有正常图像作为辅助数据进行自监督训练,并在其余工业和医疗数据集上直接进行小样本测试。为了在 VisA 本身上进行评估,则采用 MVTecAD 作为辅助训练集。值得注意的是,由于采用了自监督范式,任何无需像素级标注的图像都可以用作辅助训练数据,包括自然场景中的图像。
实验对比: 对比的SOTA方法包括: RegAD, AnomalyGPT, FastRecon, WinCLIP, APRIL-GAN, PromptAD, 其中PromptAD 采用的是类别依赖的范式,而其他方法采用的是类别泛化的范式。定量的对比结果如下:
表1
图3
实验结果显示,在7个数据集上几乎均取得了最优的性能,甚至超过了类别依赖的方法PromptAD。在1, 2, 4, 8, 16-shot时, 像素级的AP水平也达到了SOTA水平。
定性的对比结果如下:
图4
由可视化结果可以看出DictAS在工业和医疗数据集上的异常定位结果更加精确和完整。
表2
由于不需要像传统的方法一样构建记忆库检索,或者类别依赖的方法在每个目标类别中进行训练,DictAS在取得优良性能的同时也获得了更快的推理速度。
本文提出了一种新颖的类别可泛化小样本异常分割(FSAS)框架 DictAS。受人类检验员工作方式的启发,我们将 FSAS 重新表述为字典查找任务:当查询特征无法从由正常参考图像特征构建的字典中检索到时,即判定为异常。
通过在辅助数据集的已见类别上进行自监督训练,DictAS 学习到可迁移的字典查找能力,从而能够有效泛化到未见类别的 FSAS 任务中。为了进一步增强异常判别能力,本文引入了查询判别正则化(Query Discrimination Regularization),并与查询损失联合优化,使异常特征更难被字典检索。最终的异常图是通过计算查询特征与其检索结果之间的余弦距离得到的。
在七个工业和医疗数据集上的大量实验表明,DictAS 在 FSAS 性能上超越了现有的最先进方法,同时实现了最快的推理速度并保持了可比的稳定性。
来源:极市平台