选择性结构化状态空间模型综述

摘要 基础模型正为深度学习中大多数令人兴奋的应用程序提供源动力,但它们几乎普遍基于注意力架构及其核心的注意力模块。由于注意力架构的结构特点,其在长序列上的计算效率表现糟糕。因此,许多次二次时间架构,如线性注意力、门控卷积和循环模型,以及结构化状态空间模型已经被开发出来,以解决上述问题,但它们在语言处理等重要模态上的表现不如注意力架构,这些模型的一个关键弱点是它们无法执行基于内容的推理。而最新提出的选择性结构化状态空间模型对上述模型做了一些改进。本综述总结了现有的几种基础模型的优缺点,总结了SSM模型的发展历程,描述了一种新式模型,即选择性结构化状态空间模型的工作原理及其优势和缺陷。并且介绍了选择性结构化状态空间模型的应用与实验,对其的未来发展进行了展望。

关键词 SSM; 序列模型;基础模型;深度学习

1 介绍(Introduction

基础模型(Foundation models),即在海量数据上进行预训练,然后针对下游任务进行调整的大型模型。它们已成为现代机器学习的有效范式。

这些基础模型的骨干网络通常是序列模型(即模型捕捉数据中的时序关系或顺序依赖性),操作来自各种领域的任意序列输入,如语言、图像、语音、音频、时间序列和基因组学。由于先前的主流基础模型,如CNN,RNN等,都存在无法捕捉长程序列依赖关系的问题,因而现代的基础模型主要基于单一类型的序列模型:Transformer这种自注意力结构的效果主要归功于它能够在上下文窗口中密集地传递信息,使其能够对复杂数据进行建模。然而,这一特性也带来了一些根本性的缺陷:无法对有限序列外的任何内容进行建模,同时存在序列长度的二次缩放问题。为尝试解决现有基础模型架构中的诸类问题,提出了状态空间模型(State Space Models,SSM),并由此演化出了选择性结构化状态空间模型(Selective Structured State Space Models,S6,or Mamba)

为了深入了解状态空间模型的发展现状,尤其是选择性结构化状态空间模型的发展现状,本综述进行了深入地调研。调研按照时间发展顺序,从三方面展开。

在第一部分,本文对比了过去三种主流的深度学习基础模型CNN,RNN,Transformer的优势与劣势,并由这三种基础模型的劣势,提出了为什么要提出状态空间模型这种新的模型结构。

在第二部分,本文介绍了状态空间模型的基本原理,并按照时间发展顺序,介绍了由状态空间模型发展出的结构化状态空间模型和选择性结构化状态空间模型,讨论了他们相对先前的版本有了哪些进步,以及相较于之前提到的基础模型有哪些优势和缺点。

在第三部分,本文介绍了选择性结构化状态空间模型的应用,介绍了由选择性结构化状态空间模型发展出的一些主干网络,以及选择性结构化状态空间模型在各个深度学习下游领域的应用。

最后,本文对选择性结构化状态空间模型进行了总结,并尝试探讨选择性结构化状态空间模型在深度学习领域应用的未来前景。

2 过去主流深度学习基础模型的对比

2.1 循环神经网络RNN

循环神经网络,源自于1982年由Saratha Sathasivam 提出的霍普菲尔德网络,是指在全连接神经网络的基础上增加了前后时序上的关系,每一时刻的输出不仅取决于当前时刻的输入,还与上一时刻的输入相关。由于增加了上下文的联系,相较于全连接神经网络,循环神经网络可以更好地处理比如机器翻译等的与序列相关的问题。 

图1 RNN网络架构示意图

循环神经网络的问题在于:

(1)串行化训练,速度缓慢:它的网络结构决定了其是串行化的模型,由于在训练中的反向传播时需要被迫地逐个通过各个单元,以进行参数的更新,导致了其较为缓慢的训练速度(当然在推理中,此问题会被避免)。

(2)“遗忘”的存在:循环神经网络也存在着“遗忘”的问题,由于随着时间的推移,就像人类的记忆力,RNN会出现梯度消失或梯度爆炸问题,所以其并没有办法有效处理长程的依赖关系,即其面对非常长的序列时依旧无能为力。

2.2 卷积神经网络CNN

卷积神经网络利用卷积层可以高效处理输入数据,相较RNN网络,其可以进行并行训练,可以一定程度上提高训练速度,但缺少了推理方面的优化,固定的卷积核限制了推理的速度,故在速度方面较为平庸。

同时普遍来看,卷积神经网络也存在着这样的问题:
即主要关注局部,容易忽略全局信息:特征卷积核在一次卷积时只对输入数据的一部分进行运算,这使得卷积神经网络的结构主要关注局部特征,故容易忽视全局特征,故其也存在面对长程序列时的遗忘特性。

图2 CNN网络架构示意图

2.3 注意力架构(Transformer)

注意力架构相较于前两种网络,其输入与序列中每一元素都有关,同时可以进行并行训练,能够良好地捕捉长程依赖关系,这时其相较RNNCNN的显著优势。

图3 注意力网络架构示意图

但其缺点也很显而易见:

1无法对有限序列外的任何内容进行建模:自注意力机制通常在一个固定的上下文序列(如一个句子或一个图像的一部分)内操作,这意味着它只能考虑并处理这个序列内的信息。当需要理解或分析的信息分散在更大的上下文中,或者与当前序列相隔较远时,自注意力可能无法有效捕捉这些信息,因为自注意力缺乏直接的机制来处理或引入窗口之外的信息。这在处理长篇文本或大规模图像时尤其成问题,因为重要信息可能分布在整个文档或图像中。

2)序列长度的二次缩放问题:自注意力的一个关键特点是它在计算时需要比较和加权序列内的所有元素对。这意味着如果序列长度增加,所需的比较(和因此的计算资源)会以二次方的方式增加。即计算的复杂度是序列长度的平方。例如,如果序列长度加倍,那么所需的计算量将增加四倍。这种二次缩放性质使得自注意力在处理大序列或长序列时变得效率低下。

