计算效率和性能双赢！南大联合中移动发布高效多模态大模型新范式—— p-MoD

文章链接：https://arxiv.org/pdf/2412.04449
项目链接：https://github.com/MCG-NJU/p-MoD

亮点直击

• MoD机制的引入：通过选择性处理重要视觉token，显著减少计算负担，提升模型效率。
• 创新的TanhNorm与STRing设计：解决了MoD模块在多模态训练中的稳定性和性能问题，即使在有限数据下也能精准评估token的重要性。
• PRD策略：通过逐层减少保留token的比例，针对视觉token的冗余问题提供了有效的解决方案，大幅提高计算效率。
• 全面的实验验证：在多个基准测试中，p-MoD模型不仅保证了高效性，还表现出了超越传统基线的性能。

总结速览

解决的问题：

• 高计算成本：尽管多模态大语言模型（MLLMs）在多任务中表现出色，但其庞大的训练和推理成本严重阻碍了其发展。计算的主要负担来自于通过LLM处理的大量视觉token。
• 视觉token冗余：在处理高分辨率图像和视频时，视觉token数量急剧增加，导致计算成本暴涨，限制了MLLMs的进一步应用。

提出的方案：

• MoD机制：通过引入深度混合（Mixture-of-Depths, MoD）机制，选择每个Transformer层要处理的最重要的视觉token，同时跳过冗余的token，从而提高计算效率。
• 创新设计：

• TanhNorm：用于权重归一化的tanh门控方法，确保MoD模块的正确初始化，增强训练稳定性并解决推理时的数值稳定性问题。
• STRing：对称token重加权模块，利用有限的训练数据和语言监督信号来精确评估token的重要性。
• PRD策略：渐进式比例衰减（Progressive Ratio Decay, PRD）策略，基于移位余弦调度逐层减少保留token的比例，优化视觉token的冗余问题。

应用的技术：

• MoD机制：最初用于大语言模型的MoD机制被改造为适用于多模态语言模型。
• TanhNorm与STRing：用于解决MoD模块的初始化与训练稳定性问题，提升有限数据下的token重要性预测。
• PRD策略：通过逐层减少视觉token的保留比例，进一步提高计算效率。

达到的效果：

• 高效推理：在推理时，p-MoD模型的TFLOPs仅为基线模型的55.6%，KV缓存存储仅为53.8%。
• 高效训练：训练时，p-MoD仅使用77.7%的GPU时间，显著减少了计算成本。
• 优异表现：在14个基准测试中，p-MoD模型的表现与LLaVA-NeXT基线相当，甚至超越了基线，证明了高效性和性能的平衡。

方法

本节介绍 p-MoD 模型。如下图 2 所示，本文的模型由 p-MoD 层组成，这些层通过 tanh 门控权重归一化（TanhNorm）和对称 token 重加权（STRing）来升级 MoD 架构。关键的设计是渐进式比例衰减（PRD）策略，它控制不同层之间的 token 保留比例。接下来的部分中，将首先简要回顾 MoD，然后详细解释p-MoD 模型的每个组件。

回顾 Mixture-of-Depths

在经过 transformer 层处理后，选中的 tokens 会根据其对应的权重进行缩放。通过这种方式，权重预测器参与到梯度路径中，使其参数能够通过反向传播进行更新。这个过程称为 token 重加权。

在原始为 LLM 设计的 MoD 模块中， X代表文本 tokens。在本文的多模态场景中， X代表视觉 tokens。仅在视觉 tokens 上应用 MoD，因为它们占用了主要的计算负担，并且在深层中显示出高冗余性。

改进 Mixture-of-Depths 模块

与从头训练 MoD 基于 LLMs 的模型不同，将 MoD 集成到 MLLMs 中面临几个挑战。将新的 MoD 模块插入到预训练的 LLM 中可能会导致训练和推理的不稳定。相对较小的多模态数据量可能不足以训练 MoD 模块。两种增强 MoD 模块的设计，以应对这些挑战。

Tanh 门控权重归一化

