【人工智能】GlobalMamba: 用于Vision Mamba的全局图像序列化

论文总结

img

2024年10月14日第一次发布在arXiv上。

这篇文章的作者团队来自清华大学和加州大学伯克利分校,成员包括王成昆、郑文钊、周杰、陆机文。王成昆和郑文钊是论文的主要贡献者,周杰是通讯作者,负责整体研究方向和监督。

动机

现有的视觉 Mamba 模型通常通过将图像分割成局部图块并将其平铺成一维序列进行因果处理,但这种方法忽略了图像的固有二维结构和全局信息。为了解决这一问题,论文提出了GlobalMamba,该模型通过全局图像序列化方法捕捉图像的频域信息,从而增强对图像全局特征的理解。

创新点

频域转换和分割:首先通过离散余弦变换(DCT)将图像从空间域转换到频域,以获得其频谱分布。将频率范围划分成多个区间(低频到高频),每个区间分别转换回空间域,形成对应的子图像。

因果序列生成:对每个频段的子图像进行轻量化卷积处理,将其转化为因果序列,频率由低到高排列,以此捕捉图像的全局和局部特征。

全局Mamba框架:将生成的因果序列输入到基于 Mamba 的编码器中,通过多个模块进行特征提取,并在分类、检测、分割等下游任务上应用该特征。

实验结果

图像分类:在ImageNet-1K数据集上,GlobalMamba比其他基准模型在Top-1准确率上有所提升,尤其是比Vim和VMamba模型高0.2%-0.6%。

目标检测:在COCO数据集上,与VMamba相比,GlobalMamba在1x和3x训练策略下都表现更优。

语义分割:在ADE20K数据集上进行语义分割实验,GlobalMamba在平均交并比(mIoU)上表现比其他基准模型略高。

对比算法

论文主要与其他视觉 Mamba 模型(如Vim、VMamba、LocalMamba)进行了对比。

对比方法包括使用常见的二维平铺、局部窗口平铺和多方向扫描策略的视觉 Mamba 框架。

img
img
img
img

表5探究了频率段数K(代表频率划分的粒度,并直接决定了因果序列的长度)的设置,当K=4时,无论是模型大小均表现出理想的性能,进一步增大K虽然会增加序列长度,但不会带来显著的性能提升。

表6探究了因果Transformer的应用。通过将DeiT-S和Swin-T的原始自注意力机制修改为因果形式,并在ImageNet分类任务中应用,测试了提出的全局图像序列化(GIS)方法的有效性。

数据集

ImageNet-1K:用于图像分类。

COCO:用于目标检测和实例分割。

ADE20K:用于语义分割。

改进空间

由于高频成分对应较低的下采样率,仍然存在部分平铺操作。未来的工作将集中于完全避免简单的平铺操作,进一步提升因果序列的鲁棒性。

DCT将图像转换到频域后,直接通过简单的阈值将频谱划分为低频、中频和高频段。这种分割方式虽然能够捕捉到基本的频域信息,但划分标准比较固定,可能无法适应不同图像的特征分布。

当前使用轻量CNN直接对降采样后的子图生成标记(tokens),这一步也是平铺操作。虽然轻量化卷积能减少计算量,但这种方式对特征的提取较为浅层。

论文阅读

Abstract

视觉Mamba模型已展现出在视觉token数量上的线性复杂度表现,其高效性来自于对图像token的顺序处理。然而,大多数现有方法采用基于图像块的token化,并将其展开为一维序列进行因果处理,这忽略了图像的内在二维结构相关性。同时,通过局部图像块的顺序处理难以提取全局信息。本文提出了一种全局图像序列化方法,将图像转换为包含二维图像全局信息的因果token序列。我们首先使用离散余弦变换(DCT)将图像从空间域转换到频率域,并按照对应的频率范围排列像素。然后,将同频段内的每组像素重新转换回空间域,以生成一系列待token化的图像。基于所提出的全局图像序列化方法,我们构建了一个视觉Mamba模型——GlobalMamba,其因果输入格式能够更好地利用图像序列间的因果关系。大量实验验证了GlobalMamba的有效性,包括在ImageNet-1K上的图像分类、COCO上的目标检测以及ADE20K上的语义分割。

Introduction

img

Mamba模型(Gu & Dao, 2023;Lieber等,2024)由于其高效性,近期在深度学习社区备受关注。相比广泛采用的基于Transformer的架构,Mamba将计算复杂度从O(n²)降低到O(n),其中n表示输入序列的长度,这基于状态空间模型(SSMs)(Gu等,2022;2021a;b;Gupta等,2022)。Mamba通过一系列硬件友好算法(如并行扫描)加速了原先状态变量的顺序计算,提高了实践中的效率。Mamba在图像表示学习(Zhu等,2024;Ma等,2024)、视频理解(Li等,2024)和点云分析(Liang等,2024)等领域表现出竞争力和良好的应用潜力。

