参数量分析#

首先需要特别区分两个概念：

• Total Parameters：模型中所有参数的总量；
• Active Parameters：处理一个 token 时实际参与计算的参数量。

MoE 的优势正是来自二者之间的差异：

需要注意的是，MoE 并不意味着整个模型的计算量简单地变成原来的 $\frac{K}{N}$。Attention、Router 以及 Expert 之间的通信等部分仍然需要额外计算。因此，更准确地说，MoE 主要降低的是其中稀疏 Expert 部分的计算量。要理解这一改动对计算量的影响，首先需要明确每个 Expert 的规模是如何确定的。

设原 Dense FFN 的输入与输出维度均为 $d_{\text{model}}$，隐藏维度为 $d_{\text{ff}}$，即：

原始 Transformer 通常取 $d_{\text{ff}} = 4 d_{\text{model}}$。其计算量，若以 MACs 粗略衡量，约为：

将其替换为 $N$ 个 Expert 时，有两种常见方案。

路由#

为了减少每次前向传播的计算量，MoE 将原本的一个 FFN 替换为 $N$ 个独立的 Expert，每次只激活其中 $K$ 个。这一思路可追溯到稀疏门控 MoE。

例如，一个 MoE Layer 包含 8 个 Expert，每个 token 只激活其中 2 个。如果采用上一节所说的等参拆分，即 8 个 Expert 合起来才对应原来一个 Dense FFN，那么在忽略 Router、通信等额外开销，并假设每个 Expert 计算量相同的情况下，Expert 部分的计算量约为原 Dense FFN 的：

但要注意，这个结论依赖于等参拆分。如果采用扩容拆分，则激活 2 个 Expert 的计算量是原 Dense FFN 的 2 倍；此时 $\frac{2}{8}$ 只表示激活了总 Expert 容量的 $\frac{2}{8}$。因此，MoE 的稀疏激活来自：对一个 token 只计算被选中的 $K$ 个 Expert，而不是计算全部 $N$ 个 Expert。单个 Expert 是否比原 Dense FFN 小，取决于拆分方案：

• 如果采用等参拆分，每个 Expert 确实会变小，但 $N$ 个 Expert 合起来与原 Dense FFN 相当；
• 如果采用扩容拆分，每个 Expert 可以保持原规模，此时激活计算量可能大于原 Dense FFN，但相对于所有 Expert 的总容量仍然只激活了 $\frac{K}{N}$。

但是，对于一个输入 token，我们需要知道应该选择哪些 Expert，这就需要 Router。关于 Router 与稀疏门控的经典设计，可参考稀疏门控 MoE。

最简单的一种 Router 可以使用一个可学习的线性层，将输入 token 的隐藏状态 $\mathbf{x}$ 映射到 $N$ 个 Expert 对应的 routing score：

其中 $\mathbf{x}$ 是输入 token 的隐藏状态，$W_r\in\mathbb{R}^{d_{\text{model}} \times N}$ 是 Router 的可学习参数，$\mathbf{s}\in\mathbb{R}^{N}$ 是每个 Expert 对应的 routing score。接下来，Router 先对所有 $N$ 个 Expert 的 score 做 Softmax，得到每个 Expert 的 routing probability，这也是 Switch Transformer 中的定义：

记 $\mathbf{p} = (p_1, p_2, \dots, p_N)$ 为所有 Expert 的 routing probability 组成的向量，然后选出概率最高的 $K$ 个 Expert，记为集合 $\mathcal{I} = \operatorname{TopK}(\mathbf{p}, K)$，并只对被选中的 Expert 将概率归一化，得到 routing weight：

这样保证 $\sum_{i\in\mathcal{I}} w_i = 1$。最终，只有这些被选中的 Expert 会参与当前 token 的计算，并根据对应的 routing weight 对 Expert 的输出进行加权求和：

其中 $E_i$ 表示第 $i$ 个 Expert。

以上过程可以独立实现为一个 Router 模块：

1
class Router(nn.Module):
2
    def __init__(self, d_model: int, num_experts: int, top_k: int) -> None:
3
        super().__init__()
4
        self.top_k = top_k
5
        self.weight = nn.Linear(d_model, num_experts, bias=False)
6

7
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
8
        logits = self.weight(x)                                              # [T, N]
9
        probs = torch.softmax(logits, dim=-1)                                # [T, N]
10
        topk_probs, topk_idx = probs.topk(self.top_k, dim=-1)
11
        routing_weights = topk_probs / topk_probs.sum(dim=-1, keepdim=True)  # [T, K]
12
        return probs, topk_idx, routing_weights

其中 probs 保留所有 Expert 的 routing probability，topk_idx 与 routing_weights 是 Top-K 选择与归一化的结果，三者都会在后续 MoE Layer 的 dispatch 与负载均衡损失中使用。

