MoE环游记：从几何意义出发-CSDN博客

本文链接：https://blog.csdn.net/c9yv2cf9i06k2a9e/article/details/145540696

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

前两年福至心灵之下，开了一个“Transformer 升级之路”系列，陆续分享了主流 Transformer 架构的一些改进工作和个人思考，得到了部份读者的认可。这篇文章开始，我们沿着同样的风格，介绍当前另一个主流架构 MoE（Mixture of Experts）。

MoE 的流行自不必多说，近来火出圈的 DeepSeek-V3 [1] 便是 MoE 架构，传言 GPT-4 也是 MoE 架构，国内最近出的一些模型也有不少用上了 MoE。

然而，虽然 MoE 的研究由来已久，但其应用长时间内都不愠不火，大致上是从去年初的《Mixtral of Experts》[2] 开始，MoE 才逐渐吸引大家的注意力，其显著优点是参数量大，但训练和推理成本都显著低。

但同时 MoE 也有一些难题，如训练不稳定、负载不均衡、效果不够好等，这也是它早年没有流行起来的主要原因。不过随着这两年关注度的提升，这些问题在很大程度上已经得到解决，我们在接下来的介绍中会逐一谈到这些内容。

问题定义

首先要指出的是，这里会用笔者自己的一种理解思路来介绍 MoE，在必要的地方会附上相应的参考文献，但不会对 MoE 架构进行系统的追根溯源，还请读者见谅。

我们知道，Transformer 模型由 Attention 层和 MLP 层组成，MoE 替换的是模型中 MLP 层。MLP 层又分 FFN（FeedForward Network）和 GLU（Gated Linear Unit）两种，主流的是 GLU，但简单起见我们还是以 FFN 为例：

其中是输入向量（行向量），是两个参数矩阵，f 是 Element-wise 的激活函数。设 n 是一个能整除 D 的整数，那么上述可以等价地用分块矩阵写成：

其中，这里的切片按照 Python 规则来。由此可见，FFN 可以等价表示成 n 个向量之和，每个向量代表了一个小模型的输出，每个小模型计算量相同，这些小模型就是 MoE 中的“Expert”。

MoE 提出的问题是：

能否只挑 k 个向量的和来逼近 n 个向量的和呢？这样就可以将计算量降低到 k/n 了。

模长排序

这个问题其实我们在《低秩近似之路（三）：CR》已经探究过，写成数学公式是：

记，那么它又可以写成：

这个问题的精确求解是比较困难的，但有一个简单的近似解：当两两正交时，我们有：

上式最优解显然就是让模长最小的 n-k 个等于 1，这又等价于说挑出模长最大的 k 个向量来逼近 n 个向量之和。当不满足两两正交的条件时，我们依然用它来作为一个近似解。它的几何意义也很直观，模长越大的向量，在求和过程中越不容易被抵消，从而作用越突出。

此外，在《低秩近似之路（三）：CR》中我们还讨论了一种依概率采样的逼近过程，在方差最小的假设下得到的最优采样概率同样有正比于模长的特点，所以总的来说按向量模长排序是一个简单但不失有效的策略。

MoE初现

现在策略已经有了——“挑模长最大的 k 个向量”——可是细想之下我们会发现它并不实用：要挑模长最大的 k 个向量，就得把所有向量的模长都算出来，这又意味着要把所有的先算出来，可我们的原本目的却是减少的计算量！

为了解决这个矛盾，我们需要重新设计每个 Expert 模型，使得它的模长可以低成本地计算出来。什么意思呢？首先我们将归一化得到，这样每个的模长都相同了。接着我们定义：

其中是参数矩阵，是一个的激活函数，说白了这就是一个 d 维到 n 维的线性变换加激活函数，所以计算量是比较小的，这部分模型在 MoE 中被称为“Router”。

的作用是什么呢？预测每个 Expert 的模长！换言之，我们将作为第 i 个 Expert 的模长，才是完整的 Expert，它被分解为两部分：计算量比较小的模长以及计算量比较大的方向。

为了减少计算量，我们先计算出，挑出最大的k个后才去计算相应的，最后乘上并求和：

这便是 MoE 模型的基本公式。由于计算中只保留了 Top-k 部分，所以它本质上属于一种 Sparse 模型，而原本的 FFN 或者 k=n 时的模型，通常称为对应的 Dense 模型。

思路概括

不管是熟悉 MoE 还是不熟悉 MoE 的读者，可能都会对上述过程有点陌生，因为这是笔者自己闭门造车的一种 MoE 理解路线，但因为其几何意义更明确，所以本质上应该是更好理解的。

我们再来整理一下整个思路：

1. 一个常规的 Dense 模型 FFN，可以等价改写为 n 个 Expert 向量之和；

2. 为了节省计算量，我们试图挑出 k 个向量求和来逼近原本的 n 个向量之和；

3. 转化为数学问题求解后，我们发现挑选规则是模长最大的 k 个向量；

4. 直接去算 n 个 Expert 的模长然后选 k 个实际上是不省计算量的，所以要重新设计 Expert；

5. 将归一化得到，然后用另外的小模型（Router）预测模长，最终的 Expert 为；

6. 此时，我们就可以先算全体，挑出 k 个后才去计算，达到节省计算量的目的。

为何如此

可能有些读者疑问，为什么要做这个看似复杂的过程？原本的 MoE 不是挺好理解的吗？一般的 MoE 形式为：

也就是求和前少了对的归一化，此时也没有模长的意义，它纯粹是一个用来对 Expert 排序的打分模型（即 Router）。可为什么将乘到 Expert 上去就能让 Router 学会正确排序 Expert 呢？笔者发现只有《Sparse Backpropagation for MoE Training》[3] 对此给出了一个解释，但还是稍欠直观。

而在本文的几何视角下，我们会发现很多问题就“豁然开朗”了。我们将 Expert 重新参数化为后，Dense 模型对应于全体求和，而 MoE 对应于选 Top-k 后求和，这是 Dense 模型的一个有理论保证的逼近。我们没有去考虑 Router 如何选择 Expert，只是每一步都尽可能逼近 Dense 模型，这可以说是既要大参数、又要小计算量的最佳选择。

现在的几何意义是模长而不是概率，所以激活函数就没有归一化的要求了，除了 Softmax 外，像 Sigmoid、ReLU 都可以考虑使用，也可以考虑我们在《Softmax后传：寻找Top-K的光滑近似》介绍的 Top-k 光滑近似。Router 使用非归一化的激活函数，有助于避免 k > 1 时 Expert 之间的恶性竞争，有时候能取得更好的效果。

最后补充一点，我们前面定义，目的是让所有模长相同，实际操作中不是一定要 L2 Normalize，也可以是其他等价操作，比如 gamma 参数恒等于 1 的 RMS Norm，它更符合我们的输出习惯。