极市导读
本文介绍了一种名为MambaAD的新方法,它首次将Mamba模型应用于多类无监督异常检测任务。MambaAD结合了预训练的CNN编码器和由不同尺度的局部增强状态空间(LSS)模块组成的Mamba解码器,通过混合扫描(HS)编码器和状态空间模型(SSM)来建立全局联系,并在多个异常检测数据集上取得了最先进的性能。 >>加入极市CV技术交流群,走在计算机视觉的最前沿
(来自浙大,腾讯优图,南洋理工大学)
论文:https://arxiv.org/abs/2404.06564
主页:lewandofskee.github.io/projects/MambaAD/
代码:github.com/lewandofskee/MambaAD
内容总结(太长不看版)
过去基于CNN和Transformer的结构的算法被广泛应用于异常检测领域取得了一定的效果。但是CNN无法处理长距离信息的关联性,而Transformer受限于其平方级别的计算复杂度。最近基于Mamba的模型凭借着其长距离建模的出色能力与线性计算复杂度引起了广泛的研究。在本文中,我们首次将Mamba应用于多类无监督异常检测领域并提出MambaAD如图1所示包含一个预训练的CNN编码器和由不同尺度的局部增强状态空间(LSS)模块成的Mamba解码器。被提出的LSS 模块包含并行的连续混合状态空间(HSS)模块和多核的卷积操作,能够同时具有长距离的建模能力与局部信息的建模能力。HSS模块所包含的混合扫描(HS)编码器将特征图编码至5种不同的扫描方式和8种不同的扫描方向并输送至SSM中建立全局的联系。其中我们采用了Hilbert扫描方式和8种扫描方向有利于提升特征序列的建模能力。大量实验表明,我们在6种不同异常检测数据集上、7种不同的指标上取得了SoTA,证明了Mamba AD方法了有效性。
1、引言
尽管基于合成和基于特征编码的方法在AD领域取得了不错的效果,但是这些方法需要额外的设计与不可轻易扩展的框架。基于重建的方法如RD4AD和UniAD具有非常好的效果与较好的可扩展性。RD4AD提出了预训练教师模型与学生模型并在多尺度特征层面上进行异常值对比。尽管基于CNN的RD4AD在多尺度上的有着捕捉相邻的信息能力并取得了较好的性能,但是其无法建立长距离的相关性。首个多类异常检测算法UniAD是基于预训练的CNN编码器和Transformer解码器的架构。尽管Transformer有着全局建模的能力,但是由于其平方级别计算复杂度,UniAD仅在最小尺度的特征图上对比得到异常图,这无疑会减弱模型性能。
2、MambaAD贡献
最近,Mamba在大语言模型中取得了出色的表现,有着远小于transformer的线性复杂度并且能够与transformer相媲美的效果。近期大量的工作将Mamba引入视觉领域,涌现大量基于Mamba的研究工作。本文首次将Mamba引入异常检测领域,构建了MambaAD架构有着全局+局部的建模能力,并且利用它的线性复杂度在多尺度上计算异常图并且有着较低的参数量与计算复杂度。具体来说MambaAD使用金字塔结构的自编码器结构来重建多尺度的特征,通过一个预训练编码器和提出的基于Mamba结构的解码器。其中基于Mamba结构的解码器由不同尺度与数量的局部增强状态空间(Locality-Enhanced State Space) LSS module组成。LSS module由两部分组成:连续的(Hybrid State Space) HSS模块用于捕捉全局的信息和并联多核的卷积操作用来建立局部的联系。最终的输出特征不仅包含基于Mamba的全局建模能力,还包含了基于CNN的建立局部相关性的能力。所提出的HSS模块探索了5种不同的扫描方式和8种不同的扫描方向,其中的HS编码器和解码器分别将特征图编码至不同的扫描方式和方向的序列并解码。HSS模块能够增强在多个方向上的全局感受野并且所使用的Hilbert的扫描方式也更加适用于工业产品位置集中在中心区域的特点。通过在不同尺度的特征图上计算异常图并相加,MambaAD在6个不同的异常检测数据集上取得了SoTA性能并且模型参数量与计算复杂度也非常低。具体来说,我们的贡献如下:
1)我们提出了MambaAD首次使用Mamba来解决多类无监督异常检测任务,它能够在很少的模型参数量和计算复杂度上进行多尺度训练与推理。
2)我们设计了一个LSS module,连接的HSS模块和并行的多内核卷积分别提取全局特征相关性与局部信息关联性,实现全局加局部的统一建模。
3)我们探索了HSS模块即5种方法8种多方向的混合扫描方式如图3所示,能够增强复杂的异常检测图像在不同类别不同形态下的全局建模能力。
4)我们证明了MambaAD在多类异常检测任务的优越性和高效性。在6个不同的异常检测数据集上达到SOTA如表1所示,并且有着非常低的模型参数与计算复杂度见表2。
公众号后台回复“数据集”获取100+深度学习各方向资源整理
极市干货
点击阅读原文进入CV社区
收获更多技术干货