arXiv:2503.13440cs.CV2025-03ICCV被引 8

用Mamba替换部分Transformer,实现更快更省的视觉语言模型。

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

  • 混合Mamba-2与Transformer解码器,结合两者优势。
  • 推理速度提升3.6倍,显存减少27.5%且性能不降。
  • 适合追求高效推理的视觉语言模型应用者。

随着具有线性复杂度的RNN模型发展,Transformer的二次复杂度难题有望被突破。值得注意的是,新兴的Mamba-2已展现出与Transformer相当的性能,弥合了二者差距。然而,由于序列处理和梯度消失问题,RNN模型难以捕捉长程依赖,限制了上下文理解能力,导致收敛缓慢、资源消耗高,并在下游理解与复杂推理任务中表现不佳。本文提出一种混合模型MaTVLM,将预训练视觉语言模型(VLM)中的部分Transformer解码层替换为Mamba-2层。基于注意力与Mamba-2的内在关联,我们用对应注意力权重初始化Mamba-2以加速收敛。随后,采用单阶段知识蒸馏,以预训练VLM作为教师模型,将知识迁移至MaTVLM,进一步提升收敛速度与性能。此外,我们研究了不同蒸馏损失的影响。在多个基准上评估显示,MaTVLM在性能上可媲美教师模型,超越现有基于Mamba的VLM及同规模模型。尤为突出的是,其推理速度最高达教师模型的3.6倍,显存占用降低27.5%,且无性能损失。代码与模型已在http://github.com/hustvl/MaTVLM发布。

原文摘要 · Abstract (English)

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.

视觉语言模型Mamba高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。