MambaOut：我们真的需要为视觉任务引入 Mamba 吗？
计算机视觉与模式识别
2024-05-21 v3 人工智能
机器学习
摘要
Mamba 是一种具有状态空间模型（SSM）类循环神经网络的 token mixer 架构，最近被引入以解决注意力机制的二次复杂度问题，随后被应用于视觉任务。然而，与卷积和注意力基模型相比，Mamba 在视觉任务中的性能往往不佳。本文深入探讨了 Mamba 的本质，概念性地得出结论：Mamba 特别适合处理长序列和自回归特征的任务。对于视觉任务而言，图像分类不符合上述两种特征之一，我们假设 Mamba 对此任务并非必需；检测和分割任务同样不具备自回归特征，但符合长序列特征，因此我们认为仍值得探索 Mamba 在这些任务中的潜力。为经验验证我们的假设，我们构建了一系列称为 MambaOut 的模型，通过堆叠 Mamba 块但移除其核心 token mixer（SSM）来实现。实验结果强有力地支持了我们的假设。具体而言，我们的 MambaOut 模型在 ImageNet 图像分类中超越了所有视觉 Mamba 模型，表明 Mamba 确实不必要于该任务。至于检测和分割任务，MambaOut 无法匹配最新视觉 Mamba 模型的性能，这表明 Mamba 在长序列视觉任务中具有潜力。代码已公开于 https://github.com/yuweihao/MambaOut
引用
@article{arxiv.2405.07992,
title = {MambaOut: Do We Really Need Mamba for Vision?},
author = {Weihao Yu and Xinchao Wang},
journal= {arXiv preprint arXiv:2405.07992},
year = {2024}
}
备注
Code: https://github.com/yuweihao/MambaOut