基于上述各个基础模型的优劣,状态空间模型(SSM)被提出,以尝试解决上述各模型的问题。

3 状态空间模型(SSM)的发展历程

3.1 状态空间模型的基本原理

状态空间模型(SSM)的功能类似于TransformerRNN,用于处理时间序列,如文本,信号等。它在神经网络中的实现思路来源于现代控制理论,因为在神经网络方面,系统的“状态”通常是其隐藏状态(hidden state),故在连续函数情况下,如下图所示的结构可以表示为:

(其中h为隐函数,x为输入,y为输出)

图4  SSM网络架构示意图

如果输入一个连续的信号,找到状态表示h(t)在分析上是具有挑战的。此外,由于通常输入是离散信息(比如文本),因此希望有一个离散的模型。在此情况下,采用零阶保持技术(Zero-Order Held technique,ZOH),上述公式可使用固定公式:

(

表示离散化的时间间隔)

推导并离散化表示为:

(

)

对于每一个序列

,可以由上述公式递推得出如下的表达式:

从上述公式不难看出,SSM架构具有完全递推属性,

是可以由上述公式层层递推得到;同时,由于

可以视为卷积核

,而此卷积核是通过固定数值的矩阵得到的,只要能确定

,就可以进行并行运算。不难看出SSM架构同时具有RNN和CNN的特性。

图7  SSM网络架构循环表示示意图

图6  SSM网络架构卷积表示示意图

由上,有了这些表示,可以使用一个巧妙的技巧——根据任务选择表示。在训练期间,使用并行化的卷积表示,在推理阶段,使用高效的循环表示:SSM架构的优势主要有两点:

  1. 训练时选择CNN表示,并行特点:训练时可以用卷积表示实现并行计算,提高训练速度。
  2. 推理时选择RNN表示,拥有的逐个快速生成输出的特性。

这种混合表示被称为线性状态空间层(Linear State-Space Layer, LSSL),它有一个重要的属性,即线性时不变性(LTI)。这里LTI表示SSM的参数

对于所有时间步长而言都是固定的。这意味着在推理时,无论向SSM提供什么序列,

的值会保持不变。这种情况就是说,它不具备内容感知能力。这便是最起初的SSM的缺陷——无法根据内容推理。

3.2 状态空间模型的的演化和发展

3.2.1 结构化状态空间模型(Structured State Space Models,S4)

对比上文中提到的RNN,CNN,Transformer等基础模型,SSM可以说集合了RNN和CNN的优势,但相较Transformer架构,SSM在对长序列建立长程的依赖关系方面仍然有所欠缺。同时,人们希望解决能够通过SSM架构解决Transformer中存在的“无法对有限序列外的任何内容进行建模”的问题

但上文提到的SSM架构显然不具备这个能力,因为矩阵

