Mamba-R: Vision Mamba ALSO Needs Registers

这篇论文与Vision Transformers类似，发现了Vision Mamba的特征图中也存在着众多的伪影。这些伪影对应于图像低信息背景区域中出现的高范数标记，而在Vision Mamba中，这些伪影要严重得多--即使是使用微小模型，它们也广泛地存在于背景区域。为了解决这个问题，我们采用了先前的解决方案，即在Vision Mamba中引入寄存器标记。为了更好地应对Mamba块的单向推理范式，我们引入了两个关键修改：1）均匀插入寄存器到输入标记序列中，2）将寄存器用于最终的决策预测。我们将这种新的架构称为Mamba-R。定性观察表明，与普通的Vision Mamba相比，Mamba-R的特征图更加清晰，更加关注语义上有意义的区域。定量上，Mamba-R具有更强的性能和更好的可扩展性。例如，在ImageNet基准测试中，我们的基准尺寸Mamba-R获得了82.9%的准确率，显著优于Vim-B的81.8%；此外，我们首次成功地将模型扩展到大型模型（即具有341M参数），获得了83.2%的竞争准确率（如果使用384x384的输入进行微调，则为84.5%）。对下游语义分割任务的额外验证也支持了Mamba-R的有效性。

本文旨在解决Vision Mamba模型中出现的高范数令牌在低信息背景区域中出现的问题，这些问题在该模型中比Vision Transformers更为严重。

为了解决这个问题，本文提出了一种新的架构Mamba-R，通过在输入令牌序列中均匀插入寄存器令牌并将其用于最终决策预测来改善Mamba块的单向推理范式。

实验结果表明，相比于原始的Vision Mamba模型，Mamba-R的特征图更加干净，集中在语义上更有意义的区域，并且在性能和可扩展性方面表现更好。在ImageNet基准测试中，Mamba-R的基本大小达到了82.9%的准确率，显著优于Vim-B的81.8%。此外，该模型还能够成功地扩展到大型模型大小，并达到了83.2%的竞争性准确率。

与该领域的相关研究包括Vision Transformers和其他基于Transformer的视觉模型，例如DeiT、PVT和CaiT等。