深入理解 MoE (混合专家模型)：稀疏激活与高效扩展

1. MoE 的核心思想#

1.1 为什么要混合专家？#

传统的 Transformer 模型在处理每个 token 时，都会激活全部参数——无论输入是简单的 “Hello” 还是复杂的数学推导，模型都要动用”全副武装”。

这就好比：

无论客户咨询什么问题（买衣服、退款、投诉），你都让整个公司的所有部门同时参与处理
效率低下，资源浪费

混合专家模型 (Mixture of Experts, MoE) 提供了一种更聪明的解决方案：

让不同的”专家”处理不同类型的任务，只激活相关的专家，其余专家处于休眠状态。

1.2 MoE 的核心思想#

MoE 的核心思想可以用一个生活化的例子来理解：

想象一家医院：

普通医院：每个医生都要会看所有科目的病（心脏病、肝病、肺病…）
专科医院：每个医生专注于自己的领域（心脏病专家只处理心脏病）

MoE 就是深度学习领域的”专科医院”：不同的专家网络专注于处理不同类型的信息，路由器（Router）负责将任务分配给最适合的专家。

1.3 密集模型 vs 稀疏模型#

特性	密集模型 (Dense)	稀疏模型 (MoE)
激活方式	所有参数参与计算	只激活部分参数
参数量	参数量 = 计算量	参数量大，但计算量小
内存需求	全部参数常驻内存	所有专家常驻，但前向计算只需部分
计算效率	O(N)	O(1) 近似常数（相对于专家数）
代表模型	GPT-3, LLaMA	Mixtral, DeepSeek-MoE, Grok-1

1.4 MoE 发展简史#

1
1991: 混合专家模型诞生
2
  │  Jacob et al. "Adaptive Mixtures of Local Experts"
3
  │
4
2010-2020: 沉寂期
5
  │  主要用于小型模型和特定任务
6
  │
7
2021: GShard
8
  │  Google 将 MoE 引入大模型，实现 6000 亿参数
9
  │
10
2022: ST-MoE
11
  │  Stability AI 提出稳定训练的 ST-MoE
12
  │
13
2023: Mixtral 8x7B
14
  │  开源 MoE 模型引爆社区
15
  │
16
2024: DeepSeek-MoE, DBRX, Grok-1
17
  │  各大公司发布自研 MoE 模型
18
  │
19
2024-2025: 全面爆发
20
  │  Mistral Large, Qwen2, MiniMax, 等纷纷采用 MoE

2. MoE 架构详解#

2.1 整体架构#

MoE 层通常替代 Transformer 中的前馈网络 (FFN) 层：

1
┌─────────────────────────────────────────────────────────────────┐
2
│                      标准 Transformer 层                          │
3
├─────────────────────────────────────────────────────────────────┤
4
│                                                                  │
5
│  输入: Hidden States                                             │
6
│       │                                                         │
7
│       ▼                                                         │
8
│  ┌─────────────┐                                                │
9
│  │ Self-Attention │  ← 所有 token 共享                          │
10
│  └──────┬──────┘                                                │
11
│         │                                                        │
12
│         ▼                                                        │
13
│  ┌─────────────┐                                                │
14
│  │     FFN     │  ← 替换为 MoE 层                                │
15
│  └─────────────┘                                                │
16
│                                                                  │
17
└─────────────────────────────────────────────────────────────────┘
18

19
                        ↓ 替换为 ↓
20

