MOE 详解:结合 Transformer 架构的完整分析
MOE 详解:结合 Transformer 架构的完整分析
本文基于 Google 的论文《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》,结合 Transformer 架构,详细解释 MOE 的工作原理、数学公式和训练过程。
目录
一、MOE 概述
1.1 什么是 MOE?
MOE(Mixture of Experts) 是一种条件计算技术,核心思想是:
- 用多个专家网络(Experts)替换单一的前馈网络
- 使用门控网络(Gating Network)为每个输入动态选择少量专家
- 只有被选中的专家参与计算,实现稀疏激活
1.2 核心优势
| 特性 | 传统 MLP | MOE |
|---|---|---|
| 参数总量 | 受限于计算量 | 可以非常大 |
| 每次推理激活的参数 | 100% | 仅 top-k 专家 |
| 模型容量 | 有限 | 大幅增加 |
| 计算成本 | 与参数量成正比 | 相对恒定 |
示例:DeepSeek-V3
- 总参数:6710 亿
- 每次推理激活参数:370 亿(仅 5.5%)
- 专家数量:256 个
- 每个 token 激活专家:9 个(1 个共享 + 8 个选中)
1.3 发展历程
1991: Jacobs 等人提出 Mixture of Experts 概念
2017: Google 提出 Sparsely-Gated MOE(本文重点)
2022: Switch Transformer、DeepSpeed-MoE
2024: Mixtral、DeepSeek-MoE、DeepSeek-V3
二、MOE 在 Transformer 中的位置
2.1 回顾标准 Transformer Block
Input X
│
├──────────────────────────────┐
↓ │ (Residual)
Layer Norm → Multi-Head Attention │
│ │
↓ │
+ ←────────────────────────────┘
│
├──────────────────────────────┐
↓ │ (Residual)
Layer Norm → MLP (FFN) │ ← 【MOE 替换这里】
│ │
↓ │
+ ←────────────────────────────┘
│
↓
Output
2.2 MOE Transformer Block
MOE 替换 MLP 层:
Input X
│
├──────────────────────────────┐
↓ │ (Residual)
Layer Norm → Multi-Head Attention │
│ │
↓ │
+ ←────────────────────────────┘
│
├──────────────────────────────┐
↓ │ (Residual)
Layer Norm → MOE Layer │ ← 【MOE 层】
│ ┌─────────────┐ │
│ │ Gating Net │ │
│ │ ↓ │ │
│ │ [E1][E2]... │ │
│ │ [En] │ │
│ └─────────────┘ │
↓ │
+ ←────────────────────────────┘
│
↓
Output
2.3 参数对比
| 组件 | 标准 Transformer (GPT-2) | MOE Transformer |
|---|---|---|
| Embedding | 50257 × 768 | 50257 × 768 |
| Attention | 4 × 768² = 2.4M | 4 × 768² = 2.4M |
| MLP/MOE | 768 × 3072 × 2 = 4.7M | n × (768 × 3072 × 2) + 门控 |
| 激活参数 | 全部 | top-k 专家 |
三、MOE 层的结构
3.1 整体架构
Input x (768维)
│
┌──────────────┴──────────────┐
│ │
↓ ↓
┌─────────────────┐ ┌─────────────────────────┐
│ Gating Network │ │ Experts │
│ G(x) │ │ E1, E2, ..., En │
│ │ │ (n个前馈网络) │
└────────┬────────┘ └─────────────────────────┘
│ │
│ 门控权重 │ 专家输出
│ [g1, g2, ..., gn] │ [E1(x), E2(x), ..., En(x)]
│ │
└──────────────┬──────────────┘
│
↓
加权求和 (稀疏)
y = Σ G(x)_i × E_i(x)
│
↓
Output y (768维)
3.2 核心组件
3.2.1 专家网络(Experts)
每个专家是一个独立的前馈网络,结构与标准 MLP 相同:
参数:
- \(W_1^{(i)} \in \mathbb{R}^{d \times d_{ff}}\):第一层权重(768 × 3072)
- \(W_2^{(i)} \in \mathbb{R}^{d_{ff} \times d}\):第二层权重(3072 × 768)
- 每个专家约 4.7M 参数
3.2.2 门控网络(Gating Network)
门控网络决定每个输入应该由哪些专家处理:
简单 Softmax 门控:
其中 \(W_g \in \mathbb{R}^{d \times n}\),\(n\) 是专家数量。
Noisy Top-K 门控(推荐):
3.3 MOE 输出公式
稀疏性:由于 \(G(x)\) 是稀疏的(只有 top-k 个非零),我们只需计算被选中的 k 个专家。
四、前向传播详解
4.1 输入设置
假设:
- 输入 \(x \in \mathbb{R}^{768}\):一个 token 的表示
- 专家数量 \(n = 8\)
- top-k = 2(每个 token 选择 2 个专家)
4.2 Step 1: 计算门控 logits
x 形状: (768,)
W_g 形状: (768, 8)
H(x) = x × W_g
形状: (8,) ← 每个专家一个分数
示例:
H(x) = [1.2, 0.5, -0.3, 2.1, 0.8, -0.5, 1.5, 0.2]
↑ ↑
专家0 专家3
4.3 Step 2: Noisy Top-K 选择
添加噪声(训练时):
noise = StandardNormal() × Softplus(x × W_noise)
H(x) + noise = [1.25, 0.48, -0.28, 2.15, 0.82, -0.52, 1.48, 0.18]
Top-2 选择:
- 专家3: 2.15 ← 最高
- 专家0: 1.25 ← 第二高
KeepTopK(H(x), 2) = [1.25, -∞, -∞, 2.15, -∞, -∞, -∞, -∞]
4.4 Step 3: Softmax 得到门控权重
Softmax([1.25, -∞, -∞, 2.15, -∞, -∞, -∞, -∞])
= [e^1.25, 0, 0, e^2.15, 0, 0, 0, 0] / (e^1.25 + e^2.15)
= [3.49, 0, 0, 8.58, 0, 0, 0, 0] / 12.07
= [0.289, 0, 0, 0.711, 0, 0, 0, 0]
门控权重:
G(x) = [0.289, 0, 0, 0.711, 0, 0, 0, 0]
↑ ↑
专家0 专家3
4.5 Step 4: 计算被选中专家的输出
只计算 G(x)_i ≠ 0 的专家:
专家0:
E_0(x) = W_2^(0) × GELU(W_1^(0) × x + b_1^(0)) + b_2^(0)
形状: (768,)
专家3:
E_3(x) = W_2^(3) × GELU(W_1^(3) × x + b_1^(3)) + b_2^(3)
形状: (768,)
其他专家 (1,2,4,5,6,7): 不计算! ← 节省计算
4.6 Step 5: 加权求和
y 形状: (768,)
这就是 MOE 层的输出!
4.7 批处理情况
对于一个 batch 的 tokens(形状 (batch_size, seq_len, 768)):
输入: X 形状 (B, S, 768) ← B个样本,S个token
对于每个 token x_{b,s}:
1. 计算门控权重 G(x_{b,s})
2. 选择 top-k 专家
3. 计算选中专家的输出
4. 加权求和
输出: Y 形状 (B, S, 768)
注意: 不同 token 可能选择不同的专家!
4.8 前向传播流程图
输入 x (batch_size × seq_len × 768)
│
├─────────────────────────────────────────────────┐
│ │
↓ ↓
┌───────────────────────┐ ┌─────────────────────────┐
│ Gating Network │ │ 8 个专家 │
│ │ │ │
│ H(x) = x × W_g │ │ E_0: MLP (768→3072→768)│
│ + noise │ │ E_1: MLP (768→3072→768)│
│ │ │ E_2: MLP (768→3072→768)│
│ TopK + Softmax │ │ ... │
│ │ │ E_7: MLP (768→3072→768)│
└───────────┬───────────┘ └───────────┬─────────────┘
│ │
│ G(x) = [0.289, 0, 0, 0.711, 0, 0, 0, 0] │
│ │
│ 选中专家: {0, 3} │
│ │
└──────────────────┬─────────────────────────┘
│
↓
只计算 E_0(x) 和 E_3(x)
│
↓
y = 0.289 × E_0(x) + 0.711 × E_3(x)
│
↓
输出 y (batch_size × seq_len × 768)
五、反向传播详解
5.1 总体梯度流向
∂L/∂y
│
┌───────────────────────┴───────────────────────┐
│ │
↓ ↓
∂L/∂G(x) (门控梯度) ∂L/∂E_i(x) (专家梯度)
│ │
↓ ↓
∂L/∂W_g, ∂L/∂W_noise ∂L/∂W_1^(i), ∂L/∂W_2^(i)
5.2 MOE 输出公式回顾
对于稀疏情况(只有 top-k 个专家被激活):
5.3 Step 1: 对专家输出的梯度
关键:只有被选中的专家才有非零梯度!
假设 ∂L/∂y 形状: (768,)
对于专家0 (G(x)_0 = 0.289):
∂L/∂E_0(x) = 0.289 × ∂L/∂y
形状: (768,)
对于专家3 (G(x)_3 = 0.711):
∂L/∂E_3(x) = 0.711 × ∂L/∂y
形状: (768,)
对于未选中的专家 (G(x)_i = 0):
∂L/∂E_i(x) = 0 ← 没有梯度更新!
5.4 Step 2: 对门控权重的梯度
对于专家0:
∂L/∂G(x)_0 = E_0(x)^T × ∂L/∂y
形状: 标量
对于专家3:
∂L/∂G(x)_3 = E_3(x)^T × ∂L/∂y
形状: 标量
5.5 Step 3: 通过 Softmax 反向传播
Softmax 的反向传播:
注意:只有 top-k 个专家参与 Softmax,所以求和只在这些专家上进行。
5.6 Step 4: 对门控网络权重的梯度
x 形状: (768,)
∂L/∂H(x) 形状: (8,) ← 只有 top-k 个非零
∂L/∂W_g = x^T × ∂L/∂H(x)
形状: (768, 8)
5.7 Step 5: 对每个专家网络的反向传播
对于被选中的专家(如专家0),其内部的 MLP 反向传播:
专家0 的结构:
h_1 = GELU(x × W_1^(0) + b_1^(0))
E_0(x) = h_1 × W_2^(0) + b_2^(0)
反向传播:
∂L/∂W_2^(0) = h_1^T × ∂L/∂E_0(x)
∂L/∂h_1 = ∂L/∂E_0(x) × (W_2^(0))^T
∂L/∂(pre_gelu) = ∂L/∂h_1 ⊙ GELU'(pre_gelu)
∂L/∂W_1^(0) = x^T × ∂L/∂(pre_gelu)
5.8 Step 6: 对输入 x 的梯度
梯度从两条路径汇合:
从门控网络:
从专家网络:
5.9 完整反向传播流程图
∂L/∂y (768,)
│
┌─────────────────────────┴─────────────────────────┐
│ │
↓ ↓
┌─────────────────────────────────────┐ ┌─────────────────────────────────────┐
│ 对门控权重的梯度 │ │ 对专家输出的梯度 │
│ │ │ │
│ ∂L/∂G(x)_0 = E_0(x)^T × ∂L/∂y │ │ ∂L/∂E_0(x) = 0.289 × ∂L/∂y │
│ ∂L/∂G(x)_3 = E_3(x)^T × ∂L/∂y │ │ ∂L/∂E_3(x) = 0.711 × ∂L/∂y │
│ │ │ │
└──────────────┬──────────────────────┘ └──────────────┬──────────────────────┘
│ │
↓ ↓
┌─────────────────────────────────────┐ ┌─────────────────────────────────────┐
│ 通过 Softmax 反向 │ │ 对每个专家 MLP 反向 │
│ │ │ │
│ ∂L/∂H(x) = Softmax_backward(...) │ │ 专家0: ∂L/∂W_1^(0), ∂L/∂W_2^(0) │
│ │ │ 专家3: ∂L/∂W_1^(3), ∂L/∂W_2^(3) │
│ (只在 top-k 专家上计算) │ │ (只更新被选中的专家!) │
│ │ │ │
└──────────────┬──────────────────────┘ └──────────────┬──────────────────────┘
│ │
↓ ↓
┌─────────────────────────────────────┐ ┌─────────────────────────────────────┐
│ 对门控网络权重的梯度 │ │ 对输入 x 的梯度 (专家部分) │
│ │ │ │
│ ∂L/∂W_g = x^T × ∂L/∂H(x) │ │ ∂L/∂x|_exp = Σ G_i × ∂E_i/∂x │
│ │ │ │
└──────────────┬──────────────────────┘ └──────────────┬──────────────────────┘
│ │
│ │
└──────────────────┬───────────────────────┘
│
↓
∂L/∂x = ∂L/∂x|_gating + ∂L/∂x|_experts
│
↓
传递给前一层 (Attention)
六、训练中的负载均衡
6.1 问题:专家不均衡
如果不加约束,门控网络会倾向于:
- 总是选择某几个专家
- 这些专家训练更多,变得更好
- 门控网络更倾向于选择它们
- 形成恶性循环
训练过程中的不均衡示例:
训练初期:
专家使用次数: [120, 115, 118, 122, 125, 110, 108, 112]
分布相对均匀
训练后期 (无负载均衡):
专家使用次数: [450, 380, 25, 15, 10, 8, 7, 5]
只有前两个专家被使用,其他专家几乎废弃!
6.2 解决方案:辅助损失函数
6.2.1 重要性损失(Importance Loss)
定义专家 \(i\) 对一个 batch \(X\) 的重要性:
重要性损失:
其中 CV 是变异系数(标准差/均值)。
直觉:如果所有专家的重要性相等,CV = 0,损失为 0。
6.2.2 负载损失(Load Loss)
重要性相等不意味着接收的样本数相等。定义负载:
其中 \(P(x, i)\) 是在重新采样噪声时,专家 \(i\) 被选中的概率。
6.3 总损失函数
其中 \(\mathcal{L}_{\text{CE}}\) 是标准的交叉熵损失。
6.4 负载均衡的效果
| 损失配置 | Test Perplexity | CV(Importance) | CV(Load) | 最大负载/平均负载 |
|---|---|---|---|---|
| 无辅助损失 | 39.8 | 3.04 | 3.01 | 17.80 |
| \(w_{imp}=0.2\) | 35.6 | 0.06 | 0.17 | 1.47 |
| \(w_{load}=0.2\) | 35.7 | 0.22 | 0.04 | 1.15 |
| \(w_{imp}=0.1, w_{load}=0.1\) | 35.6 | 0.06 | 0.05 | 1.14 |
结论:辅助损失显著改善了模型质量(perplexity 从 39.8 降到 35.6)和负载均衡。
七、完整数值示例
7.1 设置
输入: x ∈ R^768 (一个 token 的表示)
专家数量: n = 8
每个专家: MLP (768 → 3072 → 768)
top-k: 2
7.2 前向传播
Step 1: 计算门控 logits
# x 形状: (768,)
# W_g 形状: (768, 8)
# W_noise 形状: (768, 8)
H_clean = x @ W_g # (8,)
# H_clean = [1.2, 0.5, -0.3, 2.1, 0.8, -0.5, 1.5, 0.2]
noise = torch.randn(8) * F.softplus(x @ W_noise)
# noise = [0.05, -0.02, 0.02, 0.05, 0.02, -0.02, -0.02, -0.02]
H = H_clean + noise
# H = [1.25, 0.48, -0.28, 2.15, 0.82, -0.52, 1.48, 0.18]
Step 2: Top-K 选择
# 选择 top-2
top_values, top_indices = torch.topk(H, k=2)
# top_values = [2.15, 1.48]
# top_indices = [3, 6] # 专家3和专家6
# 构建稀疏门控
H_masked = torch.full((8,), float('-inf'))
H_masked[top_indices] = top_values
# H_masked = [-inf, -inf, -inf, 2.15, -inf, -inf, 1.48, -inf]
Step 3: Softmax
G = F.softmax(H_masked, dim=-1)
# G = [0, 0, 0, 0.663, 0, 0, 0.337, 0]
# 专家3 ↑ 专家6 ↑
Step 4: 计算选中专家的输出
# 只计算专家3和专家6
E3_output = expert_3(x) # (768,)
E6_output = expert_6(x) # (768,)
# 每个专家内部:
# h1 = GELU(x @ W1 + b1) # (3072,)
# output = h1 @ W2 + b2 # (768,)
Step 5: 加权求和
y = G[3] * E3_output + G[6] * E6_output
# y = 0.663 * E3_output + 0.337 * E6_output
# y 形状: (768,)
7.3 反向传播
假设上游梯度:
dL_dy = ... # (768,) 从后续层传来
Step 1: 对专家输出的梯度
dL_dE3 = G[3] * dL_dy # = 0.663 * dL_dy
dL_dE6 = G[6] * dL_dy # = 0.337 * dL_dy
# 其他专家梯度为 0
Step 2: 对门控权重的梯度
dL_dG3 = (E3_output * dL_dy).sum() # 标量
dL_dG6 = (E6_output * dL_dy).sum() # 标量
Step 3: 通过 Softmax 反向
# 只在 top-k 专家上计算
G_topk = torch.tensor([G[3], G[6]]) # [0.663, 0.337]
dL_dG_topk = torch.tensor([dL_dG3, dL_dG6])
# Softmax 反向
s = (dL_dG_topk * G_topk).sum() # 加权和
dL_dH_topk = G_topk * (dL_dG_topk - s)
# 映射回原始索引
dL_dH = torch.zeros(8)
dL_dH[3] = dL_dH_topk[0]
dL_dH[6] = dL_dH_topk[1]
Step 4: 对门控网络权重的梯度
dL_dWg = x.unsqueeze(1) @ dL_dH.unsqueeze(0)
# 形状: (768, 8)
Step 5: 对每个专家的反向传播
# 专家3 (被选中)
dL_dW2_3 = h1_3.unsqueeze(1) @ dL_dE3.unsqueeze(0) # (3072, 768)
dL_dh1_3 = dL_dE3 @ W2_3.T # (3072,)
dL_dpre_gelu_3 = dL_dh1_3 * gelu_derivative(pre_gelu_3) # (3072,)
dL_dW1_3 = x.unsqueeze(1) @ dL_dpre_gelu_3.unsqueeze(0) # (768, 3072)
# 专家6 类似...
# 其他专家: 梯度为 0,不更新!
Step 6: 对输入 x 的梯度
# 从门控网络
dL_dx_gating = dL_dH @ W_g.T # (768,)
# 从专家网络
dL_dx_exp3 = G[3] * (dL_dpre_gelu_3 @ W1_3.T) # (768,)
dL_dx_exp6 = G[6] * (dL_dpre_gelu_6 @ W1_6.T) # (768,)
# 总梯度
dL_dx = dL_dx_gating + dL_dx_exp3 + dL_dx_exp6
7.4 参数更新统计
参数更新情况:
门控网络:
W_g: (768, 8) ← 全部更新
W_noise: (768, 8) ← 全部更新
专家网络:
专家0: 不更新 (未被选中)
专家1: 不更新
专家2: 不更新
专家3: W1_3, W2_3, b1_3, b2_3 ← 更新 ✓
专家4: 不更新
专家5: 不更新
专家6: W1_6, W2_6, b1_6, b2_6 ← 更新 ✓
专家7: 不更新
本次迭代:
- 门控网络参数: 全部更新
- 专家参数: 只有 2/8 = 25% 更新
八、MOE 的优势与挑战
8.1 优势
| 优势 | 说明 |
|---|---|
| 模型容量大 | 可以有数千亿参数 |
| 计算效率高 | 每次只激活少量专家 |
| 专家专业化 | 不同专家学习不同的模式 |
| 训练并行 | 不同专家可以在不同 GPU 上 |
8.2 挑战
| 挑战 | 解决方案 |
|---|---|
| 负载不均衡 | 辅助损失函数 |
| 通信开销 | 混合数据并行和模型并行 |
| 训练不稳定 | Noisy Top-K, 负载均衡 |
| 显存占用 | 专家参数分布在多个 GPU |
8.3 实验结果
来自原论文的结果:
| 模型 | 参数量 | ops/timestep | Test Perplexity |
|---|---|---|---|
| LSTM-2048-512 | 151M | 151M | 43.7 |
| MoE-4096-h | 4.3B | 8.9M | 34.1 |
结论:
- MOE 模型有 28倍 的参数量
- 但只有 6% 的计算量
- 同时获得了 22% 的性能提升
8.4 专家专业化示例
来自机器翻译实验的观察:
| 专家 | 专业化领域 |
|---|---|
| 专家 381 | 研究、创新相关词汇 |
| 专家 752 | 领导力、核心作用相关 |
| 专家 2004 | 速度、快速相关副词 |
专家 381 激活的上下文:
"... with researchers , ..."
"... to innovation . ..."
"... technology innovations is ..."
专家 752 激活的上下文:
"... plays a core ..."
"... plays a critical ..."
"... assume a leadership ..."
专家 2004 激活的上下文:
"... with rapidly growing ..."
"... to swiftly ..."
"... the fastest ..."
九、与标准 Transformer 的对比总结
9.1 结构对比
标准 Transformer Block: MOE Transformer Block:
Input Input
│ │
↓ ↓
Layer Norm Layer Norm
│ │
↓ ↓
Multi-Head Attention Multi-Head Attention
│ │
↓ ↓
Residual Add Residual Add
│ │
↓ ↓
Layer Norm Layer Norm
│ │
↓ ↓
┌─────────────────┐ ┌─────────────────────────────┐
│ MLP │ │ MOE Layer │
│ (768→3072→768) │ │ ┌─────────┐ ┌──────────┐ │
│ │ │ │ Gating │ │ Experts │ │
│ 一个前馈网络 │ │ │ Network │ │ E1...En │ │
│ │ │ └────┬────┘ └────┬─────┘ │
│ ~4.7M 参数 │ │ └──────┬─────┘ │
│ 100% 激活 │ │ 加权求和 (稀疏) │
└─────────────────┘ │ │
│ │ n × ~4.7M 参数 │
↓ │ top-k/n × 100% 激活 │
Residual Add └─────────────────────────────┘
│ │
↓ ↓
Output Residual Add
│
↓
Output
9.2 关键公式对比
| 组件 | 标准 MLP | MOE |
|---|---|---|
| 前向传播 | \(y = W_2 \cdot \text{GELU}(W_1 \cdot x)\) | \(y = \sum_{i \in \text{TopK}} G(x)_i \cdot E_i(x)\) |
| 参数量 | \(2 \cdot d \cdot d_{ff}\) | \(n \cdot 2 \cdot d \cdot d_{ff} + d \cdot n\) |
| 计算量 | \(2 \cdot d \cdot d_{ff}\) | \(k \cdot 2 \cdot d \cdot d_{ff} + d \cdot n\) |
| 激活比例 | 100% | \(k/n\) |
9.3 反向传播对比
| 步骤 | 标准 MLP | MOE |
|---|---|---|
| 梯度传播 | 传给所有参数 | 只传给被选中的专家 |
| 参数更新 | 所有参数 | 门控网络全部 + 选中专家 |
| 辅助损失 | 无 | 负载均衡损失 |
十、参考资料
- 原论文:Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
- Switch Transformer:Switch Transformers: Scaling to Trillion Parameter Models
- DeepSeek-MoE:DeepSeekMoE: Towards Ultimate Expert Specialization
- Mixtral:Mixtral of Experts
本文结合 Transformer 架构,详细分析了 MOE 的工作原理。MOE 通过稀疏激活实现了"大参数量、小计算量"的目标,是构建超大规模语言模型的关键技术。

浙公网安备 33010602011771号