可能是仅记住了少量先前的输入与当前输入之间的关系,特别是在循环表示的上下文中,它只回顾先前的状态。

那么不难看出,若想提高SSM架构的“记忆”能力,矩阵

尤为关键。它占据着相当重要的地位,因为矩阵

不仅参与了先前状态信息的获取,用于当前输出信息的生成,还决定了这些信息如何在时间中被传递——即我们应该更好地创建矩阵

,使其保留更多的上下文信息。

由此提出了结构化状态空间模型(S4),它不再将矩阵

初始化为随机矩阵,而是由HiPPO矩阵(High-order Polynomial Projection Operators Matrix)构建矩阵

阵。

图7  HiPPO矩阵作用示意图

HiPPO尝试将当前看到的所有输入信号压缩为系数向量,构建一个“可以很好地捕获最近的语义单元并衰减旧的语义单元”的状态表示,即 通过函数逼近产生状态矩阵

的最优解,其公式可以表示如下:

图8  HiPPO矩阵计算公式及形式

从数学上讲,HiPPO矩阵是通过跟踪Legendre多项式的系数来实现的,使用HiPPO构建矩阵

使得它能够逼近所有以前的历史。这使得在被应用于循环表示和卷积表示中时,S4结构可以处理远程依赖性。

同时为了有效避免HiPPO本身的N2尺寸带来的过多的计算量,我们可以利用低秩矩阵表示HiPPO:

其中,

是酉的,

是对角的,

这样使得S4模型是次二次时间架构,相较通过注意力矩阵压缩历史的Transformer模型,虽然捕捉长程依赖的能力 相对弱一些,但显著提高了处理数据的速度。

3.2.2 选择性结构化状态空间模型(Selective Structured State Space Models)

由上文的叙述我们可以知道,从SSM到S4模型增强了模型建立长程的依赖关系的能力,但SSM及S4模型仍然存在一个缺点——无法执行基于内容的推理。由于训练时采用了“卷积”的表示方式,故在训练后“卷积核”固定,即矩阵

的值在训练后保持不变,无论向SSM提供什么序列,

的值保持不变。这种情况就是说,模型不具备内容感知能力。

因此,由于固定的

矩阵,SSM无法执行基于内容的推理,因为它将每个语义单元都视为相等重要性。相比之下,对于Transformer来说,这些任务相对容易,因为它们根据输入序列动态地改变注意力。它可以有选择地“查看”或“关注”序列不同的部分。

图9  如果没有选择性,S4会花费相同的“精力”来处理每个单词,因为A,B,C矩阵在推理时恒定。(但训练时还是会更新A,B,C的值)

而我们希望SSM能够根据输入序列的不同重要性进行处理。

故而提出了选择性结构化状态空间模型(Selective Structured State Space Models,S6, or Mamba)。与先前的研究相比,Mamba主要有三点创新:

(1)对输入信息有选择性处理(Selection Mechanism):

图10  S4对比S6

为了能够使得模型在推理时具有内容选择性,在S4中影响输入的B矩阵、影响状态的C矩阵的大小从原来的(D,N)(D指的是输入向量的维度,比如一个颜色的变量一般有RGB三个维度,N指SSM的隐藏层维度hidden dimension,一般设的比较小,远小于L)变成了(B,L,N)(这三个参数分别对应batch size、sequence length、hidden state size),故而是根据输入数据计算的。

矩阵的大小由原来的D变成了(B,L,D),意味着对于一个 batch里的每个 语义单元 (总共有B*L个)都有一个独特的

,且每个位置的B矩阵、C矩阵、

矩阵都不相同。这意味着对于每个输入的语义单元,现在都有独特不同的B矩阵、C矩阵与之对应(即Mamba模型在推理时,可根据不同的输入数据动态计算矩阵B、C和步长Δ的值,但用于这些计算的参数(即决定如何计算这些矩阵和步长的函数或映射)是固定不变的),故而可以解决内容感知问题。

同时,虽然矩阵A没有获得数据依赖特性,但是通过SSM的离散化操作之后,会经过上述公式(至于为什么不赋予矩阵A数据依赖性后面会解释):

由于

参与了计算,使得矩阵

变成(B, L, N, D)的数据依赖的张量,即离散化之后,

的“输入数据依赖性”能够让整体的

与输入相关。

我们下面对每个变量进行具体分析:

对于步长

