Multimodal Mamba: Decoder-only Multimodal State Space Model via Quadratic to Linear Distillation

TL;DR

mmMamba tackles the quadratic complexity and cross-modality coupling in decoder-only vision-language models by distilling Transformer knowledge into a linear-complexity Mamba-2 state-space framework. The authors introduce a seeding strategy and a three-stage progressive distillation, enabling two architectures: mmMamba-linear (fully linear) and mmMamba-hybrid (interleaved Transformer/Mamba-2). Empirical results show competitive performance against both linear and quadratic baselines and substantial efficiency gains, including up to 20.6× speedups and significant GPU memory reductions at long context lengths. This approach removes dependence on pre-trained vision encoders and offers flexible deployment options for long-context multimodal understanding.

Abstract

Recent Multimodal Large Language Models (MLLMs) have achieved remarkable performance but face deployment challenges due to their quadratic computational complexity, growing Key-Value cache requirements, and reliance on separate vision encoders. We propose mmMamba, a framework for developing linear-complexity native multimodal state space models through progressive distillation from existing MLLMs using moderate academic computational resources. Our approach enables the direct conversion of trained decoder-only MLLMs to linear-complexity architectures without requiring pre-trained RNN-based LLM or vision encoders. We propose an seeding strategy to carve Mamba from trained Transformer and a three-stage distillation recipe, which can effectively transfer the knowledge from Transformer to Mamba while preserving multimodal capabilities. Our method also supports flexible hybrid architectures that combine Transformer and Mamba layers for customizable efficiency-performance trade-offs. Distilled from the Transformer-based decoder-only HoVLE, mmMamba-linear achieves competitive performance against existing linear and quadratic-complexity VLMs, while mmMamba-hybrid further improves performance significantly, approaching HoVLE's capabilities. At 103K tokens, mmMamba-linear demonstrates 20.6$\times$ speedup and 75.8% GPU memory reduction compared to HoVLE, while mmMamba-hybrid achieves 13.5$\times$ speedup and 60.2% memory savings. Code and models are released at https://github.com/hustvl/mmMamba

Figures (3)

• Figure 1: Comprehensive comparison of mmMamba.(a) Our mmMamba can build linear-complexity and hybrid decoder-only VLM by distilling the knowledge in Transformer to Mamba-2. (b) By distilling from the quadratic-complexity decoder-only VLM HoVLE, our mmMamba-linear achieves competitive performance against existing linear and quadratic-complexity VLMs with fewer parameters (e.g., 2$\times$ fewer than EVE-7B), while mmMamba-hybrid surpasses them across all benchmarks and approaches the teacher model HoVLE's performance. (c)-(d) We compare the speed and memory of mmMamba-linear and mmMamba-hybrid with the teacher model HoVLE on the same single NVIDIA 4090 GPU. mmMamba-linear maintains consistently low latency and memory usage, while mmMamba-hybrid's resource consumption scales significantly better than HoVLE. At 103K tokens, mmMamba-linear demonstrates 20.6$\times$ speedup compared to HoVLE and saves 75.8% GPU memory, while mmMamba-hybrid achieves 13.5$\times$ speedup and saves 60.2% GPU memory.
• Figure 2: Initialize Mamba-2 from Transformer. By comparing the mechanism similarity in Sec. \ref{['sec:preliminary']}, we directly inherit $\boldsymbol{W}_Q$, $\boldsymbol{W}_K$, $\boldsymbol{W}_V$, $\boldsymbol{W}_O$ parameters (blue) from trained Transformer layer and carefully initialize the extra parameters (orange) including $a$, $\boldsymbol{W}_\gamma$, $\boldsymbol{W}_{\text{conv}}$, and $\boldsymbol{W}_G$ in Mamba-2 to initially mimic the Transformer's behavior, providing a strong foundation for subsequent distillation.
• Figure 3: Progressive distillation pipeline of our mmMamba. We keep MLP layers, text and image patch embedding layers and freeze them in subsequent distillation training stages. Stage-1: Train the newly-introduced SSM-specific parameters while freezing inherited Transformer parameters in a layer-wise manner. Stage-2: Train all parameters to align Mamba's state representation with Transformer in a layer-wise manner. Stage-3: Train all the Mamba layers of the model to align the end-to-end behavior with the teacher Transformer-based VLM.