混合专家，也得术业有专攻。
对于目前的混合模态基础模型，常用的架构设计是融合特定模态的编码器或解码器，但这种方法存在局限：无法整合不同模态的信息，也难以输出包含多种模态的内容。
为了克服这一局限，Meta FAIR 的 Chameleon 团队在近期的论文《Chameleon: Mixed-modal early-fusion foundation models》中提出了一种新的单一 Transformer 架构，它可以根据下一个 token 的预测目标，对由离散图像和文本 token 组成的混合模态序列进行建模，从而在不同模态之间进行无缝推理和生成。
在约 10 万亿混合模态 token 上完成预训练之后，Chameleon 表现出了适应广泛的视觉和语言能力，能很好地处理多种不同的下游任务。Chameleon 在生成混合模态长回答任务的表现尤其亮眼，它甚至打败了 Gemini 1.0 Pro 和 GPT-4V 等商用模型。然而对于 Chameleon 这样各种模态会在模型训练的早期混合起来的模型，想要拓展它的能力，需要投入大量算力。
基于以上问题，Meta FAIR 的团队对路由式稀疏架构（routed sparse architecture）进行了一番研究探索，提出了 MoMa：模态感知型专家混合架构。
论文标题：MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
论文地址：https://arxiv.org/pdf/2407.21770
之前已有研究表明，这类架构可以有效地扩展单模态的基础模型的能力，也可以增强多模态对比学习模型的性能。但是，将其用于较早将各种模态融合的模型训练还是一个机遇与挑战并存的课题，还少有人研究。
该团队的研究基于这一洞见：不同模态具有固有的异构性 —— 文本和图像 token 具有不同的信息密度和冗余模式。
在将这些 token 整合成统一的融合架构的同时，该团队也提出通过整合针对具体模态的模块来进一步优化该框架。该团队将这一概念称为模态感知型稀疏性（modality-aware sparsity），简称 MaS；其能让模型更好地捕获每个模态的特征，同时还能通过部分参数共享和注意力机制维持强大的跨模态整合性能。
之前的 VLMo、BEiT-3 和 VL-MoE 等研究已经采用了混合模态专家（MoME/mixture-of-modality-experts）方法来训练视觉 – 语言编码器和掩码式语言建模，来自 FAIR 的研究团队更进一步将 MoE 的可用范围又推进了一步。
模型架构
早期融合
本文提出的新模型基于 Chameleon 的早期融合架构，其做法是在一个统一 Transformer 中，将图像和文本表示成一系列离散 token。Chameleon 的核心是一个基于 Transformer 的模型，其会在图像和文本 token 的组合序列上应用自注意力机制。这能让该模型捕获模态内和模态间的复杂关联。该模型的训练使用的目标是下一 token 预测目标，以自回归方式生成文本和图像 token。
在 Chameleon 中，图像的 token 化方案采用了一个学习型图像分词器，它将基于大小为 8192 的 codebook 将 512 × 512 的图像编码成 1024 个离散 token。对于文本的分词将使用一个词表大小为 65,536 的 BPE 分词器，其中包含图像 token。这种统一的分词方法可以让模型无缝处理图像和文本 token 交织错杂的任意序列。
借助这种方法，新模型继承了表征统一、灵活性好、可扩展性高、支持端到端学习这些优点。
在此基础上（图 1a），为了进一步提升早融合模型的效率和性能，该团队还引入了模态感知型稀疏性技术。
宽度扩展：模态感知型混合专家
该团队提出了一种宽度扩展方法：将模态感知型模块稀疏性集成到前向模块中，从而扩展标准混合专家（MoE）架构。
该方法基于这一洞见：不同模态的 token 有各自不同的特征和信息密度。
通过为每个模态构建不同的专家分组，可让模型开发出专门的处理路径，同时维持跨模态的信息整合能力。
图 1b 展示了这种模态感知型专家混合（MoMa）的关键组件。简单来说，先是对各个特定模态的专家进行分组，然后实现分层路由（分为模态感知型路由和模态内路由），最后选择专家。详细过程参见原论文。
总体来说，对于一个输入 token x，MoMa 模块的形式化定义为：
在 MoMa 计算之后，该团队又进一步使用了残差连接和 Swin Transformer 归一化。
Mixture-of-Depths（MoD）
之前也有研究者探索将稀疏性引入深度维度，他们的做法要么是随机丢弃某些层，要么就是使用可学习的路由器。
该团队的做法参考了第二种方法，同时整合了近期提出的混合深度（MoD）技术。有关 MoD 的更多介绍可参阅AI在线报道《DeepMind 升级 Transformer，前向通过 FLOPs 最多可降一半》。
具体而言，如下图所示，该团队的做法是在每个 MoD 层中，在混合专家 （MoE）路由之前都集成 MoD，从而确保在模态分离之前，整批数据都能应用 MoD。
推理
在推理阶段，我们不能直接使用 MoE 的专家选择路由或 MoD 的层选择路由，因为在一批数据中进行 top-k（选择前 k 个）选择会破坏因果关系。
为了保证推理的因果关系，受上述 MoD 论文的启发，研究团队引入了辅助路由器（auxiliary router），其作用是仅基于 token 的隐藏表征预测该 token 被某个专家或层选中的可能性。
升级改造（Upcycling）
在优化表征空间和路由机制方面，对于一个从头开始训练 MoE 架构，存在一个独特的难题。该团队发现：MoE 路由器负责为每个专家划分表征空间。但是，在模型训练的早期阶段，这个表征空间并非最优，这就会导致训练得到的路由函数也是次优的。
为了克服这一局限，他们基于 Komatsuzaki 等人的论文《Sparse upcycling: Training mixture-of-experts from dense checkpoints》提出了一种升级改造方法。
具体来说，首先训练一个每个模态都有一个 FFN 专家的架构。经过一些预先设定的步数之后，再对该模型进行升级改造，具体做法是：将每个特定模态的 FFN 转换成一个专家选择式 MoE 模块，并将每个专家初始化为第一阶段训练的专家。这里会在保留前一阶段的数据加载器状态的同时重置学习率调度器，以确保第二阶段的训练能使用已刷新的数据。
为了促进专家更加专业，该团队还使用了 Gumbel 噪声来增强 MoE 路由函数，从而使得新的路由器能以可微分的方式对专家进行采样。
这种升级改造方法加上 Gumbel-Sigmoid 技术，可克服学习到的路由器的局限性，从而提升新提出的模态感知型稀疏架构的性能。
效率优化
为促进 MoMa 的分布式训练，该团队采用了完全分片式数据并行（FSDP/Fully Sharded Data Parallel）。但是，相比于常规 MoE，该方法存在一些特有的效率难题，包括负载平衡问题和专家执行的效率问题。
对于负载平衡问题，该团队开发了一种平衡的数据混合方法，可让每台 GPU 上的文本 – 图像数据比例与专家比例保持一致。
对于专家执行的效率问题，该团队探索了一些策略，可帮助提升不同模态的专家的执行效率：
将各个模态的专家限制为同构的专家，并禁止将文本 token 路由到图像专家，反之亦然；
使用模块稀疏性（block sparsity）来提升执行效率；
当模态的数量有限时，按顺序运行不同模态的专家。
由于实验中每台 GPU 处理的 token 都足够多，因此即使使用多个分批次矩阵乘法，硬件利用率也不算大问题。因此，该团队认为对于当前规模的实验环境而言，按顺序执行的方法是比较好的选择。
其它优化
为了进一步提升吞吐量，该团队还采用了其它一些优化技术。
其中包括降低梯度通信量、自动化的 GPU 核融合等一般优化操作，研究团队还通过 torch.compile 实现了图优化。
此外，他们还针对 MoMa 开发了一些优化技术，包括跨不同层复用模态 token 索引，以最高效地同步 CPU 和 GPU 之间的设备。
实验
设置
实验中使用的预训练数据集和预处理过程与 Chameleon 一样。为了评估扩展性能，他们训练模型使用的 token 数量超过 1 万亿。
表 1 给出了密集和稀疏模型的详细配置情况。
不同计算层级的扩展性能
该团队分析了不同模型在不同计算层级上的扩展性能，这些计算层级（FLOPs）相当于三种大小的密集模型：90M、435M 和 1.4B。
实验结果表明，一个稀疏模型仅使用总 FLOPs 的 1/η 就能比肩同等 FLOPs 的密集模型的预训练损失（η 表示预训练加速因子）。
模态解绑
引入特定模态的专家分组可提高不同规模模型的预训练效率，这对图像模态尤其有益。如图 3 所示，使用 1 个图像专家和 1 个文本专家的 moe_1t1i 配置显著优于相应的密集模型。
扩展每个模态分组的专家数量还能进一步提升模型性能。
混合深度与专家
该团队观察到，当采用 MoE 和 MoD 以及它们的组合形式时，训练损失的收敛速度会得到提升。如图 4 所示，向 moe_1t1i 架构添加 MoD（mod_moe_1t1i）可大幅提升不同模型大小的模型性能。
此外，在不同的模型大小和模态上，mod_moe_1t1i 能媲美甚至超过 moe_4t4i，这表明在深度维度上引入稀疏性也能有效提升训练效率。
另一方面，还能看到堆叠 MoD 和 MoE 的收益会逐步下降。
扩展专家的数量
为了研究扩展专家数量的影响，该团队进行了进一步的消融实验。他们探索了两种场景：为每种模态分配同等数量的专家（平衡）以及为每种模态分配不同数量的专家（不平衡）。结果见图 5。
对于平衡的设置，从图 5a 可以看到，随着专家数量提升，训练损失会明显下降。但文本和图像损失表现出了不同的扩展模式。这表明每种模态的固有特性会导致出现不同的稀疏建模行为。
对于不平衡的设置，图 5b 比较了同等专家总数（8）的三种不同配置。可以看到，一个模态的专家越多，模型在该模态上的表现通常就越好。
升级改造
该团队自然也验证了前述的升级改造的效果。图 6 比较了不同模型变体的训练曲线。
结果表明，升级改造确实能进一步改善模型训练：当第一个阶段有 10k 步时，升级改造能带来 1.2 倍的 FLOPs 收益；而当这个步数为 20k 时，也有 1.16 倍的 FLOPs 收益。
此外，还能观察到，随着训练推进，经过升级改造的模型与从头开始训练的模型之间的性能差距会不断增大。
吞吐量分析
稀疏模型通常不能立即带来性能增益，因为稀疏模型会增加动态性和相关的数据平衡问题。为了量化新提出的方法对训练效率的影响，该团队通常控制变量实验比较了不同架构的训练吞吐量。结果见表 2。
可以看到，相比于密集模型，基于模态的稀疏性能实现更好的质量 – 吞吐量权衡，并且能随专家数量增长展现出合理的可扩展性。另一方面，尽管 MoD 变体取得了最好的绝对损失，但由于额外的动态性和不平衡性，它们的计算成本往往也更高。
推理时间性能
该团队也评估了模型在留存的语言建模数据和下游任务上的表现。结果见表 3 和 4。
如表 3 所示，通过使用多个图像专家，1.4B MoMa 1t1i 模型在大多数指标上都优于相应的密集模型，只有在 COCO 和 Flickr 上的图像到文本条件困惑度指标例外。进一步扩展专家数量也能提升性能，其中 1.4B MoE 8x 在图像到文本性能上达到了最佳。
此外，如表 4 所示，1.4B MoE 8x 这个模型还非常擅长文本到文本任务。1.4B MoMa 4t4i 在所有条件图像困惑度指标上表现最佳，而其在大多数基准上的文本困惑度也非常接近 1.4B MoE 8x。
总体而言，在混合文本和图像两种模态的数据上，1.4B MoMa 4t4i 模型的建模结果最好。
更多详细内容，请阅读原论文。