最近的研究将Mamba引入计算机视觉领域,通过将图像数据转换为一维token序列以适应其输入格式(Zhu等,2024;Liu等,2024;Huang等,2024;Yang等,2024)。具体来说,首先进行图像块嵌入,将图像转为一定分辨率的token,然后以行列顺序或局部窗口的方式在全局或局部范围内逐行、逐列展平这些token(Liu等,2024;Huang等,2024)。然而,此操作直接破坏了图像token之间的因果顺序。图像数据的空间域中的相邻区域通常编码相似的视觉信息,而空间上远离的区域则通常表现出显著差异,这一现象被称为图像的局部不变性属性。因此,简单的token展平过程可能导致原本空间上接近的图像块被置于展开序列中较远的位置,反之亦然。这种方法在Mamba框架下未能提供适当的图像建模顺序。此外,视觉Mamba中的每个图像token通常只包含局部信息,难以捕捉全局特征,从而在建模能力上存在一定的局限性。

为了解决这些问题,我们提出了GlobalMamba,一种带有全局图像序列化的改进视觉Mamba模型,如图1所示。我们首先通过离散余弦变换(DCT)将原始图像从空间域转换到频率域,以获取其频谱分布。我们将频谱划分为多个区间,从低频到高频,接着在频率域中按频段对像素分组,在分组过程中对超出指定频段的频率幅值置零。随后,我们通过逆变换将这些频谱分割回投影到空间域。每个频段分别进行token化处理,生成代表不同频段且具有广泛全局视觉感受野的token集合。我们按频率升序将这些token排列成一维因果序列,并将其输入Mamba特征提取过程。我们的GlobalMamba按照频率顺序构建因果token序列,使得模型能够以类似于人类的方式理解图像(即首先获取轮廓等低频信息,然后逐步补充细节)。GlobalMamba中的token本质上与离散频率区间相关联,增强了视觉数据光谱信息的全局封装。此外,因果序列的构建符合神经网络的频率原理,即倾向于优先拟合输入数据的低频分量,而低频信息在图像分类等视觉任务中通常起到更关键的作用。我们在多个任务上进行了大量实验以评估模型的有效性,包括ImageNet-1K上的图像分类(Russakovsky等,2015)、COCO上的目标检测(Lin等,2014)和ADE20K上的语义分割(Zhou等,2019)。与采用的基线相比,GlobalMamba在各项任务中表现出显著提升(如在ImageNet-1K上相比Vim提高了0.6%),展现了其优越性。

Related Work

Vision Mambas

卷积神经网络(CNNs)和视觉Transformer(ViTs)是计算机视觉中最常用的两类骨干网络。CNNs因其局部感受野的先验特性而在大多数视觉任务中长期担任基础骨干网络(He等,2016;Liu等,2022;Szegedy等,2015;Simonyan & Zisserman,2014)。特别是ResNet(He等,2016),通过高效的残差结构防止梯度消失问题,成为最广泛使用的卷积架构。同时,ViTs凭借其卓越的扩展能力和对多模态输入的适应性,成为一种新兴的视觉基础模型架构(Dosovitskiy等,2020;Liu等,2021;Li等,2022)。受Mamba(Gu & Dao,2023)在自然语言处理中成功应用的启发,近年来出现了一些将其应用于视觉理解任务的研究(Zhu等,2024;Liu等,2024;Huang等,2024;Yang等,2024;Hu等,2024;Patro & Agneeswaran,2024)。其中,Vision Mamba(Vim)(Zhu等,2024)率先将Mamba架构适配至计算机视觉领域的应用,通过将图像token展平为一维序列以适应其输入格式。随后,VMamba(Liu等,2024)和LocalMamba(Huang等,2024)进一步丰富了图像序列化过程,通过多方向扫描和局部窗口扫描等策略增强特征提取能力。此外,ZigMa(Hu等,2024)将Mamba架构进一步应用于视觉生成任务。然而,这些方法在图像处理中普遍需要对token进行展平处理,削弱了图像中固有的局部不变性特征。因此,生成的一维token序列缺乏必要的前后元素间和相邻token间的因果关系。此外,这些展平的token包含的多是空间上局限的信息,缺乏对全局上下文的综合把握。解决这一缺陷,增强因果关系及全局感知的保留,构成了我们所提方法的主要目标。