原始的 MoD 机制是为从头训练 MoD 基于 LLMs 的模型设计的。在多模态训练场景中，需要将新的 MoD 模块插入到预训练的原始 LLM 中。一种常见的方法是交替训练，其中在第一阶段仅训练新插入的模块，在第二阶段同时训练所有模块。然而，这种方法不适合 MoD，因为 MoD 模块在 token 重加权过程中通过其权重来缩放 tokens，而后续被冻结的 LLM 层无法处理被缩放的 tokens。本文的实验结果一致表明，训练 MoD 模块并冻结 LLM 层无法收敛。因此，唯一可行的方法是直接解决这一挑战：将 MoD 模块插入到 LLM 中，并同时训练它们。

这里  是一个超参数，用于控制归一化权重的范围。虽然很简单，但 TanhNorm 确保了以下几点：

这些特性带来了以下好处：

• 归一化的权重在训练开始时接近零，这确保了通过插入 MoD 模块后，LLM 的正确训练初始化保持不变。
• 零中心的权重分布降低了某些 tokens 在所有 MoD 层中被正（或负）权重重复缩放的风险。这确保了训练的稳定性，并减轻了推理过程中数值稳定性的问题。

上述两个好处有助于提高模型性能。在第 4.4 节中，我们验证了 TanhNorm 的有效性。

对称 Token 重加权

与用于训练 MoD 基于 LLMs 的文本数据相比，MLLMs 是在规模小几个数量级的多模态数据上训练的。训练一个权重预测器（MoD 模块中唯一可学习的部分）来准确且稳健地评估视觉 tokens 的重要性并分配相应的权重是具有挑战性的。权重预测器模块的梯度来自于 token 重加权过程。原始 MoD 仅对选中的 tokens 进行 token 重加权，如方程 2 所示。这样，语言监督信号仅监督权重预测器为选中的 tokens 分配高权重，而对跳过的 tokens 的权重预测过程没有监督（即没有梯度）。

为了充分利用有限的训练数据，通过对选中和跳过的 tokens 对称地应用 token 重加权过程来增强 MoD 机制。通过我们对称的 token 重加权（STRing）模块，方程 3 可以进一步修改为：

通过这种方式，MoD 模块可以充分利用训练过程中的语言监督信号，并在有限的训练数据上学习准确评估 token 的重要性。

渐进式比例衰减（Progressive Ratio Decay）

在通过 TanhNorm 和 STRing 模块升级 MoD 后，成功地训练了基于 MoD 的 MLLMs。然而，模型的性能与效率之间的权衡仍远未令人满意。当 token 保留比例设置低于 70% 时，模型的性能急剧下降。

原始的 MoD 基于 LLMs 在所有 MoD 层中采用固定的 token 保留比例。这一策略假设 tokens 在不同层中的冗余程度相同。在多模态场景中这一假设是不成立的。在每个 transformer 层中的自注意力作用下，视觉 tokens 会逐渐相互聚合信息，文本 tokens 也会从视觉 tokens 中获取相关信息。因此，随着层数的加深，视觉 tokens 应该变得越来越冗余。

为了验证本文的假设，在基于 MoD 的 MLLM 上进行了一系列探索性实验，在所有层中设置 70% 的 token 保留比例。如下图 3a 所示，首先将层划分为若干组。在每个实验中，仅对其中一组层应用较低的 token 保留比例（15%），而保持其他层不变，并在 14 个基准上评估该模型的表现。图 3b 中的结果强烈支持了我们的假设：将较低的 token 保留比例应用于越深的层，性能越好。这表明，视觉 tokens 在较深的层中具有更高的冗余性，这与之前的研究 [4, 45] 的观察结果一致。

本文设计了一个渐进式比例衰减（PRD）策略。如前图 2 右侧所示，PRD 按照 shifted cosine 调度逐层减少 token 保留比例。假设模型有 L 层，第 l 层的 token 保留比例可以表示为：

这里， 是一个平移因子，可以通过垂直平移余弦衰减曲线来灵活控制模型的整体计算成本。

