MaTVLM: Hybrid Mamba-Transformer for Efficient Vision-Language Modeling

TL;DR

MaTVLM addresses the inefficiency of transformer-only VLMs by hybridizing Mamba-2 with transformer decoder layers. It initializes Mamba-2 from attention and trains with a single-stage distillation against a pre-trained teacher, improving convergence and reducing compute. Empirically, MaTVLM achieves up to $3.6\times$ faster inference and $27.5\%$ lower GPU memory usage while maintaining competitive accuracy across benchmarks. This approach demonstrates a practical path to resource-efficient yet capable Vision-Language Models in real-world deployments.

Abstract

With the advancement of RNN models with linear complexity, the quadratic complexity challenge of transformers has the potential to be overcome. Notably, the emerging Mamba-2 has demonstrated competitive performance, bridging the gap between RNN models and transformers. However, due to sequential processing and vanishing gradients, RNN models struggle to capture long-range dependencies, limiting contextual understanding. This results in slow convergence, high resource demands, and poor performance on downstream understanding and complex reasoning tasks. In this work, we present a hybrid model MaTVLM by substituting a portion of the transformer decoder layers in a pre-trained VLM with Mamba-2 layers. Leveraging the inherent relationship between attention and Mamba-2, we initialize Mamba-2 with corresponding attention weights to accelerate convergence. Subsequently, we employ a single-stage distillation process, using the pre-trained VLM as the teacher model to transfer knowledge to the MaTVLM, further enhancing convergence speed and performance. Furthermore, we investigate the impact of differential distillation loss within our training framework. We evaluate the MaTVLM on multiple benchmarks, demonstrating competitive performance against the teacher model and existing VLMs while surpassing both Mamba-based VLMs and models of comparable parameter scales. Remarkably, the MaTVLM achieves up to 3.6x faster inference than the teacher model while reducing GPU memory consumption by 27.5%, all without compromising performance. Code and models are released at http://github.com/hustvl/MaTVLM.

Figures (3)

• Figure 1: Comprehensive comparison of our MaTVLM.(a) Performance comparison across multiple benchmarks. Our MaTVLM achieves competitive results with the teacher model TinyLLaVA, surpassing existing VLMs with similar parameter scales, as well as Mamba-based VLMs. (b) Speed Comparison of Token Throughput. Tokens generated per second for different token lengths. Our MaTVLM achieves a $3.6\times$ speedup compared to the teacher model TinyLLaVA. (c) GPU Memory Usage Comparison. A detailed comparison of memory usage during inference for different token lengths, highlighting the optimization advantages with a 27.5% reduction in usage for our MaTVLM over TinyLLaVA.
• Figure 2: The proposed MaTVLM integrates both Mamba-2 and transformer components. The model consists of a vision encoder, a connector, and a language model same as the base VLM. The language model is composed of both transformer decoder layers and Mamba-2 decoder layers, where Mamba-2 layers replace only attention in transformer layers, while the other components remain unchanged. The model is trained using a knowledge distillation approach, incorporating probability distribution and layer-wise distillation loss. During the distillation training, only Mamba-2 layers and the connector are trainable, while transformer layers remain fixed.
• Figure 3: We initialize certain weights of Mamba-2 from attention based on their correspondence. Specifically, the linear weights of $x, B, C$ in Mamba-2 are initialized from the linear weights of $V, K, Q$ in the attention mechanism. The remaining parameters, including $\Delta_t$ and $A$, are initialized randomly.