Paper Reading: 探索多类无监督异常检测的状态空间模型, MambaAD (NeurIPS 2024)-CSDN博客

本文链接：https://blog.csdn.net/Transfattyacids/article/details/144836683

MambaAD

MambaAD: Exploring State Space Models for Multi-class Unsupervised Anomaly Detection

MambaAD：探索多类无监督异常检测的状态空间模型， NeurIPS 2024

paper: https://arxiv.org/abs/2404.06564

code: https://lewandofskee.github.io/projects/MambaAD/

单位：浙江大学，腾讯优图，南洋理工大学

摘要

近年来，异常检测领域在基于卷积神经网络（CNN）和Transformer的技术进展中表现出了显著的效果。然而，CNN在处理长距离依赖关系时存在困难，而Transformer则因其二次方的计算复杂度而面临较大的计算负担。基于Mamba的模型因其卓越的长距离建模能力和线性计算效率，受到了广泛关注。本研究首次将Mamba应用于多类别无监督异常检测，提出了MambaAD方法。该方法由一个预训练编码器和一个包含多尺度（局部增强状态空间，Locality-Enhanced State Space, LSS）模块的Mamba解码器组成。

所提出的LSS模块通过将并行级联的（混合状态空间，Hybrid State Space, HSS）模块与多核卷积操作相结合，能够有效捕捉长距离和局部信息。HSS模块利用（混合扫描，Hybrid Scanning, HS）编码器，将特征图通过五种扫描方法和八个方向进行编码，从而通过（状态空间模型，State Space Model, SSM）增强全局联系。特别地，Hilbert扫描和八方向编码的使用显著提高了特征序列的建模能力。

在六个不同的异常检测数据集和七个评估指标上的综合实验，验证了该方法的有效性，并实现了最新的最佳（state-of-the-art）性能。这进一步证明了MambaAD方法在异常检测任务中的卓越表现。

简介

基于重构的方法，如RD4AD和UniAD，（比基于特征嵌入和合成的方法）表现出优越的性能和更好的可扩展性

如图 1 (a) ：基于 CNN 的 RD4AD 有效地捕获本地上下文，但它缺乏建立远程依赖的能力。

如图 1 (b)：UniAD是第一个多类 AD 算法，依赖于预训练的编码器和转换器解码器架构。尽管变压器具有优越的全局建模能力，但它受到二次计算复杂度的阻碍，它将UniAD限制在最小特征图上的异常检测，可能会影响其性能。

最近Mamba 在大型语言模型中展示了卓越的性能，与Transformer 相比，线性复杂度显着降低，同时保持了可比的效率。

如图1 ©：MambaAD结合了全局和局部建模能力，利用其线性复杂度来计算多个尺度的异常图。它拥有较低的参数计数和计算需求，这使得它非常适合实际应用。

在这里插入图片描述

图 1：与 (a) 基于局部 CNN 的 RD4AD [8] 和 (b) 基于全局 Transformer 的 UniAD [43] 相比，的具有线性复杂度的 MambaAD 能够整合全局和局部建模的优势，多尺度特征赋予它更精细的预测精度。

具体来说，MambaAD 采用金字塔结构的自动编码器来重建多尺度特征，利用预训练的编码器和基于 Mamba 架构的新解码器。

这种基于 Mamba 的解码器由不同尺度和数量的局部增强状态空间 (LSS) 模块组成。

每个 LSS 模块包括两个组件：
一系列用于全局信息捕获的混合状态空间 (HSS) 块和用于建立局部连接的并行多核卷积操作。
生成的输出特征将 Mamba 结构的全局建模能力与 CNN 的局部相关强度相结合。
所提出的HSS模块研究了五种不同的扫描方法和8个扫描方向，其中(混合扫描)HS编码器和解码器编码和解码特征映射分别转化为各种扫描方法和方向序列。
HSS模块在多个方向上增强了全局接受域，并利用希尔伯特扫描方法特别适合于工业产品特征的中心浓度。

