TransMamba:从Transformer到Mamba的快速通用架构适配
TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba
AI总结:
TransMamba提出两阶段框架,利用Transformer预训练知识初始化Mamba架构,并通过自适应知识蒸馏,实现跨架构迁移,在多种任务上超越基线,加速Mamba模型训练。
摘要:
基于Transformer的架构已成为单模态和多模态基础模型的支柱,这主要归功于其通过注意力机制实现的扩展性,从而产生了丰富的公开预训练模型生态系统,如LLaVA、CLIP和DeiT等。与此同时,新兴的次二次复杂度架构(如Mamba)通过线性复杂度的全局上下文建模提供了有前景的效率提升。然而,从头训练这些架构仍然资源密集(例如,在数据和计算时间方面)。受此挑战的驱动,我们探索了一种跨架构知识迁移范式,称为TransMamba,该范式促进了Transformer预训练知识的重用。我们提出了一个两阶段框架来加速基于Mamba的模型的训练,确保其在单模态和多模态任务中的有效性。第一阶段利用预训练的Transformer模型来初始化Mamba架构的关键组件。为了弥合架构和维度上的差距,我们开发了一种选择性权重子克隆策略和一种分层初始化方案,该方案优先处理前n层。在此初始化的基础上,第二阶段引入了一种自适应多方向知识蒸馏方法。该机制采用逐层自适应缩放因子,将Mamba表示与其Transformer对应物对齐,同时适应多模态Mamba架构固有的扫描顺序变化。尽管使用缩减的训练数据集和更紧凑的模型架构,TransMamba在多种基于Mamba的主干网络(如PlainMamba、Vmamba、ViM和VideoMamba)以及下游任务(如图像分类、视觉问答、文本-视频检索和多模态推理)上始终优于基线方法。所有代码和实现细节将予以发布。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。