21
┌─────────────────────────────────────────────────────────────────┐
22
│                      MoE Transformer 层                          │
23
├─────────────────────────────────────────────────────────────────┤
24
│                                                                  │
25
│  输入: Hidden States                                             │
26
│       │                                                         │
27
│       ▼                                                         │
28
│  ┌─────────────┐                                                │
29
│  │ Self-Attention │  ← 所有 token 共享                          │
30
│  └──────┬──────┘                                                │
31
│         │                                                        │
32
│         ▼                                                         │
33
│  ┌─────────────────────────────────────────────────┐            │
34
│  │                    MoE 层                        │            │
35
│  │  ┌──────────┐                                   │            │
36
│  │  │ Router    │ ← 决定每个 token 去哪个专家        │            │
37
│  │  └────┬─────┘                                   │            │
38
│  │       │                                          │            │
39
│  │       ▼                                          │            │
40
│  │  ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┐   │            │
41
│  │  │ E1  │ E2  │ E3  │ ... │ EN  │     │     │   │            │
42
│  │  │专家1│专家2│专家3│     │专家N│ 空闲│ 空闲│   │            │
43
│  │  └─────┴─────┴─────┴─────┴─────┴─────┴─────┘   │            │
44
│  └─────────────────────────────────────────────────┘            │
45
│                                                                  │
46
└─────────────────────────────────────────────────────────────────┘

2.2 数学公式#

对于输入 token $x$ ，MoE 层的输出定义为：

y = \sum_{i=1}^{N} G(x)_i \cdot E_i(x)

其中：

$N$ 是专家的总数
$E_i(x)$ 是第 $i$ 个专家网络的输出
$G(x)_i$ 是第 $i$ 个专家的门控权重（通常经过 softmax）

对于稀疏 MoE，通常只激活 top-k 个专家：

G(x) = \text{Softmax}(\text{TopK}(W_g x, k))

这里 $\text{TopK}(\cdot, k)$ 操作保留最大的 k 个值，其余设为 $-\infty$ 。

2.3 PyTorch 实现#

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
import math
5

6
class MoELayer(nn.Module):
7
    """
8
    混合专家层
9
    - num_experts: 专家数量
10
    - top_k: 每个 token 激活的专家数量
11
    - d_model: 模型维度
12
    - d_ff: FFN 隐藏层维度
13
    """
14
    def __init__(self, num_experts: int, top_k: int, d_model: int, d_ff: int):
15
        super().__init__()
16
        self.num_experts = num_experts
17
        self.top_k = top_k
18
        self.d_model = d_model
19
        self.d_ff = d_ff
20

21
        # 路由器网络
22
        self.router = nn.Linear(d_model, num_experts, bias=False)
23

24
        # 专家网络列表
25
        self.experts = nn.ModuleList([
26
            nn.Sequential(
27
                nn.Linear(d_model, d_ff),
28
                nn.GELU(),
29
                nn.Linear(d_ff, d_model)
30
            )
31
            for _ in range(num_experts)
32
        ])
33

34
    def forward(self, x: torch.Tensor):
35
        """
36
        x: (batch_size, seq_len, d_model)
37
        返回: (batch_size, seq_len, d_model)
38
        """
39
        batch_size, seq_len, d_model = x.shape
40

41
        # 重塑为 2D 以便并行处理
42
        x_flat = x.view(-1, d_model)  # (batch_size * seq_len, d_model)
43

44
        # 路由器计算：获取每个专家的原始分数
45
        router_logits = self.router(x_flat)  # (tokens, num_experts)
46

47
        # 计算概率分布
48
        router_probs = F.softmax(router_logits, dim=-1)
49

50
        # 选择 top-k 个专家
51
        top_k_probs, top_k_indices = torch.topk(
52
            router_probs, self.top_k, dim=-1
53
        )
54

55
        # 归一化 top-k 概率
56
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
57

58
        # 初始化输出
59
        output = torch.zeros_like(x_flat)
60

61
        # 对每个专家分别计算并聚合
62
        for i in range(self.top_k):
63
            expert_idx = top_k_indices[:, i]  # 每个 token 的第 i 个专家索引
64
            expert_weight = top_k_probs[:, i]   # 对应的权重
65

66
            # 按专家分组计算
67
            for e_idx in range(self.num_experts):
68
                # 找出分配给当前专家的 token
69
                mask = (expert_idx == e_idx)
70
                if mask.sum() == 0:
71
                    continue
72

73
                # 获取这些 token 的输入和权重
74
                token_indices = mask.nonzero(as_tuple=True)[0]
