Multi-Head Mixture-of-Experts

发表于 更新于 分类于 阅读次数： 本文字数： 1.8k 阅读时长 ≈ 7 分钟

Multi-Head Mixture-of-Experts[1]

作者是来自MSRA的Xun Wu等人，论文引用[1]:Wu, Xun et al. “Multi-Head Mixture-of-Experts.” ArXiv abs/2404.15045 (2024): n. pag.

Time

-2024.Apr

Key Words

  • low expert activation
  • multi-head
  • 一句话总结：类似多头注意力的操作，将输入分成多给sub-tokens，每个sub-tokens给到experts，最后将所有的输出在进行merge，还原为初始的形状，每个sub-tokens包含了不同feature space的语义信息

总结：

  1. **稀疏MoE在不增加计算成本的情况下，扩展了model的capacity，然而，它展示出了low expert activation的问题，仅有一小部分experts被激活，用于优化，导致suboptimal的性能，限制了在复杂任务中学习大量experts的有效性。在本文中，作者提出了Multi-MoE，MHMoE将每个输入的token或分成多个sub-tokens，然后这些sub-tokens被分配给多个并行的experts进行处理，无缝合成为原来的token form。以上的操作使得MH-MoE显著地提高了expert的activation，同时在不同的experts汇总，集体attend to 多个representation spaces，来加深context understanding，另外，值得注意地是: MH-MoE直接可以执行，和其它的SMoE框架解耦，使得很容易地和这些框架集成。

  1. 大的capacity models，例如LLMs和LMMs，在多个domains和task中展示出了它们的efficacy，为了进一步增强perfomance，一个可靠的方法是通过增强parameter count，扩展这些models，但是大多数是densely-activated large-capacity models(Dense Models)，用所有的参数来处理所有的inputs，这些models的size极大，降低了推理速度，限制了它们的实用性。 一个alternative是Sparse Mixtures of Experts，加速了model scalability，同时缓解了大量的计算成本，相比于Dense Model，SMoE包含了并行的FFN(称之为experts)，通过一个router，对于特定的输入tokens，激活distinct experts，产生了重要的效率提升。例如，GShard将一个Dense Model从2B扩展到了600B，比100B Dense Model的训练成本更低。最近，Mixtral的一个SMoE model，包含8个expert，超过了LLaMA-2和GPT-3.5.

    尽管它成功了，SMoE展示出了low experts activation issue，意味着在优化和推理期间，仅有一小部分的experts被激活了，其中的大部分没有用到，因此，SMoE未能利用这些experts的full expressive power，特别是这些experts的数量比较大，显著地限制了SMoE的有效性和可扩展性。

    本文的目标是实现denser expert activation，同时不增加计算成本。为了是实现这个，作者提出了MH-MoE，MH-MoE的工作流如图所示，受多头自注意力的启发，MH-MoE将每个input token划分为多个sub-tokens，将它们分给不同的experts，经expert处理之后，sub-tokens无缝地合成为原始的形状，实现了denser expert activation，绕开了额外的计算负担。

    另外，作者观察到一个有趣的现象：对于有更丰富的语义信息的tokens，从这些tokens划分来的sub-tokens更可能分配给distinct experts，作者猜想：将sub-tokens分配给distinct experts使得MH-MoE同时focus 不同expert内的多种representation spaces中的information，在vision和language patterns中，确保了对于subtle difference的更granular的理解。最后实现了更好的、finer-grained understanding ability。MH-MoE的优势有:

    • higher experts activation和better scalability: MH-MoE缓解了lower expert activate问题，增强了更大的experts利用
    • Finer-grained understanding ability: 基于token的语义的丰富程度，将sub-tokens分配给不同的experts，MH-MoE在不同的experts，能够attend to不同的representation spaces的information，实现更好的finer-grained understanding ability。
    • Seamless integration: MH-MoE的执行是直接的，和其它的SMoE的优化方法解耦，使其很容易和其他的框架集成。

  2. SMoE增强了模型的capacity，同时保持了恒定的计算需求，比densely-activate models实现了更好的performance，不同于densely-activated models，每个MoE layer包含 N个独立的FFN作为experts，还有一个gating function，用来建模一个概率分布，表示experts输出的权重，对于每个input token的hidden representation。 大多数的路由机制就是从 \(N\)个experts中选择 top-k个experts，这样的路由机制能够实现数据并行和专家并行的结合。一些工作表明：k 值越大，模型的性能越好，然而，K值增加，传统训练top-k路由的方法变得没有那么高效，本文中，作者提出了一个MH-MoE，是一个简单高效的方式，实现denser expert activation，不增加计算复杂度。

  3. MH-MoE: 输入的inputs tokens是 ：\(X \in R^{l \times d}\)，l 是tokens的数量，d表示token的维度，在MH-MoE中，每个并行的layer包含一组N个experts，每个表示为 \(f^{FFN}_i: \mathcal{R}^{d/h} \rightarrow \mathcal{R}^{d/h}\)，h表示heads的数量。完整的架构如下：X 通过一个权重为 \(W_{head}\) 的multi-head layer进行project，之后，\(\hat{X}\)中的每个token在token维度被分成 h个sub-tokens，这些sub-tokens根据原来的token sequence并行的组织在一起，构成了一个新的feature space。本文中，作者主要关注于top-k routing，有最高的top-k routing score的experts被激活， \[ \mathbf{o}_j^i = \mathbf{x}_j^i + \sum_{p \in \Phi} g\left(f_p^\text{FFN}\right) \cdot f_p^\text{FFN}\left(\mathbf{x}_j^i\right). \]

    之后，所有得到的 \(o^i_j\) 以原始sub-tokens的方式进行重新组合，然后，concatenate到一起。然后经过一个merge 操作变成original token form，然后用一个merge layer对其进行project，得到多个捕捉多个detailed information信息的集成，就是最后的MH-MoE layer的输出。

    通过以上的操作，能够有效地增加路由到特定专家的data的数量，实现denser expert activation，另外，MH-MoE layer中, input和output输出的shape是没有变化的，没有引入额外的计算，特别地，作者引入了一个超参数，来scale 每个expert的inner dimensions，旨在平衡multi-head layer和merge layer引入的参数，和原始的SMoE对齐模型的参数和计算复杂度，sub-tokens的分配使得能够在这些experts中，从多个feature space捕获semantic information，然后增强模型的能力，实现finer-grained understanding。

  4. training objectives: 为了缓解expert load 不平衡的问题，给定sub-token set \(\ddot{\mathbf{x}}\) 和sub-tokens路由到 \(p^{th}\) expert的频率 \(t_p\)，作者following 现有的工作，来计算load balancing loss \(\mathcal{L}_{\text{balance}}\): \[ \mathcal{L}_{\text{balance}} = \frac{N}{|\bar{\mathbf{X}}|} \sum_{p=1}^{N} \sum_{\mathbf{x}_j^i \in \bar{\mathbf{X}}} t_p \cdot g\left(f_p^\text{FFN}\right), \]

    \(N\) 表示experts的数量，\(\ddot{\mathbf{x}}\) 是sub-tokens的数量， \(g\left(f_p^\text{FFN}\right)\)，表示将特定的sub-token路由到 expert的gating value.

    • task specific loss: \(\mathcal{L}_{task}\) 依赖于MH-MoE 用来学习的特定的任务，在预训练English-focused Language Modeling task中，作者利用language modeling loss. 整个的training objective是这样的： \[ \mathcal{L} = \mathcal{L}_{task} + \alpha \mathcal{L}_{balance}\] \(\alpha\) 是load balancing的coefficent。

Workflow of MH-MoE \(Fig.1^{[1]}\) 对于视觉数据，路由到不同的experts的不同的heads尝试朴拙不同方面的details，对于language data，不同的heads能够捕捉

Workflow of MH-MoE \(Fig.2^{[1]}\)