TransMamba：跨架构训练实现Transformer知识向Mamba迁移，两阶段策略及多方法赋能单多模态任务 !

Transformer架构在单模态和多模态基础模型中备受青睐，因其在注意力模块上的灵活可扩展性。因此，许多预训练的Transformer模型，例如LLaVA、CLIP和DEIT，已公开可用。近期研究引入了亚二次复杂度的架构，如Mamba，这使得全局感知具有线性复杂度。然而，从零开始训练特定任务的专业亚二次复杂度架构既耗费资源又耗时。 作为推动因素，作者探索了跨架构训练的方法，将现有Transformer模型中的知识转移到替代架构Mamba中，称为TransMamba。作者的方法采用两阶段策略来加速训练新的Mamba模型，并确保其在单模态和跨模态任务中的有效性。 针对架构差异，作者在转移知识之前将中间特征投影到对齐的潜在空间中。此外，作者提出了一种名为Weight Subcloning and Adaptive Bidirectional Distillation (WSAB)的方法，用于无条件限制地进行知识转移。 对于跨模态学习，作者提出了一个跨Mamba模块，将语言 Aware 整合到Mamba的视觉特征中，增强了Mamba架构的跨模态交互能力。尽管使用不到75%的通常所需的训练数据从零开始训练，但TransMamba在多种网络架构和下游任务上，包括图像分类、视觉问答和文本视频检索，表现出了显著的优势。

1. Introduction

Transformer [53] 架构对计算机视觉领域产生了深远的影响，它们灵活的注意力模块结构被认为是其成功的关键因素之一。尽管这些架构非常流行，但 Transformer 遇到了计算复杂度问题，因为其注意力机制的计算复杂度呈二次方增长 [2]，这导致了计算和内存使用的增加。因此，这给模型优化和扩展带来了重大挑战，阻碍了它们的广泛应用。为应对这一挑战，近期的研究引入了一些亚二次的架构，例如 Mamba 和 RWKV [13, 40]。然而，为了针对各种下游任务从头开始训练专门的亚二次模型，会面临显著的计算负担，并产生更高的二氧化碳排放量。幸运的是，作者观察到许多基于 Transformer 的预训练模型，例如 LLaVA [42] 和 CLIP [42] 等已经公开可用。

自然地提出一个问题：作者是否可以从这些广泛可用的基于 Transformer 的预训练模型向亚二次模型，如 Mamba，转移知识？本文旨在研究如何利用广泛可用的基于 Transformer 的预训练模型来为亚二次模型，如 Mamba，提供更经济且高效的训练方法。

在作者的研究中，作者着重解决了两个关键挑战：

1. 跨架构学习问题：这涉及到将一个架构框架（预训练的Transformer模型）的知识适应到另一个框架（Mamba模型）。此外，作者必须确保转移的知识能够保持有效性，并增强目标模型的性能，而不损害其结构完整性。
2. 为基于SSM的模型配备跨模态交互能力：这包括开发方法以无缝集成和处理来自不同模态（例如文本和图像）的信息，从而提高基于SSM的模型在复杂任务中的通用性和应用范围。作者应该确保基于SSM的模型能够理解和利用各种类型数据之间的关系。

针对第一个挑战，作者的方法引入了一种两阶段策略，通过利用预训练Transformer模型的知识来提高基于SSM模型的训练效率和性能。作者首先使用一个简单的MLP层将不匹配的表示转换到对齐的潜在空间中，从而实现了跨架构中间层的知识蒸馏。然而，最后的Transformer层与不同 Level 的基于SSM的层之间存在一致性不一致的问题。

作者采用余弦相似度来评估各层之间的相关性，并提出自适应知识蒸馏策略，即高相似度特征接收较小的权重系数，反之亦然。为了解决双路逆Mamba处理图像信息时的不一致优化问题，作者提出了双向（前向/后向）蒸馏策略分别优化不同的处理过程。仅仅依赖于蒸馏可能不能有效地促进知识转移，特别是在大型模型中。权重复用对于确保传递知识的准确性和效率至关重要。具有类似架构的模型权重可以直接用于初始化。然而，由于架构差异，跨架构权重复用具有挑战性。受到[45]的启发，作者从预训练Transformer模型中初始化基于SSM的模型。两者之间的区别在于架构相似性和跨架构（注意力机制和SSM）。对于架构差异，除了QKV投影层外，作者重新训练其余权重。对于维度差异，作者使用部分重要权重来自预训练Transformer模型以匹配基于SSM的模型的维度。权重子克隆也有助于加速收敛。