75
                expert_input = x_flat[token_indices]
76
                expert_weight_e = expert_weight[token_indices]
77

78
                # 计算专家输出
79
                expert_output = self.experts[e_idx](expert_input)
80

81
                # 加权累加到输出
82
                output[token_indices] += expert_weight_e.unsqueeze(-1) * expert_output
83

84
        return output.view(batch_size, seq_len, d_model)

3. 路由器 (Router) 机制#

3.1 路由器的作用#

路由器是 MoE 的”大脑”，负责决定：

每个 token 应该由哪些专家处理
每个被选中的专家应该分配多大的权重

3.2 常见的路由器类型#

3.2.1 线性路由器 (Linear Router)#

最简单的路由器，只有一个线性变换：

G(x) = \text{Softmax}(W_g x)

3.2.2 带噪声的 Top-K 路由器#

1
class NoiseTopKRouter(nn.Module):
2
    """带噪声的 Top-K 路由器"""
3
    def __init__(self, d_model: int, num_experts: int, top_k: int):
4
        super().__init__()
5
        self.gate = nn.Linear(d_model, num_experts, bias=False)
6
        self.top_k = top_k
7

8
        # 可学习的噪声温度
9
        self.noise_std = nn.Parameter(torch.zeros(num_experts))
10

11
    def forward(self, x: torch.Tensor):
12
        # 原始 logits
13
        logits = self.gate(x)  # (tokens, num_experts)
14

15
        # 添加噪声（训练时使用）
16
        if self.training:
17
            noise = torch.randn_like(logits) * self.noise_std.exp()
18
            logits = logits + noise
19

20
        # Top-K 选择
21
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
22

23
        # 处理未选中专家的 logits（设为 -inf）
24
        mask = torch.zeros_like(logits).scatter_(1, top_k_indices, 1.0)
25
        logits = logits.masked_fill(mask == 0, float('-inf'))
26

27
        # Softmax
28
        return F.softmax(logits, dim=-1), top_k_indices

3.2.3 路由器类型对比#

路由器类型	优点	缺点	应用
Linear	简单高效	可能负载不均	早期 MoE
Top-K	只激活部分专家，节省计算	K 值需调优	Mixtral
Noise Top-K	探索性更强，负载均衡好	噪声需训练	Switch Transformer
Token Choice	细粒度控制	通信开销大	研究场景
Expert Choice	保证负载均衡	输出长度可变	EP-MoE

3.3 Top-K 选择策略#

Top-K 中的 K 是一个关键超参数：

K 值	特点	适用场景
K=1	最稀疏，只选 1 个专家	极端稀疏，显存敏感
K=2	平衡选择	Mixtral 8x7B
K=4+	更稳定，但更接近密集	追求性能上限

K 值选择的经验法则

计算预算有限：K=1 或 K=2
追求性能：K=4 到 K=8
专家数量多时：可以适当增大 K

4. 负载均衡 (Load Balancing)#

4.1 为什么需要负载均衡？#

如果路由器总是选择同样的几个”明星专家”，会导致：

计算瓶颈：少数专家过载，成为计算瓶颈
训练不稳定：负载不均的专家训练不均衡
资源浪费：未使用的专家浪费显存

4.2 辅助损失函数#

4.2.1 辅助负载均衡损失#

Mixtral 使用的辅助损失：

L_{balance} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i

其中：

$f_i = \frac{1}{T} \sum_{t=1}^{T} \mathbb{1}(\text{argmax}(G(x_t)) = i)$ ：被分配给专家 i 的 token 比例
$P_i = \frac{1}{T} \sum_{t=1}^{T} G_i(x_t)$ ：专家 i 的平均路由概率
$\alpha$ 是辅助损失权重（通常 0.01）

1
def load_balancing_loss(router_probs, top_k_indices, num_experts, alpha=0.01):
2
    """
3
    计算负载均衡辅助损失
4
    router_probs: (tokens, num_experts) - 路由器概率分布
5
    top_k_indices: (tokens, top_k) - 每个 token 选中的专家索引
6
    """
