微软让MoE长出多个头，大幅提升专家激活率
MH-MoE 能优化几乎所有专家，实现起来非常简单。
混合专家（MoE）是个好方法，支持着现在一些非常优秀的大模型，比如谷歌家的 Gemini 1.5 以及备受关注的 Mixtral 8x7B。
稀疏混合专家（SMoE）可在不显著增加训练和推理成本的前提下提升模型的能力。比如 Mixtral 8×7B 就是一个 SMoE 模型，其包含 8 个专家（共 7B 参数），而其表现却可以超过或比肩 LLaMA-2 70B 和 GPT-3.5。
但是，它也有两个问题。一是专家激活率低 —— 也就是搞不好会出现下图这种情况：
具体来说，就是在优化时只有一小部分专家会被激活，如图 1a 所示（8.33% 的激活率），这会导致在学习应对复杂任务的大量专家时，会出现性能次优和效果不佳的问题。
二是无法细粒度地分析单个 token 的多重语义概念，比如多义词和具有多重细节的图块。
近日，微软研究院和清华大学提出了多头混合专家（MH-MoE）。顾名思义，MH-MoE 采用了多头机制，可将每个输入 token 分成多个子 token。然后将这些子 token 分配给一组多样化的专家并行处理，之后再无缝地将它们整合进原来的 token 形式。
论文标题：Multi-Head Mixture-of-Experts
论文地址：https://arxiv.org/pdf/2404.15045
代码地址：https://github.com/yushuiwx/MH-MoE
图 2 展示了 MH-MoE 的工作流程。可以看到，当输入单个 token 时，MH-MoE 会将其分成 4 个子 token，进而激活 4 个专家，而 SMoE 仅激活 1 个专家。
如图 2 所示，分配给专家 3 和 2 的子 token 包含对图块内每个角色动作的详细理解，而分配给专家 1 和 4 的子 token 则显式地建模了错误的同源词「camera」的语义。
专家处理完成后，再将子 token 无缝地重新整合进原来的 token 形式，由此可以避免后续非并行层（例如注意力层）的任何额外计算负担，同时还集成从多个专家捕获的语义信息。
这样的操作可让 MH-MoE 从整体上关注来自不同专家内不同表征空间的信息，从而可以加深上下文理解能力，同时提升专家激活率。该项目的代码也将发布。
MH-MoE 的具有以下优势：
专家激活率更高且扩展性更好。MH-MoE 能优化几乎所有专家，从而可以缓解专家激活率低的问题并大幅提升更大专家的使用率，如图 1a 所示实现了 90.71% 的激活率，这能让模型能力获得更高效的扩展。
具有更细粒度的理解能力。MH-MoE 采用的多头机制会将子 token 分配给不同的专家，从而可以联合关注来自不同专家的不同表征空间的信息，最终获得更好更细粒度的理解能力。举个例子，如图 1b 的明亮区域所示，子 token 会被分配给更多样化的一组专家，这有助于捕获语义丰富的信息。
可实现无缝整合。MH-MoE 实现起来非常简单，而且与其它 SMoE 优化方法（如 GShard）无关，反而可以将它们整合起来一起使用以获得更好的性能。
方法
图 3 给出了 MH-MoE 的整体架构，其使用了多头机制将每个 token 分拆为子 token，然后将这些子 token 路由给不同的专家。
多头混合专家
为了能清楚说明，这里仅描述单层 MH-MoE。
首先，通过一个多头层将输入 token 序列投射成一个新序列。
之后，沿 token 维度将新序列中的每个 token 分拆为多个子 token，并根据原始 token 序列并行排布这些子 token，进而构成一个新的特征空间。
然后将所有这些子 token 输送给一个门控函数。将特定子 token 路由到第 p 个专家的门控值的计算方式为：
对于路由方法，这篇论文关注的重点方法是 top-k 路由，也就是激活路由分数最大的 k 个专家。然后让这些激活的专家处理子 token。
之后，按子 token 原来的顺序重新排布并整合所得结果。
然后，通过一个 token 合并操作将所得整合结果转换回原始 token 形式。
最后，使用一个融合层将转换后的结果投射成多个特征的有效整合形式，此时这些特征已捕获了不同专家表征空间的详细信息。这样便可得到单层 MH-MoE 的最终输出。
训练目标
MH-MoE 的训练目标是最小化两个损失：针对具体任务的损失和辅助性的负载平衡损失。
实验
实验设置
为了进行比较，该研究的实验采用了两种基准模型：(1) Dense，这是没有整合稀疏激活的并行模块（SMoE 层）的 Transformer 解码器。(2) X-MoE，基于 Chi et al. (2022) 的论文《On the representation collapse of sparse mixture of experts》提出的方法的实现。
实验中的 MH-MoE 基于 X-MoE 并使用了与其一样的设置。
实验任务有三个：以英语为中心的语言建模、多语言语言建模、掩码式多模态建模。
更多有关数据集和模型架构的设置请参阅原论文。
困惑度评估
他们在两种专家设置（8 个专家和 32 个专家）下研究了所有预训练模型和预训练任务的验证困惑度曲线。图 4 给出了困惑度趋势，表 1 是最终的困惑值。
据此可以看出：
相比于基准，MH-MoE 的困惑度总是更低，这说明其能更有效地学习；
在三个不同的设置中，MH-MoE 的困惑度是最低的；
当专家数量增多时，MH-MoE 的困惑度会下降，这说明随着专家数量增多，其表征学习能力会提升，模型也能从中受益。
这些结果表明 MH-MoE 在多种预训练范式下都有更优的学习效率和语言表征能力。
下游任务评估
为了验证 MH-MoE 的效果，该团队也为每个预训练任务执行了对应的下游任务评估。
以英语为中心的语言建模
这里使用了 9 个不同的零样本评估基准，可以评估模型解决多种不同自然语言任务的能力，比如常识推理、一般语言理解和知识理解。评估框架为 LLM Evaluation Harness。结果见表 2。
可以看到，相比于 Dense 模型，X-MoE 有明显优势，这说明较大的模型能让 SMoE 模型（如 X-MoE）受益。总体而言，MH-MoE 在所有基准上都表现最佳。
多语言语言建模
他们在跨语言自然语言推理（XNLI）语料库（14 种语言）上评估了新的多语言语言模型。评估框架依然是 LLM Evaluation Harness，同样使用了零样本设置。结果见表 3。
MH-MoE 依然表现最佳，这体现了多头机制在建模跨语言自然语言方面的有效性。
掩码式多模态建模
他们也在社区广泛使用的视觉 - 语言理解和生成基准上执行了评估，包括视觉问答、视觉推理和图像描述。评估结果见表 4。
可以看到，MH-MoE 在这三个任务上有着全面的优势。这些结果表明 MH-MoE 具有更强的视觉信息理解能力，这也验证了新提出的多头机制在捕获视觉数据中的不同语义和详细信息方面的有效性。
消融研究
为了验证 MH-MoE 各组件和参数的效果，该团队也进行了消融研究。他们研究的内容包括头的数量、多层感知器层（包括多头层和融合层）、token 拆分与融合操作、MLP 层的数量。
表 5、6、7 给出了研究结果。整体而言，MH-MoE 各组件的效果得到了验证，并且他们也得到了一些有趣的结果，比如从表 7 可以看出单层 MLP 足以实现 token 分割和融合。
分析
专家激活分析
最后该团队还通过可视化分析等方法对 MH-MoE 进行了分析。
图 5 给出了 X-MoE 和 MH-MoE 中专家激活的分布情况。
可以看到，MH-MoE 的专家激活率明显更高，并且随着头的数量 h 增大，专家激活的频率也会上升。
图 6 则对比了 X-MoE 和 MH-MoE 的可扩展性（专家数量从 8 扩展到 256）。
可以看到 MH-MoE 的优势非常明显，并且 X-MoE 的下游性能会在专家数为 64 时达到饱和，而 MH-MoE 却还能继续提升。
分析细粒度理解能力
为了进一步分析多头机制对 MH-MoE 的帮助，该团队更深入地分析了其理解多样且复杂的语义信息的能力，比如理解语言中的多义词和错误同源词（记为 PF token）以及图像中的信息丰富的区域。
对于语言数据，他们计算和比较了从 PF token 和非 PF token 拆分出的子 token 的散度层级（即这些子 token 路由到的不同专家的数量）。结果见图 7。
可以看到相比于非 PF token，PF token 的散度分布明显靠右。这说明，在 MH-MoE 的推理过程中，PF token 会将其子 token 路由到更多不同专家，从而会捕获到与非 PF token 不同的语义信息，实现更好的多义词和错误同源词建模。
对于图像数据，他们分析的是不同图块的散度层级在训练过程中的变化情况，结果见图 8。
有趣的是，可以看到随着训练步骤增多，高频纹理区域（即有丰富语义信息的区域）的散度层级会逐渐增大，而低频纹理区域的散度层级则会逐渐降低。这表明在训练过程中，MH-MoE 倾向于将具有复杂纹理的区域的 token 路由到更多不同专家，由此可让模型对该区域的语义有更细粒度的理解。
该团队也执行了复杂性和参数分析，详见原论文。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道：[email protected]
微信扫码关注该文公众号作者