Mamba-R来了!视觉Mamba也需要Registers!

点击下方卡片,关注“CVer”公众号

AI/CV重磅干货,第一时间送达

点击进入—>【Mamba/多模态/扩散】交流群

添加微信号:CVer111,小助手会拉你进群!

扫描下方二维码,加入CVer学术星球!可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料,及最前沿应用!发论文/搞科研/涨薪,强烈推荐!

016d7484e0e2e2fcd1238d1138bbbf36.png

转载自:极市平台 | 编辑:科技猛兽

在CVer公众号后台回复:Mamba,下载综述!

导读

 

与原始的视觉 Mamba 架构相比,Mamba® 的特征图看起来更加清晰,更专注在语义有意义的区域。Mamba® 也获得了更好的性能,比如 Mamba®-B 在 ImageNet 上达到了 82.9% 的精度,优于 Vim-B 的 81.8%。 

本文目录

1 Mamba®:视觉 Mamba 也需要寄存器
(来自 Johns Hopkins University, University of Florida, UC Santa Cruz)
1 Mamba® 论文解读
1.1 背景:视觉 Mamba 的特征中也存在 "伪影"
1.2 视觉 Mamba 架构
1.3 视觉 Transformer 中的特征 "伪影"
1.4 视觉 Mamba 中的特征 "伪影":高范数异常 tokens 仍旧存在
1.5 视觉 Mamba 需要寄存器
1.6 实验结果

太长不看版

ICLR24 杰出论文|解决监督与自监督视觉Transformer「伪影」问题!只需要再多来几个token

本文发现 Vision Mamba 中存在着与 Vision Transformer 中相似的 "伪影 (artifacts)" 现象[1]。这些伪影指的是一张图片中信息量较低的背景区域中出现的高范数 tokens,而且作者发现该现象在视觉 Mamba 中更加严重。在 ViT 里面只有尺寸相对比较大的模型存在 "伪影" 现象,而对于视觉 Mamba 而言,即使是小模型中也同样存在这个现象。

为了缓解这个问题,作者遵循 Vision Transformers Need Registers 的做法将 register tokens 引入 Vision Mamba 。register tokens 指的是区别于原有模型 tokens 之外的,一些新的,与输入无关的 tokens。它们会与原来的 image tokens 一起喂入 ViT 或者视觉 Mamba 模型。

为了更好地应对 Mamba Block 的单向推理范式,本文引入了2个关键的改进:

  1. 把 register tokens 均匀地插入输入 token 序列中。之前在[1]中,是把几个 register tokens 并排插入 image tokens 一侧,再把它们一起喂给 ViT。

  2. 在视觉 Mamba 模型的末端,将所有的 register tokens 进行 concat 起来之后,形成最终预测的综合图像表征,来进行预测。

作者把新的架构命名为 Mamba®。与原始的视觉 Mamba 架构相比,Mamba® 的特征图看起来更加清晰,更专注在语义有意义的区域。Mamba® 也获得了更好的性能,比如 Mamba®-B 在 ImageNet 上达到了 82.9% 的精度,优于 Vim-B 的 81.8%。

303b1342dfe3b94a201140c0fec50f01.png
图1:Mamba® 的架构:通过将与输入无关的 register tokens 均匀地插入到输入序列中来解决 Vision Mamba 的伪影问题。在最后一层,将 register tokens 的输出连接起来,形成最终预测的全局表征

1 Mamba®:视觉 Mamba 也需要寄存器

论文名称:Mamba®: Vision Mamba ALSO Needs Registers (Arxiv 2024.05)

论文地址:

https://arxiv.org/pdf/2405.14858

项目主页:

https://wangf3014.github.io/mambar-page/

1.1 背景:视觉 Mamba 的特征中也存在 "伪影"

