MoE

在上一篇文章中，我们介绍了MoE的一个几何诠释，旨在通过Dense模型的最佳逼近出发来推导和理解MoE。同时在文末我们也说了，给出MoE的计算公式仅仅是开始，训练一个实际有效的MoE模型还有很多细节补，比如本文要讨论的负载均衡（Load Balance）问题。 负载均衡，即"不患寡而患不均"，说白了就是让每个Expert都在干活，并且都在干尽可能一样多的活，避免某些Expert浪费算力。负载均衡既是充分利用训练算力的需求，也是尽可能发挥MoE大参数量潜力的需求。 问题分析 我们知道，MoE的基本形式是 $$ \boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i $$ 对于传统MoE，$\boldsymbol{\rho}$是一个概率分布（Router），$\boldsymbol{e}_i=\boldsymbol{v}_i$，$\boldsymbol{v}_i$是一个小型FFN（Expert）的输出；而对于我们上一篇推导的几何MoE，$\boldsymbol{\rho}$没有归一化的要求，它预测的是Expert的模长，而$\boldsymbol{e}_i=\boldsymbol{v}_i/\Vert\boldsymbol{v}_i\Vert$预测的是Expert的方向。 不管哪种格式的MoE，实际表现都差不多，只是理解视角的不同。但要注意，虽然MoE的公式给人的感觉是"每遇到一个Token，就去找相应的Expert来计算"，但实际训练时其实是反过来的：先给每个Expert分配好相应的算力，然后将Token分配（Route）到所属的Expert中并行计算，这也就为什么负责打分的$\boldsymbol{\rho}$被称为Router。 这样一来，如果Expert的分配不均衡，就可能出现如下局面：某些Expert（Dead Expert）几乎一直闲置，浪费算力；某些Expert要处理的Token太多，根本忙不过来，只能Token Drop（即放弃处理部分Token）。从理论上来说，出现Dead Expert意味着MoE没有达到预期的参数量，即花了大参数量的显存，结果只训出来小参数量的效果。 所以，不管是从训练还是性能角度看，我们都希望保证Expert的负载均衡。 辅助损失（Auxiliary Loss） 促进负载均衡的常规思路是添加与之相关的损失函数，我们通常称之为"Aux Loss（Auxiliary Loss）"，目前主流用的Aux Loss最早可以追溯到2020年的《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》。 介绍Aux Loss之前，我们需要先引入一些新概念。首先，我们已经提到对于一般的MoE来说，$\boldsymbol{\rho}$未必是概率分布，我们将归一化的$\boldsymbol{\rho}$记为$\boldsymbol{p}=[p_1,p_2,\cdots,p_n]$，以及它Top-$k$版为$\boldsymbol{f}=[f_1,f_2,\cdots,f_n]$，其中 $$ p_i = \frac{\rho_i}{\sum_{i=1}^n \rho_i},\qquad f_i = \begin{cases}1/k, & i\in \mathop{\text{argtop}}_k \boldsymbol{\rho} \ 0, & i\not\in \mathop{\text{argtop}}_k \boldsymbol{\rho}\end{cases} $$ 接着我们定义$\boldsymbol{P}=\mathbb{E}[\boldsymbol{p}],\boldsymbol{F}=\mathbb{E}[\boldsymbol{f}]$，这里的$\mathbb{E}$是指对所有样本的所有Token做平均。不难看出，$\boldsymbol{F}$就是Expert当前的负载分布，而$\boldsymbol{P}$则相当于$\boldsymbol{F}$的一个光滑近似。 有了这些记号，我们就可以写出Aux Loss为： $$ \mathcal{L}_{\text{aux}} = \boldsymbol{F}\cdot \boldsymbol{P} = \sum_{i=1}^n F_i P_i \tag{1} $$ ...

MoE（Mixture of Experts）架构的流行自不必多说，近来火出圈的DeepSeek-V3便是MoE架构，传言GPT-5也是MoE架构，国内最近出的一些模型（Qwen3系列相关）也有不少用上了MoE。然而，虽然MoE的研究由来已久，但其应用长时间内都不愠不火，大致上是从去年初的《Mixtral of Experts》开始，MoE才逐渐吸引大家的注意力，其显著优点是参数量大，但训练和推理成本都显著低。 但同时MoE也有一些难题，如训练不稳定、负载不均衡、效果不够好等，这也是它早年没有流行起来的主要原因。不过随着这两年关注度的提升，这些问题在很大程度上已经得到解决，我们在接下来的介绍中会逐一谈到这些内容。 问题定义 我们知道，Transformer模型由Attention层和MLP层组成，MoE替换的是模型中MLP层。MLP层又分FFN（FeedForward Network）和GLU（Gated Linear Unit）两种，主流的是GLU，但简单起见我们还是以FFN为例：$$y=f(xW^{(A)})W^{(B)}$$其中$x\in\mathbb{R}^d$ 是输入向量（行向量），$W^{(A)}\in\mathbb{R}^{d\times{D}}$, $W^{(B)}\in\mathbb{R}^{D\times{d}}$ 是两个参数矩阵，$f$是Element-wise的激活函数，设$n$是一个能整除$D$的整数，那么上面的FFN可以用分块矩阵等价： $$ \begin{equation}\boldsymbol{y} = f\big(\boldsymbol{x}\begin{bmatrix}\boldsymbol{W}^{(A)}_1 & \boldsymbol{W}^{(A)}_2 & \cdots & \boldsymbol{W}^{(A)}_n\end{bmatrix}\big)\begin{bmatrix}\boldsymbol{W}^{(B)}_1 \\ \boldsymbol{W}^{(B)}_2 \\ \vdots \\ \boldsymbol{W}^{(B)}_n\end{bmatrix} = \sum_{i=1}^n \underbrace{f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i}_{\boldsymbol{v}_i}\end{equation} $$ 其中 $W^{(A)}_i = W^{(A)}_{[:,(i-1)c:ic]}$, $W^{(B)}_i = W^{(B)}_{[(i-1)c:ic,:]}$, $c= D/n$，这里的切片按照Python规则来。由此可见，FFN可以等价表示成n个向量 $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n$ 之和，每个向量代表了一个小模型$f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i$的输出，每个小模型计算量相同，这些小模型就是MoE中的“Expert”。 MoE提出的问题是： 能否只挑k个向量的和来逼近n个向量的和呢？这样就可以将计算量降低到k/n了。 模长排序 要解决上述的问题，实质上是要解决低秩近似的问题，数学公式就是: $$\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \lambda_i \boldsymbol{v}_i - \sum_{i=1}^n\boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \lambda_i = k\end{equation}$$ 记$\gamma_i = 1 - \lambda_i$，那么它又可以写成： $$\begin{equation}\mathop{\text{argmin}}_{\gamma_1,\gamma_2,\cdots,\gamma_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \gamma_i = n - k\end{equation}$$ 这个问题的精确求解是比较困难的（NP Hard），但有一个简单的近似解：当$v_i$两两正交时，我们有 $$\begin{equation}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2 = \sum_{i=1}^n \gamma_i^2 \Vert\boldsymbol{v}_i\Vert^2 = \sum_{i=1}^n \gamma_i \Vert\boldsymbol{v}_i\Vert^2\end{equation}$$ 上式最优解显然就是让模长$\Vert\boldsymbol{v}_i\Vert$最小的$n-k$个$\gamma_i$等于1，这又等价于说挑出模长最大的$k$个向量来逼近$n$个向量之和。当$v_i$不满足两两正交的条件时，我们依然用它来作为一个近似解。它的几何意义也很直观，模长越大的向量，在求和过程中越不容易被抵消，从而作用越突出。 ...