TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba

📄 arXiv: 2502.15130v2 📥 PDF

作者: Xiuwei Chen, Wentao Hu, Xiao Dong, Sihao Lin, Zisheng Chen, Meng Cao, Yina Zhuang, Jianhua Han, Hang Xu, Xiaodan Liang

分类: cs.CV

发布日期: 2025-02-21 (更新: 2025-10-09)

💡 一句话要点

TransMamba：利用Transformer预训练知识快速适配Mamba架构，实现通用架构迁移

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Mamba Transformer 知识迁移 知识蒸馏 预训练模型 多模态学习 架构适配

📋 核心要点

  1. 现有Mamba架构训练成本高昂，从头训练需要大量数据和时间，限制了其应用。
  2. TransMamba提出一种跨架构知识迁移方法，利用预训练Transformer知识初始化Mamba模型，加速训练过程。
  3. 实验表明，TransMamba在多种Mamba变体和下游任务上均优于基线方法，验证了其有效性。

📝 摘要（中文）

基于Transformer的架构已成为单模态和多模态基础模型的主干，这主要归功于其通过注意力机制实现的可扩展性，从而产生了丰富的公开预训练模型生态系统，如LLaVA、CLIP和DeiT等。与此同时，像Mamba这样新兴的亚二次复杂度架构，通过以线性复杂度实现全局上下文建模，提供了有希望的效率提升。然而，从头开始训练这些架构仍然需要大量的资源（例如，在数据和时间方面）。为了应对这一挑战，我们探索了一种跨架构的知识迁移范式，称为TransMamba，它有助于重用Transformer预训练的知识。我们提出了一个两阶段框架来加速基于Mamba的模型的训练，确保它们在单模态和多模态任务中的有效性。第一阶段利用预训练的Transformer模型来初始化Mamba架构的关键组件。为了弥合架构和维度上的差距，我们开发了一种选择性的权重子克隆策略和一种分层初始化方案，优先考虑前n层。在这一初始化的基础上，第二阶段引入了一种自适应的多方向知识蒸馏方法。这种机制采用逐层自适应缩放因子，使Mamba表示与其Transformer对应物对齐，同时适应多模态Mamba架构固有的扫描顺序变化。尽管使用减少的训练数据集和更紧凑的模型架构进行操作，TransMamba在各种基于Mamba的主干网络（例如，PlainMamba、Vmamba、ViM和VideoMamba）和下游任务（例如，图像分类、视觉问答、文本视频检索和多模态推理）中始终优于基线方法。所有代码和实现细节都将发布。

🔬 方法详解

问题定义：现有Mamba架构虽然在效率上具有优势，但从头开始训练需要大量的计算资源和数据，这限制了其在实际应用中的推广。如何高效地训练Mamba模型，使其能够利用已有的Transformer预训练知识，是一个亟待解决的问题。

核心思路：TransMamba的核心思路是利用Transformer预训练模型中的知识来初始化Mamba模型，从而避免从头开始训练Mamba模型。通过知识迁移，可以显著减少训练时间和数据需求，同时提高Mamba模型的性能。

技术框架：TransMamba采用两阶段训练框架。第一阶段是初始化阶段，利用预训练的Transformer模型来初始化Mamba架构的关键组件，包括选择性的权重子克隆策略和分层初始化方案。第二阶段是知识蒸馏阶段，引入自适应的多方向知识蒸馏方法，通过逐层自适应缩放因子，使Mamba表示与其Transformer对应物对齐。

关键创新：TransMamba的关键创新在于其跨架构的知识迁移方法，特别是选择性的权重子克隆策略和自适应的多方向知识蒸馏方法。选择性的权重子克隆策略能够有效地将Transformer模型中的知识迁移到Mamba模型中，而自适应的多方向知识蒸馏方法能够更好地对齐Mamba和Transformer模型的表示，从而提高Mamba模型的性能。

关键设计：在初始化阶段，TransMamba优先考虑前n层的初始化，认为这些层对于模型的性能至关重要。在知识蒸馏阶段，TransMamba使用逐层自适应缩放因子来调整Mamba和Transformer模型之间的表示对齐程度，以适应不同层的特征差异。损失函数的设计也考虑了多模态Mamba架构的扫描顺序变化，以确保知识迁移的有效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TransMamba在多种基于Mamba的主干网络（如PlainMamba、Vmamba、ViM和VideoMamba）和下游任务（如图像分类、视觉问答、文本视频检索和多模态推理）中均取得了显著的性能提升，证明了其有效性。具体性能数据将在论文发布后公开。

🎯 应用场景

TransMamba具有广泛的应用前景，可应用于图像分类、视觉问答、文本视频检索和多模态推理等领域。通过利用预训练的Transformer知识，TransMamba能够加速Mamba模型的训练，降低计算成本，并提高模型性能，从而推动Mamba架构在实际应用中的普及。

📄 摘要（原文）

Transformer-based architectures have become the backbone of both uni-modal and multi-modal foundation models, largely due to their scalability via attention mechanisms, resulting in a rich ecosystem of publicly available pre-trained models such as LLaVA, CLIP, and DeiT, etc. In parallel, emerging sub-quadratic architectures like Mamba offer promising efficiency gains by enabling global context modeling with linear complexity. However, training these architectures from scratch remains resource-intensive (e.g., in terms of data and time). Motivated by this challenge, we explore a cross-architecture knowledge transfer paradigm, termed TransMamba, that facilitates the reuse of Transformer pre-trained knowledge. We propose a two-stage framework to accelerate the training of Mamba-based models, ensuring their effectiveness across both uni-modal and multi-modal tasks. The first stage leverages pre-trained Transformer models to initialize critical components of the Mamba architecture. To bridge architectural and dimensional gaps, we develop a selective weight subcloning strategy and a layered initialization scheme that prioritizes the early $n$ layers. Building on this initialization, the second stage introduces an adaptive multi-directional knowledge distillation method. This mechanism employs layer-wise adaptive scaling factors to align Mamba representations with their Transformer counterparts, while accommodating the scanning order variations inherent to multi-modal Mamba architectures. Despite operating with a reduced training dataset and a more compact model architecture, TransMamba consistently outperforms baseline approaches across diverse mamba-based backbones (e.g., PlainMamba, Vmamba, ViM and VideoMamba) and downstream tasks (e.g., image classification, visual question answering, text-video retrieval and multimodal reasoning). All code and implementation details will be released.