最近的状态空间模型 (State Space Models, SSMs) 展示了它们在序列建模任务中的巨大潜力。与 Transformer 中与序列长度呈现二次计算复杂度不同,SSM 模型具有线性计算复杂度,在处理长序列方面占有显著的效率优势。SSM 的一个典型架构是 Mamba[2],它使用选择性扫描技术以及一套硬件优化的设计。这种创新促进了线性计算复杂度和内存开销的循环模型的高效训练和推理。一些研究,比如 Jamba[3],也证明了 Mamba 架构能够实现与 Transformer 相当的很有竞争力的性能水平,尤其是在处理自然语言和音频方面。

此外,Mamba 架构也已成功应用于各种视觉任务[4][5]。这种扩展的动机主要来自处理高分辨率图像和视频所带来的计算量的挑战。这些数据类型通常会导致传统模型难以高效地处理的很长的输入序列。原因是对于较长的输入序列,卷积神经网络 (ConvNets) 等传统模型的感受野相对较小,而 Vision Transformers (ViTs) 会带来较大的计算和显存的代价。Vision Mamba (Vim) 架构有减少这些限制的潜力。它们不仅更高效地管理计算和内存需求,而且还在各种通用视觉任务 (包括分类、分割和图像生成) 中提供了强大的性能。

尽管性能具有竞争力,但本文观察表明,视觉 Mamba 内部存在着与 Vision Transformer 中相似的 "伪影 (artifacts)" 现象。"伪影" 指的是一些不正常 (范数很高) 的 tokens,具体而言这些 tokens 语义背景较少,但是却出乎意料地包含有丰富的全局信息[1]。

作者发现,在视觉 Mamba 架构中这种伪影问题不仅存在,而且实际上相比 ViT 更为严重。图1为 Vim[4]和本文 Mamba® 的特征图。如图 2 的第 2 列和第 4 列所示,"伪影" 在 Vim 的特征图中清晰可见。大量的比较大的激活值不仅涵盖了语义内容,而且还扩展到了信息量很少的背景区域。

8ea481a5a1a74fef49c8b1b7037562ea.png
图2:原始 Vision Mamba (Vim) 和本文 Mamba® 的特征图。Vim 的特征图中出现了大量的伪影,使模型难以关注图像中视觉上有意义的内容。相比之下,本文模型表现出了更加清晰的特征激活

1.2 视觉 Mamba 架构

SSM 的原始定义是一个线性时不变 (Linear Time-Invariant, LTI) 系统, 它将输入刺激 通过隐藏状态 投影到输出响应 。对于连续输入, 系统可以用一组线性常微分方程来表示, 如下所示:

ee63560aeb5e3cd329041f2be68bf8d5.png

其中, 是权重参数。

通过对这个常微分方程组进行离散化,可以使用连续时间 SSM 来处理离散输入,如语言、语音和图像像素。为此,模型可以求得解析解:

8da1ff608ba16cda29b23ae79fd36c98.png

然后用零阶保持器近似,得到离散模型如下:

195390ff1f470015adf0869f6491d878.png

其中, 为离散输入的变换参数, 为估计离散区间的可学习参数。与循环推理模式不同,这种结构化 SSM (S4) 允许通过卷积过程进行高效计算:

633c32be6cf6257667a8584ee31715a6.png

其中, 为卷积核, 卷积过程为: 。

然而,线性时间不变性的结构状态空间模型的性质极大地限制了其拟合上下文信息的能力,因此很难进行缩放,并实现与 Transformer 架构相当的性能。选择性状态空间模型,也称为 Mamba 或 S6[2],通过引入与输入有关的参数 来改进它,从而得到一个时变系统,可以对更复杂的输入进行建模。值得注意的是,通过扫描算法,Mamba 模块可以通过并行计算有效地训练和推理。

为了使 Mamba 用于视觉任务,首先通过 ViT 中的补丁嵌入将图像处理成顺序输入。然而,标准的 Mamba 是一个单向模型,其中序列中的每个 token 只能访问来自前一个 token 的信息。这一特性虽然可以很好地处理一维语言信号,但极大地限制了模型从二维视觉信号中收集上下文信息的能力。为了克服这个限制,一种常见的解决方案是重新配置 Mamba block 以进行双向扫描。具体来说,序列从开始到结束再次从结束到开始扫描一次,随后对两个扫描的输出进行平均以获得综合的表征。作者遵循这种扫描设计。

