PaperlayerView exact revision on arXiv

Source-linked AI summary

MobileMamba: Lightweight Multi-Receptive Visual Mamba Network

Haoyang He, Jiangning Zhang, Yuxuan Cai, Hongxu Chen, Xiaobin Hu, Zhenye Gan, Yabiao Wang, Chengjie Wang, Yunsheng Wu, Lei Xie

arXiv:2411.15941v1cs.CV

TL;DR

CNNs’ local receptive fields limit long-range correlations, Transformers’ quadratic complexity raises high-resolution computational overhead, and low FLOPs in lightweight Mamba models do not necessarily mean fast inference. MobileMamba combines a three-stage network and an MRFFI module with WTE-Mamba, multi-kernel depthwise convolution, and identity mapping, alongside training and testing strategies. MobileMamba† achieves ImageNet-1K Top1 accuracies of 76.9/78.9/80.7/82.2/83.3/83.6 for models ranging from 200M to 4G FLOPs, with experiments validating transferability in high-resolution downstream tasks.

  • Problem

    CNNs’ local receptive fields limit long-range correlations, Transformers’ quadratic complexity raises high-resolution computational overhead, and low FLOPs in lightweight Mamba models do not necessarily mean fast inference.

    Source: §1 · pp. 1–2

  • Method

    MobileMamba combines a three-stage network and an MRFFI module with WTE-Mamba, multi-kernel depthwise convolution, and identity mapping, alongside training and testing strategies.

    Source: §1 · p. 2, §3.2 · p. 4, §3.3 · p. 5

  • Results

    MobileMamba† achieves ImageNet-1K Top1 accuracies of 76.9/78.9/80.7/82.2/83.3/83.6 for models ranging from 200M to 4G FLOPs, with experiments validating transferability in high-resolution downstream tasks.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    Experiments on ImageNet-1K and high-resolution downstream tasks validate the method’s effectiveness, efficiency, and transferability.

    Source: §5 · p. 8

    Mamba models still need substantial improvement in CPU and edge-device acceleration.

    Source: §5 · p. 9

Abstract

from arXiv · show

Previous research on lightweight models has primarily focused on CNNs and Transformer-based designs. CNNs, with their local receptive fields, struggle to capture long-range dependencies, while Transformers, despite their global modeling capabilities, are limited by quadratic computational complexity in high-resolution scenarios. Recently, state-space models have gained popularity in the visual domain due to their linear computational complexity. Despite their low FLOPs, current lightweight Mamba-based models exhibit suboptimal throughput. In this work, we propose the MobileMamba framework, which balances efficiency and performance. We design a three-stage network to enhance inference speed significantly. At a fine-grained level, we introduce the Multi-Receptive Field Feature Interaction(MRFFI) module, comprising the Long-Range Wavelet Transform-Enhanced Mamba(WTE-Mamba), Efficient Multi-Kernel Depthwise Convolution(MK-DeConv), and Eliminate Redundant Identity components. This module integrates multi-receptive field information and enhances high-frequency detail extraction. Additionally, we employ training and testing strategies to further improve performance and efficiency. MobileMamba achieves up to 83.6% on Top-1, surpassing existing state-of-the-art methods which is maximum x21 faster than LocalVim on GPU. Extensive experiments on high-resolution downstream tasks demonstrate that MobileMamba surpasses current efficient models, achieving an optimal balance between speed and accuracy.

1. Introduction

Lightweight vision models face a trade-off between CNNs’ local receptive fields, Transformers’ quadratic cost, and Mamba models’ sometimes slow inference despite low FLOPs. MobileMamba combines a three-stage network, multi-receptive-field feature interaction, and training/testing strategies, with reported gains on ImageNet-1K and high-resolution downstream tasks.

