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

原创

已于 2024-12-31 01:06:34 修改 · 1.4k 阅读

25 ·

CC 4.0 BY-SA版权

文章标签：

#笔记 #论文笔记 #pytorch #深度学习 #计算机视觉

于 2024-12-31 00:18:15 首次发布

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解码器以及残差连接组成。这种混合扫描方法的设计使得模型能够更好地捕捉全局特征，从而提高对复杂异常图像的建模能力。