到这里，MoE 已经实现了稀疏计算，但又出现了一个新的问题：

Router 会不会把大量 token 都发送给少数几个 Expert？

如果某些 Expert 被频繁选择，而其他 Expert 几乎没有 token，那么虽然模型拥有大量参数，但实际上只有少数 Expert 得到了充分训练和使用。因此，我们还需要解决 Expert 之间的负载均衡问题。

负载均衡#

理想情况下，如果有 $N$ 个 Expert，那么一个 batch 的 tokens $T$ 应该比较均匀地分配到这些 Expert。例如 $N = 8$，$T = 1024$，如果完全均衡，那么每个 Expert 平均处理的 token 数量为：

但实际上，Router 会根据 token 的内容动态进行选择，因此不同 Expert 接收到的 token 数量可能存在较大差异。如果大量 token 都被发送给同一个 Expert，就会产生两个问题：

1. 部分 Expert 负载过高，成为整个 MoE Layer 的计算瓶颈；
2. 其他 Expert 接收到的 token 太少，导致这些参数无法得到充分训练。

因此，MoE 通常会引入额外的 Load Balancing Loss，鼓励 Router 将 token 更均衡地分配给不同 Expert。一种经典的方法来自 Switch Transformer。 对于一个 batch，可以定义 Expert $i$ 实际接收到的 token 比例为：

其中 $T$ 是 batch 中的 token 数量，$\mathbb{I}(\cdot)$ 为指示函数，当其中的条件成立时取值为 $1$，否则取值为 $0$。同时，可以计算 Router 对 Expert $i$ 的平均概率：

其中 $p_{t,i}$ 就是上一节定义的 routing probability，即 token $t$ 经 Softmax 后对 Expert $i$ 的概率。然后可以构造一个辅助损失：

其中 $\alpha$ 用于控制负载均衡损失对整体训练目标的影响。

为什么乘积 $f_iP_i$ 能鼓励均衡分配？关键在于两者的可微性不同：$f_i$ 来自 Top-K 的硬选择，对 Router 参数不可微分；$P_i$ 是 Softmax 的软概率，可微分。因此梯度只通过 $P_i$ 传播，而 $f_i$ 只作为系数携带负载信息。对 $P_i$ 求偏导：

也就是说，某个 Expert 接收的 token 越多，$f_i$ 就越大，Router 分配给它的概率就会被越强地压低。这是一个负反馈：过载的 Expert 在后续迭代中被分配的概率降低，接收的 token 随之减少，直到各 Expert 的负载趋于均衡。反过来，负载越低的 Expert 受到的抑制越弱，Router 自然会把更多概率分配给它们。

另外，系数中的 $N$ 是归一化因子：完全均衡时每个 Expert 被分配的比例为 $f_i = \frac{K}{N}$，代入可得 $\mathcal{L}_{\mathrm{aux}} = \alpha K$，损失值与 Expert 数量无关，因此 $\alpha$ 可以在不同规模的 MoE 中统一控制辅助损失的权重。

这个损失的目的并不是要求所有 Expert 学习完全相同的内容，而是避免 Router 长期将大量 token 集中发送到少数 Expert，从而改善 Expert 的利用率。

对应的辅助损失可以写为：

1
def load_balancing_loss(
2
    probs: torch.Tensor,    # [T, N]
3
    topk_idx: torch.Tensor, # [T, K]
4
    alpha: float,
5
) -> torch.Tensor:
6
    N = probs.size(1)
7
    dispatch_mask = torch.zeros_like(probs)
8
    dispatch_mask.scatter_(1, topk_idx, 1.0) # [T, N]
9
    f = dispatch_mask.mean(dim=0)            # [N]
10
    P = probs.mean(dim=0)                    # [N]
11
    return alpha * N * (f * P).sum()

其中 dispatch_mask 对应公式中的指示函数 $\mathbb{I}(\cdot)$，由离散的 topk_idx 计数得到，天然不可微分；梯度只经由 P 传回 Router 参数，与前文的分析一致。

不过，即使加入了负载均衡损失，实际路由过程中仍然可能出现某个 Expert 接收到过多 token 的情况。因此，MoE 通常还需要为每个 Expert 设置一个容量上限，即 Expert Capacity。Switch Transformer 和 GShard 都讨论了 Expert Capacity、token dispatch 与 overflow 等问题。

一个常见的计算方式为：

其中 $C$ 是每个 Expert 的容量，$T$ 是 batch 中的 token 数量，$N$ 是 Expert 数量，$CF$ 是 Capacity Factor。Capacity Factor 为 Expert 额外提供了一定的缓冲空间。例如，当 $CF = 1$ 时，每个 Expert 的容量大约等于平均分配情况下的 token 数量；当 $CF > 1$ 时，容量在平均分配量的基础上留出余量，Expert 可以接收超过平均数量的 token，从而减少因容量不足而被丢弃的 token。

