前言

动机：近期笔者完成了一项较为理论的工作，并计划将其延申出一项实际应用——Mixture of Experts (MoE)。因此近期决定较为系统的学习 MoE。幸运的是，苏神正好有一个关于 MoE 的系列，因此笔者在这里整理一份个人笔记。

继续进行 MoE 系列的学习。符号定义见前文。

均匀分布的反思

抛开效率上的需求不谈，均匀分布就一定是效果最好的方向吗？

Shared Expert

首先回顾 MoE 的基本形式：

因为在前面笔记中提及的改进 trick 和 Shared Expert 思想是正交的，这里使用最基本形式以方便读者进行理解。

Shared Expert将上式改为：

如何理解 Eq. ($\ref{eq:share-1}$)？我们强制选择 $s$ 个 Experts，然后再在 $n-s$ 个 Expert pool 中选择 $Top_{k-s}$ 个 Experts。这里的 $s$ 通常不会太大，取 1 或者 2。

Shared Expert原则上不增加模型参数量和推理成本。但即便如此，DeepSeekMoE和苏神的一些实验显示，Shared Expert依然能一定程度上提升模型效果。

通常将 Eq. ($\ref{eq:share-1}$) 写成：

这里引入 $\lambda$ 的是为了平衡 Shared Expert 和 Routed Expert。

苏神在《Muon is Scalable for LLM Training》提出，适当的 $λ$ 应使得两者在初始化阶段模长接近一致。具体来说，我们假设每个Expert在初始化阶段具有相同的模长（不失一般性，可以直接设为1），并且满足两两正交，然后假设Router的logits服从标准正态分布（即零均值、单位方差，当然如果觉得有必要，也可以考虑其他方差）。这样一来，$s$ 个 Shared Expert 的总模长就是 $\sqrt{s}$，而 Routed Expert 的总模长是 $\lambda\sqrt{\sum_{i\in \mathop{\text{argtop}}{k-s} \boldsymbol{\rho}{[s:]}} \rho_{i+s}^2}$，通过让它等于 $\sqrt{s}$ ，就可以估计出 $\lambda$：

苏神还给出了数值模拟的脚本，并给出结论：数值模拟的结果和大厂调参的结果较为一致。

import numpy as np

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def softmax(x):
    return (p := np.exp(x)) / p.sum()

def scaling_factor(n, k, s, act='softmax', renorm=False):
    factors = []
    for _ in range(10000):
        logits = np.random.randn(n - s)
        p = np.sort(eval(act)(logits))[::-1][:k - s]
        if renorm:
            p /= p.sum()
        factors.append(s**0.5 / (p**2).sum()**0.5)
    return np.mean(factors)

"""
DeepSeek-V2有n=162,k=8,s=2，Softmax激活并且没有重归一化，上述脚本的模拟结果约等于16，而DeepSeek-V2的λ正好是16

DeepSeek-V3则有n=257,k=9,s=1，Sigmoid激活且重归一化，脚本的结果大约是2.83，而DeepSeek-V3的λ则是2.5
"""

scaling_factor(162, 8, 2, 'softmax', False)
scaling_factor(257, 9, 1, 'sigmoid', True)

非均匀性

Shared Expert 思想给了我们一个回答笔记开头问题的答案：抛开效率上的需求不谈，均匀分布就一定是效果最好的方向吗？——不一定！ 因为 Shared Expert 也可以理解为某些 Expert 一定会被激活，于是整体来看，这将导致一个非均匀的 Expert 分布：

除了 Shared Expert 外，DeepSeekMoE 所提的另一个改进点是 Fine-Grained Expert，它指出在总参数量和激活参数量都不变的情况下，Expert 的颗粒度越细，效果往往越好。比如，原本是 n 选 k 的 Routed Expert，现在我们将每个 Expert 缩小一半，然后改成 2n 选 2k，那么总参数量和激活的参数量都还是一样的，但后者表现往往更好。

但要注意，Fine-Grained Expert 并非是无成本的，n 越大，Expert 之间的负载往往越不均衡，并且 Expert 之间的通信和协调成本也会增加，所以 n 也不能无限增加，有一个效果和效率都友好的舒适区间。

最优分配促进负载均衡

