2024顶会神器！Mamba最新综述！

标题：A Survey of Mamba

作者：Haohao Qu, Liangbo Ning, Rui An, Wenqi Fan, Tyler Derr, Xin Xu, Qing Li

机构：香港理工大学、范德比尔特大学

原文链接：https://arxiv.org/abs/2408.01129

1. 摘要

深度学习作为一项关键技术，在人工智能（AI）领域引发了显著变革，极大地改变了人类的生活方式。作为最具代表性的架构，Transformer赋能了众多先进模型，尤其是包含数十亿参数的大语言模型（LLMs），成为深度学习的基石。尽管取得了令人瞩目的成就，但Transformer仍面临固有局限性，特别是注意力计算的二次计算复杂度导致的耗时推理问题。最近，一种受经典状态空间模型启发的新型架构Mamba应运而生，成为构建基础模型的有前景的替代方案。Mamba在保持序列长度方面近乎线性的可扩展性的同时，提供了与Transformer相当的建模能力。这引发了越来越多的研究，积极探索Mamba在不同领域实现卓越性能的潜力。鉴于这种快速发展，迫切需要一篇系统综述来整合现有的基于Mamba的模型，以便全面理解这一新兴模型架构。因此，在本次综述中，我们对最近的与 Mamba 相关的研究进行了深入调查，涵盖三个方面：基于 Mamba 的模型进展、使 Mamba 适应不同数据的技术，以及 Mamba 可以出类拔萃的应用领域。具体而言，我们首先回顾了各种具有代表性的深度学习模型的基础知识以及 Mamba-1 和 Mamba-2 的细节作为初步介绍。然后，为了展示 Mamba 对人工智能的重要性，我们全面回顾了侧重于 Mamba 模型架构设计、数据适应性和应用的相关研究。最后，我们讨论了当前的局限性，并探讨了各种有前景的研究方向，为未来研究提供更深入的见解。推荐学习： 国内首个面向具身智能方向的理论与实战课程

2. 引言

在过去的二十年里，深度学习作为最突出的人工智能（AI）技术，在医疗保健、自动驾驶系统、推荐系统和金融服务等多个领域引发了革命性变革。这一时期见证了众多深度神经网络（DNNs）的涌现，这些网络极大地改变了人类的生活方式，为个人带来了巨大便利。其中，一个显著的例子是U-Net，这是视觉领域中的一种鲁棒的深度学习模型，广泛应用于医学成像领域，以检查MRI和CT扫描等放射学图像。其应用有助于疾病的识别和诊断，展示了在关键医疗保健领域的有效性。此外，图神经网络（GNNs）被用于处理图结构数据，以支持智能服务，如推荐系统，该系统可向用户推荐个性化的内容、产品或服务。另外，循环神经网络（RNNs）因其能够捕获对准确翻译至关重要的序列和上下文信息而被广泛用于机器翻译，从而使来自不同语言背景的人们能够有效地交流和理解彼此的想法、观点和信息。

在众多的深度学习架构中，Transformer最近脱颖而出，并在广泛的应用领域中确立了其主导地位。例如，作为最具代表性的大型基础模型，ChatGPT 等大型语言模型（LLMs）从根本上说是基于 Transformer 架构构建的。通过将其模型规模扩展到数十亿级别，并在各种数据源的混合上进行训练，这些基于 Transformer 的模型在语言理解、常识推理和内容内学习方面展现出了令人印象深刻的能力，达到了人类级别的智能。这一显著的成功得益于注意力机制，它使基于 Transformer 的模型能够专注于输入序列的相关部分，并促进更好的上下文理解。然而，注意力机制也引入了显著的计算开销，该开销随输入大小的增加而呈二次方增长，这在处理长输入时带来了挑战。例如，计算成本的快速增长使得 Transformer 在处理大量序列时变得不切实际或不可行，从而限制了其在文档级机器翻译或长文档摘要等任务中的应用。