,由数据进行计算,期望达到如果某个输入比较重要,它的步长就更长些,被重点关注;如果某个输入不太重要它的步长就相对变短,被直接忽略。即较小的步长导致忽略特定单词,而更多地使用先前的上下文。较大的步长则专注于输入的单词而不是上下文。

图11  步长选择性忽略示意图

从而对于不同的输入,达到选择性关注或忽略的目标,做到详略得当,主次分明。

而对于矩阵B,C,可理解为:通过输入数据改变B和C,可以允许模型更精细地控制是否让输入x进入状态 h,或状态h进入输出 y,所以 B 和 C 类似于 RNN 中的输入门和输出门(门控)。

(2)硬件感知的算法(Hardware-aware Algorithm):

如上文所述,由于A,B,C这些矩阵现在是动态的了,因此无法使用卷积表示来计算它们(CNN需要固定的卷积核),因此,我们只能使用循环表示,如此也就而失去了卷积提供的并行训练能力。对于S6(Mamba)而言,提出了硬件感知算法,以实现并行化训练。

图12  先前的循环计算方法

对于先前的循环计算方法来说:每个状态都是前一个状态乘以矩阵

加上当前输入乘以矩阵

的和。如:

这被称为扫描操作(scan operation),可以使用 for 循环轻松计算,然这种状态之下想并行化是不可能的(因为只有在获取到前一个状态的情况下才能计算当前的每个状态)。

而S6提出了并行扫描(parallel scan)算法使得最终并行化成为可能,其假设我们执行操作的顺序与关联属性无关。因此,我们可以分段计算序列并迭代地组合它们,即动态矩阵B和C以及并行扫描算法一起创建选择性扫描算法(selective scan algorithm)。

图13  S6提出的并行扫描算法

举例而言,S6的并行化计算如下所示:

由:

可得到:

同理有:

则如上可推导出:

由上述的推导可以看出,S6模型的并行扫描法使用了类似于数字电路中超前进位加法器的设计原理,定义了一种新的加法:

图14  一种新的加法用于并行运算,每次运算取右侧项

由于与输入

有关的矩阵

每次仅与当前输入

做运算,故每一步

的运算可以单独计算,再将他们存储起来,乘上

后,可用于后面

的运算。

像上述这样,便可一层一层地进行隐状态

的计算,最后便可并行化地得到输出。

同时,S6模型也对显存的分配有了特定的提升,利用其自身显存占用小的优势,尝试限制我们需要从DRAM到SRAM以及反之的复制内容次数——在小型但高效的SRAM与大型但效率略逊的DRAM之间的的传输速度有限。即避免一有个结果便从SRAM写入到DRAM,而是待SRAM中有一批结果再集中写入DRAM中,从而降低来回读写的次数。

具体来说:直接将SSM的参数

从慢速DRAM载到快速SRAM中。然后,在SRAM中进行离散化,得到

。接着,在SRAM中进行上面说到的并行化隐状态计算,最后与C做计算并写入输出到DRAM。注意,当输入从DRAM加载到SRAM时,中间状态不被保存,而是在反向传播中重新计算,这样比从相对较慢的DRAM中读取所有这些中间状态要便宜的多。                      

具体的效果如下图所示:

图15  希望达到的效果

(3)更简单的架构:

将大多数SSM架构比如H3的基础块,与现代神经网络比如transformer中普遍存在的门控MLP相结合,组成新的Mamba块,重复这个块,与归一化和残差连接结合,便构成了Mamba架构:

图16  Mamba块

首先进行线性投影以扩展输入嵌入,然后,在应用选择性 SSM之前先进行卷积(以防止独立的语义单元计算)

图中的选择性SSM(Selective SSM)具有以下特性:

1.Recurrent SSM通过离散化创建循环SSM

2.HiPPO对矩阵A进行初始化A以捕获长程依赖性

3.选择性扫描算法(Selective scan algorithm)选择性压缩信息

4.硬件感知算法(Hardware-aware algorithm)加速计算

最后,包含归一化层和用于选择“预测的语义单元的Softmax层。

至此,我们介绍了选择性结构化状态空间模型(S6或Mamba)的发展及详细原理的介绍。相较之前的基础模型,它实现了:

1.并行化的快速训练与循环式的快速推理相结合。

2.能够很好地完成长序列建模任务,善于捕捉长程依赖关系。

3.新的硬件算法以实现显存的合理分配,实现占用显存的降低。

