点击下方卡片,关注 「AI视界引擎」 公众号
状态空间模型(SSM)是一种用于描述和分析动态系统行为的数学模型。该模型在控制理论、信号处理、经济学和机器学习等多个领域得到了广泛应用。在深度学习领域,状态空间模型被用于处理序列数据,如时间序列分析、自然语言处理(NLP)和视频理解。
通过将序列数据映射到状态空间,可以更好地捕捉数据中的长期依赖关系。特别是现代SSM在NLP中表现出强大的表示能力,尤其是在长序列建模中,同时保持了线性时间复杂度。
值得注意的是,基于最新的状态空间模型,Mamba 将时间变化参数融合到SSM中,并制定了一种硬件感知算法以实现高效的训练和推理。
鉴于其令人印象深刻的效率和强大的长距离依赖建模能力,Mamba有望成为一种可能超越Transformer的新AI架构。
最近,一些研究试图通过将Mamba从自然语言领域扩展到视觉领域,探索Mamba在通用视觉、多模态、医学图像分析和遥感图像分析等领域的潜力。
为了全面理解Mamba在视觉领域的情况,作者进行了全面的调查并提出了一个分类学研究。
本调查关注Mamba在各种视觉任务和数据类型上的应用,并讨论了其前身、近期进展以及对广泛领域的深远影响。
由于Mamba目前正处于上升趋势,
Mamba 项目:https://github.com/lx6c78/Vision-Mamba-A-Comprehensive-Survey-and-Taxonomy。
1 Introduction
自从深度学习出现以来,卷积神经网络(CNNs)和Transformer在各个视觉领域任务中变得普遍。CNNs因其简单的架构和可扩展性而受到欢迎。
然而,Vision Transformer 的引入打破了这一局面。这归功于其自注意力机制,它能够实现全局感知场和动态权重。基础模型的出现近年来为AI带来了变革和机遇。
特别是,基于自注意力机制Transformer架构的基础模型具有广泛的应用范围和令人兴奋的性能。然而,CNNs在处理长距离依赖方面存在困难,而Transformer则承受着次要的计算复杂性的负担。
最近,状态空间模型(SSM)在捕获语言序列的动态和依赖关系方面展示了显著的有效性。[30]引入了一种结构化状态空间序列模型(S4),专门设计用于以线性复杂度建模远程依赖,为自然语言理解的进一步发展带来了新的动力。名为Mamba 的新变种的的成功已被证明能够在实际数据上达到更高的性能,序列长度可达百万 Level 。这使其成为近期热门话题。
Mamba结合了选择性扫描(S6)并享受快速推理,具有与序列长度线性扩展的优点。
其计算优势比Transformer快五倍以上。
随着自然语言和计算机视觉逐渐形成融合趋势,将性能卓越的大型语言模型(LLM)Mamba扩展到计算机视觉是一个吸引人的方向。Vim [136] 是第一个基于纯SSM的模型,用于处理密集预测任务,并且是SSM首次应用于视觉通用 Backbone 网络。
该框架解决了Mamba在处理图像序列时遇到的两个挑战:建模的单向性和位置感知的缺乏。
VMamba 提出了一种跨扫描策略,以弥合1D数组与2D序列扫描之间的差距。Mamba-ND [54] 旨在将Mamba架构扩展到多维数据上的各种视觉任务,包括1D、2D和3D数据,以及更多其他视觉任务。
这些工作探索了Mamba的潜力,并在视觉领域显示出希望。
在Mamba架构的各个方面已经进行了大量工作,以探索SSM在垂直领域视觉任务中的应用。
特别是由于Mamba在计算效率和长距离依赖建模能力方面的卓越表现,许多针对医学图像分析的垂直领域Mamba迅速涌现,以及渐进学习,医学图像分割,U-Mamba 首次将U-net与Mamba结合作为混合CNN-SSM模型处理高分辨率医学图像数据。
这种方法在效率和性能方面优于最先进的CNN和基于Transformer的分割网络。此后,针对医学图像分析提出了一系列Mamba变体。此外,社区还提出了一些扩展,以探索Mamba在遥感图像处理、分析和理解方面的能力边界。
作者的调查与现有的关于Mamba的调查不同,展示了其独特的优势。
首先,作者采用清晰且精确的Mamba变体分类法,系统地组织和总结了Mamba在计算机视觉及其垂直领域的发展和应用程序。
与当前的评论相比,作者的分类法更具针对性和读者友好性,便于研究行人更容易理解和导航相关内容。其次,作者为每个子类别提供了深刻且合理的分类法,确保全面覆盖不同的垂直领域和问题。
这种分类法不仅加深了读者对每个类别特征和应用场景的理解,而且帮助研究行人在特定领域中找到所需的信息。第三,作者确保对每种方法的原理和技术细节进行详尽的解释,旨在加深读者对其意图的理解。
作者的评论文章不仅提供了概念解释,还提供了具体实例和实验结果,以帮助读者理解和应用这些方法。此外,本文还包括了两部分内容,分别是“数据类型”和“亮点”,它们提供了对各个类别方法补充性的分类和解释。这些补充性内容为研究行人提供了更多参考和选择,帮助他们快速定位并理解与Mamba建模相关的最新进展。
2 Formulation of Mamba
状态空间模型(SSMs)[30, 31, 34, 57, 78, 28]利用中间状态变量来实现序列到序列的映射,这使得能够处理长序列。基于结构化状态空间的模型,如S4 [30],使用低秩修正来调节模型的某些参数,使得能够稳定对角化,并将SSMs简化为已充分研究的Cauchy核。这种方法解决了先前SSM模型计算和内存需求过大的问题。S4相较于先前的序列建模方法是一个显著的改进。Mamba [28]通过引入选择机制和硬件感知算法,根据输入序列参数化SSMs,从而改进了SSMs。这解决了离散模态问题,在处理语言和基因组学中的长序列时,实现了更直接的加工处理。
State Space Models
SSMs的主要表示是连续时间表示,如图1所示,它由四个参数(,,,)定义。这种表示将一组随时间变化的输入通过一个隐藏状态转换为一系列输出。映射过程可以表达为以下形式:
具体来说,等式1中的第一个方程,即状态方程,通过将矩阵与输入相乘,然后是矩阵与先前的状态相乘。参数是演化参数,它存储了由系数矩阵表示的所有先前的历史信息,并确定了先前隐藏状态对用于更新下一个时刻隐藏状态空间状态的隐藏状态的影响程度。
投影参数矩阵决定了输入对隐藏状态的影响程度。等式1中的第二个方程,即输出方程,描述了隐藏状态如何通过投影参数矩阵转换为输出,以及输入如何通过矩阵影响输出。通常,参数提供了从输入到输出的直接信号,也称为跳过连接,可以通过省略来进行简化,即假设。
如图1所示,S4受到用于1D序列的连续系统的启发,并使用采样时间尺度参数将连续SSM转换为离散SSM。因此,从连续信号输入到输出的映射被转换为序列到序列的映射:,这有助于处理离散标记输入,如图像和文本。
这个序列到序列过程有两个阶段。第一阶段是离散化,它为启发式门控机制提供了原则性的基础。连续参数通过采样时间尺度参数转换为离散参数。这种转换通常使用零阶保持(ZOH)来完成,得到离散公式:
在第一阶段,模型通过状态空间表示来编码输入序列:
其中 , 和 。
第二阶段是计算阶段,模型通过全局卷积来计算输出:
其中 是输入序列 的长度,而 是一个结构化的卷积核。
由于当输入为常数时,三个离散参数 , 和 是常数,卷积核 是静态的,因此避免了重新计算输入到状态和状态到输出的过程,并且可以像卷积神经网络一样并行训练以实现高效计算。为了充分利用卷积和循环模式的优势,SSMs 在训练阶段使用卷积模式进行并行训练,并在推理阶段使用循环模式进行有效的自回归推理。这是因为生成下一个时间步的输出只需要当前时间步的状态,而不是整个输入历史。然而,在循环模式下,循环神经网络(RNNs)倾向于随着时间的推移而忘记某些信息。
为了解决这个问题,线性状态空间层(LSSL)[31] 使用 HiPPO 连续时间记忆理论[29],试图将所有当前观察到的输入信号压缩成一个系数向量。它使用矩阵 构建一个状态表示,该状态表示能够很好地捕获最近的标记并衰减较旧的标记。由 HiPPO 构建的 矩阵包含 个多项式压缩器,用于压缩连续增长的历史数据,并在训练过程中调整系数的具体值。然而,连续增加的顺序导致了维度的爆炸,因此作者将来自方程式 1 的状态方程得到的隐藏状态通过线性组合转换为输出,即使用输出方程的投影参数矩阵 。
此外,为了解决由于与矩阵 的重复矩阵乘法导致的离散时间SSM的复杂性问题,S4提出了一种实用的解决方案。它通过利用一个结构化结果来增强矩阵 ,从而简化了SSMs。具体来说,作者采用了正常加低秩(NPLR)方法对HiPPO矩阵进行分解。这种方法允许稳定的对角化,并显著简化了Cauchy核的计算,使得该方法更具实用性和在现实世界场景中的应用性。
Selective State Space Models (S6)
- 选择性机制。 由于先前的线性时不变状态空间模型缺乏内容意识,无法基于输入内容实现选择性任务,这导致S4对所有的标记给予相同的关注。然而,在现实中,标记的重要性是不同的,其重要性程度会随着训练过程的动态变化。因此,在重要内容上投入更多努力,并动态调整重要性 Level 以匹配复杂的输入内容更为有效。
基于上述考虑,Mamba将选择性机制合并到状态空间模型中,以获得选择性状态空间模型(S6)。具体来说,为了操作具有批大小,长度和通道的输入序列,Mamba旨在为每个通道独立应用SSM。在Mamba [28]中,矩阵,矩阵和S4中的都是输入的函数,允许模型根据输入自适应地调整其行为。
将选择机制融入模型后的离散化过程如下:
其中 , 和 。 和 是将输入投影到维空间的线性函数,而 将隐藏状态维度线性地投影到所需的维度,与RNN门控机制相连。通过上述计算,参数 ,, 成为了长度为的输入函数,将时不变模型转换为时变模型,从而实现选择性。
的大小已从更改为,这意味着对于批次中的每个标记(总共有个),都有一个唯一的用于输入数据依赖性,以及更细粒度的控制功能。的步长越大,模型越关注输入,而不是存储的状态。相反,步长越小,模型将越忽略特定的输入,从而更多地关注存储的状态。
参数和成为S4上的依赖于输入数据的函数,从而允许更精细地控制输入是否到达状态,或者状态是否到达输出。参数并未成为数据依赖的,但在SSM的离散化操作之后,可以通过的数据依赖性使与输入相关。同时,由于参数的维度为,它在每个SSM维度中都有不同的角色,从而实现对之前所有内容的精确泛化,而不仅仅是简单的压缩。
- 硬件感知状态扩展。 选择性机制克服了先前线性时不变状态空间模型的限制。然而,时间变化的特性带来了计算上的挑战。在Mamba中,由于此时的输入和输出不是简单的静态映射关系,不能使用固定卷积核的高效卷积计算。因此,研究行人设计了一种循环模式下的硬件感知并行算法。
模型通过扫描而不是卷积来计算。Mamba的扫描算法避免了RNN无法并行执行的陷阱。由于在扫描操作中需要先前的状态来计算新状态,直接使用循环计算不能实现并行扫描。研究行人发现,每个状态实际上是之前状态的整个压缩,即可以直接使用之前的状态来计算新状态。那么可以假设操作执行的顺序与关联属性无关。因此,Mamba通过分段计算序列并迭代结合输入数据依赖的参数来实现选择扫描算法。
另一方面,由于GPU具有许多处理器并且能够执行高度并行的计算,Mamba利用GPU的HBM和快速的SRAM IO,通过使用 Kernel 融合来避免频繁将SRAM写入HBM。具体来说,Mamba在更高速度的SRAM内存中执行离散化和递归操作,然后将输出写回HBM。当从HBM将输入加载到SRAM时,中间状态不被保留,而是在反向传播中重新计算。
如图2所示,Mamba将SSM的基本块与当今神经网络中普遍使用的MLP块相结合,形成了一个新的Mamba块,这个块被堆叠并与归一化和残差连接结合,形成了Mamba网络架构。
Discussion and Summary
选择机制使Mamba具有线性计算复杂度和长距离依赖建模能力,而硬件感知的状态扩展使其在内存使用上高效。凭借这两项关键技术,Mamba在各种应用中显示出超越先前状态空间模型的巨大潜力。
3 Mamba in General Vision Tasks
本节回顾了Mamba及其变体在通用视觉中的应用,包括高级/中级视觉、低级视觉和3D视觉。在以下小节中,作者将介绍为每个任务重新设计的Mamba变体。考虑到在视觉任务中选择性扫描策略的重要性,作者在图3中进一步总结了现有的经典2D扫描机制。
Vision Backbone with Mamba
语言建模领域的成功激励研究行人设计通用且高效的视觉基础架构,使用先进的Mamba模型。由于视觉数据的非因果性质,直接将Mamba应用于图像块和平面图像将不可避免地导致接受域受限,因为无法估计与未扫描块的关系。V Mamba [68] 将这个问题称为方向敏感问题。在这种情况下,方法需要考虑数据的 spatial 结构和全局相关性,而不仅仅是数据的顺序。如图3(b)所示,设计了一个跨扫描模块(CSM)来遍历空间域,并使用四向扫描策略将任何非因果视觉图像转换为一系列顺序块。具体来说,图像块是
图4:视觉状态空间(VSS)块(图像来自[68])。
Mamba作为一个基础架构,也可以服务于其他方法,如多模态信息处理、扩散模型、域泛化、视觉-语言模型等。遵循Mamba原则,如何与这些方法协同工作,仍需要进一步探索。
参考
[1].Vision Mamba:.
点击上方卡片,关注 「AI视界引擎」 公众号