在实验中，证明了使用 PRD 策略的 p-MoD 模型在相同计算预算下，显著优于使用固定保留比例的模型。它还优于其他类型的比例调度器。

实验

设置

模型
为了评估 p-MoD 的有效性，选择了两个代表性的开源 MLLM：LLaVA-1.5  和 LLaVA-NeXT，作为实验中的基线模型。将对应的 p-MoD 模型命名为 p-MoD-LLaVA-1.5 和 p-MoD-LLaVA-NeXT。LLaVA-1.5 将输入图像调整为固定分辨率，这与其 CLIP 视觉编码器的训练设置对齐，将图像编码为 576 个 token。LLaVA-NeXT 将高分辨率图像分成多个片段，这些片段由视觉编码器独立处理。这一策略增强了视觉感知能力，但也导致了更多的视觉 token（最多 2880 个）和更高的计算成本。

训练数据
遵循基线方法的训练设置。LLaVA-1.5 和 LLaVA-NeXT 都使用 LCS-558K 作为预训练数据。在指令调优阶段，LLaVA-1.5 使用 LLaVA-mix-665。原始的 LLaVA-NeXT 模型使用了 760k 数据进行微调，但 LLaVA 团队发布的 LLaVA-NeXT 训练数据包含 779k 数据。为了公平比较，我们使用 LLaVA 团队发布的 779k 数据和官方 LLaVA-NeXT 代码库训练 LLaVA-NeXT 模型。

基准
在 14 个基准上进行了全面的实验：GQA、OK-VQA、AI2D 和 ScienceQA-IMG 是传统的视觉问答基准；DocVQA、TextVQA、ChartQA 和 InfographicVQA 聚焦于细粒度的视觉问答；POPE 评估 MLLMs 的幻觉；SEED-Bench（图像）、RealWorldQA、MME、MMBench 和 MMMU 是为 MLLMs 定制的综合性基准。为了确保本文的结果可以方便地复现，我们使用 LLaVA 团队开发和维护的 lmms-eval 评估框架评估本文的模型。

评估指标
本文报告所有基准的性能，并结合反映训练和推理效率的多种指标。我们报告了 TFLOPs，它衡量推理计算复杂度，以及 KV 缓存存储，这是推理过程中的主要内存瓶颈。此外，我们还报告了模型训练中的 GPU 小时数和推理延迟。

实现细节

主要结果

将 p-MoD 集成到 LLaVA-1.5 和 LLaVA-NeXT 基线模型中，并在 14 个基准上进行了全面评估，如下表 1 和表 2 所示。与基线模型相比，p-MoD-LLaVA-1.5 和 p-MoD-LLaVA-NeXT 在所有基准上都实现了可比甚至更好的性能，同时显著节省了推理 TFLOPs 和 KV 缓存存储。

表 1 包含了一般的 VQA 和 MLLM 基准，衡量了 MLLM 在各个方面的综合能力。我们的 p-MoD-LLaVA-NeXT 模型在显著减少推理 TFLOPs 和 KV 缓存存储的同时，达到了可比的性能。我们的 p-MoD-LLaVA-1.5 模型甚至超过了 LLaVA-1.5 基线模型，推理 TFLOPs 为 58.7%，KV 缓存存储为 53.8%。

在表 2 中，值得注意的是，像 DocVQA、ChartQA、TextVQA 和 InfoVQA 这样的文本丰富和基于图表的基准需要细粒度的视觉感知和推理能力。只有在模型成功识别出与答案相关的区域并且这些区域占整个图像的比例较小时，才能给出正确的答案。值得一提的是，p-MoD-LLaVA-NeXT 在 DocVQA 和 InfoVQA 上几乎没有性能下降，在 ChartQA 上甚至提高了 0.2% 的准确率，同时在时间和内存效率上也取得了显著改进。

以上结果表明，本文的方法在保持性能的同时，显著提高了效率。

性能与效率的权衡

总之，表 3 展示了一个趋势：随着计算成本的降低，p-MoD 的性能线性下降，下降幅度较浅，提供了卓越的性能-效率权衡。

消融研究

