Mamba-R: 视觉Mamba也需要寄存器。
0. TL;DR
这篇论文针对 Vision Mamba 模型中存在的特征图伪影问题，提出了 Mamba® 架构。通过在输入序列中均匀插入寄存器token，并在最终预测时回收利用这些寄存器，Mamba® 有效减少了特征图中的伪影，提升了模型性能和可扩展性。
1. 背景介绍
近年来，状态空间模型（SSMs）在序列建模任务中展现出巨大潜力。与 Transformer 的二次计算复杂度相比，SSMs 具有线性计算复杂度，能够更高效地处理长序列。Mamba 架构作为 SSMs 的一种实例，通过选择性扫描技术和硬件优化设计，实现了高效的训练和推理。它不仅在自然语言和音频处理任务中表现出色，还被成功扩展到视觉任务领域。然而，Vision Mamba 在处理视觉输入时暴露出特征图伪影问题，这些伪影对应于图像低信息背景区域中出现的高范数token，严重影响了模型的特征提取和优化稳定性。
2. Mamba® 模型
在 Vision Mamba 模型中，存在与 ViT 类似的特征伪影问题。这些伪影表现为在低信息背景区域中出现大量高范数token，它们在推理过程中倾向于丢弃局部信息以保留全局特征，从而影响特征图质量。研究表明，Vision Mamba 的伪影问题比 ViT 更为严重，主要体现在：
- 伪影在背景区域更为普遍，且具有更高的范数值。
- 即使在小型 Vision Mamba 模型中也存在伪影，并且随着模型尺寸的增加而加剧。
- 这些伪影对模型的特征提取和优化稳定性构成挑战，可能导致训练不稳定和扩展困难。
为解决 Vision Mamba 的特征伪影问题，论文提出了 Mamba® 架构，通过在输入序列中均匀插入寄存器token，并在最终预测时回收利用这些寄存器，具体包括以下两个关键修改：
- 稀疏分布的寄存器token：在输入序列中均匀插入寄存器token，以增强与局部token的有效交互，帮助捕获和保留不同位置的重要语义信息。
- 寄存器头部用于最终预测：在最终预测时，将寄存器token的输出通过线性层降维后拼接成一个全局表示，而不是像 ViT 那样丢弃寄存器。这种设计借鉴了多头注意力机制中拼接的思路，更有效地保留了所有寄存器的信息。
3. 实验分析
在 ImageNet 数据集的图像分类实验上，Mamba® 模型展现了显著的性能提升和可扩展性：
- 相比纯 Mamba 架构 Vim，Mamba® 在所有尺寸模型上均表现出更好的性能和可扩展性。
- 与混合架构（2D 卷积 + Mamba）相比，Mamba® 在参数量和准确率之间取得了更好的平衡。
- 与传统卷积网络和 Vision Transformers 相比，Mamba® 在准确率和效率方面具有竞争力。
在 ADE20K 数据集上，Mamba® 模型在语义分割任务中也展现了优越性能：
- 相比基于 ResNet 的模型，Mamba® 在参数量相近的情况下取得了更高的 mIoU。
- 与基于 Vision Transformers 的模型相比，Mamba® 在语义分割任务中表现出色，证明了其在视觉任务中的广泛适用性。
插入寄存器通常能带来性能提升，且存在一个最佳的寄存器数量。例如，在 Base 模型中，使用 12 个寄存器时性能最佳，过多或过少的寄存器都会导致性能下降。
均匀插入寄存器并将其全部用于最终预测的设计，相比将寄存器集中在序列某一端或仅使用部分寄存器的设计，能取得更好的性能。这表明寄存器在序列中的稀疏分布有助于 Vision Mamba 的特征提取。
论文展示了不同模型的特征图，直观地呈现出 Mamba® 相比原始 Vision Mamba 在减少伪影、聚焦语义有意义区域方面的优势。
此外，寄存器令牌有时能显示出不同的特征模式，突出图像中的不同对象或语义元素。