在第二个挑战中，现有基于SSM的模型缺乏跨模态交互能力[9, 27]，这对其应用于各种多模态场景造成了重大限制。受跨注意力机制[5]的启发，作者提出了一种跨Mamba模块，用于将多模态理解能力整合到SSM中。通过堆叠Mamba和跨Mamba模块，新型Mamba模型得以有效处理视觉-语言任务，并且保持了高效性和性能。

本工作主要选用Mambas作为SSM-Based模型的实现。作者采用两阶段方法，将预训练Transformer的知识转移到Mamba及其变体中。作者的方法速度快且具有普适性，能够提升Mamba的训练效果，并适用于多种架构。此外，作者通过包括图像分类、视频检索和视觉问答在内的详尽实验设置验证了所提方法的泛化能力。

1. 作者提出了一种快速且通用的两阶段框架，该框架将现有预训练Transformer模型的知识转移到基于SSM的新模型中，从而在较低的成本下提升训练效率和后续性能。
2. 重量子克隆与自适应双向蒸馏（WSAB）：作者使用重量子克隆技术实现跨架构转移，高效利用预训练知识。
3. 此外，作者还提出了一种适用于Mamba架构的自适应前向/后向蒸馏方法，并赋予Mamba多模型交互的能力。
4. 全面验证：作者在广泛的 Backbone 网络架构和多种应用场景下验证了所提出的方法，包括视觉问答、视频检索和图像分类。

2. Related Work

Transformers 变换了范式，在视觉任务中展示了强大的感知能力，并已成为视觉领域最有前景的基础模型之一。单模态任务：早期基于ViT的模型通常需要大规模的数据集进行训练，并且具有相对简单的架构。随后，DeiT 采用了一些训练技术来解决优化过程中遇到的问题，研究逐渐倾向于在网络设计中融入视觉感知的归纳偏置。例如，社区提出了分层ViT[8, 11, 35]，逐步降低 Backbone 网络的特征分辨率。此外，其他研究则利用卷积神经网络(CNN)的优势，引入卷积操作[6, 8, 54] 或者设计结合了CNN和ViT模块的混合架构[8]。多模态任务：CLIP 利用多模态预训练重新定义分类任务为目标检索，从而推动了跨领域应用的发展。LLaVA 将CLIP与大语言模型相结合，通过对生成的视觉-语言指令数据进行端到端微调，取得了优秀的多模态指令数据集性能。然而，注意力机制[2]在图像Token长度上表现出二次复杂度，这导致了对诸如目标检测[60]、语义分割[51]等密集预测任务带来了显著的计算开销。这一局限性限制了Transformers的应用效果。

状态空间模型 状态空间模型（SSMs）已被证明在通过状态空间变换捕捉语言序列的动力学和依赖关系方面非常有效。结构化的状态空间序列模型（S4）[20, 21, 49]专门设计用于处理长距离依赖关系，并具有线性复杂度。在引入S4之后，提出了更多相关模型，如S5 [48]、H3 [17] 和GSS [39]。

Mamba通过结合数据依赖的状态空间层和被称为并行扫描的选择机制（S6）[19]而脱颖而出。与依赖于具有二次复杂度的注意力机制的Transformer模型相比，Mamba在处理长序列时具有线性复杂度，表现出色。在计算机视觉领域，SSM最初被应用于像素级图像分类，而S4则用于管理电影剪辑分类中的长程时间依赖关系。此外，Mamba的潜力激发了大量研究，展示了其在诸如目标检测[60]和语义分割[38]等视觉任务中的优越性能以及更高的GPU效率。不同于以往工作，作者的TransMamba旨在以跨架构的方式探索利用预训练Transformer模型的知识来构建具有Mamba架构的新模型的潜力。

迁移学习 多种方法提出了将Transformer的知识转移到卷积神经网络(CNN)的方法。[33]引入了交叉注意力投影和组线性投影，以在两个投影特征空间中对齐学生的特征与教师模型。[18]提出了一种名为跨模型知识蒸馏(CMKD)的新方法用于音频分类，其中使用CNN和音频光谱 Transformer (AST)模型相互作为教师来训练学生模型。[24]讨论了在特征空间中丢弃与架构相关的信息，以防止学生模型受到不相关信息的干扰。最近，在自然语言处理(NLP)领域的一些研究[1, 55]聚焦于从Transformer向Mamba的知识转移过程。