7
    tokens = router_probs.size(0)
8

9
    # 计算每个专家被选中的频率
10
    # f_i: 被分配给专家 i 的 token 比例
11
    f_i = torch.zeros(num_experts, device=router_probs.device)
12
    for i in range(top_k_indices.size(1)):
13
        indices = top_k_indices[:, i]
14
        counts = torch.bincount(indices, minlength=num_experts).float()
15
        f_i += counts / tokens
16

17
    # P_i: 每个专家的平均路由概率
18
    P_i = router_probs.mean(dim=0)
19

20
    # 负载均衡损失
21
    loss = alpha * num_experts * torch.sum(f_i * P_i)
22

23
    return loss

4.2.2 Expert Capacity 限制#

另一种方法是限制每个专家的最大容量：

1
class CapacityBoundedRouter(nn.Module):
2
    """带容量限制的路由器"""
3
    def __init__(self, d_model: int, num_experts: int, top_k: int,
4
                 capacity_factor: float = 1.25):
5
        super().__init__()
6
        self.gate = nn.Linear(d_model, num_experts, bias=False)
7
        self.top_k = top_k
8
        self.capacity_factor = capacity_factor
9

10
    def forward(self, x: torch.Tensor):
11
        batch_size, seq_len = x.shape[:2]
12
        num_tokens = batch_size * seq_len
13

14
        # 计算容量
15
        capacity = int(num_tokens * self.capacity_factor / self.num_experts)
16

17
        # 路由器
18
        logits = self.gate(x.view(-1, x.size(-1)))
19
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
20

21
        # 容量检查
22
        # 对于每个专家，限制分配的 token 数量
23
        outputs = []
24
        dispatch_counts = torch.zeros(self.num_experts, device=x.device)
25

26
        for k in range(self.top_k):
27
            expert_ids = top_k_indices[:, k]
28
            for e in range(self.num_experts):
29
                mask = expert_ids == e
30
                if mask.sum() > capacity:
31
                    # 保留概率最高的 token
32
                    expert_logits = top_k_logits[:, k][mask]
33
                    keep_indices = expert_logits.topk(capacity)[1]
34
                    full_indices = mask.nonzero(as_tuple=True)[0][keep_indices]
35
                    # 处理超容情况
36
                    ...
37

38
        return ...

4.3 辅助损失可视化#

1
负载均衡前：              负载均衡后：
2

3
专家使用分布              专家使用分布
4
│                          │
5
│ ████                     │    █
6
│ ████                     │ █  █
7
│ ████ ████                │ █  █  █
8
│ ████ ████ ████           │ █  █  █  █
9
├───────────────           ├───────────────
10
 E1 E2 E3 E4               E1 E2 E3 E4
11

12
专家1、2过载               负载更加均衡
13
专家3、4空闲

5. 分布式训练策略#

5.1 专家并行 (Expert Parallelism)#

在大规模 MoE 训练中，专家被分布到不同的 GPU 上：

1
┌─────────────────────────────────────────────────────────────────┐
2
│                        4 GPU 专家并行                            │
3
├─────────────────────────────────────────────────────────────────┤
4
│                                                                  │
5
│  GPU 0          GPU 1          GPU 2          GPU 3             │
6
│  ┌───────┐      ┌───────┐      ┌───────┐      ┌───────┐         │
7
│  │Expert1│      │Expert2│      │Expert3│      │Expert4│         │
8
│  │Expert5│      │Expert6│      │Expert7│      │Expert8│         │
9
│  └───────┘      └───────┘      └───────┘      └───────┘         │
10
│       ↑              ↑              ↑              ↑            │
11
│       └──────────────┴──────────────┴──────────────┘            │
12
│                        All-to-All 通信                           │
13
│                                                                  │
14
└─────────────────────────────────────────────────────────────────┘

5.2 通信模式#

1
def expert_parallel_forward(alltoall_enabled=True):
2
    """
3
    专家并行的前向传播
4
    """