下面我们将介绍选择性结构化状态空间模型的实际效果。

4 选择性结构化状态空间模型的应用

4.1 选择性结构化状态空间模型的主干网络

像这部分我们主要分两部分介绍基于S6模型的主干网络:

  1. ViM(Vision Mamba):ViM是一种基于Mamba的架构,它的工作原理如下:

首先将输入图像分成小块,然后将小块投影成token。然后将这些token输入 ViM 编码器。然后序列中的每个标记经历两个单独的线性变换。

该算法在向前和向后两个方向上处理这些转换后的token,模拟双向神经网络层。与用于建模文本序列的 Mamba 模型不同,ViM 编码器可以向前和向后处理标记序列。对于每个方向,该过程涉及应用一维卷积,然后应用 Sigmoid 线性单元 (SiLU) 激活函数。

ViM的主要贡献是提供了双向扫描技术,为Mamba结构不适用于2D图像的问题提出了解决方法——双向模型可以更丰富地理解图像上下文,这是准确图像分类和分割的关键因素。

  1. VMamba:VMamba-Tiny的架构如下图所示。首先使用一个Stem节点将输入图像分割成多个patchs,类似于 ViTS,但没有将patchs进一步展平成1D序列,这种修改保留了图像的2D结构。然后,堆叠多个 VSS 块,保持相同的维度,构成“Stage1”。之后通过patch合并操作对特征映射进行下采样,构建分层特征表示。随后再堆叠多个下采样和VSS块,由此往复,创建更多Stage。这样就构成了一个类似CNN和ViT的基础模型,生成的架构可以在实际应用中作为其它视觉模型的替代品。

总的来说,该架构能够在保持全局感受野和动态权重的同时,将计算复杂度降低到线性,通过引入交又扫描模块(CSM),VMamba有效地解决了因ViTS计算复杂度高而带来的性能限制,特别是在处理高分辨率图像时。

图17  ViM和VMamba

  1. Mamba ND:Mamba ND致力于将Mamba的输入扩展到包括图像和视频在内的多维数据。

它主要解决数据缺乏预定义的排序,同时具有固有的空间维度所带来的挑战。考虑到将数据平坦化为1D序列的大量可能方式,Mamba ND仅包括通过沿其维度轴在向前或向后方向上平坦化数据的扫描排序。

然后,它将作为由1D Mamba层组合的Mamba ND块以交替顺序应用于序列。

Mamba ND的作者进行了广泛的实验来探索排序的不同组合。此外,他们将输入数据的一维划分为多个排序,采用不同的Mamba层排列,并将序列分解为更小的序列。结果表明,Mamba层链和简单的交替方向排序实现了优越的性能。MambaND区块的最终设计如图所示。

图18  Mamba-ND

简单来说了,Mamba-ND 通过简单地跨层交替序列排序,将 SSM 扩展到更高的维度。与 Transformer 相比,该模型在某些任务中能够在更低的参数数量下实现更强的性能。

4.2 选择性结构化状态空间模型的各领域应用

(1)图像

在图像方面,S6模型,即Mamba架构可完成分类,分割,生成,图像恢复等等任务。
与基于Transformer的方法相比,S6模型往往能以显著更少的参数,展示相差不大的优越性能。同时,基于Mamba的架构对较长token序列的可扩展性使得它们可以被用于高分辨率图像(例如,全幻灯片图像和遥感图像)和高维图像(例如3D医学图像)的分析以用于识别目的。

近期,Mamba架构也被广泛应用于几个低级别的任务,包括图像去雾、曝光校正、泛锐化、超分辨率、去噪、医学图像重建和水下图像增强。

(2)视频
视频理解是计算机视觉研究的基本方向之一。它的主要目标是有效地掌握长上下文中的时空表示,而这正中选择性结构化状态空间模型下怀,它也因此在这一领域表现出色,在保持线性复杂性和实现有效的长期动态建模之间实现了平衡。

这种创新方法促进了其在各种视频分析任务中的广泛采用,如视频目标分割、视频动作识别、视频生成和表示学习。

(3)多模态

多模态任务在CV领域发挥着至关重要的作用,因为它们有助于整合各种信息源,丰富视觉数据的理解和分析。这些任务的目标是聚合多种模态,包括文本和视觉信息、具有附加组件(如深度或热图像)的RGB图像以及各种形式的医学成像数据。