[55]提出利用大型 Transformer 模型的注意力层的线性投影权重进行少GPU资源条件下的跨架构蒸馏，实现了性能与大型 Transformer 模型相当的结果。[1]认为Transformer和自适应统计建模(Adaptive Statistical Models, SSMs)都通过不同的混叠矩阵形式作用于Token序列，并提出了一种渐进式蒸馏策略，通过匹配SSM的不同粒度 Level 来蒸馏Transformer架构。[30]采用简单的蒸馏损失来利用教师网络丰富的表示知识。据作者所知，目前在视觉和多模态领域对Transformer到Mamba的知识转移的研究较少。图像信息的引入复杂化了Mamba结构，使知识转移变得更加困难。本文关注如何快速将Transformer的知识转移到Mamba在视觉和多模态领域的应用上。

3. Method

作者的研究旨在探讨二次架构与次二次架构之间跨架构学习的可能性。具体来说，作者以Mamba作为案例研究，目标是在更经济、高效和稳健的方式下，将预训练的Transformer模型的知识转移到Mamba模型中。本节首先解释SSM的核心原理，随后概述作者提出的训练方法TransMamba。接着，作者深入探讨TransMamba中处理单模型任务和跨模型任务的策略，并详细描述整个Pipeline。

3.1. Preliminary

状态空间模型（SSMs）基于连续系统构建，通过一个隐藏状态将1维函数或序列转换为。形式上，SSMs使用以下常微分方程（ODE）来描述输入数据：

其中， 表示系统的演化矩阵，而 和 分别表示投影矩阵。这种连续的微分方程通过离散化在现代状态空间模型（SSM）中进行近似。Mamba 是连续系统的一种离散版本，引入了一个时间尺度参数 ，用于将连续参数 转换为其离散对应物 。这种转换通常通过零阶保持（ZOH）方法实现，定义如下：

其中包含 (N) 个特征向量时，输出序列 (Y) 可以表示为：

这可以表示为：

3.2. TransMamba

特征校准。首先，需要使Transformer的潜在空间与Mamba的潜在空间对齐。考虑两个潜在空间：（学生模型），（教师模型），目标是将和转移到相似的空间，中。然而，两种模型的特征维度可能会有所不同，Mamba的维度可能小于预训练的Transformer。作者首先通过零填充的方法使Mamba模型的特征维度与Transformer模型的维度相匹配，同时保持其内在结构。一个简单而有效的MLP层被用来实现Mamba特征维度与Transformer特征维度的对齐。这一过程为后续的知识迁移提供了坚实的基础。

自适应双向知识蒸馏。知识蒸馏被用来将基于Transformer的知识转移给mamba模型。在作者的方法中，考虑了一个采用Transformer架构的教师模型和一个采用Mamba架构的学生模型。TransMamba的目标是从预训练的Transformer模型向新型的Mamba基模型转移知识。虽然受到了蒸馏方法的启发，但简单的方法仅利用教师模型的任务逻辑值或特征输出作为监督信号来调整学生模型的结果，但作者发现这种方法效果有限，如表5所示。作者认为原因可能是两模型框架之间的差异。直接通过最后一层特征对模型进行约束可能导致中间特征出现严重不一致，从而无法使Mamba的整体优化方向与Transformer保持一致。因此，作者选择在整个层之间进行余弦相似度的知识蒸馏，并仅使用教师网络的最后一层作为监督信息。

然而，直接优化中间层特征可能会导致某些层与教师网络过度对齐，而其他层则不足对齐。为了解决这个问题，作者提出了一种自适应优化方法，根据特征之间的变化相似度为不同的特征分配不同的权重。作者重用了之前计算的特征相似值来计算总的相似度。然后，为每个层的特征分配了一个权重系数。这样可以确保在所有层上更均衡地优化特征一致性。

其中 是超参数， 表示 Softmax。对于图像任务，双向 Mamba 必要时需要解决 Mamba 无法有效处理全局信息的问题。作者简化了双向 Mamba 的输出形式如下：

双向计算将原始矩阵 从下三角矩阵转换为密集矩阵。

与标准的Transformer形式 ( Y = S V = (S X)V ) 相比，其中 ( S = \operatorname{softmax}(Q K^{\mathrm{T}}/\sqrt{D}) )，双向形式 ( Y = C(M X) ) 包含重复的对角元素，并在VideoMamba [37]中有所展示。如果作者使用普通的Transformer特征来对齐带有重复对角元素的Mamba特征，可能会导致某些矩阵特征的过度优化或不足优化。因此，作者提出了一种双向蒸馏过程以避免这一问题。作者将正向和反向的SSM蒸馏过程分开。对于正向过程，作者直接使用Transformer对齐后的输出特征作为监督信号。