5
    # 步骤 1: 本地计算路由器
6
    local_router_output = compute_router(local_tokens)
7

8
    # 步骤 2: All-to-All 通信（如果启用了跨设备通信）
9
    if alltoall_enabled:
10
        # 将 token 分发到对应专家所在的设备
11
        dispatched_tokens = all_to_all(local_router_output, experts_per_device)
12
    else:
13
        dispatched_tokens = local_router_output
14

15
    # 步骤 3: 在各个设备上计算对应的专家
16
    expert_outputs = []
17
    for local_expert in local_experts:
18
        expert_outputs.append(compute_expert(dispatched_tokens))
19

20
    # 步骤 4: 汇总输出
21
    output = all_to_all(expert_outputs)  # 恢复原始顺序
22

23
    return output

5.3 通信优化#

优化技术	描述	效果
All-to-All	跨设备 token 分发	必要开销
通信重叠	计算与通信流水线	隐藏延迟
梯度累积	增大 batch 而不减损	减少通信频率
量化通信	FP16/INT8 传输	减少带宽

6. 主流 MoE 模型#

6.1 Mixtral 8x7B#

Mixtral 是由 Mistral AI 发布的开源 MoE 模型，开创了开源 MoE 的先河。

特性	参数
专家数量	8
每个 token 激活专家数	2
总参数量	~46.7B
激活参数量	~12.9B
上下文长度	32K
许可证	Apache 2.0

架构特点：

每个 Transformer 层都有 MoE 层
使用稀疏 MoE，2/8 的专家被激活
支持高质量代码和指令跟随

1
# Mixtral 的 MoE 层结构（简化）
2
class MixtralSparseMoeBlock(nn.Module):
3
    def __init__(self, hidden_size, num_experts=8, top_k=2):
4
        super().__init__()
5
        self.num_experts = num_experts
6
        self.top_k = top_k
7

8
        self.gate = nn.Linear(hidden_size, num_experts, bias=False)
9

10
        # 每个专家是一个完整的 FFN
11
        self.experts = nn.ModuleList([
12
            SwigluFFN(hidden_size)  # SwiGLU 激活函数
13
            for _ in range(num_experts)
14
        ])
15

16
    def forward(self, x):
17
        # ... 标准的 Top-K MoE 前向传播
18
        pass

6.2 DeepSeek-MoE#

DeepSeek AI 提出的 DeepSeek-MoE 在细粒度专家和共享专家方面进行了创新。

特性	DeepSeek-MoE 16B	DeepSeek-MoE 67B
总专家数	64	64
激活专家数	2	2
共享专家	2	2
总参数量	16.4B	67.0B
激活参数量	2.4B	5.8B

创新点：

细粒度专家分割：将专家分割成更小的单元，增加灵活性
共享专家：某些专家始终被激活，捕获共同知识
专家级词汇分离：不同专家专注于不同类型的词汇

6.3 DBRX#

Databricks 的 DBRX 是一个强大的开源 MoE 模型。

特性	参数
专家数量	16
每个 token 激活专家数	4
总参数量	132B
激活参数量	36B
上下文长度	32K

6.4 模型对比#

模型	总参数量	激活参数量	专家数	Top-K	开源
Mixtral 8x7B	46.7B	12.9B	8	2	✅
Mixtral 8x22B	141B	39B	8	2	✅
DeepSeek-MoE 16B	16.4B	2.4B	64	2	✅
DeepSeek-V2	236B	21B	128	6	✅
DBRX	132B	36B	16	4	✅
Grok-1	314B	~79B	8	2	✅
Qwen1.5-MoE	109B	20B	64	8	✅
MiniMax-01	456B	45B	256	8	✅

7. 专家 specialty 现象#

7.1 什么是专家 specialty？#

研究表明，MoE 中的专家会自然地分化出不同的”专业领域”：

专家类型	典型行为	示例
语法专家	处理词法、句法结构	Python、括号匹配
语义专家	处理词汇含义	同义词、反义词
领域专家	专注特定知识领域	数学、医学、法律
格式化专家	处理特定输出格式	JSON、XML、Markdown
代码专家	处理编程逻辑	函数调用、变量引用

