点击下方卡片，关注「AI视界引擎」公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
随着具有线性复杂度的RNN模型的发展，Transformer的二次复杂度挑战有望被克服。值得注意的是，新兴的Mamba-2已经展示了竞争力的表现，将RNN模型和Transformer之间的差距拉近了。
然而，由于顺序处理和梯度消失问题，RNN模型在捕捉长范围依赖方面存在困难，限制了其上下文理解能力。这导致了收敛速度慢、资源需求高以及在下游理解和复杂推理任务中的表现不佳。
在此项工作中，作者通过用Mamba-2层替换预训练的大语言模型（VLM）中的一部分Transformer解码器层，提出了一种混合模型MaTVLM。利用注意力机制与Mamba-2之间的内在关系，作者使用相应的注意力权重初始化Mamba-2以加速收敛。
随后，作者采用单阶段蒸馏过程，使用预训练的大语言模型作为教师模型，将其知识转移到MaTVLM上，进一步提高收敛速度和性能。此外，作者研究了训练框架内不同蒸馏损失的影响。
作者在多个基准测试上评估了MaTVLM，结果显示其在教师模型和其他大语言模型上的性能相当，但在Mamba基的大语言模型和同等参数量级的模型上表现更优。令人惊奇的是，MaTVLM在不牺牲性能的情况下比教师模型快至多3.6倍，并且减少了27.5%的GPU内存消耗。代码和模型已发布。
- Introduction
近年来，大型跨模态模型（VLMs）取得了迅速进展。这些模型主要基于Transformer架构构建。然而，由于Transformer在序列长度上的二次复杂度，这使得VLMs在训练和推理时计算密集。最近，一些RNN模型作为Transformer的潜在替代方案出现，并提供了与序列长度线性相关的扩展能力。值得注意的是，Mamba 在长距离序列任务中表现出色，其计算效率超过了Transformer。
多项研究[29, 37, 43, 63, 66, 67]探索了将Mamba架构整合进VLMs的方法，通过用基于Mamba的大语言模型（LLM）替换基于Transformer的大语言模型。这些工作展示了与Transformer相比具有竞争力的表现，并在推理速度上取得了显著提升。然而，这些方法也存在一些局限性：
(1) Mamba采用顺序处理方式，无法像Transformer那样捕捉全局上下文，从而限制了这些VLMs在复杂推理和问题解决任务中的性能[53, 56]；
(2) Mamba的顺序处理特性导致长序列训练期间梯度传播效率低下，在从头开始训练VLMs时收敛速度较慢，这使得高计算成本和大量的训练数据成为这些VLMs的重大 Bottleneck ；
(3) 当前这些VLMs的训练方案较为复杂，需要多阶段训练才能达到最优性能。这一过程既耗时又耗费大量计算资源，使得基于Mamba的VLMs难以扩展到更广泛的应用中。
为了解决上述问题，作者提出了一种名为Mamba-Transformer 视觉-语言模型（MaTVLM）的新颖方法，该模型结合了Mamba-2和Transformer组件，力求在计算效率与整体性能之间取得平衡。首先，注意力机制与Mamba本就有紧密联系，移除注意力机制中的Softmax操作会将其转换为一个线性RNN，显示出其结构上的相似性。作者将在第3.2节详细分析这种关系。此外，将Mamba应用于大语言模型（LLMs）的研究[48, 49]表明，混合使用Mamba的模型在某些任务上优于纯Mamba和纯Transformer Baseline 模型。受这一连接和实验发现的启发，将Mamba与Transformer组件相结合是一个有前景的方向，既能提高推理能力又能保持计算效率。具体来说，作者以TinyLLaVA [64]为基础视觉-语言模型，并用Mamba解码器层部分替代其Transformer解码器层，而其他部分则保持不变。
为了在最大限度地提高MaTVLM性能的同时尽量减少训练成本，作者提出从预训练的基础大语言模型中提炼知识。首先，作者根据第3.2节的描述，使用相应的注意力权重初始化Mamba-2，这对于加快Mamba-2层的收敛速度非常重要。此外，在提炼训练过程中，作者采用概率分布损失和逐层提炼损失来引导学习过程，使得只有Mamba-2层可训练，而保持Transformer层固定。值得注意的是，与大多数需要复杂多阶段训练的大语言模型不同，作者的方法涉及一个单一阶段的提炼过程。
尽管采用了简化训练方法，作者的模型在多个基准测试中展示了全面的表现，如图1所示。与教师模型TinyLLaVA相比，它表现出竞争性的效果，并且在基于Mamba的VLM和具有相似参数规模的其他基于Transformer的VLM中表现更优。此外，作者的模型通过提高3.6倍的速度和减少27.5%的内存使用，进一步突显了其实用优势，从而验证了其在实际应用中的有效性。这些结果强调了作者方法的有效性，为作者未来在模型开发和优化方面的进展提供了前景。
综上所述，本文做出了三项重要贡献：
作者提出了一种新的混合视觉语言模型架构MaTVLM，有效地将Mamba-2和Transformer组件结合在一起，在计算效率与高性能之间取得平衡。·作者提出了一种新颖的一阶段知识蒸馏方法，用于Mamba-Transformer混合视觉语言模型。
通过利用预训练的知识，作者的方法加速了收敛过程，提升了模型性能，并加强了视觉语言理解能力。·作者的方法在保持 Baseline VLM竞争力的同时，显著实现了3.6倍的推理速度提升和27.5%的内存使用率降低。此外，该方法在多个基准测试中优于基于Mamba的VLM以及具有相似参数量的现有VLM。
- Related Work
2.1. Efficient VLMs
近年来，高效且轻量级的视觉语言模型（VLMs）取得了显著进展。为了提高效率，开发了多种面向学术研究的VLMs，例如TinyLLaVA-3.1B [64]、MobileVLM-3B [14] 和LLaVA-Phi [68]。与此同时，通过利用高分辨率图像和长上下文文本的大规模数据集，像Qwen2.5-VL-3B [5]、InternVL2.5-2B [10]等面向商业应用的模型也取得了出色的性能。
作者的工作优先考虑效率和资源约束，而非大规模、以商业为导向的训练。与以往的方法不同，通过整合Mamba-2 [17]，作者的方法在显著降低计算需求的同时实现了竞争性的性能，因此非常适用于资源受限的环境部署。
2.2. Structured State Space Models
结构化状态空间模型（S4）在序列长度上可以线性扩展。Mamba[23]引入了选择性的状态空间模型（SSM），而Mamba2[17]在此基础上通过将SSM与注意力机制变体关联起来进行改进，实现了2至8倍的速度提升，并且性能与Transformer相当。基于Mamba的视觉语言模型（VLM）[29, 43, 63, 69]主要用预训练的Mamba-2语言模型完全替代了基于Transformer的大语言模型（LLMs），同时实现了竞争性的性能和增强的计算效率。
作者的研究创新性地将 Mamba 和Transformer融合到VLMs中，结合了两者的优势，而不是完全用Mamba-2取代Transformer。通过采用混合方法并引入单阶段蒸馏策略，作者增强了模型的表现能力，提高了效率，并在保持计算效率以实现实际部署的同时，优于之前的基于Mamba的VLMs，在性能上取得了卓越的表现。
2.3. Hybrid Mamba and Transformer
近年来的研究，如MambaInLlama [49] 和MOHAWK [6] 证明了混合Mamba-Transformer架构在大语言模型中的有效性，实现了效率和性能的显著提升。此外，MambaVision [27] 将这种混合方法扩展到了视觉模型中，引入了一种基于Mamba-Transformer的 Backbone 网络，在图像分类和其他视觉相关任务上表现出色，展示了将SSM与 Transformer 结合的潜力。
与以往关于大语言模型或视觉Backbone的研究不同，作者的工作将Hybrid Mamba-Transformer扩展到了多模态语言模型，并设计了一个简洁的架构，其中包含一个高效的单阶段蒸馏策略，从而提高收敛速度、减少推理时间并降低内存消耗，以便于实际部署。
2.4. Knowledge Distillation
近年来，大语言模型（LLM）的知识蒸馏逐渐受到关注，而视觉语言模型（VLM）的知识蒸馏研究仍较为有限[20, 52, 55]。DistillVLM[20] 使用均方误差（MSE）损失对齐注意力和特征图，MAD[52] 对齐视觉和文本的 Token，LLAVADI[55] 强调了联合 Token 和 Logit 对齐的重要性。
在这些进展的基础上，作者将知识蒸馏集成到一种混合Mamba-Transformer框架中，并采用单阶段蒸馏策略，从基于Transformer的教师模型转移知识。这提高了收敛性，增强了性能，并降低了高效VLM部署的计算成本。
- Method
大型多模态视觉-语言模型（VLMs）处理更长的序列，导致训练和推理速度较慢。正如之前提到的，Mamba-2 架构表现出线性扩展，并且与Transformer相比提供了显著更高的效率。为了充分利用这些优势，作者提出了一种结合了Mamba-2 和Transformer组件的混合VLM架构MaTVLM，旨在在计算效率与最佳性能之间取得平衡。
3.1. Mamba Preliminaries
Mamba [23] 主要建立在结构化状态空间序列模型（S4）之上，如公式1所示，这是一种在深度学习中用于序列建模的近期发展，与RNN、CNN以及经典的状态空间模型有着紧密的联系。
Mamba 引入了选择性状态空间模型（Selective SSMs），如公式 2 所示，与标准的线性时不变（LTI）形式 1 不同，使其能够在每个时间步长上选择性地关注或忽略输入。在诸如语言处理等信息丰富的工作中，它的性能已被证明优于 LTI SSMs，特别是在状态大小
增加的情况下，能够处理更多的信息容量。
Mamba-2 [17] 在 Mamba 的选择性 SSMs 中引入了状态空间对偶（SSD）框架，该框架通过不同形式的结构半可分离矩阵分解建立了 SSMs 与各种注意机制之间的理论联系。借助这一框架，Mamba-2 实现了2至8倍的计算加速，同时保持与 Transformer 相当的竞争性能。
3.2. Hybrid Attention with Mamba for VLMs
如图2所示，MaTVLM基于预训练的语言模型构建，包括一个视觉编码器、一个连接器和一个语言模型。语言模型最初由Transformer解码层组成，其中一些解码层在作者的模型中被Mamba-2解码层替换。这种替换仅修改了注意力机制至Mamba-2，而其他组件保持不变。根据配置的比例（例如，
,
），作者将Mamba-2解码层均匀分布。由于Mamba-2与注意力机制共享某些连接，部分权重可以从原始Transformer层初始化，具体细节如下。
正式地，对于输入序列
中的
，transformer解码器层中的注意力定义如下：
其中，
是输入嵌入的维度，而
、
和
是可学习的权重。
移除公式3中的softmax操作后，注意力变为：
上述结果可以以下述线性RNN的形式重新表述：
将公式5与公式2进行比较，作者可以观察到它们之间的如下映射关系：
因此，作者像图3所示的那样，使用Transformer层的相应权重初始化Mamba-2层的上述权重，其余权重则随机初始化。除了Mamba2层外，所有其他权重与原始Transformer的一致。
3.3. Knowledge Distilling Transformers into Hybrid Models
为了进一步提升MaTVLM的表现，作者提出了一种知识蒸馏方法，该方法将Transformer层的知识转移到Mamba-2层。作者将预训练的VLM作为教师模型，将作者的MaTVLM作为学生模型。接下来，作者将介绍蒸馏策略。
概率分布蒸馏 首先，作者的目标是尽量减小模型之间概率分布之间的距离，具体来说，就是模型在应用softmax函数之前的logits输出之间的距离。这种做法在知识蒸馏中被广泛采用，因为通过使模型的输出分布对齐，可以使学生模型从教师模型的预测中获得更加细腻的理解。为了实现这一点，作者使用带有温度缩放因子的Kullback-Leibler（KL）散度作为损失函数。温度因子调整了概率分布的平滑程度，使得学生模型能够捕捉到经过软化处理的教师模型分布中的细微差别。损失函数定义如下：
软化的概率 (P_{t}(i)) 和 (P_{s}(i)) 分别通过将教师模型和学生模型的logits应用温度缩放的softmax函数计算得到：
其中，
是温度缩放因子，较高的温度会产生更软的分布；
是教师模型的logit（即未经过Softmax的输出），而
则是学生模型对应的logit。
逐层蒸馏 为了确保学生模型中的每个Mamba层与教师模型中的相应层对齐，作者采用了一种逐层蒸馏策略。具体而言，当给定相同的输入时，该方法通过最小化学生模型中Mamba层的输出与其教师模型中相应Transformer层之间的L2范数来实现这种对齐。这些输入来自教师模型的前一层，从而保证了上下文的一致性和连续性。通过对中间特征表示的对齐，学生模型可以更有效地复制教师模型的层级特征提取过程，从而提升其整体性能。假设学生模型中Mamba层的位置为
，则这种对齐的损失函数定义如下：
其中，
和
分别表示教师模型和学生模型在第
层的输出。
序列预测损失 最后，除了上述提到的蒸馏损失外，作者还计算学生模型输出的序列预测与ground truth之间的交叉熵损失。这种损失用于引导学生模型学习正确的序列预测，这对于模型在下游任务上表现出色至关重要。该损失函数定义如下：
其中，
是 GT 序列，
是学生模型预测的序列。
单阶段蒸馏训练 为了充分利用所提出的蒸馏方法的互补优势，作者将概率分布损失、层间蒸馏损失和序列预测损失整合到一个统一框架中进行单阶段蒸馏训练。在训练过程中，作者设置它们各自的权重如下：
其中，
、
、
是超参数，用于控制每个损失项相对重要性的权重。
作者将进行一系列实验，以彻底探究这三种损失函数的个体贡献及其相互作用。通过分别分析它们单独和组合的效果，作者旨在深入了解每种损失函数如何影响学生模型的学习过程、中间表示的质量以及最终预测的准确性。这将帮助作者理解每种损失函数在提升模型整体性能方面所扮演的具体角色，并确保选择的损失函数有效地参与到优化过程中。
这个单阶段框架高效地结合了两个知识蒸馏目标和一个预测任务，使得学生模型能够无缝地从不同的损失组件中获取梯度。统一的损失函数不仅在训练过程中加速了收敛，还确保了学生模型可以从层次特征蒸馏和全局预测对齐中受益。此外，所提出的方法具有灵活性，并可以根据特定任务的要求通过调整损失组件的权重轻松适应各种神经网络架构和任务。通过将这两个蒸馏目标和预测任务整合到一个框架中，学生模型MaTVLM实现了显著的性能提升，同时保持了计算效率，使其在实际应用场景中具有高度适用性。
- Experiments
4.1. Implementation Details
作者在实验中选择了TinyLLaVA-Phi-2-SigLIP-3.1B [64]作为教师视觉语言模型。该模型的视觉编码器是SigLIP模型[61]，在包含4亿参数和分辨率为
的WebLi数据集[9]上预训练而成。语言模型部分为Phi-2 [31]，包含27亿个参数。在MaTVLM中，作者用Mamba-2层分别替换掉教师模型中
、
和
的Transformer解码器层，并确保其分布均匀。如图2所示，学生模型的可训练参数仅为Mamba-2层和连接器。
在训练过程中，损失函数超参数设置为
和
，表示概率分布损失和层间蒸馏损失被赋予了相同的权重，而序列预测损失未被采用。作者使用批量大小为 64，并使用 AdamW 优化器进行模型优化，其中权重衰减率为 0.01，动量参数分别为
和
。学习率设置为
，并遵循一个包含预热稳定衰减阶段的学习率调度方案，其中预热和衰减阶段分别占总训练步骤的 10%。根据 TinyLLaVA 的配置，作者采用 ShareGPT4V [8] 的 SFT 数据集进行训练。该数据集用由 GPT-4V [1] 生成的 23K 高质量图像-文本对替换了 LLaVA-Mix-665K [39] 数据集中与图像描述任务相关的图像-文本对，从而提高了数据质量。
4.2. Main Results
性能比较 如表1所示，作者展示了多种基于视觉的语言模型(VLMs)在多个基准上的性能对比，包括MME [21]、MMBench [62]、TextVQA [44]、GQA [30]、MM-Vet [59]、ScienceQA [41]、POPE [36]、MMMU [60] 和VQAv2 [22]。作者在表中列出了具有不同架构代表性的VLMs [4, 7, 15, 33-35, 39, 58, 65]。为了进行比较，作者特别突出显示了与MaTVLM参数量相似的模型[14, 38, 55, 64, 68]，以及那些整合了Mamba [29, 43, 63] 的模型。首先，与教师模型TinyLLaVA [64]相比，MaTVLM在MME上的表现提升了17.6点。在所有基准测试中，性能下降幅度均控制在2.6点以内，仅在MMBench和ScienceQA上的下降幅度大于4.9点，显示出该模型的竞争优势。基于基于Transformer的大语言模型(LLMs) [3, 13, 47]，MaTVLM在大多数基准测试上与之相匹敌，在POPE和MMMU上分别取得了0.2点和1.6点的性能提升。此外，相比于参数量相似的其他VLMs，作者的MaTVLM几乎在所有基准测试中都表现出更优的表现，在MME上的显著改进达到87.7点，在TextVQA上的改进则为7.0点。最后，相比于集成Mamba的VLMs，作者的MaTVLM在大多数基准测试中表现出最佳性能，在TextVQA中只落后于Cobra [63]和ML-Mamba [29]约0.2点。总之，这些结果突显了MaTVLM在各种不同基准测试中持续展现出稳健且竞争力强的性能，证明了其有效性及在实际应用中的强大潜力。
推理速度对比
作者使用NVIDIA GeForce RTX 3090对作者的MaTVLM与教师模型TinyLLaVA的推理速度进行了评估。如图1(b)所示，在相同的生成Token长度设置下，MaTVLM相比TinyLLaVA在FlashAttention2 [16, 18]的支持下实现了最高达
倍的更快推理速度。换句话说，随着生成Token长度的增加，MaTVLM与TinyLLaVA之间的推理时间差距持续扩大。此外，Mamba-2层更高的混合比例也进一步提升了推理速度。这表明，在推理过程中，作者的MaTVLM具有更优的效率，使其更适合实际应用。
内存使用量比较 作者进一步将作者提出的MaTVLM与TinyLLaVA在NVIDIA GeForce RTX 3090上进行了GPU内存使用量的比较。如图1(c)所示，相较于TinyLLaVA，作者的MaTVLM在Token长度为16,384时显著降低了内存占用，峰值降低率达27.5%。值得注意的是，当Token长度达到32,768时，TinyLLaVA遇到了内存溢出错误，而作者的MaTVLM则依然可以正常运行。这种内存使用量的降低归功于MaTVLM优化的架构，该架构有效地平衡了计算效率和性能，使其更适合部署在资源受限的设备上。
4.3. Ablation Study
Mamba-2 混合比例 作者通过改变 Mamba-2 层的比例（
，
和
），并评估其在八个跨模态基准上的性能影响。如表2 所示，
的比例获得了最高的平均分数，比
的比例高出了 2.2 分，表明过多的 Mamba-2 层可能削弱了全局依赖关系建模。
的比例在 MM-Vet 和 MMMU 上表现最佳，但整体上略逊一筹，得分比
的比例低 0.5 分。这一发现突显了平衡 Mamba-2 和 Transformer 层的重要性，以优化不同任务的性能。尽管较高的 Mamba-2 比例（
）可能会限制模型捕捉远程依赖关系的能力，较低的比例（
）则保留了 Transformer 的优势，但可能无法充分利用 Mamba-2 的优点。
Mamba-2 混合层位置 作者进一步研究了 Mamba-2 层在模型中的位置对性能的影响。具体来说，作者在四种配置中将transformer解码器层替换为 Mamba-2 层：全部置于开头、全部置于中间、全部置于结尾以及均匀分布。值得注意的是，在全部置于结尾的配置中，无法有效实现知识蒸馏，导致响应不一致。如表3所示，在所有基准测试中，均匀分布配置的性能最高，分别比全部置于开头和全部置于中间的配置高出1.9和3.4分。这些结果强调了在整个模型中均匀集成 Mamba-2 层以优化性能的重要性。
蒸馏损失 如第3.2节所述，作者采用三种蒸馏损失：概率分布损失
、层间蒸馏损失
以及序列预测损失
。为了研究每种损失对MaTVLM性能的影响，作者在表4中进行了消融实验。初始阶段，作者分别使用这三种损失。结果显示，
显著提高了性能，平均得分提高了5.8分。加入
后，性能进一步提升，在结合
和
时达到了最高的平均得分，比单独使用
高出0.9分。这表明概率分布对齐和层间特征匹配都有助于知识转移。然而，当重新引入
与
和
一起时，观察到有轻微1.0分的下降，这表明直接来自
的监督可能干扰了蒸馏过程。
- Limitations
尽管MaTVLM具有许多优势，但它也存在一些局限性。虽然使用预训练的注意力权重初始化Mamba-2有助于加速收敛，但可能未能充分利用其隐式状态表示。这可以通过定制化的初始化策略来改进，例如梯度匹配或额外的预训练阶段。
此外，由于本研究中可用的有限GPU资源，作者没有探索模型在更大规模上的性能表现。未来的工作可以利用更多的计算资源，系统地研究最佳的Mamba-2集成比例，并在更大的VLM上进行混合实验以评估可扩展性和性能。解决这些挑战将进一步提高大型VLM中混合架构的效率和适用性。
参考
[0]. MaTVLM: Hybrid Mamba-Transformer for Efficient Vision-Language Modeling .
点击上方卡片，关注「AI视界引擎」公众号