中山大学&华为诺亚方舟实验室等提出TransMamba：从Transformer到Mamba的高效跨架构知识迁移

最新推荐文章于 2025-08-18 11:07:02 发布

AI启智汇

最新推荐文章于 2025-08-18 11:07:02 发布

阅读量1.3k

点赞数 40

CC 4.0 BY-SA版权

文章标签： transformer 架构深度学习计算机视觉人工智能分类图像处理

本文链接：https://blog.csdn.net/helloaiti/article/details/146511944

“TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba”

论文提出了TransMamba方法，通过两阶段策略将预训练Transformer模型知识转移到Mamba模型，有效提升训练效率和性能，在多种任务上表现优异。

论文地址：https://arxiv.org/pdf/2502.15130

代码将会开源

研究动机

Transformer 架构因注意力模块的灵活可扩展性，在单模态和多模态基础模型中备受青睐，有大量预训练模型可用。但其注意力机制的二次计算复杂度带来高昂计算成本和内存消耗，限制了模型优化与扩展。Mamba 等亚二次架构虽能以线性复杂度实现全局感知，但为特定任务从头训练这类模型既耗费资源又耗时。因此，文章探索跨架构训练，将Transformer模型中的知识转移到Mamba模型，以实现更高效且经济的训练。

主要贡献

-快速通用框架：提出两阶段框架，将预训练Transformer模型知识转移到新的基于状态空间模型（SSM）的Mamba模型，低成本提升训练效率和性能。

-权重子克隆和自适应双向蒸馏（WSAB）：权重子克隆用于跨架构知识转移，自适应双向蒸馏为Mamba架构设计，赋予其多模型交互能力。

-全面验证：在多种骨干网络架构和应用中验证了方法的有效性，涵盖视觉问答、视频检索和图像分类任务。

方法

-初步知识：状态空间模型（SSMs）基于连续系统，通过隐藏状态将1D函数或序列进行转换，使用常微分方程描述输入数据。Mamba 是连续系统的离散版本，引入了选择性扫描机制（S6），通过与输入数据相关联的参数，以线性复杂度处理长序列。

-TransMamba

特征校准：通过零填充和MLP层，使Mamba模型的特征维度与Transformer模型匹配，对齐两者的潜在空间，为知识转移奠定基础。

自适应双向蒸馏：传统基于任务logits或特征输出的蒸馏方法效果不佳。文章提出基于余弦相似度的跨层知识蒸馏，仅用教师网络最后一层作为监督信息，并根据特征相似度自适应分配权重，优化特征一致性。针对双向Mamba处理图像信息时的问题，提出双向蒸馏，分别优化前向和后向SSM蒸馏过程，避免矩阵特征的过度或欠优化。

权重子克隆：由于Transformer和Mamba架构存在差异，跨架构权重继承困难。文章提出权重子克隆技术，针对结构差异，修改Mamba框架并使用Transformer模型参数初始化部分参数；针对维度差异，选择重要参数初始化Mamba模型，加速模型收敛。

下游任务：通过图像分类（单模态任务）和视觉问答、视频检索（多模态任务）验证 TransMamba的有效性。在多模态任务中，引入CrossMamba模块增强图像和文本的交互，同时针对模型训练中的问题，采用不同的初始化策略和蒸馏方法。

-论文提出的两阶段框架：旨在将预训练Transformer模型的知识高效转移到基于状态空间模型（SSM）的 Mamba 模型中，提升Mamba模型的训练效率与性能，具体如下：

第一阶段：特征校准。

维度匹配：考虑到Transformer模型与Mamba模型的特征维度可能不同，且Mamba 模型维度可能小于预训练的 Transformer 模型。首先对 Mamba 模型进行零填充，使其特征维度与 Transformer 模型一致，同时保持其底层结构不变。

潜在空间对齐：利用简单且高效的 MLP 层，将Mamba模型的特征维度与 Transformer 模型进行对齐，使两者的潜在空间相似。这一操作是后续知识转移的基础，通过将两个模型的特征映射到相似空间，为知识蒸馏创造有利条件，让Mamba模型能够更好地接收来自 Transformer 模型的知识。

第二阶段：自适应双向蒸馏与权重子克隆。

自适应双向蒸馏：跨层知识蒸馏：采用基于余弦相似度的知识蒸馏方法，在模型的所有层进行操作。计算教师模型（Transformer 架构）与学生模型（Mamba 架构）各层特征之间的余弦相似度，以衡量特征的相似程度。仅将教师网络的最后一层作为监督信息，避免直接使用最后一层特征约束导致的中间特征不一致问题，确保Mamba模型的优化方向与Transformer模型更为契合。

自适应双向蒸馏：自适应权重分配：为解决直接优化中间层特征可能出现的某些层过度对齐或对齐不足的问题，根据特征相似度计算每个层特征的权重系数。相似度高的特征，分配较低的权重系数；相似度低的特征，分配较高的权重系数。通过这种自适应的权重分配方式，实现对所有层特征一致性的更平衡优化，使 Mamba 模型在学习 Transformer 模型知识时，能够更合理地调整自身参数。

自适应双向蒸馏：双向蒸馏策略：针对图像任务，双向Mamba对于处理全局信息具有必要性，但它的结构与标准Transformer不同，直接用Transformer特征对齐可能导致优化问题。因此，文章提出双向蒸馏过程，将前向和后向SSM蒸馏过程分开。前向过程中，直接使用对齐后的Transformer输出特征作为监督信号；后向过程中，将 Transformer 的特征进行反转，使其与Mamba矩阵特征对齐后作为监督信号。通过这种方式，避免因结构差异造成的优化问题，提高知识转移的效果。

权重子克隆：由于Transformer和Mamba架构在注意力和状态空间模型机制上存在差异，跨架构权重继承面临挑战。文章提出权重子克隆技术，针对结构差异，在Mamba框架中添加MLP层并修改RMS 层为Layer Norm，然后使用 Transformer 模型参数初始化除 SSM 外的所有部分参数。针对维度差异，根据神经元的重要性选择参数，具体做法是在使用部分原始数据对预训练模型进行微调后，选取梯度更新后权重变化较小的参数来初始化Mamba模型，加速模型收敛。

实验

-图像分类：在CIFAR - 100、ImageNet - subset和ImageNet1K数据集上，与普通Mamba模型相比，TransMamba 表现更优。如TransMamba - P比Mamba - P的准确率高出 2.83%。从训练过程的损失函数收敛图和准确率图来看，TransMamba收敛速度更快，准确率更高，在仅使用50% 数据时，PMamba模型就几乎能达到最佳性能，证明了Transformer模型的知识能显著缩短Mamba模型的训练过程。

-视觉问答：在保持总层数不变的情况下，用Mamba替换原Transformer结构得到的 Trans - LLaVA模型，参数数量更少。与LLaVA - 3.2 - 1B相比，Trans - LLaVA在GQA、VQA、VisWiz等任务上表现更出色，且在多个指标上接近3B模型的性能。从一些定性结果示例可以看出，Trans - LLaVA能有效理解用户问题并准确回答。