MoE 与 LLM

MoE

深度学习时代，尤其是在 NLP 大模型大行其道的今天，MoE（Mixture-of-Experts）是指这样一种结构：

在 MoE 模型中，输入数据被 Router（路由）选择性地依照一定规则分配给最适合处理它的 Expert（专家），最终推理结果由专家输出混合形成。因此整个结构被称为混合专家。

可能有人在这种结构中看到了机器学习领域集成方法的影子。二者的确相似，但目的上有着区别。

集成方法利用了稀疏的不同「视角」来增强模型表现。其每种方法可能较简单（过程、参数），但由于更易于训练、从更多角度上反映了数据特性，因此集成后达到了极强的性能。各挑战赛的 Top 方法多数为集成方法，如果你有所了解，GPT 4 同样为集成模型^[1]。

不过 MoE 的出发点似乎略有不同。在大量的 LLM 实践中，人们意识到了 Transformer + 超多参数 力大砖飞的暴力效果，在克服一系列问题后诞生了 GPT 这种怪物。于是一个问题诞生了：怎么在有限制的计算资源下，继续增加参数并获得更好的性能？当前 LLM 的训练和推理成本已经有点无法承受，继续密集地增加参数并非可行，并且在训练难度上也会遇到问题。因此，MoE 这种稀疏结构理念再次出现在了 LLM 的 Transformer 结构中：

• 更易训练，更快推理
• 较低的硬件需求
• 推理性能展现出一定优势

MoE 的结构及 Transformer 集成

一个 MoE 的结构主要包含：

• 稀疏的 MoE 组件： 也就是多个专家形成的集合体。最简单的形式是一个 FFN，我们会在下文中继续解释。
• 路由/门：有选择性的将输入分配给专家。如何路由是 MoE 的关键。

首次应用：GShared

MoE 首次被塞到 Transformer 中是在 GShared中。GShard 在编码器和解码器中使用 top-2 Router，将 FFN 层替换为 MoE。

在 MoE 于 Transformer 的首次应用中讨论了以下几个关键问题：

• 路由方法：每次选择 top-1 expert，并依据概率选择 top-2 expert。
• 专家容量：限制每个专家能处理的输入数量。一个是出于防止某个 expert 被过度使用，另一方面则是技术性问题（张量的维度需要被预先定义）。

经典结构：Switch Transformer

一种较经典的结构是用 MoE 替换 Transformer 中每个 Block 的 FFN 结构后形成的新模型 Switch Transformer^[8]，如下图所示。相比 GShared 来说，Switch Transformer 的训练更容易。

• 路由方法：只使用 top-1 expert。减少计算量。为了方式 Router 塌陷，加入了 Auxiliary Loss。
• 专家容量：$\frac{\text{tokens}}{\text{experts}} \cdot \alpha$，其中 $\alpha$ 是 capacity factor。

除此之外，另有论文认为 MoE 可替换 Transformer 中的 Attention 结构^[9]。

• 细粒度调整：Attention 的微调可以更细粒度地控制模型的行为，其直接影响输入权重，而 FFN 更多涉及线性转换
• 效果显著：在某些任务中可以带来更显著的性能提升。Attention 机制在处理上下文和长程依赖关系上起关键作用
• 稳定性：在 Attention 上微调带来更好的模型稳定性，尤其是在大规模模型和复杂任务下

另外，该文章引入了一种新的 Router loss，z-loss，用于保证 Router 的稳定性。TODO

MoE 的优势就在于这种 Expert + Router 结构带来的稀疏性（当然在性能上也有竞争优势）。MoE 推理时每次只激活少量的参数，从而在进一步扩大模型参数量的同时控制住了计算量。

Router

Router 作为分配输入的主要结构，对整个网络有着关键性作用。记 Router 为 $[g_1(x),g_2(x),\dots,g_N(x)]$，则一种 MoE 网络可以简单描述为

$y=\sum g_i(x)f_i(x)$