贡献
• 提出了 MambaAD，它创新地应用 Mamba 框架来解决多类无监督异常检测任务。这种方法以最小的模型参数和计算复杂度实现多尺度训练和推理。
• 设计了一个局部增强状态空间 (LSS) 模块，包括级联基于 Mamba 的块和并行多核卷积，提取全局特征相关性和局部信息关联，实现了全局和局部模式的统一模型。
• 探索了混合状态空间 (HSS) 块，包括五种方法和八个多方向扫描，以增强跨各种类别和形态的复杂异常检测图像的全局建模能力。
• 展示了 MambaAD 在多类异常检测任务中的优越性和效率，在六个具有七个指标的不同 AD 数据集上实现了 SoTA 结果，同时保持了非常低的模型参数和计算复杂度。

背景知识

State Space Models，SSM状态空间模型
状态空间模型 (SSM) 因其在处理长序列建模方面的有效性而备受关注。具体来说，结构状态空间序列 (S4, structure state-space sequence) 通过对角结构的参数化有效地对远程依赖关系 (LRD) 进行建模，解决了以前工作中遇到的计算瓶颈。在S4的基础上，提出了许多模型，包括S5、H3，特别是Mamba。Mamba 在 S4 中引入了与数据相关的选择机制，该机制提供了与 CNN 或 Transformer 不同的新范式，保持了长序列处理的线性可扩展性。

在这项工作中，开发了 MambaAD 来利用 Mamba 的远程建模能力和线性计算效率进行多类无监督异常检测。这种方法创新地将 SSM 的广泛建模能力与 CNN 的详细局部建模程序相结合，从而实现无缝集成。
（有兴趣的友友可以了解一下：RNN–>SSM–>S4–>Mamba，这里不做科普（因为我看了一个下午Mamba也还是一知半解，就不误人子弟了orz

方法

在这里插入图片描述

图 2：所提出的 MambaAD 概述，它使用金字塔自动编码器框架通过所提出的高效和有效的局部增强状态空间 (LSS) 模块重建多尺度特征。具体来说，每个 LSS 由：1）级联混合状态空间（HSS）块来捕获全局交互； 2）并行多核卷积操作来补充局部信息。聚合多尺度重构误差作为推理的异常图。

DWConv: depth-wise Conv block
右上角是图b

从图1(a)中可以看到，MambaAD的整体架构由三部分组成：预训练的CNN编码器、Half-FPN瓶颈层和Mamba解码器。预训练的编码器负责提取图像的多尺度特征，这些特征随后被送入Half-FPN瓶颈层进行融合。瓶颈层的输出接着被送入Mamba解码器，该解码器由多个Locality-Enhanced State Space (LSS)模块组成，这些模块在不同尺度上进行操作。

4 个 Mamba 解码器，深度配置为 [3,4,6,3]。最终的损失函数是在三个尺度上跨特征图计算的均方误差 (MSE) 的总和

LSS由 HSS块(捕获全局交互) 和并行多核卷积运算(补充局部信息)组成

LSS作用： 重构多尺度特征
参考图2(b)，LSS模块是MambaAD的核心组件之一。每个LSS模块由多个Hybrid State Space (HSS)块和多核卷积操作组成。HSS块负责捕获全局信息，而多核卷积操作则用于捕获局部信息。这种设计使得LSS模块能够有效地结合Mamba的全局建模能力和CNN的局部特征关联能力。具体来说，输入特征Xi首先被分成全局特征Gi和局部特征Li，其中Gi通过HSS块处理，而Li则通过多核卷积层处理。最终，通过通道维度的拼接和1x1的卷积操作，恢复原始输入特征的维度，得到该模块的输出Xo。