1. Introduction

  • CNNs’ local effective receptive fields limit long-range correlations, while Transformers’ quadratic complexity raises overhead, especially for high-resolution downstream inputs.The introduction also notes that CNN-based methods can improve downstream performance by increasing computational load.Source: §1 · pp. 1–2
  • Lightweight Mamba models capture long-range dependencies with linear computational complexity, but low FLOPs do not necessarily translate into fast inference, and current models show slow speed and poor performance.Figure 2 presents the accuracy-versus-speed comparison for Mamba-based methods.Source: §1 · p. 2, Figure 2 · p. 2
  • MobileMamba uses three stages and an MRFFI module that combines global WTE-Mamba features, multi-kernel local convolution, and identity mapping to reduce redundancy while integrating receptive fields.The module is described as enhancing high-frequency edge-detail extraction; knowledge distillation, extended training, and normalization-layer fusion are also used to improve performance or inference speed.Source: §1 · p. 2
  • MobileMamba† reaches ImageNet-1K Top-1 accuracies of 76.9/78.9/80.7/82.2/83.3/83.6 across 200M–4G FLOPs, surpassing existing CNN, ViT, and Mamba-based methods.The reported comparison appears in Figure 1.Source: Figure 1 · p. 1, §1 · p. 2
  • High-resolution downstream results include PSPNet improvements of +7.2 and +0.4 over MobileNetv2 and Mobile-ViTv2, respectively, using 8.5% and 11.2% of their FLOPs.The passage also reports results for RetinaNet, SSDLite, DeepLabv3, Semantic FPN, and PSPNet.Source: §1 · p. 2

2. Related Work

Lightweight visual models trade receptive-field range against computational cost, motivating CNN, ViT, and hybrid designs. State-space models offer linear-complexity long-range modeling, and visual Mamba variants explore scanning strategies, though the cited related work identifies a remaining performance gap.

2.1. Lightweight Visual Models

  • CNN-based lightweight models reduce computation with depthwise separable convolution and cheap operations, but their local receptive fields contrast with ViTs’ global modeling and quadratic cost.Lightweight ViTs and hybrid CNN-ViT methods seek to retain global receptive fields while reducing overhead or combining local and global information.Source: §2.1 · pp. 2–3

2.2. State Space Models

  • Visual SSM work includes bidirectional Mamba blocks, VMamba’s Cross-Scan, EfficientVMamba’s skip-sampling Efficient Scan, and LocalVim’s local-window scanning, yet the cited work says no lightweight Mamba network surpasses existing CNN and ViT methods.The paper frames its exploration around better performance, lower computational complexity, and faster inference.Source: §2.2 · p. 3

3. Methdology

MobileMamba combines a three-stage network with an MRFFI module and training and testing strategies. Its design choices target faster inference alongside improved classification performance and multi-receptive-field feature processing.

3.1. Coarse-Grained Design of MobileMamba

  • The three-stage network is faster than the compared four-stage design and improves both Top-1 and Top-5 accuracy by +0.4↑.The comparison uses similar-throughput conditions on ImageNet-1K.Source: §3.1 · p. 3

3.2. Fine-Grained Design of MobileMamba

  • WTE-Mamba combines bidirectional global modeling with wavelet-based local processing to enhance high-frequency edge extraction.Haar wavelets produce features at different frequency scales; local convolution and inverse wavelet transformation restore the feature-map size, and the local output is added to the global Mamba output.Source: §3.2 · p. 4
  • MK-DeConv splits local channels across different kernel sizes and concatenates their outputs to provide interaction across receptive fields.The design is described as extracting local information with varying effective receptive fields.Source: §3.2 · p. 4
  • Identity mapping on the remaining channels reduces high-dimensional feature redundancy and unnecessary computation.The paper presents this as a way to improve processing speed.Source: §3.2 · pp. 4–5
  • At similar FLOPs, Mamba outperforms the compared RNN alternatives in performance and throughput, with RWKV6 second to Mamba.TTT and xLSTM are reported to trail Mamba and RWKV6 on both measures.Source: §3.2 · p. 5

3.3. Training and Testing Strategies

  • Training uses soft knowledge distillation and extends the small model’s training from 300 to 1000 epochs; inference fuses batch normalization into preceding layers.The authors report that the small model had not fully converged under conventional 300-epoch training, while normalization fusion reduces the number of inference layers.Source: §3.3 · p. 5

4. Experiments

MobileMamba is evaluated on ImageNet-1K classification and high-resolution detection, instance segmentation, and semantic segmentation, with ablations examining its design and efficiency. The experiments report strong GPU throughput but also identify weaker CPU and mobile-device efficiency.

4.1. Implementation Details

  • ImageNet-1K models are trained from scratch for 300 epochs; enhanced models use 1000 epochs and knowledge distillation, while downstream GPU throughput is measured on one Nvidia L40S.The baseline classification resolution is 2242, and the enhanced-model teacher is TResNet-L.Source: §4.1 · p. 5

