Note

MoE 如何扩大参数量并控制推理计算量？

介绍专家网络、路由器、Top-K 稀疏激活与加权输出。

混合专家模型（MoE）是如何在不显著增加推理成本的情况下，有效扩大模型参数规模的？请简述其工作原理。

混合专家模型（Mixture of Experts, MoE）是一种模型架构，它的核心思想是通过 “稀疏激活”（Sparse Activation） 的策略，来解决模型规模与计算成本之间的矛盾。它允许模型拥有巨大的总参数量，但在处理任何一个输入时，只动用其中一小部分参数，从而在不显著增加推理成本（FLOPs）的情况下，大幅提升模型容量。

工作原理如下：

  1. 用“专家”替换FFN层：

    • 在标准的Transformer架构中，计算量最大的部分之一是前馈网络（Feed-Forward Network, FFN）层。
    • MoE架构将模型中的部分或全部FFN层替换为MoE层。一个MoE层由两部分组成：
      • N个“专家”（Experts）： 每个专家本身就是一个独立的、规模较小的FFN。
      • 1个“门控网络”或“路由器”（Gating Network / Router）： 这是一个小型的神经网络，通常是一个简单的线性层。

  2. 动态路由决策：

    • 当一个词元（token）的向量表示来到MoE层时，它首先被送入路由器。
    • 路由器的作用是 “决策” ，判断这个token应该由哪些专家来处理最合适。它会输出一个包含N个分数的向量，代表该token与N个专家的“匹配度”。

  3. Top-K稀疏激活：

    • 路由器输出的分数经过Softmax归一化后，系统并不会激活所有的专家。相反，它只选择分数最高的Top-K个专家（K通常很小，比如1或2）。
    • 这就是“稀疏激活”的关键：对于每一个token，只有极少数（K个）专家被激活并进行计算，其余的（N-K个）专家则完全不参与，不产生任何计算成本。

  4. 加权输出：

    • 被选中的K个专家分别对输入的token向量进行处理，得到K个输出向量。
    • 最终的输出是这K个输出向量的加权和，权重同样由路由器的输出分数决定。

如何实现“参数大但成本低”？

  • 假设一个模型有8个专家（N=8），并且每次只激活2个（K=2），如Mixtral-8x7B模型。
  • 总参数量： 模型的总参数量是所有共享部分（如注意力层）的参数量，加上所有8个专家的参数量之和。这使得模型的总参数规模可以非常大（例如达到47B）。
  • 推理成本： 但在进行一次前向传播（推理）时，对于任意一个token，实际参与计算的只有共享部分和被激活的2个专家。因此，其计算量（FLOPs）约等于一个规模小得多的“稠密”模型（例如一个13B的模型）。
  • 结论： MoE成功地将总参数量（代表模型的知识容量）和单次推理的计算量（代表模型的速度和成本）解耦，从而实现了“用小模型的成本，获得大模型的知识”。