对于反向过程，作者将Transformer的特征反转以与Mamba矩阵的特征对齐。

因此，对于任何任务而言，总的损失函数如下所示：

权重子克隆。通常，具有相同架构的模型可以继承权重。然而，由于结构和维度的不同，跨架构的权重继承极具挑战性。为了探索如何完全继承Transformer架构的知识，作者提出了一种基于Mamba模型的权重子克隆技术。与Transformer和Mamba架构的主要区别在于它们的注意力机制（attn）和状态空间模型（SSM）。对于结构差异，作者在现有的Mamba框架中增加了MLP层，并将RMS层修改为LayerNorm。除了SSM部分外，作者使用Transformer模型参数初始化所有其他部分的参数。然而，在LLaVA模型中，用现有Mamba模型初始化SSM结构对确保训练稳定起着重要作用。对于维度差异，作者根据神经元重要性的显著性选择更多重要的参数，从而更好地进行模型初始化。具体来说，在使用一部分原始数据微调预训练模型之后，相比于变化较大的权重，变化较小的模型权重对于当前任务更为关键。因此，作者使用在梯度更新后权重变化最小的参数来初始化Mamba模型。

下游任务。在作者的工作中，包含两种类型的任务以验证TransMamba的有效性，包括一个单模态任务：图像分类，以及两个跨模态任务：视觉问答和视频检索。

单模态任务：对于图像分类任务，作者采用三种目前最先进的（SOTA）Mamba架构模型——Vmamba[34]、PlainMamba[58]和VisionMamba[59]作为学生模型。作者使用预训练的（ImageNet1k/21k）Transformer DeiT模型[52]展示了从头训练一个Mamba模型的过程。

多模态任务：使用TransMamba构建大型多模态模型的一个关键挑战在于，新的Mamba架构可能会破坏预训练模型的参数分布，导致文本和图像表示之间存在不一致性。重新进行类似于LLaVA的预训练无效，因为预训练的知识已经受损。因此，在TransMamba与LLaVA结合的架构中，应当增强图像与文本之间的交互。为此，作者引入了CrossMamba来解决这一问题。首先简化Mamba的计算公式，，与标准Transformer形式相比。然后，作者将，，。此外，作者也设置了，使其与具有类似的模态输入方式。这样，CrossMamba可以有效促进文本与图像之间的交互。在基于Mamba训练大型模型时，会随机出现梯度发散的现象，这在[50, 61]的NLP研究中也有观察到。适当的初始化对于确保稳定的训练至关重要。作者使用权重子克隆法对SSM以外的参数进行了初始化。但在训练过程中，模型表现出收敛困难的问题。为缓解这一问题，作者采用了来自NLP领域预训练的Mamba模型对SSM参数进行了初始化[61]。此外，作者发现卷积层的初始化显著影响后续训练。因此，这些参数被用标准正态分布进行初始化。这些策略还与前述的蒸馏方法相结合。对于视频检索任务，SSM参数则使用VideoMamba [31]进行初始化。

4. Experiments

4.1. Experimental Setup and Details

在本节中，作者简要介绍了相应模型的数据集、实现细节以及评估指标。更多详细信息详见附录中的数据集部分。单模型实验：对于图像分类，作者在CIFAR100 [29]、ImageNet100 [10] 和 ImageNet1000 [10] 这三个基准数据集上进行了实验。多模型实验：对于视觉问答，作者在论文中使用了LLaVA-1.5-pretrain数据集中50万条通用图例样本和LLaVA-1.5-finetune数据集中60万条通用图例和对话样本。此外，作者还在论文中使用了两个视频-文本数据集，即MSR-VTT [57] 和 DiDeMo [26]。

实施细节。对于图像分类任务，作者基于VMamba [34]、PlainMamba [58] 和ViM [59] 构建了代码库。具体来说，作者使用AdamW [36] 优化器和学习率为5e-04训练所有Mamba模型共计300个epoch。对于视觉问答任务，作者采用预训练的CLIP-ViT-L/14 [43] 作为视觉编码器，并使用两层MLP作为 Projector 。教师和学生模型均利用LLaMA-3.2 [14] 家族构建其基础模型。具体而言，教师模型采用3B参数配置，而学生模型则基于LLaMA-3.2-1B构建0.6B参数规模的基础模型。作者首先使用LLaVA训练方法训练LLaMA-3B模型，以获得LLaVALLaMA3.2-3B模型作为教师模型。随后，作者仅使用0.6百万条通用标题样本训练作者的Mamba模型。作者将批量大小设置为128，并使用Adam优化器和2e-5的学习率。在整个过程中，作者在每个epoch使用16块V100 GPU进行训练。对于视频检索任务，作者使用AdamW优化器 [36] 训练所有Mamba模型共计5个epoch。学习率设置为1e-4并采用余弦衰减策略。作者使用预训练的CLIP4Clip作为冻结的教师模型。批量大小设置为128，并在4块A100 GPU上进行训练。