因果序列建模

递归神经网络(RNNs)(Jordan,1997;Hochreiter & Schmidhuber,1997;Cho,2014)是深度学习领域中最早的架构范式,具有天然捕捉顺序因果关系的能力。它们以序列数据作为输入,并沿序列进展递归地进行运算,各节点通过链状结构连接。因此,RNNs特别适用于时间序列和自然语言等具备时间因果性的样本。Mamba(Gu & Dao,2023)也具有类似RNN的中间隐状态变量,其状态变量之间的迭代方式同样遵循时间序列。因此,在没有因果顺序的情况下直接使用Mamba对视觉token建模缺乏合理性。因果序列建模也存在于Transformer的解码器部分(Kim等,2018)。当前,大型语言模型普遍采用仅包含解码器的架构,通过下一token预测来提取因果输入序列的特征(Radford,2018;Radford等,2019;Brown,2020;Touvron等,2023a;b;Dubey等,2024),适用于语言理解和生成任务。然而,将仅解码器架构直接应用于视觉分类任务时,效果并不理想,其准确率低于具备全局注意力交互的模型(Chen等,2020)。此外,Tian等(Tian等,2024)将原始的下一token预测转换为下一尺度预测,以增强序列间的因果性,从而提升视觉生成任务的质量。本文通过频率分割加强图像序列间的因果关系,以增强其在后续建模过程中的兼容性。

频率分析

频率分析在深度学习和计算机视觉领域展示出巨大的发展潜力。大量研究探讨了频率原理,认为神经网络在学习过程中倾向于优先拟合数据中的低频信号(Xu等,2019;2024;Luo等,2019)。这些研究利用频率原理来执行深度学习的解释性分析,并指导相应的训练过程。此外,一些工作利用频率分析来促进视觉任务的实际应用(Liang等,2023;Xu等,2020;Qin等,2021;Rao等,2023;Xie等,2021)。例如,Xu等(Xu等,2020)发现CNN对低频通道具有较高的敏感性,通过频域中的特征选择策略缓解了由于空间下采样导致的信息损失。Rao等(Rao等,2023)构建了GFNet,在频域内以对数线性复杂度建模长时间空间依赖性。本文利用频率划分构建视觉token序列,使得Mamba的建模能够遵循从低到高频的因果顺序,从而在一定程度上缓解图像局部不变性的破坏。每个图像token也可更专注于其相应频带内的全局信息,这为以往仅包含局部信息的视觉模型提供了更优的替代方案。

Proposed Approach

在本节中,我们首先简要介绍Mamba的基础知识。随后,详细阐述频率分割的具体原理和操作流程。最后,我们概述GlobalMamba,并进行相应的分析。

预备知识

img

Mamba模型通过将参数从时间不变转向时间可变来优化其特征表示能力。尽管时间不变参数有助于训练效率,但它对不同时间实例的输入缺乏特异性区分,从而限制了模型的特征表达能力。具体来说,Mamba使用不同的线性变换矩阵从输入中获得参数B和C,同时通过线性变换和相应的激活函数来确定参数∆。然而,时间可变参数的引入阻碍了模型转化为卷积形式进行并行训练。为此,Mamba引入了多种硬件优化算法来实现加速,例如并行扫描技术。因此,Mamba在保证训练效率的同时,将时间复杂度约束在O(n),相比复杂度为O(n²)的transformer模型具有一定的优势。

基于频率的全局图像序列化

img
图2详解

图2展示了GlobalMamba模型的频域全局标记化过程,分为以下几个关键步骤:

  1. 离散余弦变换 (DCT):原始图像 ( x(i, j) ) 通过离散余弦变换转换到频域,得到频谱图 ( F(u, v) )。在频域中,低频分量代表图像的整体轮廓信息,而高频分量表示更细致的局部特征。
  2. 频率分段 (Frequency Segmentation):在频谱图上,将频率划分为多个频段(如图中的 ( f_0, f_1, f_2, f_3 )),从低频到高频进行分割。每个频段对应不同的图像细节层次,通过分段可以有选择性地保留或提取不同频率的特征。
  3. 逆离散余弦变换 (IDCT):对每个频段的频谱部分进行逆变换,将其转换回空间域。这会得到多个图像子块 ( x_k(i, j) ),每个子块只包含某一频段的信息,表现出不同层次的细节。
  4. 降采样 (Downsampling):为了减少计算量和加速处理,对每个频段的子图像进行降采样。这一步会保留主要的结构信息,同时降低图像分辨率。
  5. 轻量化标记生成 (Lightweight Tokenizer):将降采样后的子图像输入一个轻量卷积神经网络(CNN),生成因果序列的标记。这些标记代表了图像在不同频率层次上的特征。
  6. 因果标记序列 (Causal Tokens):最终生成的因果标记序列按照频率顺序(从低频到高频)排列。模型可以利用这些序列化的标记在下游任务中进行进一步的处理,既保留了图像的全局特征,又能有效地捕捉到局部细节。

这一方法的关键在于通过频域分割和标记序列化,GlobalMamba能够有效地表示图像的全局结构与局部细节,并提高计算效率。

img

DCT后的频谱图表现出明显的低频系数在左上象限聚集,而高频成分则散布在右下角。同时,频谱图相对于主对角线呈现出对称性。考虑到频率成分的层次组织,我们将频谱图划分为离散的频率段,遵循从低频到高频的进程,并与主对角线垂直对齐。具体而言,设K为要划分的频率段数。我们不均匀地将主对角线划分为K个段,以考虑频率分布的非均匀性,使得第k个划分点距离左上角的距离为整个对角线长度的1/(2K−k)。在每个划分点处,相对于主对角线绘制一条垂直线。连续垂直线之间的间隔定义了每个频率段的频谱域,封装了该段内的相应频率分布。我们将与每个划分点对应的最大频率记为fk,划分后的频率带可以表示为(0, f1, ..., fK)。

img

全局Mamba

img
图3详解

图3展示了GlobalMamba的整体框架,分为几个主要部分,从图像的全局标记化到特征提取和分类过程。

  1. 全局标记化 (Global Tokenization)
    首先,图像通过图2所示的全局标记化模块被转换为因果序列标记(Causal Tokens),这些标记代表了不同频段的图像特征信息,作为模型的输入。
  2. Vision Mamba块 (Block of Vision Mambas)
    接下来,这些因果标记会经过多个Vision Mamba块进行特征提取。每个Vision Mamba块由以下几部分组成:

标准化 (Normalization):对输入数据进行标准化,以稳定模型训练过程并提升性能。
SSM块 (SSM Block):SSM(Structured State-space Model)块是用于捕捉全局上下文信息的模块,能够更好地理解图像的整体结构。
MLP层 (MLP Layer):多层感知机(MLP)层用于非线性映射,进一步提取特征。
在Vision Mamba块的结构中,标准化和SSM块的输出通过加法运算相结合,形成了特征提取的逐层堆叠模式。这些块会多次重复,以增强模型对特征的表达能力。

  1. 下采样或恒等映射 (Downsampling or Identity)
    对于某些金字塔架构(如VMamba),可能需要在不同层次进行下采样,以构建多尺度特征。在图中,"Downsampling or Identity"模块表示可以选择对特征进行下采样以实现特征金字塔,或者保持恒等映射(不改变特征维度)。
  2. 分类器或下游任务头 (Classifier or Downstream Head)
    经过多层Vision Mamba块后,提取的最终特征输入到分类器或下游任务头。这个模块用于具体的任务,例如图像分类、目标检测或语义分割。

总结
该框架通过全局标记化捕捉了图像的频域信息,并使用多层Vision Mamba块逐层提取特征,最终得到用于下游任务的高层次表示。整个结构设计旨在高效捕捉图像的全局和局部信息,并能适应多尺度任务需求。

img
img

此外,通过GlobalMamba获取的tokens本质上编码了更多的全局信息,特别是在低频频谱段中。当k ≥ 4时,生成的tokens数量为单一,从而使得该单个token能够代表该频率带的全局空间特征。同时,我们的方法遵循从低频到高频的因果顺序,明确地增加了整个token序列中低频信息的比例。这种方法与人类视觉理解过程及神经网络的频率先验原则一致,后者通常优先学习低频特征,以确保在拟合高频部分以获取详细信息之前对全局信息的全面理解。值得注意的是,低频成分往往对任务理解所需的解释能力产生主导影响。

Experiments

在本节中,我们进行了广泛的实验来证明Global-Mamba的有效性。我们最初在ImageNet-1K上进行图像分类训练,然后将预训练模型转移到下游任务,如目标检测和语义分割。传统上,我们提供了一系列烧蚀研究用于对比分析和调查。我们所有的实验都是在8个RTX 3090显卡上进行的。

图像分类

