用Transformer预训练知识加速Mamba模型训练,性能更优。
TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba
- 将Transformer预训练权重迁移至Mamba架构,分两阶段初始化与蒸馏。
- 在小数据和紧凑模型下,多任务表现超越基线方法。
- 适合想快速部署高效Mamba模型的研究者与工程师。
基于Transformer的架构已成为单模态与多模态基础模型的核心,得益于其通过注意力机制实现的可扩展性,已形成丰富的预训练模型生态(如LLaVA、CLIP、DeiT等)。与此同时,新兴的次二次复杂度架构如Mamba,通过线性复杂度实现全局上下文建模,展现出显著的效率优势。然而,从头训练这些架构仍需大量资源(如数据与时间)。为应对这一挑战,我们提出跨架构知识迁移范式TransMamba,实现Transformer预训练知识的复用。设计两阶段框架:第一阶段利用预训练Transformer模型初始化Mamba关键组件,提出选择性权重子克隆策略与分层初始化方案,优先优化前n层;第二阶段引入自适应多方向知识蒸馏机制,采用逐层自适应缩放因子对齐Mamba表示与Transformer对应结果,兼顾多模态Mamba架构的扫描顺序差异。即使在更小训练数据集和更紧凑模型架构下,TransMamba在多种Mamba骨干网络(如PlainMamba、Vmamba、ViM、VideoMamba)及下游任务(图像分类、视觉问答、文本-视频检索、多模态推理)中均持续优于基线方法。所有代码与实现细节将公开。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。