然而,实现多模态目标的一个重大挑战在于有效地捕捉不同模式之间的相关性。Mamba的捕捉关系能力同样也能胜任此任务。

最近,有几种方法将Mamba架构用于许多多模态任务,包括多模态大语言模型、多模态配准、参考图像分割、时间视频基础、语义分割、运动生成和医学应用。

(4)点云

点云是一种基本的三维表示,它提供具有三维坐标的连续空间位置信息。点云的内在无序性和不规则性一直是三维视觉中的一个挑战。受Mamba的线性复杂性和全局建模能力的启发,在点云处理领域也诞生了几种基于SSM的通用主干,在此就不多过赘述。

图18  可以使用选择性结构化状态空间模型处理的代表性输入数据

4.3 选择性结构化状态空间模型的实验对比

图19  分类任务对比与分割任务对比

(1)单/多标签分类:上图展示了S6,即选择性结构化状态空间模型在图像识别和分割领域的对比实验。其中(a, b, c)基于视频的识别,(d)基于图像的识别,(e, f, g)基于医学图像的分割实验结果。

对于单标签分类问题,在广泛使用的ImageNet-1K数据集上计算现有作品的准确率。如图所示,我们可以发现,基础版本VMamba和Mamba2D在ImageNet1K数据集上取得了更好的结果,top-1的准确率分别为83.2%和83%。

我们也很容易发现目前基于Mamba(S6)的视觉模型都是微小的、小的或基础的版本,很少预训练一个大型或巨大版本的Mamba网络。

总体性能与一些基于Transformer的模型相当,但仍然不如ImageNet分类数据集上的最先进的模型。

对于多标签分类,选择行人属性识别(Pedestrian Attribute Recognition, PAR)任务,在PA100K和PETA数据集上进行实验。

实验采用ViT-S和基于Mamba的网络虚拟机作为主干。遵循基于视觉语言融合的PAR框架VTB,该框架以行人图像和属性集为输入,并预测每个属性的逻辑分数。

从实验结果可以发现,基于ViM-S的PAR模型在PETA数据集上达到81.08/73.75/80.91/84.96/82.52,低于Transformer。而在PA100K数据集上达到80.41/78.03/85.39/88.37/86.39。这些结果明显优于基于ViT-S的模型,但仍然明显低于基于Transformer网络开发的PAR算法。例如,基于vitb的VTB在PETA和PA100K数据集上达到了85.31/79.60/86.76/87.17/86.71,83.72/80.89/87.88/89.30/88.21。

图20  多标签分类对比

(2)视觉目标跟踪

在本小节中,将Mamba与Transformer进行比较,并将基于CNN的主干用于基于OSTrack的跟踪任务。具体来说,基于CNN的跟踪器有TrDiMP、ToMP50、DiMP50、PrDiMP、KYS和ATOM;基于Transformer的跟踪器是HDETrack、AiATrack、STARK、TransT、MixFormer和SimTrack。

为了实现公平的比较,在一个大规模的基于事件的跟踪数据集EventVOT上训练和测试这些跟踪器,该数据集分别包含841、18和282个视频。注意,比较中使用了三种广泛使用的评估指标,包括成功率(SR)、准确率(PR)和归一化准确率(NPR)。

结果如下:

图21  视觉追踪对比

从图中可以发现,使用Mamba骨干网替换ViT时,性能略有下降,但同时带来了参数数量的巨大减少(仅4.1M)。因此,我们可以得出结论——基于事件的跟踪对于Mamba网络来说,是一个有前途的选择。

(3)像素级图像分割

最近,Mamba网络在医学图像分割中得到了广泛的应用。例如,基于swing - transformer的模型SwinUNet在MRI心脏数据集中获得了89.33/99.57/88.46 (Dice, IoU, Accuracy)。

而相比之下,基于Mamba的UNet实现了类似甚至更好的分段结果,如Mamba-UNet,半Mamba-UNet和弱Mamba-UNet。这些结果充分证明了Mamba结构在医学图像分割中的有效性。

(数据在上述大图上)

(4)由图生成文字

对于图像到文本的生成,选择X射线医学的X射线报告生成任务图像作为输入,生成医疗报告。

在实验中,选择R2GenGPT6作为基线,并在IU-Xray数据集上评估其性能[224]。R2GenGPT由视觉编码器(Swin Transformer)、线性层和大型语言模型(llama-2-7B-chat)组成。训练方法包括最初冻结语言模型,然后对视觉编码器和线性层进行微调。