7.2 专家分化的可视化#

1
Tokens: "def calculate_fibonacci(n):"
2

3
专家激活分布：
4
┌──────────────────────────────────────────────────┐
5
│ 专家1 (代码结构)    ████████████████████ 0.45      │
6
│ 专家2 (语法)        ████████████████ 0.30         │
7
│ 专家3 (函数调用)    ██████████ 0.15               │
8
│ 专家4 (数学逻辑)    ██████ 0.08                   │
9
│ 专家5-8            ██ 0.02                       │
10
└──────────────────────────────────────────────────┘
11

12
Tokens: "The capital of France is"
13

14
专家激活分布：
15
┌──────────────────────────────────────────────────┐
16
│ 专家6 (地理知识)    ████████████████████ 0.55      │
17
│ 专家7 (命名实体)    ███████████████ 0.28          │
18
│ 专家1 (代码结构)    ██ 0.02                       │
19
│ ...                  ██ 0.15                      │
20
└──────────────────────────────────────────────────┘

7.3 路由器的路由模式#

路由器不仅关注 token 本身，还会考虑上下文：

1
# 路由器可能关注的信息
2
router_context = {
3
    "current_token": "France",        # 当前 token
4
    "surrounding_tokens": "of capital",  # 周围 tokens
5
    "position": "in sentence",       # 位置信息
6
    "task_type": "knowledge_qa",     # 任务类型
7
}

8. MoE 的挑战与解决方案#

8.1 主要挑战#

挑战	描述	影响
负载不均	少数专家被频繁选中	训练不稳定，资源浪费
通信开销	跨设备通信	分布式训练效率
显存占用	所有专家常驻	单卡显存压力大
训练不稳定	稀疏更新的梯度问题	模型收敛困难

8.2 解决方案#

8.2.1 辅助损失优化#

1
class AuxiliaryLossWithImportancePenalty(nn.Module):
2
    """
3
    带重要性惩罚的辅助损失
4
    论文: ST-MoE
5
    """
6
    def __init__(self, num_experts, alpha=0.01, beta=0.01):
7
        super().__init__()
8
        self.num_experts = num_experts
9
        self.alpha = alpha  # 原始平衡损失权重
10
        self.beta = beta    # 重要性惩罚权重
11

12
        # 可学习的专家重要性因子
13
        self.expert_importance = nn.Parameter(torch.ones(num_experts))
14

15
    def forward(self, router_probs, top_k_indices):
16
        # 原始平衡损失
17
        balance_loss = load_balancing_loss_base(router_probs, top_k_indices)
18

19
        # 专家重要性正则化
20
        # 鼓励专家重要性因子接近均匀分布
21
        importance_penalty = self.beta * torch.var(self.expert_importance)
22

23
        # 总损失
24
        return self.alpha * balance_loss + importance_penalty

8.2.2 专家卸载策略#

1
class ExpertOffloading:
2
    """
3
    专家卸载：不在每个 GPU 上保留所有专家
4
    """
5
    def __init__(self, num_experts, num_devices, experts_per_device=2):
6
        self.num_experts = num_experts
7
        self.num_devices = num_devices
8
        self.experts_per_device = experts_per_device
9

10
        # 专家到设备的映射
11
        self.expert_to_device = {
12
            i: i % num_devices
13
            for i in range(num_experts)
14
        }
15

16
    def get_experts_for_device(self, device_id):
17
        """获取某个设备上的专家"""
18
        return [
19
            e for e, d in self.expert_to_device.items()
20
            if d == device_id
21
        ]
22

23
    def forward(self, token_expert_assignments):
24
        """根据 token 分配决定哪些专家需要加载"""
25
        # 检查远程专家是否需要加载
26
        remote_experts = set()
27
        for token_assigned in token_expert_assignments:
28
            remote_experts.update(token_assigned)
29

30
        return remote_experts