Hybrid State Space (HSS) 块的扫描方法
图3(a)展示了HSS块采用的Hilbert（希尔伯特）扫描方法，该方法包含8个不同的扫描方向。HSS块的设计灵感来源于空间填充曲线，用于在编码器和解码器中以多种扫描方法和方向处理特征图。这些扫描方法包括横扫（Sweep）、扫描（Scan）、Z-order、Zigzag和Hilbert曲线等。HSS块通过Layer Normalization (LN)、Linear Layer、depth-wise卷积、SiLU激活函数、HS编码器、状态空间模型（SSMs）、HS解码器以及残差连接组成。这种混合扫描方法的设计使得模型能够更好地捕捉全局特征，从而提高对复杂异常图像的建模能力。

前置设定

状态空间模型SSM 受控制系统的启发，通过隐藏状态 $h(t)\in R^N$ 将一维刺激 $x(t)\in R^L$ 映射到响应 $y(t)\in R^L$ ，其表达式为线性常微分方程(ODEs):

其中，状态转移矩阵

S4和Mamba利用零阶保持和时间尺度参数∆将连续参数A和B从连续系统转换为离散参数A和B:

在离散化后，离散化模型公式可以表示为：
最后，从全局卷积的角度来看，输出可以定义为：

其中 * 表示卷积操作，L 是序列 x 的长度，是结构化卷积核。

MambaAD整体概述

提出了一种用于多类异常检测的MambaAD框架，如图2(a)所示。它由三个主要组件组成：预训练的基于 CNN 的编码器、Half-FPN 瓶颈和基于 Mamba 的解码器。在训练期间，编码器以三个不同的尺度提取特征图，并将它们输入到 H-FPN 瓶颈进行融合。然后将融合的输出输入Mamba 解码器，深度配置为 [3,4,6,3]。最终的损失函数是在三个尺度上跨特征图计算的均方误差 (MSE) 的总和。

在 Mamba 解码器中，引入了局部增强状态空间 (LSS) 模块。LSS 可以配置不同的阶段 $M_i$ ，其中每个阶段都表示模块内混合状态空间 (HSS) 块的数量 N。在本实验中，使用 $M_i$ = 3 和 $M_i$ = 2 的 LSS 作为主要模块。

在这里插入图片描述

LSS模块通过HSS块处理输入 $x_i$ ，HSS块捕获全局信息，两个不同的深度卷积层捕获局部信息。原始输入特征维度通过串联和卷积操作恢复。（图里误把b写成c了）

所提出的HSS块具有混合扫描(HS)编码器和HS解码器，它容纳了五种不同的扫描方法和8个扫描方向。
在这里插入图片描述

LSS

局部增强状态空间模块， Locality-Enhanced State Space Module

Transformers 在全局建模和捕获长期依赖关系方面表现出色，但在处理高分辨率特征时往往会忽略局部语义信息并表现出较高的计算复杂度。相反，CNN 通过从相邻位置捕获信息但缺乏远程建模能力来有效地对局部语义进行建模。为了解决这些限制，提出了图 2 (b) 中的 LSS 模块，它结合了基于 Mamba 的级联 HSS 块进行全局建模和并行多核深度卷积操作，用于局部信息捕获。

在这里插入图片描述

具体来说，对于输入特征 $x_i \in R^{H \times W \times C}$ ，全局特征Gi进入HSS块，局部特征Li通过卷积网络进行。全局特征 $G_i$ 通过一系列N个HSS块来获得全局信息特征 $G_o$ 。

其中 n ∈ N 是 HSS 块的数量。在本研究中，主要使用N = 2和N = 3，进一步消融实验见第4.3节。局部特征Li由两个并行的深度卷积Conv块处理，每个块由一个1 × 1Conv块、一个k × k的深度Conv块和一个另一个1 × 1 Conv块组成。

局部特征Li由两个并行的深度Conv块处理，每个块由一个1×1Conv块、一个k×k的深度Conv块和一个另一个1×1 Conv块组成。

其中，k表示深度卷积的核大小。在本实验中，使用了k = 5和k = 7，进一步的消融实验将在第4.3节中进行探讨。每个卷积模块包括一个二维卷积层（Conv 2D）、一个二维实例归一化层（Instance Norm 2D）以及一个SiLU激活函数，如图2(d)所示。局部和全局特征通过沿通道维度的拼接操作进行聚合。该模块的最终输出 $X_o$ 是通过一个1 × 1的二维卷积层将通道数恢复为与输入相匹配，并结合残差连接而得到的。