1.3 视觉 Transformer 中的特征 "伪影"

在 ViT 中,可以通过可视化 Self-attention block 中的激活分数来获得可解释的特征图。在经过恰当的预训练之后,理想的情况是注意力图会在前景物体处表现出较高的注意力得分,并且在北京出表现出较低的得分。但是在[1]中观察到的现象是:特征图中经常会出现大量的异常值,这些异常值的位置一般对应于低信息的背景区域,但却表现出异常高的注意力分数。这些异常值叫做 "伪影"。"伪影" 的特点是:

  • 位置一般对应于低信息的背景区域

  • 表现出异常高的注意力分数。

  • 一般具有很高的范数值。

  • 倾向于丢弃局部信息。

  • 携带的全局信息多。

更多细节可以参考原文[1]和:

ICLR24 杰出论文|解决监督与自监督视觉Transformer「伪影」问题!只需要再多来几个token

1.4 视觉 Mamba 中的特征 "伪影":高范数异常 tokens 仍旧存在

本文发现 Vision Mamba 模型也存在类似的问题。首先,通过计算原始 Vision Mamba 的全局输出和局部输出之间的 距离,作者观察到背景区域中的大量激活值,如图 2 所示。对其归一化后的进一步分析表明,这些背景激活也表现出较高的归一化值,类似于 ViT 中观察到的伪影。

例如,通过可视化图 3 中原始视觉 Mamba 的局部输出的归一化, 可以观察到背景中存在大量高范数 tokens,甚至模糊了前景和背景区域之间的区别。在数量上,作者在图 4(a) 中绘制了原始视觉 Mamba 的范数分布,它清楚地显示了许多具有高归一化异常值,证实了与[1]中中讨论的 ViT 中发现的一致性。

9081832f032849f28f47cad31d47bb88.png
图3:视觉 Mamba 不同层的局部图像 tokens 的 L2 范数。结果表明,与高范数标记相关的大量伪影出现在低信息区域,因此很难将主要对象与背景区分开来
50c59cee8fe89abffe59ddf0d9de2ecf.png
图4:不同层中局部输出的 L2 归一化值的分布,结果定量地表明 Mamba® 有效减少了高范数异常值的数量

同样值得注意的是,Vision Mamba 中的这些伪影也类似于 ViT 的 "伪影",倾向于携带全局信息。如图 5 所示,Vision Mamba 只需使用前 5% 的高范数 tokens 的均值作为全局特征就能获得 81.0% 的 ImageNet 精度,仅比对所有局部 tokens 进行池化低 0.1% 的精度。使用前 10% 或者 20% 的高范数 tokens 进一步使模型能够匹配使用全局池化的精度。相比之下,依赖剩余的 80% 低范数 tokens 会导致性能下降到 79.3%。

7d68a055e31d4393737e64c6d300ece6.png
图5:Vim-B 使用不同特征的 ImageNet 精度。使用一小部分高范数 tokens 来进行预测的精度明显高于低范数 tokens 的精度

然而,不同的是,作者观察到视觉 Mamba 中的伪影问题比 ViT 更严重:这些伪影在背景区域看起来更为普遍,并且表现出比 ViT 中观察到的更高的归一化值。如图 4(a) 所示,异常值 tokens 的平均范数随着层深度的增加而迅速增加,在第 23 层达到 4000 以上。与较浅特征中低于 100 的范数相比,这些极高的范数伪影很容易影响特征提取,并对模型优化提出了重大挑战,这可能会解释 Vision Mamba 中遇到的不稳定性问题和缩放困难。此外,ViT 中的 "伪影" 主要出现在较大尺寸的模型中,但是视觉 Mamba 中作者观察到小模型中也存在 "伪影" 现象。