4.2. MobileMamba on ImageNet-1K Classification

  • MobileMamba-T2 exceeds SHViT-S1 by +0.8 in Top-1, T4 exceeds VRWKV-T by +1 in Top-1 with 33% of its FLOPs, and S6 gains 1.5 Top-1 over EfficientVMamba-T while reducing FLOPs by 18.5%.The comparison spans six model scales on ImageNet-1K.Source: §4.2 · pp. 5–6, Table 3 · p. 6
  • Increasing B1 input resolution produces B2 and B4 models with higher classification results and fewer FLOPs than other models, while training strategies raise T2† Top-1 by +3.3 and Top-5 by +1.7.The B2 and B4 models increase resolution without adding network depth or width.Source: §4.2 · p. 6, Table 3 · p. 6

4.3. MobileMamba on Downstream Tasks

  • On COCO detection and instance segmentation, B1 improves SSDLite by +2 over EMO-1M at 3202 resolution and gives RetinaNet ×4.3 higher GPU throughput than EfficientVMamba-T with +2.1 mAP b.For Mask R-CNN, B1 has +57% throughput over EMO-5M, with +1.3 mAP b and +1.0 mAP m.Source: §4.3 · pp. 6–7, Table 5 · p. 6, Table 4 · p. 6
  • On ADE20K, MobileMamba-B4 gains +1.3 mIoU over EMO-2M with DeepLabv3 and +2.4 with PSPNet; Semantic FPN gains +2.1 over EMO-5M using 22% of its FLOPs.For Semantic FPN, it also gains +1.1 mIoU over EdgeViT-XS using 20% of its FLOPs.Source: §4.3 · p. 7, Table 6 · p. 7

4.4. Extra Ablation and Explanatory Analysis

  • Compared with the baseline, MobileMamba improves Top-1 by +0.9 and Top-5 by +0.6 while increasing throughput by +729 images per second.The ablations also find that reducing scan directions can raise throughput with a slight performance decrease, while R=2 and D=1 outperform R=1 and D=16 at the same scan count.Source: §4.4 · pp. 7–8, Table 8 · p. 8
  • MobileMamba models average ×3.5 faster GPU throughput than EfficientVMamba, but lag Transformer-based models in CPU throughput and iPhone latency.The authors attribute the CPU and mobile gap to Mamba engineering implementation shortcomings.Source: §4.4 · pp. 7–8, Table 7 · p. 7
  • Knowledge distillation raises Top-1 by +2 and Top-5 by +0.7, and extending training to 1000 epochs adds +0.7 and +0.5, respectively; normalization fusion increases test speed by ×1.2%.The resulting 652M-FLOP model reaches 80.7 Top-1 and 95.2 Top-5, exceeding a 1080M-FLOP model without those strategies.Source: §4.4 · p. 8, Table 8 · p. 8

5. Conclusion

The conclusion presents MobileMamba as a framework balancing performance and efficiency, with MRFFI and training/testing strategies as key components. Experiments support its effectiveness, efficiency, and transferability, while CPU and edge-device acceleration remain areas for improvement.

5. Conclusion

  • MobileMamba is designed to balance performance and efficiency, and its effectiveness and transferability are validated on ImageNet-1K and high-resolution downstream tasks.The conclusion also says its training and testing strategies further enhance performance and efficiency.Source: §5 · p. 8
  • MRFFI enhances perception across receptive fields while preserving high-frequency features and inference efficiency.The paper describes MRFFI as supporting larger effective receptive fields and fine-grained high-frequency edge extraction.Source: §5 · p. 8
  • The authors identify CPU and edge-device acceleration as areas needing substantial improvement for Mamba models.Future work will focus on inference capabilities across devices, particularly efficiency.Source: §5 · p. 9

Supplementary Material Overview

The supplementary material provides further MobileMamba analyses, downstream-task results, and implementation materials.

Supplementary Material Overview

  • It includes detailed fine-grained design, kernel-size, DropPath, and effective receptive field analyses on ImageNet-1K.It also covers pretrained models at different resolutions for downstream tasks.Source: Supplementary Material Overview · p. 12
  • Additional results cover object detection and semantic segmentation, including Mask R-CNN experiments on ADE20K.The material also reports object detection results using different frameworks on MS-COCO 2017.Source: Supplementary Material Overview · p. 12
  • The supplementary Codes folder contains model training and testing code and model log files.Source: Supplementary Material Overview · p. 12