Tanh-gated 权重归一化通过与不同的权重归一化方法进行比较来验证 TanhNorm 的有效性，如下表 4 所示。首先，确认使用没有权重归一化模块会导致训练过程中出现溢出，如第一行所示。然后，探讨 TanhNorm 的两个重要特性：

1. 零中心化归一化：我们将 TanhNorm 的零中心化特性与 Softmax 归一化函数进行比较，包括标准和偏移形式：
   从上表 4 的最后一行可以看出，TanhNorm 显著优于原始和偏移 softmax 函数。这证明了零中心化是 TanhNorm 强大性能的关键，因为尽管偏移 softmax 函数与 TanhNorm 具有相同的范围，但由于缺乏零中心化特性，表现较差。

对称 token 重新加权

对称 token 重新加权（STRing）策略的影响，如前表 4 中的最后两行所示。STRing 在 DocVQA 和 RWQA 基准测试中分别提高了 1%，且平均性能提高了 0.3%。这表明 STRing 通过更有效地利用语言监督信号来训练一个更好的权重预测器，尤其是在训练数据有限的情况下。

渐进比例衰减

为了验证渐进比例衰减（PRD）策略的有效性，进行了四个不同调度函数的实验，这些函数控制解码器层中的 token 保留比例。这些函数包括：

• 常数函数：对所有层使用相同的比例。
• 交替函数：在标准 transformer 层和具有较低保留比例的 MoD 层之间交替，原 MoD 论文中采用此方法。
• 阶梯衰减函数：以阶梯方式减少保留比例。
• 线性衰减函数：逐渐减少保留比例。

在本文的实验中，将所有层的平均保留比例固定在约 54%。下表 5 中的结果得出以下结论：

• 非衰减函数，如常数和交替函数，显著低于衰减函数。这证实了视觉 token 在较深层中存在更高的冗余性，支持渐进衰减的必要性。
• 连续衰减函数，如余弦和线性衰减，显著优于不连续的阶梯衰减函数。在这些方法中，余弦衰减函数取得了最佳结果，证明了我们采用的 shifted cosine 调度在优化层间 token 保留比例方面的有效性。

可视化

下图 5 展示了 p-MoD-LLaVA-NeXT 在不同层次上选择的 tokens。通过渐进比例衰减策略的引导，MoD 层处理的 tokens 数量逐层逐渐减少。从中可以观察到，选中的视觉 tokens 在图像中逐步集中于具有丰富语义信息的关键区域。具体而言，在上面的图中，位于不同位置的文本和图形所对应的视觉 tokens 几乎都在第 26 层以较低的比例被选中。在下面的图中，选中的 tokens 渐渐集中于标尺上的度量标记和其他物体。这些结果表明，我们的模型通过保留有信息的 tokens 有效地压缩了视觉信息，从而在显著降低训练和推理成本的同时，保持了良好的性能。

结论与未来工作

本文探讨了通过 Mixture-of-Depths 机制构建高效的 MLLMs（多模态语言模型）。提出命名为 p-MoD的模型，具有三个关键设计。首先，应用了 tanh-gated 权重归一化机制，以确保预测 token 权重在训练初始化和推理过程中保持稳定。其次，采用了对称 token 重新加权策略，充分利用语言监督信号，通过有限的训练数据训练出更优的权重预测模块。最后，模型采用了渐进比例衰减策略，逐层减少 token 保留比例。p-MoD 模型在 14 项不同基准测试中，取得了与基线模型相当或更优的结果，并且在训练和推理效率上有了显著提升。希望 p-MoD 能成为未来高效 MLLMs 研究的强大且高效的基线。

本研究的一个限制是 p-MoD 仅在 LLaVA-1.5 和 LLaVA-NeXT 模型上进行了实验，这些模型专注于单图像理解任务。认为p-MoD 在处理更多视觉 token 的任务中，如高分辨率、多图像和长视频理解任务中，可能会取得更显著的结果。将把 p-MoD 应用到其他视觉任务中的探索留待未来的研究。

本文转自AI生成未来 ，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/DhCPVbneK0adJVPx15Q3HQ​​