这些观察结果进一步表明,伪影问题对于视觉 Mamba 模型至关重要。

1.5 视觉 Mamba 需要寄存器

遵循在 ViT 中去除伪影的解决方案,作者通过在视觉 Mamba 中引入 register tokens 来解决这个问题。作者将新的架构命名为 Mamba®。与之前仅在输入序列一端附加 register tokens 不同,作者在整个序列中更密集地分布 register tokens。作者认为这样可以:

  1. 更好地解决普遍分布的 "伪影" 问题。

  2. 由于 Mamba 架构的单向性质,有助于捕捉视觉 Mamba 中经常遗漏的全局表征。

作者遵循原始的视觉 Mamba[4]的主干架构,其中输入图像首先被分解为一系列不重叠的 patches,然后馈入一堆双向 Mamba block。基于这种简单的架构,作者做了以下2个简单而有效的修改来构建 Mamba®。

稀疏分布的 register tokens

Mamba® 的输入序列由 patch embedding 产生的 𝑚 个 image tokens 和在它们之间均匀插入的 𝑛 个 register tokens 组成。不像是 token 与位置无关的 Self-attention 模块,在 Mamba 模型里面,register tokens 的位置很重要。放置得合适就有利于 register tokens 与 image tokens 的有效交互。对于循环 Mamba 模型,稀疏分布的 register tokens 有助于捕获和保留不同位置的重要语义信息。作者实证性地证明了这种范式有利于性能的提升。

用于最终预测的 register head

ViT 在最终预测期间简单地丢弃了 register tokens, 与之不同的是, 作者观察到将它们作为全局表示可以为视觉 Mamba 带来显著改进。具体来说, 给定 𝑛 个 𝑑 维的 register 向量, 作者首先应用线性层将其维数降低 𝑟 倍, 然后将它们拼接成一个维度为 𝑛×𝑑/𝑟 的单个向量, 将其称为 register head。使用拼接操作而不是求均值是由 self-attention 中的多头机制启发的, 其中拼接更有效地保留来自所有 head 的信息。下图6 是 Mamba®的详细配置。

2f773994c154a5e4bfb51e5ce47f44b9.png
图6:Mamba® 模型的配置

此外,如图 7 所示,作者观察到在某些情况下,register 可以有趣地显示不同的特征模式,突出显示场景中的不同对象或语义元素,这是没有明确优化的一个有趣的方面。

02ea67b89688098f56a1ab759dcc9734.png
图7:不同 register tokens 的特征图。register 有时可以关注到图像的不同部分或者语义

1.6 实验结果

图像分类实验结果

如图 9 所示,Mamba® 在 ImageNet 上展示了强大的性能。与现有的纯 Mamba 架构 Vim 相比,Mamba® 展示出了显著的改进,Tiny 模型优于 Vim 1.3%,Small 模型高出 0.6%。而且与 Vim 相比,Mamba® 在可扩展性方面表现出显着的增强:作者训练了一个 99M 参数的 Base 模型,实现 82.9% 的精度,和一个 341M 的 Large 模型,实现了 83.2% 的精度。通过在 384×384 输入分辨率图片进行微调,可以将性能进一步提升为 84.5%,优于 ImageNet 分类中所有先前的 Mamba 模型。

3df40babb797b1226bcb42c1a9b9d088.png
图9:ImageNet 图像分类实验结果

语义分割实验结果

如图 10 所示,Mamba® 在 ADE20k 数据集上始终表现出卓越的语义分割性能。与 Vim-S 相比,Mamba®-S 实现了 0.4% mIoU 的改进。Mamba®-B 模型达到了 47.7% 的 mIoU,显着优于类似大小的 DeiT-B 模型 2.2% mIoU。此外,Mamba®-L (具有 377M 参数) 在分割任务中也显示出巨大的可扩展性,在 ADE20k 基准测试中实现了 49.1% mIoU。