8.2.3 训练稳定性技巧#

技巧	描述	论文
Router Z-loss	惩罚路由器 logits 的高值	Z-loss
专家丢弃	随机丢弃部分专家	Switch Transformer
梯度削波	限制梯度范数	通用
专家容错	处理缺失专家的梯度	Expert Choice

9. MoE 的变体与进阶#

9.1 Soft MoE (Google)#

Soft MoE 用加权平均替代离散的专家选择：

y_i = \sum_{j} \frac{\exp(-\|x_i - v_j\|^2 / \sigma^2)}{\sum_k \exp(-\|x_i - v_k\|^2 / \sigma^2)} \cdot W_j x_i

特点：

连续化路由，完全可微
无需负载均衡
但计算量与密集模型相当

9.2 Expert Choice MoE#

每个专家选择自己负责的 token：

1
class ExpertChoiceMoE(nn.Module):
2
    """
3
    专家选择 MoE
4
    每个专家选择对自己最重要的 token
5
    """
6
    def __init__(self, num_experts, expert_capacity):
7
        super().__init__()
8
        self.num_experts = num_experts
9
        self.expert_capacity = expert_capacity
10

11
        self.gate = nn.Linear(d_model, num_experts, bias=False)
12

13
    def forward(self, x):
14
        tokens, seq_len, d_model = x.shape
15

16
        # 每个 token 对每个专家的分数
17
        expert_scores = self.gate(x)  # (batch, seq, num_experts)
18

19
        # 每个专家独立选择 top-k tokens
20
        # 注意：这里需要转置以按专家维度选择
21
        expert_outputs = []
22
        for e in range(self.num_experts):
23
            scores = expert_scores[:, :, e]  # (batch, seq)
24
            _, top_indices = torch.topk(scores, self.expert_capacity, dim=1)
25

26
            # 收集这些 token
27
            # ... 处理边界情况
28
            expert_outputs.append(expert_output)
29

30
        # 组合输出（需要填充到原始长度）
31
        return ...

9.3 Multi-head MoE#

将 MoE 扩展到多头：

1
Multi-head MoE 架构：
2

3
           ┌──────────────────────────────────────┐
4
           │         Multi-head Router            │
5
           │  ┌─────────┐ ┌─────────┐ ┌─────────┐ │
6
           │  │ Head 1  │ │ Head 2  │ │ Head H  │ │
7
           │  │Router   │ │Router   │ │Router   │ │
8
           │  └────┬────┘ └────┬────┘ └────┬────┘ │
9
           └───────┼───────────┼───────────┼───────┘
10
                   │           │           │
11
        ┌──────────┴───┐ ┌─────┴─────┐ ┌──┴──────────┐
12
        │ Expert Pool 1 │ │ Expert 2  │ │ Expert Pool H│
13
        │ E1 E2 E3 E4  │ │ E5 E6 E7 E8│ │ E9 E10 E11 E12│
14
        └──────────────┘ └──────────┘ └──────────────┘
15
                   │           │           │
16
                   └───────────┴───────────┘
17
                            │
18
                            ▼
19
                     融合输出

9.4 层级 MoE#

在不同层级使用不同粒度的 MoE：

层级	MoE 配置	理由
低层 (1-12)	细粒度专家 (64+)	捕获基础特征
中层 (13-24)	中等专家 (16-32)	捕获语义模式
高层 (25-32)	少量专家 (8-16)	捕获高层语义

10. MoE 与 Transformer 的结合#

10.1 MoE Transformer 完整层#

1
class MoETransformerLayer(nn.Module):
2
    """
3
    完整的 MoE Transformer 层
4
    """
5
    def __init__(self, d_model, num_heads, num_experts, top_k, d_ff):
6
        super().__init__()
7

8
        # Self-Attention
9
        self.self_attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
10
        self.attn_norm = nn.LayerNorm(d_model)
11

12
        # MoE FFN
13
        self.moe = MoELayer(
14
            num_experts=num_experts,
15
            top_k=top_k,
16
            d_model=d_model,
17
            d_ff=d_ff
18
        )
