lewandofskee / MambaAD

[NeurIPS 2024] Official implementation of MambaAD: Exploring State Space Models for Multi-class Unsupervised Anomaly Detection.
123 stars 3 forks source link

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

[Haoyang He1*](https://scholar.google.com/citations?hl=zh-CN&user=8NfQv1sAAAAJ), [Yuhu Bai1*](https://scholar.google.com/citations?hl=zh-CN&user=ucCvgooAAAAJ), [Jiangning Zhang2](https://zhangzjn.github.io), [Qingdong He2](https://scholar.google.com/citations?hl=zh-CN&user=gUJWww0AAAAJ), [Hongxu Chen1](https://scholar.google.com/citations?hl=zh-CN&user=uFT3YfMAAAAJ) [Zhenye Gan2](https://scholar.google.com/citations?user=fa4NkScAAAAJ&hl=zh-CN), [Chengjie Wang2](https://scholar.google.com/citations?hl=zh-CN&user=fqte5H4AAAAJ), [Xiangtai Li3](https://lxtgh.github.io/), [Guanzhong Tian1](https://scholar.google.com/citations?hl=zh-CN&user=0q-7PI4AAAAJ), [Lei Xie1](https://scholar.google.com/citations?hl=zh-CN&user=7ZZ_-m0AAAAJ) 1College of Control Science and Engineering, Zhejiang University, 2Youtu Lab, Tencent, 3Nanyang Technological University, Singapore [[`Paper`](https://arxiv.org/pdf/2404.06564.pdf)] [[`Project Page`](https://lewandofskee.github.io/projects/MambaAD/)] Our MambaAD is based on [ADer](https://github.com/zhangzjn/ADer).

Congratulations! Our MambaAD has been accepted at the NeurIPS 2024 conference!

Abstract

Recent advancements in anomaly detection have seen the efficacy of CNN- and transformer-based approaches. However, CNNs struggle with long-range dependencies, while transformers are burdened by quadratic computational complexity. Mamba-based models, with their superior long-range modeling and linear efficiency, have garnered substantial attention. This study pioneers the application of Mamba to multi-class unsupervised anomaly detection, presenting MambaAD, which consists of a pre-trained encoder and a Mamba decoder featuring Locality-Enhanced State Space (LSS) modules at multi-scales. The proposed LSS module, integrating parallel cascaded (Hybrid State Space) HSS blocks and multi-kernel convolutions operations, effectively captures both long-range and local information. The HSS block, utilizing (Hybrid Scanning) HS encoders, encodes feature maps into five scanning methods and eight directions, thereby strengthening global connections through the (State Space Model) SSM. The use of Hilbert scanning and eight directions significantly improves feature sequence modeling. Comprehensive experiments on six diverse anomaly detection datasets and seven metrics demonstrate SoTA performance, substantiating the method's effectiveness.

Overview

accuracy


🛠️ Getting Started

Installation

📜 Multi-class Results on Popular AD Datasets

Subscripts I, R, and P represent image-level, region-level, and pixel-level, respectively.

MambaAD Results

Method mAU-ROCI mAPI mF1-maxI mAU-ROCP mAPP mF1-maxP mAU-PROR Download
MVTec-AD 98.6 99.6 97.8 97.7 56.3 59.2 93.1 log & weight
VisA 94.3 94.5 89.4 98.5 39.4 44.0 91.0 log & weight
Real-IAD 86.3 84.6 77.0 98.5 33.0 38.7 90.5 log & weight
Uni-Medical 83.7 80.1 82.0 96.9 45.4 47.3 87.5 log & weight
COCO-AD 63.9 56.2 63.2 69.3 16.9 22.2 40.5 log & weight
MVTec-3D 86.2 95.8 92.8 98.6 37.5 41.1 93.6 log & weight

Citation

If you find this code useful, don't forget to star the repo and cite the paper:

@article{he2024mambaad,
      title={MambaAD: Exploring State Space Models for Multi-class Unsupervised Anomaly Detection}, 
      author={Haoyang He and Yuhu Bai and Jiangning Zhang and Qingdong He and Hongxu Chen and Zhenye Gan and Chengjie Wang and Xiangtai Li and Guanzhong Tian and Lei Xie},
      journal={arXiv preprint arXiv:2404.06564},
      year={2024},
}

Acknowledgements

We thank the great works ADer, VMamba for providing assistance for our research.