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 路由到更多分歧大师，由此可让模型对该地区的语义有更细粒度的明白。
该团队也执行了复杂性和参数剖析，详见原论文。