对应的容量计算：

1
def expert_capacity(T: int, num_experts: int, capacity_factor: float) -> int:
2
    return math.ceil(T / num_experts * capacity_factor)

这种被丢弃的情况称为 token overflow：如果某个 Expert 接收到的 token 数量超过了它的容量，超出部分的 token 将无法被该 Expert 处理。不同 MoE 实现对于 overflow token 的处理方式有所不同，因此 Expert Capacity 也是实际 MoE 实现中的重要问题。

到这里，一个基础的 MoE Layer 就包含了几个核心组件：

但此时又会出现一个新的问题：

如果我们希望不同 Expert 学习不同的知识，那么如何让 Expert 真正形成专业化分工？

如果仅仅增加 Expert 数量并进行 Top-K Routing，不同 Expert 之间仍然可能学习到大量重复的知识。这也是后续 DeepSeekMoE 希望解决的问题。

代码实现#

在代码实现中，可以先实现一个最基础的 Top-K MoE Layer。

整体流程可以抽象为：

对于输入 $X\in\mathbb{R}^{T\times d_{\text{model}}}$，Router 首先计算 $S = XW_r$，得到 $S\in\mathbb{R}^{T\times N}$。然后对每个 token 的 score 做 Softmax 得到概率，进行 Top-K 选择并对选中项归一化，得到对应的 Expert index 和 routing weight。之后，将 token 分发到对应的 Expert 中进行 FFN 计算：

其中 $\mathcal{I}_t$ 为第 $t$ 个 token 被选中的 Expert 集合，$w_{t,i}$ 为对应的 routing weight。最后再将不同 Expert 的结果按照原 token 的顺序重新组合。

基于前文实现的 Router、expert_capacity 与 load_balancing_loss，可以组装出一个完整的 MoE Layer：

1
class MoELayer(nn.Module):
2
    def __init__(
3
        self,
4
        d_model: int,
5
        d_ff: int,
6
        num_experts: int,
7
        top_k: int,
8
        alpha: float = 0.01,
9
        capacity_factor: float = 1.0,
10
    ) -> None:
11
        super().__init__()
12
        self.num_experts = num_experts
13
        self.top_k = top_k
14
        self.alpha = alpha
15
        self.capacity_factor = capacity_factor
16

17
        self.router = Router(d_model, num_experts, top_k)
18

19
        # 等参拆分（扩容拆分时为 d_ff）
20
        expert_hidden = d_ff // num_experts
21

22
        self.experts = nn.ModuleList([
23
            nn.Sequential(
24
                nn.Linear(d_model, expert_hidden),
25
                nn.GELU(),
26
                nn.Linear(expert_hidden, d_model),
27
            )
28
            for _ in range(num_experts)
29
        ])
30

31
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
32
        # x: [T, d_model]
33
        probs, topk_idx, routing_weights = self.router(x)
34
        capacity = expert_capacity(x.size(0), self.num_experts, self.capacity_factor)
35

36
        y = torch.zeros_like(x)
37
        for i in range(self.top_k):
38
            expert_ids = topk_idx[:, i]                              # [T]
39
            weights = routing_weights[:, i]                          # [T]
40
            for e in range(self.num_experts):
41
                token_ids = (expert_ids == e).nonzero().squeeze(-1)[:capacity]
42
                if token_ids.numel() > 0:
43
                    expert_out = self.experts[e](x[token_ids])
44
                    y[token_ids] += weights[token_ids].unsqueeze(-1) * expert_out
45

46
        aux_loss = load_balancing_loss(probs, topk_idx, self.alpha)
47
        return y, aux_loss

每个 Expert 在每个 Top-K 位置最多接收 capacity 个 token，超出的 token 会被直接丢弃，即前文所说的 overflow。返回的 aux_loss 即负载均衡辅助损失，训练时可加入总的优化目标。

在实际的大规模 MoE 模型中，还涉及跨设备的 token dispatch 与通信，因此实际实现会比上述过程复杂得多。Switch Transformer 的实现就需要显式处理 token dispatch 和 overflow 等问题。

细粒度与共享专家#

传统 MoE 通常将 FFN 划分为 $N$ 个较大的 Expert，然后每个 token 激活其中 $K$ 个。DeepSeekMoE 首先进一步将原来的 Expert 划分成更多、更小的 Expert：假设原来的 MoE 有 $N$ 个 Expert，每个 token 激活 $K$ 个 Expert，DeepSeekMoE 将每个 Expert 进一步划分，使 Expert 数量增加到原来的 $m$ 倍，即 $N \rightarrow mN$；每个 token 激活的 Expert 数量也相应增加到 $mK$，即 $K \rightarrow mK$。