简单来讲，Router 决定了各个专家网络 f 在最终输出的占比。不过为了达到稀疏性和低硬件需求，MoE 在推理中不使用所有 Expert，而是选取权重最高的 k 个。这样即使网络很重，每次参与计算的参数也仅占一小部分。

$k > 1$ 的另一种考量是为了防止 Router 塌陷到只选择一个 Expert 的情况。

Router 依什么路由

虽然有人（比如我）可能期望 Router 在设计上有着更具逻辑的路由结果，例如按照语言或某种更具可解释性的要素将输入转发给不同 Expert，但事实好像并非如此。取决于对 Router 的正则化方式及数据分布，Encoder 和 Decoder 结构上能观察到一些可解释的相对分工，例如 Encoder 在 NLP 任务中能观测按词性的分工^[2]。

在各 LLM 中目前实践的 Router 结构十分简单，记 $g(x)=[g_i(x)]_i$

$g(\mathbf x)=\mathrm{Softmax}(\mathbf W^\top \mathbf x)$

一句话，就是一个简单的全连接 + Softmax 激活。在得到权重后，依权重选取少量专家参与计算。

负载均衡

实际中为了避免 Router 对路由的选择出现坍缩（只选 1 个），可能会额外叠加一些噪声

$g(\mathbf x)=\mathrm{Softmax}(\mathbf W^\top \mathbf x)+ N \cdot \mathrm{Softmax}(\mathbf W_n^\top \mathbf x)$

或者在训练上增加正则项，称为 Load Balancing Loss：

$\begin{aligned} L&=\alpha \cdot N \sum_{i=1}^N r_iP_i \\ r_i&=\frac 1 T \sum_{x \in \mathcal B} \left[ i=\arg\max_i g(x) \right] \\ P_i&=\frac 1 T \sum_{x \in \mathcal B} g_i(x) \end{aligned}$

简单来讲，$r_i$ 为路由至专家 i 的输入占总输入的比例，$P_i$ 是专家 i 对所有输入的选择性总和，L 要求 Router 的输入分配及权重分配应接近平均，即要求 Router 平等对待所有专家。

在理想情况下，r 和 P 都应该是 1/N，因此在 L 中乘了一个 N 用于归一化。

除此之外，Expert 在实际设计中还涉及了「容量」问题，即负载均衡的基础上外加的硬上限。（TODO）

Token-Choice & Expert-Choice

但输入的负载均衡还有另一种解决方案^[6]。重新审视 Top k 的选择过程，我们还有其他方案。在本节中我们不区分「输入」和「Token」。

1. Token-Choice: 计算每个 Token 对 Expert 的权重，然后每个 Token 依次依照 Top 1 分配至对应专家。如果专家容量已满，则该 Token 被舍弃；亦会出现专家容量不满的情况。
2. Expert-Choice：计算每个 Token 对 Expert 的权重，然后每个专家依次选取 Top k（容量）个 Token 处理。可能出现某 Token 不被选择的情况；亦会出现 Token 被选择多次的情况。

需要主要的这个问题只在存在容量上限时才存在。在采用第二种方案时，就不再存在均衡问题了——每个专家必然都能接收到满额的 Token。考虑到对容量的高利用，以及 Token 能够被多次（不）选择，这种方案具有更好的灵活性。

另一种 Router 方案

有一种我感觉较为有趣的 Router 方案，它借助 autoencoder 的特性来计算权重^[3]。

如果一个 autoencoder 编解码后的数据与输入越接近，就说明输入数据的分布越匹配这个 autoencoder。如果将其作为 Router 使用，就能选择同分布 Expert 来处理数据了。

这种 Router 方案更重，对 Expert 的职责划分十分明显，可解释性较强。原文将该设计应用在持续学习，用于匹配输入数据分布，整体十分合适。但它在 LLM 上的用武之地有些局限，Token 粒度的 Routing 似乎与数据分布牵扯不上什么关系。在更大尺度上，也许它能用在 LLM 的大粒度持续学习上。

SoftMoE

