Mamba-Reg: Vision Mamba Also Needs Registers.

Area of Science:

• Computer Vision
• Deep Learning Architectures

Background:

• Vision Mamba, like Vision Transformers, exhibits artifacts in feature maps, particularly high-norm tokens in background areas.
• These artifacts are more severe in Vision Mamba, appearing even in small models and activating extensively in background regions.

Purpose of the Study:

• To mitigate feature map artifacts in Vision Mamba.
• To introduce an improved architecture, Mamba®, by modifying register token integration for Mamba's unidirectional inference.

Main Methods:

• Introduced register tokens into Vision Mamba, a known solution for artifact mitigation.
• Implemented two key modifications: even insertion of registers throughout the token sequence and recycling registers for predictions.
• Developed the Mamba® architecture incorporating these modifications.

Main Results:

• Mamba® demonstrates cleaner feature maps with improved focus on semantically meaningful regions compared to vanilla Vision Mamba.
• Achieved higher accuracy on ImageNet (83.0% for Mamba®-B vs. 81.8% for Vim-B).
• Successfully scaled to a 341M parameter model, achieving competitive accuracies (83.6% and 84.5% for different input sizes).

Conclusions:

• Mamba® effectively reduces artifacts and enhances feature map quality in Vision Mamba.
• The proposed modifications lead to superior performance and better scalability.
• Mamba® shows strong efficacy in downstream tasks like semantic segmentation.