19
        self.moe_norm = nn.LayerNorm(d_model)
20

21
        # 可选：共享专家（DeepSeek 风格）
22
        self.shared_expert = SharedExpertFFN(d_model, d_ff)
23
        self.use_shared_expert = True
24

25
    def forward(self, x, mask=None):
26
        # Self-Attention with residual
27
        attn_out, _ = self.self_attn(x, x, x, attn_mask=mask)
28
        x = self.attn_norm(x + attn_out)
29

30
        # MoE FFN with residual
31
        moe_out = self.moe(x)
32
        x = self.moe_norm(x + moe_out)
33

34
        # 可选：添加共享专家
35
        if self.use_shared_expert:
36
            shared_out = self.shared_expert(x)
37
            x = x + shared_out
38

39
        return x

10.2 与 SwiGLU 的结合#

Mixtral 使用 SwiGLU 激活函数：

1
class SwigluFFN(nn.Module):
2
    """
3
    SwiGLU 前馈网络
4
    用于 Mixtral 的专家
5
    """
6
    def __init__(self, d_model, d_ff, dropout=0.0):
7
        super().__init__()
8
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
9
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
10
        self.w3 = nn.Linear(d_model, d_ff, bias=False)
11
        self.dropout = nn.Dropout(dropout)
12

13
        # SwiGLU 激活
14
        self.act = nn.SiLU()
15

16
    def forward(self, x):
17
        # SwiGLU: gate(x) * x (with different projection)
18
        return self.dropout(self.w2(self.act(self.w1(x)) * self.w3(x)))

11. 核心公式总结#

MoE 层输出：

y = \sum_{i=1}^{N} G(x)_i \cdot E_i(x)

稀疏门控（Top-K）：

G(x) = \text{Softmax}(\text{TopK}(W_g x, k))

负载均衡损失：

L_{balance} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i

其中 $f_i$ 是专家 $i$ 被选中的频率， $P_i$ 是平均路由概率。

SwiGLU 激活：

\text{SwiGLU}(x) = \text{SiLU}(W_1 x) \odot (W_3 x)

12. 总结与展望#

MoE 的核心优势#

参数量大但计算高效：参数量随专家数线性增长，但激活参数量保持相对稳定
稀疏激活：每个 token 只激活少量专家，计算量远小于密集模型
专业分工：不同专家可以学习不同类型的知识或任务
可扩展性：可以轻易增加专家数量而无需显著增加计算成本

当前挑战#

负载均衡：需要精心设计辅助损失或容量机制
通信开销：分布式训练中的跨设备通信
显存压力：所有专家需要常驻显存
训练稳定性：稀疏更新可能导致梯度不稳定

未来方向#

动态专家：根据输入动态调整专家数量或结构
细粒度专家：更细粒度的专家分割以提高灵活性
多模态 MoE：统一处理文本、图像、音频等多种模态
硬件协同设计：针对 MoE 特点优化的硬件架构

MoE 正在成为大模型扩展的关键技术，它让我们能够以更低的计算成本训练和部署更大规模的模型。随着研究的深入和硬件的进步，MoE 有望在未来的 AI 系统中发挥更加重要的作用。

参考资料#

Shazeer, N., et al. (2017). “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.” ICLR.
Lepikhin, D., et al. (2020). “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” ICLR.
Fedus, W., et al. (2022). “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” JMLR.
Zoph, B., et al. (2022). “ST-MoE: Designing Stable and Transferable Sparse Expert Models.” arXiv.
Jiang, A., et al. (2024). “Mixtral of Experts.” arXiv.
Dai, D., et al. (2024). “DeepSeekMoE: Towards Ultimate Expert Specialization in Large-Scale Mixture-of-Experts Language Models.” arXiv.
Zhou, J., et al. (2024). “DBRX: A Open Large Language Model with 132B Parameters.” Databricks Blog.
DeepSeek-AI. (2024). “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model.” arXiv.