我们在ImageNet-1K数据集(Russakovsky等,2015)上评估了GlobalMamba在分类任务中的性能。该数据集包含超过128万张训练样本,涵盖1000类类别,验证集包含50,000张图像。我们采用了Vision Mamba(Vim)(Zhu等,2024)和VMamba(Liu等,2024)作为基线,并保持数据增强和优化器选择的一致性。根据参数规模,我们将模型分为GlobalMamba-M(Mini)、GlobalMamba-T(Tiny)、GlobalMamba-S(Small)和GlobalMamba-B(Base),如表1所示。训练轮数设为300,学习率调整使用余弦调度。我们对比了参数相似的方法,并提供Top-1准确率和FLOPs指标。实验结果如表2所示,其中标记为 * 的GlobalMamba模型表示在Vim上应用的简单结构,其他模型则表示在VMamba上应用的金字塔结构。结果显示,GlobalMamba在准确率方面持续优于基线方法。例如,在VMamba-S和VMamba-B模型上,我们的方法分别提高了0.3%和0.2%的分类准确率,展示了GlobalMamba方法的有效性。此外,GlobalMamba在token序列长度稍微扩展的情况下FLOPs略有增加,这将在第3.3节中进一步分析。

img

目标检测

我们在MSCOCO2017数据集(Lin等,2014)上进行了目标检测和实例分割的评估。该数据集包含超过118,000张训练图像、5,000张验证图像和超过40,000张测试图像。我们使用Mask-RCNN作为检测器,并在MMDetection代码库(Chen等,2019)中执行了1x和3x的训练调度。对比结果如表3所示。我们发现,基于SSM的方法在相似参数下优于视觉Transformer,并且GlobalMamba在不同模型规模和训练设置下均优于VMamba。例如,GlobalMamba-S在1x和3x调度下的box AP分别比VMamba-S高出0.3和0.2,而在mask AP上分别高出0.2和0.1。

img

语义分割

我们采用ADE20K数据集(Zhou等,2019)验证了GlobalMamba在语义分割中的有效性。该数据集包含20,210张训练图像、2,000张验证图像和3,000张测试图像,共有150个不同的语义类别。实验采用MMSegmentation框架(Contributors,2020)中的UPerNet(Xiao等,2018)作为分割器,训练调度为160k,结果如表4所示。结果显示,GlobalMamba在mIoU(SS)和mIoU(MS)指标上均优于VMamba。例如,GlobalMamba-S在mIoU(SS)上超越VMamba-S基线0.3,证明了我们提出框架的优越性。

img

实验分析

因果顺序

img

GlobalMamba的低到高频的因果建模顺序是其先验知识。为验证这一顺序的合理性和有效性,我们对比了高到低频的频率划分和随机频率区间选择的性能。三种频率划分的方法及其性能对比如图4所示。随机选择频率划分范围会对模型分类准确性产生不利影响,而高到低频序列的性能增益明显低于GlobalMamba采用的低频优先准则。

频率分段数

GlobalMamba通过多段频率划分生成相应的因果序列,因此频率段数K是一个关键因素,代表频率划分的粒度,并直接决定了因果序列的长度。为此,我们研究了不同划分数对模型性能的影响,并提供了直接复制和扩展Vim序列长度的性能对比(见表5)。首先,直接复制Vim中的token未能带来性能提升,甚至降低了原始模型的准确性。此外,我们观察到随着K值从2增加到6,分类性能先上升后趋于稳定。特别地,当K=4时,无论是模型大小均表现出理想的性能,进一步增大K虽然会增加序列长度,但不会带来显著的性能提升。因此,我们在主要实验中将K设为4。

因果Transformer的应用

除了Vision Mamba,解码器结构的Transformer也具备因果建模的能力。因此,我们通过将DeiT-S和Swin-T的原始自注意力机制修改为因果形式,并在ImageNet分类任务中应用,测试了提出的全局图像序列化(GIS)方法的有效性(见表6)。无论是简单结构还是金字塔类型的因果Transformer结构,性能的持续提升表明了GlobalMamba的灵活性和优越性。

img

Conclusion

在本文中,我们提出了GlobalMamba作为一种高效的视觉骨干网络用于表示学习。我们采用离散余弦变换(DCT)在频域中进行频带排列,构建从低频到高频的因果图像序列。我们进一步确保了与低频成分相关的token序列能够提取图像中的全局信息,从而显著增强了对视觉数据的全局理解。我们在多种视觉任务上验证了GlobalMamba的有效性,并进行了深入的消融研究以进行详细分析和对比。

局限性

由于高频成分对应较低的下采样率,仍然存在部分平坦化操作。未来的工作将集中于完全避免这种简单的平坦化操作,以获得更加稳健的因果序列。

© 版权声明
THE END
点赞15 分享
评论 抢沙发
头像
欢迎您留下宝贵的见解!
提交
头像

昵称

取消
昵称表情代码图片

    暂无评论内容