Elamur Research⌕
☰
home / papers / 2024 / Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

Paper · 2024 · Lianghui Zhu, Bencheng Liao, Qian Zhang et al.

Summary

Vision Transformers give each patch global, input-dependent context, but self-attention becomes slow and memory-hungry on high-resolution images with many tokens. The authors ask whether attention is needed at all. They propose Vision Mamba (Vim), a plain, non-hierarchical backbone built only from state space model (SSM) blocks derived from Mamba. Vim splits an image into patches, projects them linearly, adds position embeddings, and runs the token sequence through blocks that scan both forward and backward. This addresses two weaknesses of Mamba for images: it reads in only one direction, and it does not know where each token sits. Vim is pretrained with supervision on ImageNet and then used as a backbone for COCO detection and instance segmentation and ADE20K semantic segmentation. It beats DeiT on these tasks while being much cheaper at high resolution. This matters because Vim keeps the sequence-style, modality-agnostic design of ViT without its quadratic cost, which makes it a candidate backbone for high-resolution and multimodal foundation models.

Key points

  • On batch feature extraction at 1248×1248 (6084 tokens per image), Vim runs 2.8× faster than DeiT and uses 86.8% less GPU memory.
  • Vim's computation grows sub-quadratically and its memory grows linearly with sequence length. The authors present it as the first pure-SSM model to handle dense prediction.
  • Bidirectional design ablation:
    • A vanilla unidirectional Mamba block scores 73.2 top-1 on ImageNet and 32.3 mIoU on ADE20K.
    • The default, which adds a backward SSM preceded by a backward Conv1d, reaches 73.9 top-1 and 35.9 mIoU.
  • Classification head ablation on ImageNet-1K:
    • A class token in the middle of the sequence gives 76.1 top-1.
    • A head class token gives 75.2, a double token 74.3, mean pooling 73.9, and max pooling 73.4.
  • Detection and instance segmentation follow the ViTDet framework on COCO 2017. Code and models are released at github.com/hustvl/Vim (ICML 2024).

Interesting ideas

  • I1. Add a reverse scan instead of only flipping data. Randomly reversing sequences or alternating block directions was weaker than giving each block its own backward SSM branch. Any causal, recurrent model used for non-causal understanding tasks could benefit from this, for example long-text encoders.
  • I2. Place the summary token where recurrence sees the most. In a left-to-right recurrent model, a token at the start sees nothing. Putting the class token in the middle lets it gather context from both halves, and on ImageNet it also matches the tendency for objects to be centred. Where to insert special tokens is worth testing whenever a recurrent model must produce one pooled representation.
  • I3. Ablate on both global and dense tasks. The unidirectional block looked fine for classification but lagged clearly on segmentation, and one bidirectional variant showed the opposite pattern. Judging an architecture change by a single global metric can hide weaknesses that only appear in per-token tasks.
  • I4. Benchmark efficiency as input length grows. Reporting speed and memory across increasing resolutions, up to the point where the baseline runs out of memory, makes the real scaling advantage visible. This is a useful protocol for any comparison of long-sequence models.
  • I5. Avoid domain-specific priors to stay modality-agnostic. Vim deliberately skips 2D windows and pyramids and treats the image as a plain token sequence. This keeps the backbone compatible with text-style models and future multimodal pretraining, such as CLIP-style or masked modeling.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Huazhong University of Science and TechnologyWuhan · China59 works · 8,310 citations🔬Beijing Academy of Artificial IntelligenceBeijing · China46 works · 8,682 citations

Also on the byline: Horizon Robotics.

 Fields, topics and concepts

📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksComputer Vision and Pattern Recognition · 3,486 works

Cite

Cite as Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu and Xinggang Wang. 2024. Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model. arXiv:2401.09417.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2401.09417 · arXiv 2401.09417

@misc{zhu2024vision,
  title = {{Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model}},
  author = {Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu and Xinggang Wang},
  year = {2024},
  doi = {10.48550/arxiv.2401.09417},
  eprint = {2401.09417},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2401.09417}
}

References

