点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
近年来,状态空间模型(SSM),尤其是Mamba,在视觉建模方面展现出卓越的性能,并具备高效的模型特性。然而,由于序列模型中存在三个主要约束,Mamba在视觉任务中的应用表现不佳:
1)因果计算无法获取全局上下文;
2)在计算当前隐藏状态时存在长距离遗忘问题;3)由于输入序列的转换,导致空间结构建模能力较弱。为解决这些问题,作者研究了一种简单而强大的视觉任务 Adapter Adaptor,用于Mamba模型,该 Adapter 由两个功能模块组成:Adaptor-T和Adaptor-S。
在求解SSM的隐藏状态时,作者应用轻量级预测模块Adaptor-T,通过选择一组可学习的位置作为记忆增强,以缓解长距离遗忘问题。此外,作者利用由多尺度扩张卷积核组成的Adaptor-S,增强空间建模能力,并将图像归纳偏置引入特征输出。
这两个模块均能通过增强不可访问的特征来扩大因果计算中的上下文建模范围。作者探索了Mamba-Adaptor的三种应用方式:
1)作为各种视觉任务的通用视觉 Backbone ;
2)作为提升预训练 Backbone 性能的增强模块;
3)作为高效微调模块,用于迁移学习任务。大量实验验证了Mamba-Adaptor在三种场景下的有效性。
值得注意的是,作者的Mamba-Adaptor在ImageNet和COCO基准测试中取得了当前最佳性能。
- 引言
结构化状态空间模型(SSMs)[13-15, 38]因其高效建模长序列的能力,在语言任务中广受欢迎。近年来,在视觉Transformer[7, 33, 39, 49, 53, 55, 64, 67]取得巨大成功后,SSMs,特别是Mamba变体[13],被应用于视觉任务。虽然基于Mamba的视觉模型在多种下游任务[62, 63, 65, 75],如密集预测[24, 32, 72, 78]、图像合成[6, 36]和生成任务[8, 50, 66]中展现出潜力,但Mamba模型在视觉领域的进一步应用和适配仍基本未得到探索。
使用案例1:作为 Backbone 网络从头开始训练
使用2:继续作为助推网络进行训练
使用3:迁移学习作为 Adapter 网络尽管视觉Mamba模型在长token长度方面具有优越的线性计算复杂度,但在各种视觉任务中并未达到最佳性能[30, 43, 61, 76],这源于SSM机制中的三个内在约束。首先,SSM的内在RNN[20]机制通过先前存储的记忆(也称为隐藏状态)来递归计算下一个时间步的输出。因此,因果计算流程无法访问未来的输出状态,而2D视觉数据属于非因果风格。现有的视觉Mamba[78]模型通过双向扫描策略部分解决了这一问题,该策略将序列顺序反转以建模全局上下文。其次,SSM中的递归计算流程在序列长度增加时保持复杂度恒定,而先前隐藏状态的影响则逐渐衰减。因此,记忆衰减导致长程遗忘现象,进一步削弱了SSM的全局上下文建模能力。第三,当前的视觉Mamba模型[32, 78]严重依赖于将2D视觉数据转换为1D序列格式,这破坏了结构图像中的空间依赖关系。尽管存在专门的扫描策略,如四向[32]和局部窗口扫描[24, 72],可以部分保留空间局部性,但视觉Mamba模型仍然面临从序列建模中学习2D依赖关系的问题。
经过详细分析,作者将所提及的三种约束分为两类,即时间衰减和空间局部性。因此,作者为Mamba模型引入了一种简单而强大的视觉任务 Adapter (如图2所示),该 Adapter 由两个功能模块组成:Adaptor-T和Adaptor-S。Adaptor-T执行记忆保留,以处理每个隐藏状态的远程遗忘问题。值得注意的是,手工记忆选择在寻找最易遗忘的状态方面存在不足,从而限制了记忆保留的效果。因此,作者采用了一种轻量级可学习记忆选择方案,该方案动态选择一组最易遗忘的状态。为了增强空间建模,作者直接将多尺度深度卷积滤波器应用于二维特征输出的函数,从而有效地引入了图像归纳偏差。此外,将作者的 Adapter 集成到高度优化的Mamba求解器中至关重要。为此,作者将求解器与单位矩阵解耦,并采用高度高效的矩阵乘法,从而使Adaptor-S和Adaptor-T能够无缝地插入Mamba求解器中。作者还设计了Mamba-Adaptor的两种插入形式,即用于迁移学习的并行形式和用于从头开始训练的顺序形式。
如图1所示,作者的Mamba-Adaptor具有多种用途,包括作为通用视觉 Backbone 网络、用于进一步提升性能的增强模块,以及用于迁移学习的有效Adaptor。在分类和密集预测任务中,Mamba-Adaptor相较于近期的基于Mamba、CNN和基于transformer的方法,实现了更优越的性能。Mamba-Adaptor在多种迁移学习任务中也表现出具有竞争力的结果。总体而言,本研究的贡献如下:
作者介绍了一种即插即用的Mamba-Adaptor,它能够在各种视觉任务中提升Mambal模型的性能,且计算开销可接受。作者的Mamba-adaptor由两个功能模块组成,即Adaptor-T和Adaptor-S,它们有效解决了时间衰减和空间局部性问题。Mamba-Adaptor不仅可以作为图像分类和预测任务的通用视觉 Backbone 网络,还可以在迁移学习任务设置中发挥作用。大量的实验结果表明,作者的MambaAdaptor在多种视觉下游任务中表现出优越的性能,并在Mamba Baseline 模型上取得了显著的改进。
- 相关工作
状态空间模型。状态空间模型(SSMs)[13, 15, 38, 44],源自线性控制系统[60],提供了一种具有线性计算复杂度的有效框架,使其非常适合长序列处理。语言任务[37, 54]首先应用基于SSM的方法进行长距离上下文建模。随后,提出了一系列改进,包括结构优化[15]、关联扫描[44]和硬件加速技术[14],以使SSM与现代深度学习架构[10, 17, 18, 27, 40, 48, 69, 73]相结合。最近出现的Mamba[13]引入了输入特定参数化和并行扫描(S6),进一步确立了SSM作为Transformer有力竞争者的地位。这些创新促成了SSM在视觉任务中的广泛应用,尤其是在具有更优线性计算复杂度的Transformer替代方案中[5, 33, 51, 55-57]。先驱视觉模型如ViM[78]和VMamba[32]通过双向和交叉扫描策略将2D图像转换为序列格式,在视觉表征中应用Mamba方案。后续视觉方法通过更专用的扫描模式[24, 41, 68, 72]增强Mamba扫描的视觉能力以保留空间局部性,以及混合Transformer-Mamba架构[16, 36]作为强大的补充。尽管取得了这些进展,视觉SSM模型在克服长距离建模问题和处理2D图像数据的序列转换限制方面仍面临挑战,这阻碍了它们在需要特定视觉先验和能力的视觉任务中的性能。作者的 Adapter 旨在克服视觉Mamba方法中长距离遗忘问题和序列转换的不足。
视觉 Adapter 。 Adapter 的概念源于语言领域[22],它作为附加模块用于针对特定目的微调大语言模型[21, 46]。在视觉领域, Adapter 通常作为可训练部分,将视觉基础模型[4, 34, 42, 58](如CLIP[42])的大规模知识迁移到各种下游任务[9, 45, 47, 74]中。在迁移学习任务[22, 25, 59]中,Adaptformer[2]使用额外的分支对MLP[52]进行微调,以高效地微调预训练的视觉Transformer的下游任务。VPT[25]通过添加额外的视觉 Token 作为 Prompt 输入,高效地微调基础模型。Mamba模型的迁移学习[59]仍有许多未探索的空间。作者开发了针对Mamba的微调模块,以适应迁移学习中的下游任务。 Adapter 在视觉中的另一个用途是作为补充模块来增强基础模型的能力,例如增强普通视觉Transformer中的层次特征表示[7]。像ViTDet[29]这样的先驱工作实现了特征金字塔模块,以适应从视觉Transformer提取的普通特征进行密集预测任务。视觉 Adapter 的另一个代表性工作是ViT-adaptor[3],它设计了一个基于卷积的网络作为基于注意力的网络的补充分支,为图像任务引入了必要的归纳偏差。基于Mamba的架构[13, 15, 44]最初是为语言建模设计的,忽略了视觉先验。作者引入了视觉先验并克服了Mamba的限制,这可以作为一个加速模块。作者的工作探索了Mamba Adapter 的两种用途,它作为迁移学习中针对Mamba模型的专用 Adapter 设计,以及作为一个加速模块来进一步提升性能。
- 方法
在本节中,作者首先重述第3.1节中SSM的预备知识。然后,作者介绍所提出的Mamba-Adaptor的公式化描述(第3.2节)以及两个关键子模块:Apdator-T(第3.3节)和Adaptor-S(第3.4节)。此外,作者还展示了Mamba-Adaptor的实际实现方法,以及如何将 Adapter 集成到Mamba结构中(第3.5节)。
3.1. 预备知识
状态空间模型。状态空间模型(SSMs)通常用于分析序列数据,并对连续线性时不变(LTI)系统的输入变换进行建模。在状态空间模型中,一个输入序列
通过隐藏状态
作为中间值来计算响应
。求解状态空间模型的过程可以表示为线性常微分方程(ODEs)的组合:
其中
,
,
,以及
是线性状态转换和输出方程的转换矩阵。
SSM的离散化。为了在深度学习框架内有效解决LTI系统中的SSM[13],将连续时间方程转换为离散格式至关重要。在LTI分析中,实现这一目标常用的方法是零阶保持器(ZOH)离散化[28]。通过将输入分解为小的时间间隔,并将这些间隔内的输入视为常数,可以将连续时间系统转换为离散时间系统。ZOH有效地将SSM中的连续时间参数
在固定的采样时间间隔
上转换为它们的离散等价形式。
选择性扫描机制。为了解决线性时不变(LTI)状态空间模型(SSM)在捕捉上下文信息方面的局限性,Mamba [13] 引入了一种参数化方法,该方法具有一个输入相关的选择机制,称为S6 [15]。这允许模型根据特定输入选择相关特征,从而增强其对各种输入上下文的响应。在选择性SSM的背景下,时变权重参数给高效计算隐藏状态带来了挑战,因为卷积是静态的,并且不是设计用来处理动态权重的。Mamba [12] 采用输入相关的方案,从输入序列
中生成参数
、
和
。通过参数的离散化,顺序状态
可以以递归的方式计算:
公式1中的隐藏状态
可以使用前一个隐藏状态
和当前特征输入
递归计算。然后,作者可以从隐藏状态中获得最终输出
响应
仍然可以使用关联扫描算法 [44] 以线性复杂度高效计算。硬件优化技巧也被应用于加速求解SSM [13, 15]。
视觉Mamba的序列变换。视觉Mamba方法[24, 32, 78]需要先将二维图像数据
转换为序列格式
。为了获取全局上下文,双向扫描是基本策略,而鼓励采用多条扫描路径以获得更强的上下文建模能力。对于多条扫描序列,SSM独立处理每一条并最终合并输出。在此,作者将视觉Mamba中的序列变换和SsM求解过程命名为SS2D方案。
3.2. Mamba-Adaptor的公式化表述
作者的Mamba-Adaptor直接作用于SSM的方程3和方程4,旨在提升其内在建模能力。作者探索了求解SSM的两个主要步骤的潜在局限性,即方程3中的隐藏状态计算和方程4中的输出计算。
隐藏状态。对于索引
在方程 3 中的隐藏状态
是先前隐藏状态
和当前输入
的转换值的总和。隐藏状态是递归计算的,从空状态
开始。因此,在求解方程 3 时,距离较远的隐藏状态对当前隐藏状态
是不可访问的,导致长距离遗忘问题。 Aware 到这个问题后,对非邻近区域的隐藏状态进行加权聚合可以有效地缓解遗忘:
其中
是选定的可访问隐藏状态的索引集,
是用于求和的相应权重参数。公式 5 描述了 Adaptor-T 的核心原理,该原理处理了 Mamba 隐藏状态中的时间衰减问题。
最终输出。在公式4中,最终转换后的输出
是通过转换后的隐藏状态
与相应的输入嵌入
的求和得到的。然后,最终的序列输出
被 Reshape 为2D格式,以获得用于下一层的最终输出
。尽管
通过递归隐藏状态计算包含了长距离依赖关系,但在最终输出中,2D图像数据的空间局部性被很大程度上忽略了。为了克服这个问题,作者在 Reshape 的2D格式输出
上进行操作,通过卷积风格聚合来获得空间依赖关系,这可以有效地引入图像归纳偏差:
卷积算子在一个固定的扫描局部窗口中聚合加权输出
。对于在二维空间坐标
中的每个输出
,空间局部性增强如下:
其中
通常表示
局部邻域嵌入,
是相应的权重参数。通过增加核大小
可以扩大局部邻域。公式 6 和公式 7 构成了作者的 Adaptor-S 的核心方案,该方案在 Mamba 的输出空间局部性上运行。
3.3. Adaptor-T
Adaptor-S的目标是对每个隐藏状态
进行记忆保持。由于方程3是随意计算的,每个隐藏状态只能访问到比当前状态序号小的先前状态。粗暴地聚合所有先前状态可能效果不佳,因为先前隐藏状态的遗忘程度各不相同。采用固定的人工设计模式来选择相邻隐藏状态,类似于卷积滤波器,也是不足够的,因为衰减的隐藏状态可能会被大量忽略。
可学习记忆选择。为处理这一问题,作者提出了一种可学习记忆选择方案,该方案使用一个极其轻量级的预测层来选择一组最容易遗忘的隐藏状态。作者采用一个极其轻量级的线性层
来预测
个坐标
,以选择每个隐藏状态
的遗忘状态。每个遗忘状态
的系数也通过一个线性层
和SoftMax函数[52]进行预测:
因此,最有可能被遗忘的内存状态是动态确定的,从而消除了手动选择的需求。
多序列时间保留。类似于多头注意力机制[7]和多组卷积滤波器,总
序列被分别预测并聚合:
对于双向扫描[78],序列
的总数为2,需要预测
个坐标。因此,SSM中的每个隐藏状态都通过预测的最遗忘状态得到加强,类似于语言处理中的记忆保留方案。
3.4. Adaptor-S
遵循Adaptor-S原则,作者采用深度卷积滤波器从二维空间维度聚合输出特征y。深度卷积轻量级,带来可接受的计算开销。它以通道维度聚合特征,从而保持输出特征形状不变。
多尺度空间聚合。具体来说,为了扩大空间依赖性,作者使用具有多尺度膨胀因子的多个深度卷积滤波器。在空间坐标
处最终输出特征嵌入
的空间局部性可以表示如下:
在Adaptor-S中,
是从深度卷积滤波器
的邻近区域选择的索引。这种空间聚合过程为顺序特征表示带来了更多的结构归纳偏差,极大地增强了Mamba中的空间局部性。
3.5. 实现
作者描述了如何使用高度优化的硬件模块实现Mamba-Adaptor,并引入了两种针对不同任务设置进行 Adapter 插入的集成方法。
实际应用。将MambaAdaptor集成到高度优化的SSM实现中对于有效性至关重要。作者将求解过程分为两部分:一个原始优化的Mamba算子和高效的矩阵乘法。作者的AdaptorT和Adaptor-S可以轻松集成到Mamble层中。具体而言,作者手动生成单位矩阵和零矩阵分别用于
和
,以替换原始参数。因此,高度优化的Mamba求解器[13]可以计算隐藏状态而不是最终输出,这是Adaptor-T可以插入的位置。然后,作者使用矩阵乘法并结合Adaptor-S来计算最终输出,以增强空间建模。整个过程可以用算法1中的伪代码描述。
层插入如图4所示,作者引入两种形式将作者的 Adapter 与视觉Mamba模型集成,即顺序插入和并行插入。对于从头开始训练的图像分类和密集预测任务,作者采用顺序形式。对于微调预训练模型的迁移学习任务,作者采用并行形式。并行设计通过独立分支保留原始预训练特征,更适合迁移学习任务中的微调。
权重共享系数。在迁移学习和增强器设置中,这些设置需要零权重初始化和较少的微调参数以实现高效率,作者为Adaptor-S采用权重共享系数策略。具体而言,作者对每个隐藏状态初始化固定形状的系数参数
,从而降低了线性层
的模型成本。
- 实验
作者通过三个主要的实验设置评估了Mamba-Adaptor的有效性。第4.1节的实验评估了视觉 Backbone 网络在图像分类和密集预测任务中的性能。第4.2节的实验将Mamba-Adaptor用作一个增强模块,以提升预训练网络的表现。第4.3节的实验专注于迁移学习,其中MambaAdaptor作为一个高效的微调模块,用于下游识别任务。最后,作者在第4.4节展示了消融研究。
4.1. 视觉 Backbone 网络
模型架构。作为视觉 Backbone 网络,Mamba-Adaptor遵循广泛采用的分层 Backbone 架构[17, 23, 32, 33, 55, 70]。它主要由一个图像块嵌入层、一个包含所提出的Mamba-Adaptor的视觉Mamba层(形成四个模型阶段)以及一个用于分类和密集预测的 Head 网络组成。每个模型阶段包含一个下采样层,该层将特征图的大小减半,同时将通道维度增加一倍。在本实验中,作者引入了MambaAdaptor的两个变体,具体如下:
Mamba-Adaptor-b1具有48个通道维度和四个模型阶段中的层数量
;MambaAdaptor-b2具有96个通道维度和四个模型阶段中的层数量
。
在ImageNet-1K上的图像分类。ImageNet1K [43] 数据集包含128万张用于训练的图像和5万张用于验证的图像。作者使用相同的训练设置 [32, 33],使用训练集从头开始训练Mamba-Adaptor的两个变体,并通过报告验证集上的Top-1准确率来评估它们的性能。作者在ImageNet-1K数据集上的图像分类中评估了VMamba的性能,将其与基准方法的比较结果展示在表1中。具有相似的计算FLOPs,作者的MambaAdaptor-b1达到了78.4的Top-1准确率,比LocalViM-T高2.2%,比Vim-T高2.3%。此外,Mamba-Adaptor-b2在更大规模上也保持了其性能优势。例如,Mamba-Adaptor-b2达到了82.9的Top-1准确率,超过Swin-T高2.6%,超过VMamba-T高0.2%。
COCO数据集上的目标检测。COCO [30]是一个目标检测和实例分割数据集,包含118K张训练图像和5K张验证图像。作者使用Mamba-Adaptor作为MaskRCNN检测框架 [18]中的 Backbone 网络,以评估yinqing-Mamba-Adaptor_2505的有效性。作者遵循SwinT [33]中类似的训练策略,该策略使用在ImageNet-1K数据集上预训练的权重。MSCOCO上的结果如表2所示。作者的MambaAdaptor在框平均精度
和 Mask 平均精度
方面均表现出优越性,在
和A
训练计划下均如此。使用
微调计划和相似的计算FLOPs,Mamba-Adaptor-b1实现了最佳的目标检测mAP,达到
,比EffVMamba-S [41]高
(
),比PVT-T高
(
)。在公平比较下,Mamba-Adaptor-b2实现的实例分割mAP在
方面比 Baseline VMamba-T高
,在
方面高2.3。Mamba-Adaptor-b2在目标检测和实例分割任务中均显示出对改进视觉Mamba模型LocalVMamba-T的优势。Mamba-Adaptor展示了其在与密集预测相关的下游任务中取得出色性能的潜力。
4.2. 图像识别的Booster网络
增强器、 Baseline 和基准。作者选择在ImageNet [43]分类任务上预训练的VMambaT/S/B作为 Baseline 模型,并为其配备所提出的MambaAdaptor,同时保持可接受的计算开销。作者继续在ImageNet训练集上对配备Mamba-Adaptor的VMamba进行额外的10个epoch的微调。
Booster结果。表3中的结果表明,作者的Mamba-Adaptor能够进一步提升预训练的VMamba[32] Baseline 的性能。具体而言,在增加不到3.2%的模型参数和6.1%的FLOPS的情况下,MambaAdaptor可以将VMamba-T的top1准确率提升0.1%。对于具有较大模型规模的VMamba-B,其性能提升在top1准确率上达到了0.2%。值得注意的是,booster训练仅需10个额外轮次,这表明作者的Mamba-Adaptor作为booster模块具有很高的效率。
4.3. 转移学习 Adapter 网络
Adapter 、 Baseline 和基准。作者将Mamba-Adaptor与三种常用的微调方法进行比较。(1) 线性 Prob :在主干网络顶部添加一个额外的线性层,并对该线性层进行调优以进行评估。(2) 完全微调:作者解冻所有模型参数,并将它们一起进行训练。(3) 视觉 Prompt 微调(VPT)[25]:将额外的视觉 Token 作为 Prompt 进行微调(作者添加了128个额外的 Token )。作者选择具有三种变体(即tiny/small/base)的VMamba [32]来评估 Adapter 的性能。作者还在图4中提供了ViT-Base [7] Baseline 网络进行比较。根据[2, 25]中的设置,作者选择了三个图像识别基准:CIFAR-100 [26]包括50,000张训练图像和10K张验证图像,分辨率为
像素,跨越100个标签;街景房屋数字(SVHN)[11]数据集,用于数字分类,包含超过600K张 Token 图像,其中73K用于训练,26K用于测试,还有531K张额外的训练图像;食品-101 [1]数据集包含101种食品类别,总计101K张图像。
迁移学习结果。表3报告了在三个视觉识别基准上的结果。与其他两种微调方法线性 Prob 和VPT[25]相比,MambaAdaptor在不同数据集上保持了在Top 1分类精度方面的优越性能。具体而言,在CIFAR-100基准上,Mamba-Adaptor-T/S/B保持了90%的性能,使用的模型参数分别少于5.56%/9.25%/7.13%。相比之下,全微调方法将所有模型参数都纳入训练。线性 Prob 和VPT[25]方法在三个视觉识别基准[1, 11, 26]上的性能与全微调之间存在较大差距,验证了yinqing-Mamba-Adaptor_2505的有效性。
4.4. 消融实验
Adaptor-T/-S在分类和预测任务上的表现。表6展示了两个功能模块在不同配置下的消融研究,即Adaptor-T的静态/可学习状态选择和Adaptor-S的多尺度因子。与没有MambaAdaptor的 Baseline 模型
相比,所有其余设置(
至
)在分类和密集预测任务中均有不同程度的提升。值得注意的是,可学习状态选择方案在ImageNet1k上比静态选择方案高出0.3%的top-1准确率,在COCO目标检测的
上高出0.6%,这验证了作者动态选择方案的有效性。作者还观察到,由于层次化表示,向Adaptor-S添加额外的深度卷积核时,性能提升了0.7%(在
中)。
关于迁移学习。如表6所示,初始化方法对迁移学习任务有显著影响。与随机初始化相比,零初始化在三个基准测试中的top1准确率上显著提高了14.9%/13.6%/13.2%(
vs.
)。并行插入形式也对迁移学习场景下的性能有显著贡献,这一点可以通过
和
的比较得到验证。潜在的原因是零初始化和并行插入设计保留了原始特征。这也对booster模块设计有影响,因为预训练权重包含了ImageNet [43]的大规模先验知识,而额外的Mamba-Adaptor模块参数量较小。Adaptor-T的权重共享系数策略在降低模型参数增加的同时提升了最终性能。
关于泛化能力。作者还进行了迁移学习实验以展示泛化能力。表7显示yinqing-Mamba-Adaptor_2505在具有双向扫描的Vision Mamba [31]上也表现出良好的结果。
- 结论与局限性
本文提出了一种即插即用的Mamba-Adaptor,该模块具有三大主要应用。Mamba-Adaptor可作为图像分类和密集预测任务的通用视觉 Backbone 网络,显著提升 Baseline 性能。此外,它还可以作为增强模块插入预训练的视觉 Backbone 网络中,通过扩展训练提升性能。
在迁移学习场景下,Mamba-Adaptor是一个高效的模块,能够使预训练的Mamba Backbone 网络适应各种下游视觉识别任务。大量的实验结果验证了Mamba-Adaptor的有效性。作者希望Mamba-Adaptor为视觉领域中的Mamba方案提供一种富有洞见的解决方案。然而,将Mamba-Adaptor作为视觉 Backbone 网络进行扩展仍是一个开放问题,其应用于其他变种的视觉Mamba是未来工作的重点。
点击上方卡片,关注「AI视界引擎」公众号