A.1. Fine-Grained design analysis

The supplementary analyses examine channel ratios, dimensions, depth, convolution kernels, and DropPath to guide MobileMamba's design choices.

A.1. Fine-Grained design analysis

  • The design maps 10% of channels in the last two stages to reduce computational load, selects dimensions {192, 384, 448}, and chooses depth {1, 2, 2} after throughput-performance trade-offs.Increasing depth significantly reduces throughput; an extreme one-layer-per-stage design had higher throughput but poor performance.Source: §A.1 · p. 12, Table A1 · p. 12

A.2. Effect of kernel sizes

  • Reducing kernel size as feature-map scale decreases across stages improves performance, while using the same kernel size across stages yields similar results.Source: §A.2 · p. 12, Table A4 · p. 13

A.3. Effect of DropPath

  • For B1, DropPath rate 0.03 gives the best performance, raising Top-1 by 0.2 over no DropPath; higher rates add no improvement.The shallow T2, T4, and S6 models do not use DropPath.Source: §A.3 · p. 13, Table A5 · p. 13

A.4. Visualization of the ERF of different model methods

At a fixed 224×224 input resolution, MobileMamba shows a larger and more intense effective receptive field (ERF) at each stage than the compared lightweight models.

A.4. Visualization of the ERF of different model methods

  • MobileMamba exhibits a larger and more intense ERF at each stage than the compared state-of-the-art methods.The comparison includes MobileNet, EfficientViT, and EMO.Source: §A.4 · p. 13
  • The ERF visualizations use a fixed 224×224 input resolution.Source: §A.4 · p. 13
  • MobileMamba and EfficientViT use three stages, whereas MobileNet and EMO use the traditional four-stage approach.Source: §A.4 · p. 13

A.5. Analysis of Pre-trained Models with Different Resolutions for Downstream Tasks

The section compares downstream performance from two otherwise identical pre-trained MobileMamba variants that differ in ImageNet-1K pre-training resolution. The preferred variant depends on the task: B1 for object detection and B4 for semantic segmentation.

A.5. Analysis of Pre-trained Models with Different Resolutions for Downstream Tasks

  • MobileMamba-B1 outperforms MobileMamba-B4 on all reported metrics for SSDLite, RetinaNet, and Mask RCNN object detection.The variants differ only in ImageNet-1K pre-training resolution: 256 for B1 and 512 for B4.Source: §A.5 · p. 13
  • The compared pre-trained variants, MobileMamba-B1 and MobileMamba-B4, differ only in ImageNet-1K pre-training resolution: 256 versus 512.The section presents ERF visualizations alongside the downstream resolution analysis.Source: §A.5 · p. 13, Figure A1 · p. 13
  • MobileMamba-B4 outperforms MobileMamba-B1 on all reported metrics for DeepLabv3, Semantic FPN, and PSPNet semantic segmentation.The authors attribute this pattern to B4's higher segmentation accuracy and B1's stronger semantic information, and select the corresponding variant as each task's backbone.Source: §A.5 · p. 13

B.1. Detailed Object Detection Results

The detailed object-detection results cover MobileMamba on MS-COCO 2017 using SSDLite, RetinaNet, and Mask R-CNN.

B.1. Detailed Object Detection Results

  • The section identifies SSDLite, RetinaNet, and Mask R-CNN as the object-detection methods covered by its detailed results.Source: §B.1 · p. 13
  • Table A6 reports MobileMamba object-detection results using SSDLite and RetinaNet on MS-COCO 2017.The caption marks 512×512-resolution results with a dagger.Source: §B.1 · p. 13, Table A6 · p. 14
  • Table A7 reports MobileMamba Mask R-CNN object-detection results on MS-COCO 2017.Source: §B.1 · p. 13, Table A7 · p. 14

B.2. Detailed Semantic Segmentation Results

The detailed semantic-segmentation results evaluate MobileMamba on ADE20K using DeepLabv3, Semantic FPN, and PSPNet.

B.2. Detailed Semantic Segmentation Results

  • The section also identifies SegFormer among the methods included in its detailed ADE20K semantic-segmentation results.Source: §B.2 · p. 14
  • Table A8 reports MobileMamba semantic-segmentation results on ADE20K using DeepLabv3, Semantic FPN, and PSPNet.Source: Table A8 · p. 14, §B.2 · p. 14

Related papers

Loading 2411.15941v1…