15b380218cf3d66ff5b9aab04cced00c.png
图10:ADE20k 语义分割实验结果

消融实验:register tokens 的数量

作者探索了 register tokens 的数量对模型性能的影响。如图 11 所示,插入 register tokens 通常会导致一致的性能提升。此外,可以观察到简单地增加输出维度对性能几乎没有好处。例如,通过将 Vim-Base 的 hidden dim 从 768 维升至 2304,准确率仅提高 0.1%。此外,可以观察到使用 12 个 register token 是最佳点。如果继续增加 register tokens,最终的聚合特征维度很高,性能会饱和甚至可能下降 (比如最终维度升高到 4608)。

6f6a8b33441107783b28285c84edcc1a.png
图11:使用 Mamba®-Base (d=768) 对 register 进行消融实验。当 n=r=1 时,等价于具有 [cls] token 的视觉 Mamba (Vim) 模型

消融实验:register tokens 的设计

然后作者探索了 register tokens 的设计方案,即均匀地插入 register tokens 并使用它们进行最终的预测。结果如图 12 所示。首先,注意到性能对 register tokens 很敏感。例如,在序列的开头放置所有 register tokens 会导致性能下降 0.8% (82.1% vs. 82.9%)。同样,把 register tokens 放置在序列中间也会有 0.3% 的性能下降,表明 register tokens 的稀疏分布有利于视觉 Mamba 的特征提取。这些明显的性能差距突出了在 image tokens 之间均匀插入 register tokens 的必要性,因为 Mamba 的递归性使其对其在输入序列中的 token 位置敏感。

85c5bef9eccf01bcf2b70329c1d5f30c.png
图12:register 位置与预测范式的消融实验结果。Final Prediction 的含义是如何计算全局特征。R1 only 是指其中一个 register 并丢弃其他。Reduce and concat 是默认设置,利用线性层来减少 register 的维度并将它们连接为全局表征

此外,本文研究表明,registers 在提高视觉 Mamba 架构的性能方面起着重要的作用。使用均匀分布的 register tokens 的默认方法并在预测时使用这些 register tokens,得到的精度为 82.9%,超过了仅使用1个 register 的方法 (仅 R1,其余 tokens 被丢弃) 1.2%。这些结果证实 registers 构成了 Vision Mamba 架构的重要组成部分。

参考

  1. ^abcdefVision Transformers Need Registers

  2. ^abMamba: Linear-Time Sequence Modeling with Selective State Spaces

  3. ^Jamba: A hybrid transformermamba language model

  4. ^abcVision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

  5. ^VMamba: Visual State Space Model

在CVer公众号后台回复:Mamba,下载综述!

 
 

何恺明在MIT授课的课件PPT下载

 
 

在CVer公众号后台回复:何恺明,即可下载本课程的所有566页课件PPT!赶紧学起来!

ECCV 2024 论文和代码下载

在CVer公众号后台回复:ECCV2024,即可下载ECCV 2024论文和代码开源的论文合集

CVPR 2024 论文和代码下载

在CVer公众号后台回复:CVPR2024,即可下载CVPR 2024论文和代码开源的论文合集

Mamba、多模态和扩散模型交流群成立

 
 
扫描下方二维码,或者添加微信号:CVer111,即可添加CVer小助手微信,便可申请加入CVer-Mamba、多模态学习或者扩散模型微信交流群。另外其他垂直方向已涵盖:目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注:研究方向+地点+学校/公司+昵称(如Mamba、多模态学习或者扩散模型+上海+上交+卡卡),根据格式备注,可更快被通过且邀请进群

 
 
▲扫码或加微信号: CVer111,进交流群
CVer计算机视觉(知识星球)来了!想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料,欢迎扫描下方二维码,加入CVer计算机视觉(知识星球),已汇集上万人!

▲扫码加入星球学习
 
 
▲点击上方卡片,关注CVer公众号
整理不易,请赞和在看
  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值