MoE 基础 | GJJ’s Blog

type

Post

status

Published

date

Apr 21, 2026

slug

blog-109

summary

一些 MoE 相关的学习笔记

从几何角度看MoE的原理

我们以 FFN 为例子，来介绍大模型的的 MLP 层：

其中是输入向量，, 是参数矩阵，假设是可以整除的整数，上述公式可以等价地用分块矩阵写成：

由此可见，FFN 可以等价地表示成个向量之和，每个向量表示了一个小模型的输出，每个小模型计算量相同，每个小模型就是一个 Expert。

“能否只挑个向量的和来逼近个向量的和？这样就可以将计算量降低到了”

优化问题：

直接求解较为困难，我们可以考虑取两两正交时的解作为正式解： “挑模长最大的个向量”。

为了避免模长的计算，我们需要设计一个小的模型“Router”来计算模长，获得前大小模长的向量后，再计算对应的方向：

这就是 MoE 模型的基本公式。

MoE的负载均衡

我们希望每个 Expert 都在干活，并且尽量干一样的活，避免某些 Expert 浪费算力，充分发挥训练算力的寻求。

促进负载均衡的常规思路是添加与之相关的损失函数，我们通常称之为“Aux Loss（Auxiliary Loss）”。

// 待拓展

LOSS-Free 的引入

Aux Loss固然简单直观，但它也有一个明显的缺点——权重不好调——调低了无法促进均衡，调高了容易损害LM Loss，所以业界一直有寻找替代方案的尝试。

DeepSeek 引入了一种非常优雅的做法：

其中，是输入无关的向量，训练后就保持不变，它仅仅参与分配过程而非 MoE 的前向计算。

我们先定义 :

在偏置下 Expert 当前的分布。其中我们的目标均匀分布定义为：，负载均衡就相当于最小化

这个目标是不可导的，我们使用 STE 解决这个问题：找一个可导且和具有相同增减趋势的量作为的光滑近似：

它的梯度是：

因此使用 SGD 更新就是：

采用SGD 来更新的符号梯度下降版本就是：

将固定专家数扩展为可变专家数

在先前介绍的MoE架构中，我们习惯了为每个 Token 分配固定数量K的专家（Top-k）。但稍微动脑筋想想就会发现：Token 与 Token 之间的难度是不平等的。 一个简单的字可能只需要 1 个专家顺手处理，而一个复杂的数学逻辑可能需要 4 个专家协作。如果我们给所有 Token 都分 k=2，那对简单的 Token 是浪费，对难的 Token 是抠门。

这部分接着上一部分进行扩展，利用 Loss-Free 方案中的偏置项（Bias），把 MoE 变成动态专家系统。

从 Top-K 到 Argwhere

让我们首先回顾 MoE 的基本形式：

是专家打分，几何意义上来说是预测模长，是专家的输出；

为了解决负载均衡，引入偏置项 :

上一节内容讲到，是作为相对大小调节的相对大小，整体有一个均值在这个过程中是不起作用的，这个冗余度将在这一节中被使用。

为了实现动态数量，我们干脆去掉 Top- 的硬性约束，使用门槛：

只要预测专家得分和补偿得分的和大于 0，专家就被激活。这样每个词选中的专家数就不再固定。

数学实现

我们的训练目标依旧是找到合适的，而的优化目标有两个：

负载均衡：保证流量平均分配给每个专家。

预算控制：考虑第一个优化目标，如果不加限制，那么，即选中所有专家，这也是一种“均衡”，计算量会变得特别大。所以要让平均激活数维持在。

定义指示函数

记其期望值（Batch 平均值）为 ;

实际分布：

目标分布：

负载均衡实现：让趋近于

为了不影响整体水平，而且本身这个冗余度在负载均衡中没有影响，所以上边的更新规则可以先去均值：

这样，我们就将的均值这个自由度让出来给后续的预算控制。

此时，实际的平均激活数是，而目标是。如果实际激活数大于目标，则说明太大，此时需要我们调小：

将上述过程整理，得到最终的更新公式：

如果我们希望总预算大，就可以让整体的变大，这样的概率就越大，预算就越多。

显然地，如果我们希望预算只是不超过，而并非希望其等于，则可以让最小为 0，这样当预算不超过时不因为预算问题对进行更新：

简化

苏神在实验中发现，其实还可以更简单。如果我们让每个专家的“绝对平均负载”都去逼近，那么均衡和预算就自动合二为一了：

定义目标，更新规则可以简化为：

根据苏神的实验，这些优化过程效果上大同小异，但是最后简化后的更新过程在负载均衡和预算控制两个指标在训练前期抖动大得多。

同时，苏神还在这里提到了一个技巧：使用 RMS Norm 来替换 sign 函数，增加稳定性。在这个任务中使用 RMS Norm 替换后更新的数量级相同，同时 RMS Norm 保留了相对大小：

初始化

最后问题来到了一个“有意思但是不算十分关键的问题”：的初始化。