SoftMoE: Soft Mixture of Experts for Efficient and Effective Fine-tuning of Large Language Models

https://arxiv.org/pdf/2308.00951

这是某种比较新的 Router 方案。

简单来讲，这种方法不使用 top-k 这种人工设计的启发式硬机制从所有 token 中挑选被某专家处理的子集，而是将所有 token 混合为新的 token 输入给 expert。这个混合权重与 token 及 slot 编号相关。

因此它规避了很多问题，例如上文提到的负载失衡，以及没提到的 token 丢失、随机路由等问题。同时，它的路由过程全部可微。

不过在实际情况中，LLM 常见的自回归过程导致 token 并非一开始就全部已知，所以 router 也需要某种 mask，在训练需要额外注意。

这篇文章 diss 了很多使用 auxiliary loss 平衡路由的工作，说他们和随机路由没什么区别。这个可以关注一下。

Expert

LLM 中的 Expert 就是 Transformer 中的 FFN。

这只是 MoE 的基本结构。使用其他结构实现 Expert 同样可取。

在实践中，随 expert 的数量增加，模型表现上升逐渐趋缓。常见超参为 256/512，但相同的趋势同样出现在小数量 expert 时。

---

预训练 MoE 模型的微调价值

无论是微调 LLM 还是更具体到微调 MoE 结构的 LLM。我们都有很多办法：

首先是全量微调 MoE，以下经验来自于 huggingface 的 blog^[10]。在我看来这些经验比较符合直觉。

• MoE 相比密集模型更容易过拟合。
• MoE 在重推理上更容易拟合，在知识型任务上表现更好。
• 微调剩余部分比微调 MoE 更容易得到效果。（但我觉得这可能和参数量更有关）
• MoE 微调往往倾向于更小的 batch (65k, 262k, 1M) 和更大的 lr (1e-4, 5e-4, 1e-3)。

更加有效、稳定的参数微调方法就相对较多：

• Adapter：在网络中引入少量可训练参数。例如在 Transformer 的 attention 后和/或 FFN 后引入 MLP。
• Prompt Tuning：微调优化几个新的 token。另一个相似的概念是 Stable Duffusion 微调的 Text Inversion。
• Prefix Tuning：在我看来也是 Prompt Tuning 的一种。
• LoRA：同样是向网络中插入少量可训练参数，只是通过低秩分解进一步降低了参数量。

在 Prompt 上也有不同的编纂策略：

首先是喜闻乐见的 Prompt Learning，也就是提示学习、提示工程。这一般不涉及到微调参数，只是利用模型原本的知识，发挥 zero-shot 能力。

更加复杂有效的微调方式还有 Instruction Tuning，也就是指令微调。这种方式用更加具体的指令指出模型需要完成的任务，指导模型做出选择。这种微调方法一般也伴随着模型的微调。

在只使用 Prompt Learning 的时候，微调模型极容易造成过拟合，降低模型的综合表现。

后续的文章，我们将简单讨论一些具体的 LoRA MoE 方法。

MixLoRA

某种 MoE+LoRA 的排列组合。论文给出的动机

• MoE 在多任务上有更好的推理性能
• LoRA 在微调和资源上有更好的表现

然后它们就被缝合到一起了：使用 LoRA 构建 expert，形成 MoE 结构。

• 每个专家有一个参数共享的 FFN 和一个 LoRA。有点 get 不到这个 FFN 的目的是什么，但看起来是 pretrained 模型的一部分。
• LoRA MoE 替换了原有的 FFN；LoRA 替换了原有的 attention。
• 使用了 load balancing loss。

MoRAL

一篇声称自己解决大模型 Lifelong Learning 的文章。

• 使用了 LoRA 结构的 MoE
• 替换了 Transformer 的 attention
• Top-2 Router

但是读完之后发现 Lifelong Learning 像是个故事。对该问题有着大量的描述，但方法相比之下却缺少针对 Lifelong Learning 的特化。只能理解成利用了 LoRA、MoE 本身的某些特性。