将Swin Transformer替换为ViM模型(基于Mamba或是S6),并将结果与其他方法进行比较。

由于这两个模型都使用预训练的组件,Vision Mamba在BLEU-4和ROUGE-L评分方面表现出比Swin Transformer模型更优越的性能。

图22  由图生成文字对比

(5)人/车辆鉴定

如下图所示,对人重识别和车重识别两个重识别(re-ID)任务进行了实验。对于人员重识别,使用了四个广泛使用的数据集,包括MSMT17, Market1501, DukeMTMC和Occluded-Duke数据集。

这些数据集是从不同的场景中捕获的,来自摄像机重叠覆盖的监视系统,这些系统具有跨时间跨度,遮挡和背景干扰等挑战。 对于车辆重识别,使用VeRi-776和VehicleID数据集进行实验验证。与行人样本不同,观察视点的变化也会给车辆带来显著的外观差异,因此车辆数据集额外提供视点标签来标记车辆样本的不同视点。

对于上述数据集,我们使用累积匹配特征(CMC)曲线和平均平均精度(mAP)作为评价指标。

从实验结果可以看出,与需要复杂模块设计的基于CNN的模型相比,简单的Mamba网络已经具有有效性。

即使与DeiT、ViT等复杂度较高的模型相比,虽然ViM(Vision Mamba)训练参数较少,在VehicleID数据集上也显示出了有效性。

相比之下,VMamba的交叉扫描机制不依赖于Transformer的结构(例如,位置嵌入和类标记),在Market1501、DukeMTMC和VeRi-776数据集上取得了可比较的结果。

因此,我们可以期望在未来有更多基于Mamba的研究适用于重新识别任务。

图23  重识别任务对比

总的而言,基于选择性结构化状态空间模型(S6,or Mamba)架构的模型在许多任务上的表现能与基于Transformer的模型分庭抗礼,甚至优于Transformer架构的模型。

同时,他们仅需要相对Transformer架构很少的参数量和相对简单的网络结构。对于选择性结构化状态空间模型(S6,or Mamba)架构,我们可以认为择性结构化状态空间模型(S6,or Mamba)未来可期。

5 总结

5.1对于选择性状态空间网络模型的一些问题与挑战:
(1)可扩展性和稳定性:目前,Mamba架构在应用于ImageNet等大规模数据集时可能会有稳定性隐患。Mamba在扩展到更广泛的网络配置时不稳定的根本原因尚不清楚。

这种不稳定性经常导致Mamba框架内的梯度消失或爆炸,这阻碍了其在大规模视觉任务中的部署。

(2)因果关系问题:鉴于Mamba模型最初是为因果序列数据设计的,将其选择性扫描技术应用于非因果视觉数据带来了重大挑战。

目前的解决方法是通过采用ViM中提出的双向扫描等扫描技术来解决这一问题,其中向前和向后扫描都被用来相互补偿感受野中单向扫描的固有限制。然而,这样的扫描方法并不能从根本上解决问题,并且还会带来其他挑战。

(3)空间信息:由于Mamba架构中选择性扫描技术固有的1D特性,在将其应用于2D或更高维度的视觉数据时带来了挑战,因为它可能导致关键空间信息的丢失。为了解决这一限制,当前的方法通常从各个方向展开图像块,从而允许跨多个维度的空间信息的集成。与上述解决问题的方式一样,多个方向展开图像块也会带来其他问题。

(4)冗余和计算:如前所述,双向扫描方法和多个扫描方向的使用会导致显著的信息冗余和计算需求的增加。这些可能会降低模型性能,并降低Mamba次二次复杂度的优势。

5.2未来方向

A.数据

(1)数据效率:考虑到Mamba的计算成本与CNN相当,即使不依赖大规模数据集,它也具有提供最佳性能的巨大潜力。这一属性让Mamba有潜力完成各种下游任务/多任务和涉及预训练模型自适应的任务。

(2)高分辨率数据:由于SSM的架构在理论上简化了计算复杂性,同时更能捕捉长程依赖关系。因此其有效处理高分辨率数据(如遥感和全切片图像)或长期序列数据(如长期视频帧)的潜力具有相当大的价值。

(3)多模态数据:Mamba模型在处理扩展序列方面的熟练程度大大拓宽了其在多模态学习中的适用性。