这样做并不会简单地增加计算量，因为单个 Expert 的规模也相应减小。在保持激活参数量大致相近的情况下，更多、更小的 Expert 可以提供更加灵活的组合方式。例如，原本一个 token 只能从 $\{E_1,E_2,\cdots,E_N\}$ 中选择 $K$ 个较大的 Expert；细粒度划分后，它可以从更多的小 Expert $\{E_1,E_2,\cdots,E_{mN}\}$ 中选择 $mK$ 个。这使得不同 token 可以组合出更加丰富的 Expert 子集，从而减少不同 Expert 之间的知识冗余，并促进 Expert specialization。

但是，如果所有 Expert 都依赖 Router 进行选择，那么某些基础的、所有 token 都需要的知识也可能被重复学习。因此，DeepSeekMoE 又引入了 Shared Expert。Shared Expert 不参与普通的 Top-K Routed Expert 竞争，而是对所有 token 都进行计算，用于学习不同 token 之间更加通用、共享的知识。

于是，整个结构可以理解为：

因此，DeepSeekMoE 可以看作是在基础 MoE 的基础上进一步回答了两个问题：

• 如何让 Expert 更加细粒度，从而形成更加灵活的专业化分工？
• 如何将通用知识从 Routed Experts 中分离出来，减少 Expert 之间的重复学习？

这也构成了 DeepSeekMoE 相较于传统 MoE 的核心设计思想。

代码实现#

DeepSeekMoE 的代码实现可以在基础 MoE 上继续扩展，核心变化是：

1. 将 Routed Expert 的隐藏维度设置得更小，使 Expert 数量更多、粒度更细；
2. 增加一个或多个 Shared Expert，并让所有 token 都经过 Shared Expert；
3. 将 Routed Expert 的输出与 Shared Expert 的输出相加。

对应的输出可以写为：

其中 $\mathcal{S}$ 为 Shared Expert 集合，$\mathcal{I}_t$ 为第 $t$ 个 token 被 Top-K 选中的 Routed Expert 集合。

在前述 MoELayer 的基础上扩展：

1
class DeepSeekMoELayer(nn.Module):
2
    def __init__(
3
        self,
4
        d_model: int,
5
        d_ff: int,
6
        num_routed_experts: int,
7
        num_shared_experts: int,
8
        top_k: int,
9
        alpha: float = 0.01,
10
        capacity_factor: float = 1.0,
11
    ) -> None:
12
        super().__init__()
13
        self.num_routed_experts = num_routed_experts
14
        self.top_k = top_k
15
        self.alpha = alpha
16
        self.capacity_factor = capacity_factor
17

18
        self.router = Router(d_model, num_routed_experts, top_k)
19

20
        expert_hidden = d_ff // num_routed_experts
21

22
        def make_expert() -> nn.Module:
23
            return nn.Sequential(
24
                nn.Linear(d_model, expert_hidden),
25
                nn.GELU(),
26
                nn.Linear(expert_hidden, d_model),
27
            )
28

29
        self.routed_experts = nn.ModuleList(
30
            [make_expert() for _ in range(num_routed_experts)]
31
        )
32
        self.shared_experts = nn.ModuleList(
33
            [make_expert() for _ in range(num_shared_experts)]
34
        )
35

36
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
37
        # x: [T, d_model]
38
        probs, topk_idx, routing_weights = self.router(x)
39
        capacity = expert_capacity(x.size(0), self.num_routed_experts, self.capacity_factor)
40

41
        shared_out = torch.zeros_like(x)
42
        for expert in self.shared_experts:
43
            shared_out += expert(x)
44

45
        routed_out = torch.zeros_like(x)
46
        for i in range(self.top_k):
47
            expert_ids = topk_idx[:, i]                              # [T]
48
            weights = routing_weights[:, i]                          # [T]
49
            for e in range(self.num_routed_experts):
50
                token_ids = (expert_ids == e).nonzero().squeeze(-1)[:capacity]
51
                if token_ids.numel() > 0:
52
                    expert_out = self.routed_experts[e](x[token_ids])
53
                    routed_out[token_ids] += weights[token_ids].unsqueeze(-1) * expert_out
54

55
        aux_loss = load_balancing_loss(probs, topk_idx, self.alpha)
56
        return shared_out + routed_out, aux_loss

可以看到，DeepSeekMoE 的 dispatch 流程与基础 MoE 完全一致，差异只在两点：Expert 被切得更小、更多，以及多了一组对所有 token 恒激活的 Shared Expert。实际模型中还会有更精细的配置，例如 DeepSeek-V2 将 gating 改为 sigmoid 打分后对选中项归一化，DeepSeek-V3 又在 Router 中引入 bias 项实现无辅助损失的负载均衡，但核心结构仍然是“细粒度 Routed Expert + Shared Expert”。