MoE-LLaVA: Mixture of Experts for Large Vision-Language Models

发表时间：6 Jul 2024

论文链接：https://arxiv.org/pdf/2401.15947 作者单位：Peking University

Motivation：最近的进展表明，扩展大型视觉语言模型 (LVLM) 有效地提高了下游任务的性能。然而，现有的缩放方法使所有模型参数在计算中的每个标记都处于活动状态，这带来了大量的训练和推理成本。然而，直接应用MoE来训练稀疏的LVLM是具有挑战性的。我们观察到，同时将LLM转换为LVLM和稀疏化模型会导致显著的性能下降。

解决方法：经过多次尝试，我们发现适当的初始化对于稀疏化 LVLM 至关重要，因此，我们引入了一种简单而有效的三阶段训练策略 MoE-Tuning（也就是说之所以三阶段训练，是为了更好的初始化参数）。在这项工作中，我们提出了一种简单而有效的 LVLM 训练策略 MoE-Tuning。该策略创新性地解决了多模态稀疏学习中性能下降的常见问题，从而构建了一个具有大量参数但计算成本恒定的稀疏模型。 此外，我们提出了MoE-LLAVA，这是一种基于moe的稀疏LVLM架构，在部署过程中，它只通过路由器唯一地激活前k个专家，使剩余的专家inactive，保持silence。

MoE-Tuning 包括三个阶段：

1. 在第一阶段，只训练 MLP。

2. 在第二阶段，除了视觉编码器 (VE) 之外，所有参数都经过训练。

3. 第三阶段，FFN 用于初始化 MoE 中的专家，只训练 MoE 层。对于每个 MoE 层，每个令牌只激活两个专家，而其他专家保持沉默。

MoE-LLAVA 仅在大约 3B 稀疏激活参数的情况下，在各种视觉理解数据集上表现出与 LLAAVA-1.5-7B 相当的性能，甚至在对象幻觉基准测试中甚至超过了 LLAVA-1.5-13B。如图：

实现方式：Training framework and strategy：MoE-LLAVA 采用三阶段训练策略。

(a) 阶段1，我们只训练 MLP 将 LLM 适应视觉输入。 数据：For the first stage of pretraining, we use the pretrained data of LLaVA 1.5-558k 说明：阶段 1 的目标是让视觉 token 适应到 LLM，赋予 LLM 看懂图片的实体有哪些的能力。MoE-LLaVA 采用一个 MLP 来将图片 token 投影到 LLM 的输入域，这意味着一个个图片小 patch 被 LLM 当作伪文本 token。在这个阶段，LLM 被训练学会描述图片，理解更高层次的图片语义。在这个阶段 MoE 层不会被应用到 LVLM 中。

(b) 阶段2，训练 LLM 后端支持多模态理解能力，不涉及 MoE 层。 数据：For the second stage, we collect datasets from MIMIC-IT, LRV , SViT and LVIS to provide a robust initialization for MoE-LLaVA. 说明：用多模态的指令数据来微调是提高大模型能力和可控性的关键技术，并且在这个阶段 LLM 被调整为有多模态理解能力的 LVLM。在这个阶段该研究加入更复杂的指令，包含图片逻辑推理，文字识别等高级任务，要求模型有更强的多模态理解能力。通常来说，稠密模型的 LVLM 到此就训练完成，然而研究团队发现同时将 LLM 转为 LVLM 和把模型稀疏化是具有挑战的。因此，MoE-LLaVA 将使用第二阶段的权重作为第三阶段的初始化以降低稀疏模型学习的难度。

(c)阶段3， 在这个阶段，我们复制 FFN（阶段2训练过，阶段3还训练） 的权重来初始化每个专家。 数据：For the third stage, we utilize the same data pipeline as LLaVA-mix-665k。 说明：MoE-LLaVA 将 FFN 复制多份作为专家集合的初始化权重。当视觉 token 和文本 token 被喂入 MoE 层时，router 会计算每一个 token 和专家们的匹配权重，然后每个 token 会被送入最匹配的 top-k 个专家进行处理，最后根据 router 的权重加权求和汇聚成输出。当 top-k 个专家被激活时，其余的专家保持非活动状态，这种模型构成了具有无限可能的稀疏通路的 MoE-LLaVA。

架构：在LLAVA 1.5 (Liu et al.， 2023b)之后，我们利用CLIP-Large (Radford et al.， 2021)作为视觉编码器，MLP由两个线性层组成（与LLAVA 1.5一样都用的MLP替换LLaVA中的线性层），GELU激活函数。除非另有说明，MoE-LLAVA 使用 MoE 层交替替换 FFN，这意味着 MoE 层的数量是层总数的一半。

图c中的expert在哪里？

ST (Switch Transformer): Switch Transformer中，模型的每一层都是一个专家网络的集合，输入数据会被动态地路由到不同的专家进行处理。ST在Transformer模型中用一个稀疏的Switch前馈网络（FFN）层（浅蓝色）替换Transformer中存在的密集FFN层。该层独立地对序列中的标记进行操作，然后路由到多个FFN专家中。switch FFN层返回所选FFN的输出，然后乘以路由器阈值，然后进行合并。参考链接：https://blog.csdn.net/chattyfish/article/details/131908504

也就是说图c中的不同的FFN就是expert：

实验：Image Understanding Evaluation，Object Hallucination Evaluation， MoE-LLaVA 的多模态理解能力评估。

由于 MoE-LLaVA 是第一个基于 LVLM 搭载 soft router 的稀疏模型，因此该研究将之前的模型归纳为稠密模型。研究团队在 5 个图片问答 benchmark 上验证了 MoE-LLaVA 的性能，并报告了激活的参数量和图片分辨率。与 SOTA 方法 LLaVA-1.5 相比，MoE-LLaVA-2.7B×4 展现了强大的图片理解能力，在 5 个 benchmark 上性能非常接近 LLaVA-1.5。其中 MoE-LLaVA 用 3.6B 的稀疏激活参数在 SQAI 上超过了 LLaVA-1.5-7B 1.9%。值得注意的是，由于 MoE-LLaVA 的稀疏结构，只需要 2.6B 的激活参数量就能全面超过 IDEFICS-80B。

结论：（主要由消融实验得出）

1. MoE-LLaVA 这一稀疏模型可以用更少的激活参数达到和稠密模型相当甚至超过的性能。

2. MoE-LLaVA 的 yes ratio 占比处于较均衡状态，这表明稀疏模型 MoE-LLaVA 能够根据问题做出正确的反馈。

3. MoE-LLaVA 的专家们学到了某种特定的模式，它能够按照一定的规律进行专家分工。

4. 研究发现 text 和 image 的路由分布极其相似，例如当专家 3 在 17-27 层努力工作时，它所处理的 text 和 image 的占比是相似的。这表明 MoE-LLaVA 对于模态并无明显的偏好。

5. 研究团队发现对于某个未见的 text token 或 image token，MoE-LLaVA 始终偏向于派发专家 2 和 3 来处理模型深层。专家 1、4 倾向于处理初始化的 token。这些结果能够帮助我们更好地理解稀疏模型在多模态学习上的行为，并探索未知的可能。

局限：虽然 MoE-LLAVA 展示了有竞争力的能力，但训练稳定性存在一些困难，尤其是 16 位浮点精度。 Future：由于存在多个专门针对不同能力的专家，MoE-LLAVA 可以轻松扩展以处理其他任务，例如检测、分割、生成或处理更多模式，例如视频、深度和热量。