(4)上下文学习:在深度学习的动态环境中,上下文学习已经发展到包含越来越复杂和新颖的方法,以解决NLP、CV和多模态领域的复杂任务。这种方法上的进步对于突破现有深度学习框架的极限至关重要。Mamba模型凭借其精通上下文建模能力和捕获长程依赖关系的能力,在上下文学习应用程序中显示出更深入的语义理解和增强性能的潜力。

B.算法

(1)扫描技术:选择性扫描机制是Mamba模型的核心组成部分,最初针对1D因果序列数据进行了优化。为了解决视觉数据固有的非因果性质,许多现有方法被提出,例如ViM采用的双向扫描。此外,为了捕获2D或高维视觉数据中固有的空间信息,当前的方法通常是扩展扫描方向。

尽管这些调整还算有效,但也带来了计算冗余等等问题。对于SSM架构而言,仍然迫切地需要更具创新性的扫描方案,以更加有效地利用高维非因果视觉数据的全部潜力。

(2)融合技术:使Mamba模型适应视觉任务往往会引入冗余,使扫描输出特征的有效融合成为进一步探索的重要领域。此外,计算机视觉的基础模型各有其独特的优势;例如,细胞神经网络固有地捕捉归纳偏差,如翻译等变,而ViT以其强大的建模能力而闻名。探索融合这些不同网络架构提取的特征以最大限度地发挥其优势的方法是另一个可以研究的方向。

(3)计算效率:由于Mamba在处理高维信息时,通常需要在多个路径中扫描,将其用于视觉任务会导致计算需求增加。因此,在开发更高效、更有效的视觉Mamba模型方面有着重要的研究机会。

5.3 结论

选择性结构化状态空间模型,即S6模型,又称为Mamba,已迅速成为一种变革性的长序列建模架构。

它实现了:

(1)并行化的快速训练与循环式的快速推理相结合。

(2)HiPPO矩阵的提出,使其能够很好地完成长序列建模任务,善于捕捉长程依赖关系。

(3)新的硬件算法以实现显存的合理分配,实现速度提升。

(4)更换S4中的卷积计算模式,扩充用于计算的矩阵的维度,使其能够针对输入数据,有选择地进行处理。

同时,它在能够被应用于多个类别的任务,并且在许多任务上,它能够用相对Transformer更少的计算量,与基于Transformer的模型所表现出的性能分庭抗礼,甚至优于Transformer架构的模型。

本文对选择性结构化状态空间模型进行了全面的综述。首先,本文分析了过去三种主流的深度学习基础模型CNN,RNN,Transformer的优势与劣势;然后介绍了状态空间模型的基本原理及其发展过程,优缺点;最后介绍了选择性结构化状态空间模型的应用与相关实验,对比了其与各主流模型的性能。

相信在未来,选择性结构化状态空间模型能够得到更加长足的发展,也能够逐渐解决本文提及的诸多问题,成为一个较为成熟的基础模型。

参 考 文 献

  1. Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces[J]. arXiv preprint arXiv:2312.00752, 2023.
  2. Liu Y, Tian Y, Zhao Y, et al. Vmamba: Visual state space model[J]. arXiv preprint arXiv:2401.10166, 2024.
  3. Zhu L, Liao B, Zhang Q, et al. Vision mamba: Efficient visual representation learning with bidirectional state space model[J]. arXiv preprint arXiv:2401.09417, 2024.
  4. Xu R, Yang S, Wang Y, et al. A survey on vision mamba: Models, applications and challenges[J]. arXiv preprint arXiv:2404.18861, 2024.
  5. Wang X, Wang S, Ding Y, et al. State space model for new-generation network alternative to transformers: A survey[J]. arXiv preprint arXiv:2404.09516, 2024.
  6. A. S. Sharma, D. Atkinson, and D. Bau, “Locating and editing factual associations in mamba,” 2024.
  7. Yu W, Wang X. MambaOut: Do We Really Need Mamba for Vision?[J]. arXiv preprint arXiv:2405.07992, 2024.
  8. Gu A, Goel K, Ré C. Efficiently modeling long sequences with structured state spaces[J]. arXiv preprint arXiv:2111.00396, 2021.
  9. Wang J, Zhu W, Wang P, et al. Selective structured state-spaces for long-form video understanding[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 6387-6397.
  10. Li S, Singh H, Grover A. Mamba-ND: Selective State Space Modeling for Multi-Dimensional Data[J]. arXiv preprint arXiv:2402.05892, 2024.

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值