另一个值得注意的是其所采用的新指标 RA：

$\mathrm{RA}=\frac{\mathrm{F1}w_0+\cos(\mathrm{EMB}(a),\mathrm{EMB}(G_t))w_1}{w_0+w_1}$

其中两个 EMB 计算了模型输出和 gt 的嵌入。个人有些担心这个指标，但由于论文没有公开代码，所以就先到这了。

PESC

Parameter-Efficient Sparsity Crafting from Dense to Mixture-of-Experts for Instruction Tuning on General Tasks

在不显著降低模型性能的前提下，通过稀疏化技术减少指令微调中的计算开销，从而提高训练和推理的效率。

• MoE 结构，虽然没有使用 LoRA，但实际上它的 expert 也是某种稀疏化结构
• top-2
• 使用了 load balancing loss

整个 MoE 很常规，上面都有提及。可以视作 MoE + Instruction Tuning 的实践经验。

LoRAMOE

LoRAMoE: Alleviate World Knowledge Forgetting in Large Language Models via MoE-Style Plugin

微调模型能增强在下游任务上的表现，但会导致模型的灾难性遗忘。这篇文章使用 LoRA+MoE 的特性来抵抗部分遗忘，类似于 MoRAL 想要做的事情。

除了常规的 MoE 结构外，论文提出了一个新的 loss，称为 Localized Balancing Constraint，让 expert 出现对不同数据分布的合理分工（保留旧知识、学习新知识）。

定义 $\mathbf Q$ 为重要性矩阵，反映了某专家对某数据的关注程度：

$\mathbf Q_{n,m}=\sum_{x \in T_m}G(x)_n$

其中 $G(x)_n$ 是 Router 对专家 $n$ 的分配权重。

定义 $\mathbf I$ 为相关性矩阵，反映了某专家与某数据间的关联程度：

$\mathbf I_{n,m}=\begin{cases} 1+\delta & \text{专家 n 应处理数据 m} \\ 1-\delta & \text{专家 n 不应处理数据 m} \end{cases}$

其中 $\delta$ 是一个超参，用于调整专家的关注程度。

最终 loss 为

$\mathcal L=\frac{\sigma^2(\mathbf Q \cdot \mathbf I)}{\mu(\mathbf Q \cdot \mathbf I)}$

即，要求专家对于自己该处理的数据和不该处理的数据间有着明显的选择性。

MoLA

这篇文章主要研究了 LoRA MoE 在 LLM 中的数量该如何设置。

• 同样是常规的 LoRA MoE 结构
• top-2 router
• 使用了 load balance loss 做均衡

一路看下来也能发现这是相对常见的 MoE 实践了。文章的主要贡献在于发现低层 MoE 的专家更加冗余，高层则更加稀疏。所以在专家数量有限的时候，应该在高层引入更大的 MoE，形成一种倒三角的分布。

这篇文章也试图讨论对于持续学习的贡献，但还是简单利用了网络结构上的特性，可以拿来用但并似乎没有继续跟随研究的价值。

MoELoRA

MoELoRA: Contrastive Learning Guided Mixture of Experts on Parameter-Efficient Fine-Tuning for Large Language Models

文章针对 PEFT 微调 LLM 的共性问题（微调模型太难、参数稀少又限制灵活性）提出了 MoE+LoRA。

• 一样的结构
• top-k router
• load balance loss

一个额外的贡献是对 router 训练的对比学习正则项，看起来还挺有意思。为了让专家间形成区别、router 形成选择性，将同一个专家的输出互为正例，不同专家间的输出互为反例，从而在专家中产生区别。

$\mathcal L_{E_i}=-\sum_{q\not = k_+} \log \frac{\exp(q \cdot k_+/\tau)}{\sum_{k \in E(x)} \exp (q \cdot k / \tau)}$

其中 $E_i$ 是第 $i$ 个专家，$q$ 为专家 $i$ 输出的某个样本。整体是一个比较常见的对比学习 loss。这个贡献与其结构上并没有什么太大的关联，而是 MoE 会遇到的惯例问题，所以也可以单独拿出来使用。

