TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba

TL;DR

TransMamba tackles the resource-intensive training of sub-quadratic Mamba architectures by enabling rapid knowledge transfer from Transformer pre-trained models. It introduces a fast two-stage framework: selective subcloning to initialize Mamba from Transformer weights and adaptive multi-directional distillation to align representations across different scanning orders and depths. Empirical results across image classification, VQA, text-video retrieval, and multimodal reasoning show TransMamba outperforms baselines with substantially less data and maintains efficiency across PlainMamba, VMamba, ViM, VideoMamba backbones. This work bridges the efficiency of Mamba with the representational power of Transformers, offering a scalable path for robust multimodal foundation models.

Abstract

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.

Figures (6)

• Figure 1: (a) Comparisons of training cost (e.g., (left)), $CO_2$ emissions (e.g., (middle)) on uni-modal tasks among different methods. Compared to Mamba, TransMamba uses less data and requires shorter training time. (b) Comparison of data consumption, parameter count, and model performance across multi-modal task (e.g., (right)). TransMamba achieves performance comparable to existing models while requiring substantially less training data and fewer parameters, as measured by the POPE metric.
• Figure 2: Overview of the main components of our TransMamba. It supports uni-modal, multi-modal, and video inputs through a two-stage cross-architecture transfer process. The middle stage illustrates selective weight subcloning, enabling partial inheritance of Transformer-based parameters. The right stage depicts the adaptive multi-directional feature distillation, which combines layer-wise weighting with direction-aware distillation tailored to the multi-directional scanning structures in visual Mamba architectures.
• Figure 3: Overview of the Selective Subcloning Mechanism.
• Figure 4: Overview of the Adaptive Multi-Direction Distillation.
• Figure 5: Accuracy and loss curves under different architectures on image classification (e.g., PMamba and VMamba). (a)(b) depict the accuracy cures of PMamba and VMamba on ImageNet$^S$, highlighting the substantially accelerated convergence of the proposed TransMamba (e.g., TransPMamba-B achieving up to 2.7× speed-up) relative to baseline Mamba variants, while also surpassing both the original Mamba and the teacher model (e.g., DeiT-Pretrain) in final accuracy. (c) presents the loss curve of PMamba, further confirming the efficiency of TransMamba in optimization dynamics.