前文已经给出了两种实现负载均衡的主流思路：分别是经典的 Aux Loss 以及由 DeepSeek 提出的 Loss-Free 方案。接下来探讨第三种思路：最优分配。

首先定义：

其中，$x_{i,j}=1$ 表示第 $i$ 个 token 选中了第 $j$ 个专家，所以 $\sum_j x_{i,j} = k$ 表示第 $i$ 个 token 只分配给 $k$ 个专家；而 $\sum_i x_{i,j} = \frac{mk}{n}$ 则表示，所有 $m$ 个 token 均匀激活 Expert，即：总次数除以总共 $n$ 个专家。

由于整数规划问题（$x_{i,j}\in {0, 1}$）比较困难，考虑 Eq. ($\ref{eq:target}$) 的松弛版本，即 $x_{i,j}\in [0,1]$：

得到其拉格朗日函数：

利用拉格朗日乘子法得到：

这里对 Eq. ($\ref{eq:relax-max-min}$) 分情况讨论：

• $x_{i,j}$ 不在可行域之内，此时不满足约束条件；因此 $\min$ 操作调节 $\alpha_i, \beta_j$ 的时候会得到负无穷；
• $x_{i,j}$ 在可行域之内，此时满足约束条件；因此 $\min$ 操作调节 $\alpha_i, \beta_j$ 是作无用功；

最终在 $\max_{x_{i,j}\in[0, 1]}$ 的时候肯定会取 $x_{i,j}$ 在可行域内，满足约束条件的情况，也就是等价于原问题 Eq. ($\ref{eq:relax}$)。拉格朗日乘子法运用时，本质是将利用 max-min 嵌套的形式将约束合并进入拉格朗日函数。

鉴于这里，目标函数和可行域都是凸的，因此可以利用 Sion Minimax Theorem 交换 $\min$ 和 $\max$ 操作顺序：

这里对 Sion Minimax Theorem 的运用可参考笔者的另一篇博客。这里直接运用该定理更加方便简洁；当然，遵循通常的最优化课程，通过 “凸优化问题+Slater条件 $\Rightarrow$ 强对偶” 的性质进而进行 min-max 顺序的交换也是可行的。

在处理内层 $\max$ 的优化问题时，将优化变量之外的变量视为常量，即上述式子中的 $\alpha_i, \beta_j$；此时 $\max$ 的求解是显然的：

当 $s_{i,j} - \alpha_i - \beta_j = 0$ 时，理论上 $x^\star_{i,j}$ 任取 $[0,1]$ 均可；但是当我们取 $0$ 或者 $1$ 的时候，就会发现：$x^\star_{i,j}\in {0,1}$！这意味着，虽然我们是针对松弛过后的优化问题进行求解，但是最优解也是松弛前原问题的解，二者完全等价。

将内层优化问题的最优解 $x^\star_{i,j}\in {0,1}$ 代入 Eq. ($\ref{eq:relax-min-max}$) 得到：

下面针对该优化问题，苏神采用了交替优化的方法来进行求解。这里我对 “交替优化” 的理解非常浅薄；因此尚无法很好总结归纳背后 (是否存在抑或仅是工程上的求解) 的来龙去脉，如：交替优化求解的等价性等。这里苏神将交替优化的方法写得非常清晰，感兴趣的读者请移步苏神博客进行更进一步了解。

总之，经过交替优化的方法， 我们可以得到：对任意给定 $i$，$α^∗i$ 是将全体 $s{i,j}−β_j$ 从大到小排列后的第 $k+1$ 个元素。类似地，固定 $α_i$ 求 $β_i$ 的结果是：对任意给定 $j$，$β^∗j$ 是将全体 $s{i, j}-β_j$ 从大到小排列后的第 $mk/n+1$ 个元素。

苏神之后巧妙使用“分位数（Quantile）”的概念，将 “$n$ 个数中的第 $k+1$ 大元素、$m$ 个数中的第 $mk/n+1$ 大元素” 统一起来，它们实际上都是各自维度的 “$1−k/n$ 分位数”。Numpy、Torch 等数值框架都有 “quantile” 函数实现，利用它能避免对数据进行全排序操作，节省一些复杂度。

Reference

[1] MoE 系列笔记 v1

[2] MoE环游记：5、均匀分布的反思

[3] DeepSeekMoE

[5] MoE环游记：6、最优分配促均衡