MoCLE

尝试解决 Multi-task Instruction Tuning 时所遇到的不均衡问题。显然，方法就是 LoRA 低参数微调 + MoE 多任务适配啦……

• 类似的网络结构
• router 策略不太一样：1 个始终激活的 expert + top-1
• 向 router 添加了噪声

它认为多任务导致的 sub-optimal 由任务间的差异过大导致了 negative transfer。所以将任务根据 Instruction 嵌入的距离聚类后，作为相似任务簇塞给一个 expert 学习。

但与此同时，文章又搞了一个 universal expert，而且在消融中还提升了非常高的性能，比聚类的消融还高，这稍微有点奇怪。

另外，我觉得这种聚类非常的 hard，不如之前的 importance loss 或者 contrastive loss。

MOELoRA

依然是尝试解决 Multi-task 时所遇到的不均衡问题，而且还是专门针对医疗任务。显然，方法就是 LoRA 低参数微调 + MoE 多任务适配……

• 相同的结构
• 它的 router 有点不一样：是按照 sample -> task -> weight 得到最后结果。也就是说，相同任务的 sample 将得到相同的 weight。

MoA

Mixture-of-LoRAs: An Efficient Multitask Tuning for Large Language Models

依然是尝试解决 Multi-task 时所遇到的不均衡问题，依然是 LoRA 低参数微调 + MoE 多任务适配。

• 相同的结构。添加在 FFN 和 Attention 上。
• top-1 router，在任务级别上路由

它与其他工作的不同可以体现在：

• 首先在各个任务上分别训练 expert
• 然后再将所有任务的数据混合在一起，训练 router

似乎可以当作微调 Qwen 的经验来看待。

---

MoE 基础模型：Mistral 8x7B

Mistral 8x7B 是 Mistral 7B 的 MoE 变体，在大幅降低计算量的同时保持了与 Llama 70B、GPT-3.5 类似的性能。

该模型的 MoE 结构是上文提及的使用最广泛的一种：将 FFN 替换为 MoE 结构，Router 选择 top-2 专家。

与 Mistral 7B 一脉相承的，还有分组注意力和滑动窗口。TODO

https://github.com/HIT-SCIR/Chinese-Mixtral-8x7B

---

Benchmark

TODO

---

Appendix

MoD

近期 Google 又发了一篇新东西 MoD^[6]，提出了一种在深度上的 Mixture 策略。

说实话，看到下面这张图时有种强烈的既视感。

它的想法就是让 Router 决定是否要跳过某些层，从而更加灵活地控制计算量。在原文中给出 MoD + MoE 的集成方式

1. 将 MoD 套在 MoE 外侧。
2. 在 MoE 中加入空专家。这是一种对原模型影响极小的修改方式，但能达到的效果似乎十分有趣。

原文中讨论了很多对计算量的考量。后续有机会再看。

视觉 MoE

似乎并没有太多的工作关注 MoE 在视觉上的应用，大家还是更喜欢 NLP 上的 Transformer。Google 曾经搞过一点 Vision MoE^[7]，将 ViT 的部分 FFN 换成了上文提到的 MoE 结构。

有机会时再仔细看一看。

---

需要注意的是，MoE 实际会遇到十分多的技术性问题^[2]，本文没有讨论。

1. https://twitter.com/pommedeterre33/status/1671263789914677248 ↩
2. ST-MoE: Designing Stable and Transferable Sparse Expert Models ↩
3. Expert Gate: Lifelong Learning with a Network of Experts ↩
4. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity ↩
5. https://zhuanlan.zhihu.com/p/672002822 ↩
6. Mixture-of-Depths, Dynamically allocating compute in transformer-based language models ↩
7. Scaling Vision with Sparse Mixture of Experts ↩
8. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity ↩
9. ST-MoE: Designing Stable and Transferable Sparse Expert Models ↩
10. https://huggingface.co/blog/moe ↩