如果我们偷懒直接把 b 全设为 0，且 Router 用的激活函数是 Sigmoid：

问题来了： Sigmoid(0)=0.5。这意味着在训练刚开始时，每个专家都有 50% 的概率被选中。如果你有 32 个专家，平均每个词会激活 16 个！

后果： 预算 k 可能只有 4，这会导致初始阶段出现海量的 Token Drop（计算量溢出被迫丢弃），模型收敛会非常慢，甚至直接崩掉。

数学假设：为了给 b 一个完美的初始值，我们需要做点简单的统计假设：

假设输入向量 x 经过了 RMSNorm，满足均值为 0、方差为 1。

假设 Router 的权重 WR 初始化方差为 σ2。

那么 Router 的原始输出（Logits）近似服从正态分布 N(0,σ2d)，其中 d 是隐藏层维度。

二分法寻找“起跑线”

我们的目标是：找到一个初始的 binit，使得在训练第一步时，激活的专家数正好接近我们的预算 k:

总结

从上一章对 DeepSeek 的 “Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts”这篇工作的解读，再到这一部分对这一工作的探索，确实让人感觉到一些无比“优雅”且充满美感的工作。

均匀分布的反思

其实看到这里，如果我们稍微思考一下，就会发现，其实现实时间中的大部分情况都是“非均匀的”，而我们之前一直在强调“负载均衡”。可是均衡的负载在除了效率方面，一定是最好的吗？

Shared Expert

MoE 的基本形式如下：

Loss-Free 将替换成，在动态专家系统中，我们进一步将其推广为。

如果负载完全均衡，是否意味每个专家都成了“全才”或“庸才”？

Shared Expert:

Router Expert: 班里的语数英，政史地老师，各司其职，只有在上课时才出现；

Shared Expert: 班主任，不管上什么课都在；

数学形式如下：

压缩共性：等号后的第一项是必然选中的部分，用 DeepSeek 的话来说就是把所有 Token 都要用到的公共知识塞进共享专家；

减少冗余：让剩下的路由专家专心地去学习哪些偏门的专业知识；

几何意义：路由专家只需要学习“残差”。从几何上来看，让专家向量之间更容易满足正交假设，提高参数效率。

比例因子

实际上，前边的式子是一个理想的过程：每个Router Expert 前边有一个预测模长。如果直接相加，会导致“力量失衡”。为了让两个部分在初始化部分的模长一样，引入一个比例因子：

Fine-Grained Expert

除了 Shared Expert 以外，DeepSeekMoE提出的另一个改进点是Fine-Grained Expert：在总参数量和激活参数量都不变的情况下，Expert 的粒度越细，效果往往越好。

例如，选的系统就比选的效果要好。原论文的说法是丰富了 Expert 的组合多样性：

虽然总参数量没变，但是组合的可能性爆炸式增长。

最优分配

目前我们已经整理的负载均衡的具体实现一共有两种方案：Aux Loss 和 Loss-Free 方案。下边讨论第三种思路：最优分配。它将负载均衡视作等式约束下的线性规划问题。

问题描述：假设有个 Token，第个 Token 的 Rouer 对个 Expert 的打分记做：

那么共有个分数。目标是基于这个分数指定一个分配方案来决定每个 Token 应该激活哪些 Expert。

这在数学上是一个典型的带约束的线性规划问题：

上式属于整数优化问题，难以求解，考虑其松弛版本：

其中是分配系数，是原始得分。这样就是一个有界区域内的一个线性规划问题。

进一步，我们考虑约束优化问题的 max-min 形式，即拉格朗日乘子法：

该形式和原问题松弛版本的形式等价：如果后两项括号内的形式不满足等于 0，则在后的步骤可以取负无穷。

上边的目标关于都是线性的，同时的可行域为凸集，满足 Minimax theorem 条件，交换和的顺序后，整理一下，我们得到：

注意到，当前优化目标的第一项我们是可以进行分析的：

当时， = 1;

当时， = 0;

当时，的值无所谓;

也就是说，最后的优化目标和一开始非松弛版本的优化目标是一致的，二者完全等价。

将代入优化目标，得到，优化目标可以简化为：

我们使用交替最小化思路求解：先固定求，再固定求，交替进行。由于和有明显的对偶性，因此这两个步骤是在求解同一个问题。

我们先来看在固定来求，问题等价于：

每一项是单独累加起来的，所以将其划分为个独立的子优化问题：

将按照从大到小排列成，第大的元素为，假设我们已知，那么目标函数变成：

这一部分个人理解：
在确定的情况下，的目的是寻找最优偏置项，让模型可以刚好选出个专家。
在这个优化目标下，第一项希望尽可能大，第二项尽快能小，我们就希望在这个对的线性规划问题中，直线刚刚好卡在边界处。当的斜率是正的时候，我们要尽可能的减小，如果斜率是负的，我们要尽可能的增加，如果等于 0 时，函数变成常数。