最近，一种有前景的架构——结构化状态空间序列模型（SSMs）已经出现，它能够高效地捕获序列数据中的复杂依赖关系，成为 Transformer 的强大竞争对手。这些模型受到经典状态空间模型的启发，可以看作是循环神经网络和卷积神经网络的融合。它们可以通过递归或卷积操作高效地进行计算，实现与序列长度呈线性或接近线性的扩展，从而显著降低计算成本。更具体地说，作为最成功的SSM变体之一，Mamba在保持与序列长度呈线性扩展性的同时，实现了与Transformer相当的建模能力，从而成为研究焦点。Mamba首先引入了一个简单而有效的选择机制，该机制允许模型通过基于输入的SSM参数化来过滤掉不相关信息，同时无限期地保留必要且相关的数据。然后，Mamba提出了一种硬件感知算法，该算法通过扫描而不是卷积来递归地计算模型，从而在A100 GPU上实现了高达3倍的计算加速。如图1所示，Mamba对复杂且冗长的序列数据强大的建模能力以及接近线性的扩展性，使其作为一种新兴的基础模型，有望革新计算机视觉、自然语言处理、医疗保健等多个领域的研究和应用。例如，Zhu等人提出的Vim模型在提取高分辨率图像特征时比DeiT快2.8倍，并节省了86.8%的GPU内存。Dao和Gu揭示了SSM与注意力变体之间的联系，并提出了一种新的架构来优化选择性SSM，在语言建模任务上实现了2-8倍的速度提升。

受到 Mamba 强大的长序列建模能力和其卓越效率的启发，大量文献涌现出来，这些文献主要关注于在各种下游任务中运用和改进 Mamba。鉴于与 Mamba 相关研究的显著增加，对现有文献进行全面回顾并探讨未来研究的潜在方向至关重要。因此，在本次调查中，我们从多个角度对 Mamba 进行了全面回顾，旨在为初学者提供对 Mamba 内部工作原理的基本理解，同时帮助经验丰富的从业者了解其最新发展。具体来说，本次调查的组织结构如下：第 2 节回顾了包括 RNNs、Transformers 和状态空间模型在内的各种代表性深度神经网络的背景知识，而 Mamba 的详细情况将在第 3 节中介绍。随后，我们在第 4 节中从模块设计、扫描模式和内存管理的角度总结了基于 Mamba 的研究的最新进展。然后，第 5 节介绍了将 Mamba 应用于各种数据（包括顺序数据和非顺序数据）的技术。此外，第 6 节介绍了 Mamba 模型的代表性应用，而第 7 节则提出了面临的挑战和未来方向。最后，我们在第 8 节中总结了整个调查。

与我们的调查同时，还发布了几个相关的调查，这些调查纯粹关注状态空间模型和 Vision Mamba。与这些调查不同，本文以与 Mamba 相关的研究为中心。它从一个新颖的视角系统地分析了现有文献，以探索 Mamba 架构的演变以及基于 Mamba 的模型中所使用的数据适应方法。

3. Mamba

为了克服传统状态空间模型（SSM）在上下文感知能力方面的不足，文献提出了Mamba作为一种潜在的替代方案，有望成为一般序列基础模型的骨干。最近，Mamba-2提出了结构化空间状态对偶性（SSD），该理论框架稳健地连接了结构化SSM和各种形式的注意力机制，使我们能够将原本为Transformer开发的算法和系统优化转移到SSM上。在本节中，我们将对Mamba和Mamba-2进行简洁明了的介绍。

传统的SSM在文本和其他信息密集型数据的建模方面表现出有限的有效性，阻碍了它们在深度学习领域的进步。为了赋予SSM以Transformer的建模能力，Gu和Dao基于结构化状态空间模型提出了三种创新技术，即基于高阶多项式投影算子（HiPPO）的记忆初始化、选择机制和硬件感知计算，如图3所示。这些技术旨在增强SSM在长距离线性时间序列建模方面的能力。特别是，初始化策略建立了一个连贯的隐藏状态矩阵，有效促进了长距离记忆。然后，选择机制使SSM能够获取内容感知的表示。最后，Mamba设计了两种硬件感知计算算法，即并行关联扫描和内存重新计算，以提高训练效率。

4. Mamba模型的进步

状态空间模型（State Space Models）和曼巴（Mamba）模型最近已被研究并有望成为基础模型骨干的替代方案。尽管曼巴模型在自然语言处理方面表现出色，但它仍面临一些挑战，如记忆丧失、跨多样化任务的泛化能力不足，以及相较于基于Transformer的语言模型，捕捉复杂模式的能力较差。为了克服这些挑战，已经投入了大量努力来改进曼巴架构。现有研究主要集中在修改块设计、扫描模式和内存管理三个方面。本节将从这三个方面介绍几项关键技术，并在表1中总结了相关研究。