80 references, 23 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Bao, H., Dong, L., Piao, S., and Wei, F. Beit: BERT pre-training of image transformers. In ICLR, 2022. URL https://openreview.net/forum?id=p-BhZSz59o4. ↗
  2. Baron, E., Zimerman, I., and Wolf, L. 2-d ssm: A general spatial layer for visual transformers. arXiv preprint arXiv:2306.06635, 2023. — [arXiv:2306.06635](https://arxiv.org/abs/2306.06635) find
  3. Bavishi, R., Elsen, E., Hawthorne, C., Nye, M., Odena, A., Somani, A., and Taşırlar, S. Introducing our multimodal models, 2023. URL https://www.adept.ai/blog/fuyu-8b. ↗
  4. Cai, Z. and Vasconcelos, N. Cascade r-cnn: High quality object detection and instance segmentation. TPAMI, 2019. find
  5. Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., and Joulin, A. Emerging properties in self-supervised vision transformers. In ICCV, 2021.
  6. Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
  7. Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mohiuddin, A., Kaiser, L., Belanger, D. B., Colwell, L. J., and Weller, A. Rethinking attention with performers. In ICLR, 2021. URL https://openreview.net/forum?id=Ua6zuk0WRH. ↗
  8. Dai, Z., Liu, H., Le, Q. V., and Tan, M. Coatnet: Marrying convolution and attention for all data sizes. NeurIPS, 34, 2021. find
  9. Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In CVPR, 2009. find
  10. Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
  11. Ding, J., Ma, S., Dong, L., Zhang, X., Huang, S., Wang, W., Zheng, N., and Wei, F. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02486, 2023.
  12. Ding, X., Zhang, X., Han, J., and Ding, G. Scaling up your kernels to 31x31: Revisiting large kernel design in cnns. In CVPR, 2022. find
  13. Dong, X., Bao, J., Chen, D., Zhang, W., Yu, N., Yuan, L., Chen, D., and Guo, B. Cswin transformer: A general vision transformer backbone with cross-shaped windows. In CVPR, 2022. find
  14. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020. d’Ascoli, S., Touvron, H., Leavitt, M. L., Morcos, A. S., Biro…
  15. Fang, J., Xie, L., Wang, X., Zhang, X., Liu, W., and Tian, Q. Msg-transformer: Exchanging local spatial information by manipulating messenger tokens. In CVPR, 2022. find
  16. Fang, Y., Wang, W., Xie, B., Sun, Q., Wu, L., Wang, X., Huang, T., Wang, X., and Cao, Y. Eva: Exploring the limits of masked visual representation learning at scale. In CVPR, 2023. find
  17. Fu, D. Y., Dao, T., Saab, K. K., Thomas, A. W., Rudra, A., and Re, C. Hungry hungry hippos: Towards language modeling with state space models. In ICLR, 2023. URL https://openreview.net/forum?id=COZDy0WYGg. ↗
  18. Ghiasi, G., Cui, Y., Srinivas, A., Qian, R., Lin, T.-Y., Cubuk, E. D., Le, Q. V., and Zoph, B. Simple copy-paste is a strong data augmentation method for instance segmentation. In CVPR, 2021. find
  19. Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. find
  20. Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021a. find
  21. Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A., and Ré, C. Combining recurrent, convolutional, and continuous-time models with linear state space layers. In NeurIPS, 2021b. find
  22. Gu, A., Goel, K., Gupta, A., and Ré, C. On the parameterization and initialization of diagonal state space models. In NeurIPS, 2022. find
  23. Gupta, A., Gu, A., and Berant, J. Diagonal state spaces are as effective as structured state spaces. In NeurIPS, 2022. find
  24. He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In CVPR, 2016.
  25. Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. Densely connected convolutional networks. In CVPR, 2017. find
  26. Islam, M. M. and Bertasius, G. Long movie clip classification with state-space video models. In ECCV, 2022. find
  27. Islam, M. M., Hasan, M., Athrey, K. S., Braskich, T., and Bertasius, G. Efficient movie scene detection using state-space transformers. In CVPR, 2023. find
  28. Jia, C., Yang, Y., Xia, Y., Chen, Y.-T., Parekh, Z., Pham, H., Le, Q., Sung, Y.-H., Li, Z., and Duerig, T. Scaling up visual and vision-language representation learning with noisy text supervision. In ICML, 2021.
  29. Kalman, R. E. A new approach to linear filtering and prediction problems. 1960. find
  30. Kenton, J. D. M.-W. C. and Toutanova, L. K. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, 2019.
  31. Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. In ICLR, 2020. URL https://openreview.net/forum?id=rkgNKkHtvB. ↗
  32. Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. In NeurIPS, 2012.
  33. LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998. find
  34. Li, J., Li, D., Xiong, C., and Hoi, S. Blip: Bootstrapping language-image pre-training for unified vision-language understanding and generation. In ICML, 2022a.
  35. Li, J., Li, D., Savarese, S., and Hoi, S. Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. arXiv preprint arXiv:2301.12597, 2023.
  36. Li, Y., Cai, T., Zhang, Y., Chen, D., and Dey, D. What makes convolutional models great on long sequence modeling? In ICLR, 2022b. find
  37. Li, Y., Mao, H., Girshick, R., and He, K. Exploring plain vision transformer backbones for object detection. In ECCV, 2022c. find
  38. Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L. Microsoft coco: Common objects in context. In ECCV, 2014.
  39. Liu, H., Li, C., Wu, Q., and Lee, Y. J. Visual instruction tuning. arXiv preprint arXiv:2304.08485, 2023.
  40. Liu, S., Chen, T., Chen, X., Chen, X., Xiao, Q., Wu, B., Kärkkäinen, T., Pechenizkiy, M., Mocanu, D., and Wang, Z. More convnets in the 2020s: Scaling up kernels beyond 51x51 using sparsity. arXiv preprint arXiv:2207.03620, 2022a. — [arXiv:2207.03620](https://arxiv.org/abs/2207.03620) find
  41. Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L., Wang, Y., Ye, Q., and Liu, Y. Vmamba: Visual state space model. arXiv preprint arXiv:2401.10166, 2024. find
  42. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
  43. Liu, Z., Mao, H., Wu, C.-Y., Feichtenhofer, C., Darrell, T., and Xie, S. A convnet for the 2020s. In CVPR, 2022b. find
  44. Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In ICLR, 2019. find
  45. Ma, J., Li, F., and Wang, B. U-mamba: Enhancing long-range dependency for biomedical image segmentation. arXiv preprint arXiv:2401.04722, 2024. — [arXiv:2401.04722](https://arxiv.org/abs/2401.04722) find
  46. Mehta, H., Gupta, A., Cutkosky, A., and Neyshabur, B. Long range language modeling via gated state spaces. In ICLR, 2023. URL https://openreview.net/forum?id=5MkYIYCbva. ↗
  47. Nguyen, E., Goel, K., Gu, A., Downs, G., Shah, P., Dao, T., Baccus, S., and Ré, C. S4nd: Modeling images and videos as multidimensional signals with state spaces. In NeurIPS, 2022. find
  48. Qin, Z., Yang, S., and Zhong, Y. Hierarchically gated recurrent neural network for sequence modeling. In NeurIPS, 2023. URL https://openreview.net/forum?id=P1TCHxJwLB. ↗
  49. Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In ICML, 2021.
  50. Radosavovic, I., Kosaraju, R. P., Girshick, R., He, K., and Dollár, P. Designing network design spaces. In CVPR, 2020.
  51. Rao, Y., Zhao, W., Zhu, Z., Lu, J., and Zhou, J. Global filter networks for image classification. Advances in neural information processing systems, 34:980–993, 2021. find
  52. Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556, 2014.
  53. Smith, J. T., De Mello, S., Kautz, J., Linderman, S., and Byeon, W. Convolutional state space models for long-range spatiotemporal modeling. In NeurIPS, 2023a. find
  54. Smith, J. T., Warrington, A., and Linderman, S. Simplified state space layers for sequence modeling. In ICLR, 2023b. URL https://openreview.net/forum?id=Ai8Hw3AXqks. ↗
  55. Strudel, R., Garcia, R., Laptev, I., and Schmid, C. Segmenter: Transformer for semantic segmentation. In ICCV, 2021. find
  56. Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language modelss. arXiv preprint arXiv:2307.08621, 2023.
  57. Szegedy, C., Liu, W., Jia, Y., Sermanet, P., Reed, S., Anguelov, D., Erhan, D., Vanhoucke, V., and Rabinovich, A. Going deeper with convolutions. In CVPR, 2015. find
  58. Tan, M. and Le, Q. Efficientnet: Rethinking model scaling for convolutional neural networks. In ICML, 2019.
  59. Tan, M. and Le, Q. Efficientnetv2: Smaller models and faster training. In ICML, 2021. find
  60. Tolstikhin, I. O., Houlsby, N., Kolesnikov, A., Beyer, L., Zhai, X., Unterthiner, T., Yung, J., Steiner, A., Keysers, D., Uszkoreit, J., et al. Mlp-mixer: An all-mlp architecture for vision. In NeurIPS, 2021. find
  61. Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. Training data-efficient image transformers & distillation through attention. In ICML, 2021a.
  62. Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. Training data-efficient image transformers & distillation through attention. In ICML, 2021b.
  63. Touvron, H., Bojanowski, P., Caron, M., Cord, M., El-Nouby, A., Grave, E., Izacard, G., Joulin, A., Synnaeve, G., Verbeek, J., et al. Resmlp: Feedforward networks for image classification with data-efficient training. TPAMI, 2022. find
  64. Wang, J., Sun, K., Cheng, T., Jiang, B., Deng, C., Zhao, Y., Liu, D., Mu, Y., Tan, M., Wang, X., et al. Deep high-resolution representation learning for visual recognition. TPAMI, 2020a. find
  65. Wang, J., Yan, J. N., Gu, A., and Rush, A. M. Pretraining without attention. arXiv preprint arXiv:2212.10544, 2022. find
  66. Wang, J., Zhu, W., Wang, P., Yu, X., Liu, L., Omar, M., and Hamid, R. Selective structured state-spaces for long-form video understanding. In CVPR, 2023a. find
  67. Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020b. find
  68. Wang, W., Xie, E., Li, X., Fan, D.-P., Song, K., Liang, D., Lu, T., Luo, P., and Shao, L. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In ICCV, 2021. find
  69. Wang, W., Dai, J., Chen, Z., Huang, Z., Li, Z., Zhu, X., Hu, X., Lu, T., Lu, L., Li, H., et al. Internimage: Exploring large-scale vision foundation models with deformable convolutions. In CVPR, 2023b. find
  70. Wang, W., Ma, S., Xu, H., Usuyama, N., Ding, J., Poon, H., and Wei, F. When an image is worth 1,024 x 1,024 words: A case study in computational pathology. arXiv preprint arXiv:2312.03558, 2023c. find
  71. Wu, H., Xiao, B., Codella, N., Liu, M., Dai, X., Yuan, L., and Zhang, L. Cvt: Introducing convolutions to vision transformers. In ICCV, 2021. find
  72. Xiao, T., Liu, Y., Zhou, B., Jiang, Y., and Sun, J. Unified perceptual parsing for scene understanding. In ECCV, 2018a. find
  73. Xiao, T., Liu, Y., Zhou, B., Jiang, Y., and Sun, J. Unified perceptual parsing for scene understanding. In ECCV, 2018b. find
  74. Xie, S., Girshick, R., Dollár, P., Tu, Z., and He, K. Aggregated residual transformations for deep neural networks. In CVPR, 2017. find
  75. Xing, Z., Ye, T., Yang, Y., Liu, G., and Zhu, L. Segmamba: Long-range sequential modeling mamba for 3d medical image segmentation. arXiv preprint arXiv:2401.13560, 2024. — [arXiv:2401.13560](https://arxiv.org/abs/2401.13560) find
  76. Yan, J. N., Gu, J., and Rush, A. M. Diffusion models without attention. arXiv preprint arXiv:2311.18257, 2023. — [arXiv:2311.18257](https://arxiv.org/abs/2311.18257) find
  77. Yang, J., Li, C., Zhang, P., Dai, X., Xiao, B., Yuan, L., and Gao, J. Focal self-attention for local-global interactions in vision transformers. arXiv preprint arXiv:2107.00641, 2021. — [arXiv:2107.00641](https://arxiv.org/abs/2107.00641) find
  78. Yu, W., Luo, M., Zhou, P., Si, C., Zhou, Y., Wang, X., Feng, J., and Yan, S. Metaformer is actually what you need for vision. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10819– 10829, 2022. find
  79. Zhou, B., Zhao, H., Puig, X., Xiao, T., Fidler, S., Barriuso, A., and Torralba, A. Semantic understanding of scenes through the ade20k dataset. IJCV, 2019. find
  80. A. Visualization GT Figure 5: Visualization comparison of DeiT-Ti (Touvron et al., 2021b) and our Vim-Ti on the Cascade Mask R-CNN (Cai & Vasconcelos, 2019) framework. Thanks to the long-range context learning of SSM, we can capture the very large object in the image, which the DeiT-Ti counterpart f… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhu-2024-vision-mamba-efficient-visual-representation/