在这里插入图片描述

HSS

混合状态空间块， Hybrid State Space Block

附Mamba简化后的架构图（HSS比Mamba块就多了HS）
>
（Mamba架构）简化的模块设计将H3模块（这是大多数SSM架构的基础）与现代神经网络中常见的MLP模块相结合。与交替使用这两个模块不同，只是简单而统一地重复使用Mamba模块。与H3模块相比，Mamba将第一个乘性门替换为激活函数；与MLP模块相比，Mamba在主分支中添加了一个SSM（状态空间模型）。对于激活函数𝜎，采用了SiLU/Swish激活函数

在这里插入图片描述

HSS模块设计用于混合方法和混合方向的扫描与融合，如图2©所示。HSS模块主要由以下部分组成：层归一化（Layer Normalization, LN）、线性层、深度卷积（depth-wise convolution）、SiLU激活函数、混合扫描（Hybrid Scanning, HS）编码器、状态空间模型（State Space Models, SSMs）、HS解码器
以及残差连接。

在这里插入图片描述

HS=混合扫描方法+混合扫描方向

混合扫描方法

受空间填充曲线[45]的启发，如图3所示，本研究探索了五种不同的扫描方法:(I) Sweep， (II) Scan， (III) Z阶， (IV) Zigzag， (V) Hilbert，以评估它们对SSM建模能力的影响。最终选择希尔伯特扫描方法对特征序列中的局部和全局信息进行卓越的编码和建模，特别是在减轻对远程依赖建模的挑战方面。消融研究将给出进一步的实验结果。假设 A 是一个矩阵，AT 是 A 的转置，Alr 是 A 的左右反转，Aud 是 A 的自上而下的反转。希尔伯特曲线可以通过n阶希尔伯特矩阵得到:

在这里插入图片描述

其中
和 $E_n$ 是 n 阶的全一矩阵

在这里插入图片描述

图 3：混合扫描方向和方法。(a) HS编码器和解码器采用8个扫描方向的希尔伯特扫描方法。(b) 其他四种扫描方法进行比较。

混合扫描方向

基于之前扫描方向的设置，本研究支持八种基于Hilbert曲线的扫描方向：（i）正向，（ii）反向，（iii）宽-高（wh）正向，（iv）宽-高反向，（v）旋转90度正向，（vi）旋转90度反向，（vii）宽-高旋转90度正向，以及（viii）宽-高旋转90度反向，如图3(a)所示。多种扫描方向增强了特征序列的编码和建模能力，从而能够处理各种类型的异常特征。进一步的消融实验将在第4.3节中进行探讨。

HS编码器的目标是根据不同的扫描方法和方向组合并编码输入特征，在将其输入到SSM（状态空间模型）之前，以增强特征向量的全局建模能力。随后，HS解码器将SSM输出的特征向量解码回原始输入特征的方向，最终输出通过特征求和获得。

实验

设置

数据集: MVTec-AD, VisA, Real-IAD, MVTec-3D, Uni-Medical, COCO-AD
指标:

上述七个评估指标（表示为 mAD）的平均值来表示模型的综合能力

细节:

image size: 256 × 256
特征提取器预训练的 ResNet34
解码器: 与ResNet34 的等效深度的 Mamba 解码器用作训练的学生模型。在 Mamba 解码器中，第二个 LSS 模块中的级联 HSS 块数设置为 2，而所有其他 LSS 模块采用 3 个级联 HSS 块。

本实验采用希尔伯特扫描技术，利用八个不同的扫描方向。采用AdamW优化器，学习率为0.005，衰减率为1 × 10−4。该模型在单个NVIDIA TESLA V100 32GB GPU上进行多类设置的训练周期为500 epoch。在训练期间，不同尺度的 MSE 之和被用作损失函数。在测试阶段，利用不同尺度的余弦相似度之和作为异常图。