评价指标。对于分类任务，作者采用顶精度作为评估指标。对于视觉问答任务，作者在GQA[28]、VizWiz[23]、TextVQA[47]和MME[16]上进行实验。对于文本-视频检索任务，作者使用召回率和Mean R来评估模型性能。

4.2. Main Results

图像分类。在表1中，作者展示了针对分类任务的各种Mamba设置的层配置、维度和参数规模。与以往工作中描述的单一PlainMamba配置[58]不同，作者在新的Mamba模型中将PlainMamba扩展为三个尺度：PMambaT、PMamba-S和PMamba-B。ViM、VMamba-T和VideoMamba的层配置与文献[59]、[34]和[31]中所述相同。

作者在表2中展示了CIFAR-100、ImageNet子集和ImageNet1K的实验结果。DeiT-预训练表示在使用ImageNet-2012数据集预训练后，对CIFAR或ImageNet-100进行微调的模型。从表中可以看出，与vanilla Mamba相比，TransMamba超越了最近提出的VMamba、ViM和PlainMamba。例如，TransMamba-P在精度上比Mamba-P高出2.83%。这些结果证实了来自Vision Transformers (ViT)的知识可以有效地转移到Mamba中，从而提升了Mamba后续的表现。

为了验证TransMamba的效率，作者在Mamba训练过程中展示了损失函数收敛图和准确率图（如图3所示）。结果显示，TransMamba表现出更快的收敛速度和更高的准确率，证实了第一章中讨论的训练效率。

视觉问答 在这个实验中，作者保持总层数不变，并用Mamba替代了原有的Transformer结构。Trans-LLaVA 的参数量更少。如表3所示，相比LLaVA-3.2-1B，Trans-LLaVA 在GQA、VQA 和 VisWiz 上取得了更好的性能，并几乎与3B模型持平。由于资源限制，作者使用了一个3B模型作为教师模型。尽管如此，作者的模型在多个指标上仍取得了出色的结果。

作者展示了几个例子来说明Trans-LLaVA的定性结果。如图4所示，Trans-LLaVA能够有效地理解用户的提问，并给出准确的回答。

视频检索 在表4中可以看到，作者的所提出的方法TransMamba在视频领域两个数据集上的表现优于Mamba方法。例如，在MSRVTT数据集上，R1指标的准确率达到了41.6%，比Mamba高出了0.5个百分点，这表明了作者提出的TransMamba在视频检索中的有效性。

4.3. Ablation Studies

在本部分中，作者主要进行了分类任务的实验。更多实验详见附录。蒸馏策略：传统的logit蒸馏策略和特征蒸馏方法对Mamba的学习效果有限，这从表5中可以看出。这表明WSAB能够利用Transformer层获得的丰富表示，从而实现更全面的知识迁移，并提升Mamba模型的性能。

数据规模的影响：在表6中，作者展示了不同数据规模下模型的准确性。在PMamba模型中，仅使用数据的几乎就可以达到最佳性能，这强烈表明Transformer模型生成的丰富表示显著缩短了Mamba模型的训练过程。

5. Conclusion

在本研究中，作者进行了全面的分析和实验，以阐明关于TransMamba模型的以下洞见：1) 两阶段框架有效地将现有预训练Transformer模型的知识转移到SSM基模态上。2) 权重子克隆高效地将Transformer的知识转移至Mamba架构。自适应双向知识蒸馏过程根据不同层之间的相似性分配权重，确保各层优化的均衡。

同时，双向知识蒸馏避免了单向知识蒸馏中存在的过度或不足优化问题。此外，跨Mamba模块有效地将多模态理解能力整合到SSM中，从而在视觉语言任务中实现稳健的表现。在多个数据集上的广泛实验表明，作者的方法表现出优越的性能与效率。此外，作者的算法可以无缝适应其他基于SSM的模型，确保其广泛的应用价值和影响力。

参考

[0]. TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba .