Abstract
Transformers have been favored in both uni-modal and multi-modal foundation models for their flexible scalability in attention modules. Consequently, a number of pre-trained Transformer models, e.g., LLaVA, CLIP, and DEIT, are publicly available. Recent research has introduced subquadratic architectures like Mamba, which enables global awareness with linear complexity. Nevertheless, training specialized subquadratic architectures from scratch for certain tasks is both resource-intensive and time-consuming. As a motivator, we explore cross-architecture training to transfer the ready knowledge in existing Transformer models to alternative architecture Mamba, termed TransMamba. Our approach employs a two-stage strategy to expedite training new Mamba models, ensuring effectiveness in across uni-modal and cross-modal tasks. Concerning architecture disparities, we project the intermediate features into an aligned latent space before transferring knowledge. On top of that, a Weight Subcloning and Adaptive Bidirectional distillation method (WSAB) is introduced for knowledge transfer without limitations on varying layer counts. For cross-modal learning, we propose a cross-Mamba module that integrates language awareness into Mamba's visual features, enhancing the cross-modal interaction capabilities of Mamba architecture. Despite using less than 75% of the training data typically required for training from scratch, TransMamba boasts substantially stronger performance across various network architectures and downstream tasks, including image classification, visual question answering, and text-video retrieval. The code will be publicly available.
Abstract (translated)
Transformer模型因其在注意力模块中的灵活可扩展性而在单模态和多模态基础模型中受到青睐。因此，已经有许多预训练的Transformer模型公开可用，如LLaVA、CLIP和DEIT等。最近的研究引入了次二次架构（例如Mamba），这种架构能够以线性复杂度实现全局感知能力。然而，从头开始专门为某些任务训练此类次二次架构既耗时又耗费资源。 为此动机，我们探索了一种跨架构训练方法，将现有Transformer模型中的知识转移到替代架构Mamba上，并将其命名为TransMamba。我们的方法采用两阶段策略来加速新的Mamba模型的训练过程，在单模态和跨模态任务中都能保证有效性。针对架构差异，我们在转移知识之前将中间特征投影到对齐的潜在空间内。在此基础上，我们引入了一种无限制层计数的知识传输方法——权重子克隆与自适应双向蒸馏法（WSAB）。此外，为了增强Mamba架构在跨模态任务中的交互能力，我们提出了一种交叉Mamba模块，该模块将语言感知整合到Mamba的视觉特征中。 尽管训练TransMamba使用的数据量不到从头开始训练所需的75%，但其在各种网络架构和下游任务（包括图像分类、视觉问答以及文本-视频检索）中的表现依然显著优于其他模型。我们将公开提供相应的代码。
URL
https://arxiv.org/abs/2502.15130