Comparison with SoTAs on Different AD datasets

在这里插入图片描述

表 1：多类设置的不同 AD 数据集的定量结果。

在这里插入图片描述

图4:MVTec和VisA数据集上像素级异常分割的定性可视化。

Ablation and Analysis

在这里插入图片描述

表 2：预训练主干和 Mamba 解码器深度的消融研究。

在这里插入图片描述

表 3：不同扫描方法和方向的消融研究。

不同扫描方法和方向的有效性与效率比较。在初始阶段，对五种不同的扫描方法进行了对比，并采用了8方向扫描，结果如表3所示。结果表明，除Hilbert扫描外，其余四种方法在图像层面取得了与索引1-4类似的结果，但略逊于Hilbert扫描技术。随后，进一步研究了不同扫描方向数量的影响。随着扫描方向数量的增加（索引5-7），图像级别的指标逐渐提高，而像素级别的指标保持不变。这表明，增加扫描方向的数量可以增强SSM的全局建模能力，从而减少图像级别误分类的可能性。

最终，分析发现，在保持总扫描方向为八个的情况下，结合多种扫描技术（索引8-12）导致了性能下降。这种下降可能是由于所采用的扫描方法之间存在较大的差异而引起的。因此，选择了Hilbert扫描方法，因为它在实际工业产品中表现出更好的适应性，并且有助于降低误报的可能性。

在这里插入图片描述

表 4：LSS 模块设计的消融研究。

不同LSS设计的有效性比较。如表4所示，LSS模块的设计主要关注三个不同的设计方向：LSS模块中HSS块的数量（即Mi的值）、并行多核卷积模块的配置，以及深度卷积的核大小选择。

首先，对比了Mi = 1的情况，即每个LSS模块仅包含一个HSS块，并针对仅使用深度卷积和两端添加1 × 1卷积的配置进行了不同核大小的实验。随后，对比了没有残差连接的配置与具有相同设置但带有残差连接（Mi = 1）的结果。最后，在其他条件一致的情况下，分析了Mi = 2和Mi = 3下的表现。结果显示，当Mi = 1时，无论是否存在残差连接，其效果都逊于Mi = 2和Mi = 3。此外，在仅使用深度卷积块且Mi = 1的情况下，比较不同核大小的并行深度卷积发现，较小的卷积核（如k = 1）会显著降低性能。因此，后续的比较实验集中于较大的卷积核。

在没有残差连接的情况下，某些指标可能优于带残差连接的配置，但由于训练时间较长且收敛困难，这些配置未被用于进一步实验。在Mi = 2和Mi = 3的配置中，发现，结合1 × 1卷积的深度卷积块表现最佳。此外，核大小为k = 5和k = 7更适合提取局部特征并建立局部信息关联。因此，在本研究中，选择了包含Mi = 2和Mi = 3的HSS块数量，采用核大小为k = 5和k = 7的并行深度卷积块，并在卷积块前后搭配1 × 1卷积进行优化。

在这里插入图片描述

表 5：SoTA 方法的效率比较。

conclusion

本文介绍了MambaAD，这是首次将Mamba框架应用于异常检测（AD）的研究。MambaAD由一个预训练的编码器和一个Mamba解码器组成，并在不同尺度和深度上引入了一个新颖的LSS模块。LSS模块结合了顺序排列的HSS模块和并行多核卷积网络，将Mamba框架的全局建模能力与基于CNN的局部特征关联相结合。HSS模块利用HS编码器，将输入特征编码为五种扫描模式和八个方向，从而有助于对工业产品的中心位置特征序列进行建模。

在六个不同的异常检测数据集和七个评估指标上的广泛实验表明，该方法在实现最新的最优性能（SoTA）方面具有显著的效果。

更广泛的影响：本研究是将Mamba应用于异常检测的首次尝试，为未来的研究奠定了基础。希望所提出的MambaAD能够为异常检测领域中轻量化和高效模型的设计提供启发。