Chaos and Order

💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

引言
MoE 架构的历史与发展
Sparse MoE 的数学基础
Switch Transformer 分析
Mixtral 8x7B 架构详解
DeepSeek-V2/V3 的创新：DeepSeekMoE
路由策略对比
训练稳定性技巧
推理优化
主要 MoE 模型对比
综合实现示例：自定义 MoE Transformer Block
结论与未来展望
参考资料

引言

大规模语言模型（LLM）时代，无限扩大模型参数量在训练成本和推理成本两方面都会遇到瓶颈。Dense Transformer 对每一个输入 token 都会激活全部参数，因此参数量增加时，运算量（FLOPs）也会成比例增加。Mixture of Experts（MoE）架构通过条件计算（conditional computation）解决了这个问题。核心思路是：全部参数中只根据输入激活一部分专家（expert），在保持模型容量大幅增长的同时，让实际运算量维持恒定。

自 2017 年 Shazeer 等人在《Outrageously Large Neural Networks》论文中提出 Sparsely-Gated MoE 以来，从 2021 年 Google 的 Switch Transformer，到 2023 年 Mistral 的 Mixtral 8x7B，再到 2024 年的 DeepSeek-V2/V3，MoE 架构经历了飞速发展。2025 年，Meta 的 Llama 4 采用了 MoE，DeepSeek-R1 在 V3 架构之上把推理能力发挥到极致，在全球范围内引发关注。

本文从 MoE 架构的数学基础出发，深入分析各主要模型的设计理念、路由策略的对比分析、训练稳定性技巧，以及推理优化，力求达到论文级别的详细程度。

MoE 架构的历史与发展

MoE 的概念最早由 Jacobs 等人在 1991 年的论文《Adaptive Mixtures of Local Experts》中提出。早期做法很简单，就是用一个门控网络对多个专家网络的输出做加权求和。

现代 MoE 的关键转折点可以整理如下。

2017 年：Shazeer 等人提出基于 LSTM 的 Sparsely-Gated MoE，用 4096 个专家实现了千亿参数级模型
2021 年：Google 的 Switch Transformer 把路由简化为 Top-1，实现了 1.6 万亿参数模型
2022 年：Google 的 ST-MoE（Stable and Transferable MoE）把训练稳定性技巧系统化
2022 年：Expert Choice Routing 论文提出由专家选择 token 的反向路由
2023 年：Mixtral 8x7B 以 Top-2 路由和 SwiGLU 专家开启了开源 MoE 时代
2024 年：DeepSeek-V2 引入 Fine-Grained Expert 与 Auxiliary-Loss-Free 策略
2024 年：DeepSeek-V3 以 671B 参数（激活 37B）达到最先进性能
2025 年：Llama 4 Scout（16 个专家，109B/17B 激活）让 Meta 也采用了 MoE

Sparse MoE 的数学基础

基本公式

MoE 层的输出定义如下。

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

其中 $x$ 是输入 token 的隐藏表示， $N$ 是专家数量， $E_i$ 是第 $i$ 个专家网络， $g(x)_i$ 是门控函数分配给第 $i$ 个专家的权重。

Sparse Gating 函数

Shazeer（2017）提出的 Noisy Top-K 门控函数如下。

g(x) = \text{Softmax}(\text{TopK}(H(x), k))

H(x)_i = (x \cdot W_g)_i + \epsilon \cdot \text{Softplus}((x \cdot W_{noise})_i)

其中 $W_g$ 是门控权重矩阵， $W_{noise}$ 是噪声权重矩阵。TopK 运算只保留前 $k$ 个值，其余设为 $-\infty$ ，使其在 Softmax 之后变为 0。

PyTorch 实现：基础 Sparse Gating

import torch
import torch.nn as nn
import torch.nn.functional as F

class TopKGating(nn.Module):
    """Noisy Top-K Gating mechanism for MoE."""

    def __init__(self, input_dim: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = nn.Linear(input_dim, num_experts, bias=False)
        self.noise = nn.Linear(input_dim, num_experts, bias=False)

    def forward(self, x: torch.Tensor):
        # x shape: (batch_size, seq_len, input_dim)
        logits = self.gate(x)  # (batch, seq, num_experts)

        # Training noise for exploration
        if self.training:
            noise = torch.randn_like(logits) * F.softplus(self.noise(x))
            logits = logits + noise

        # Top-K selection
        top_k_logits, top_k_indices = logits.topk(self.top_k, dim=-1)
        # (batch, seq, top_k)

        # Sparse softmax: only over selected experts
        top_k_gates = F.softmax(top_k_logits, dim=-1)

        return top_k_gates, top_k_indices

Switch Transformer 分析

核心创新：Top-1 路由

2021 年 Fedus 等人发表的 Switch Transformer 的核心创新是Top-1 路由。以往的研究普遍认为，稳定训练至少需要激活两个以上的专家，而 Switch Transformer 证明了每个 token 只精确选择一个专家也已经足够。

g(x) = \text{Softmax}(x \cdot W_r), \quad i^* = \arg\max_i g(x)_i

路由后的输出，就是门控概率与对应专家输出的乘积。

y = g(x)_{i^*} \cdot E_{i^*}(x)

架构特点

Switch Transformer 把 T5 架构的 FFN（Feed-Forward Network）层替换成了 MoE 层。每个 MoE 层最多可以配置 2048 个专家，借此实现了 1.6 万亿参数规模的模型。Top-1 路由把通信成本削减一半，路由运算本身也随之简化。

性能

使用 64 个专家的 Switch Transformer，在相同运算量的基准下，预训练速度达到了 T5-Base 的 7 倍。这是因为模型容量在增长的同时，每个 token 的运算量维持不变。

PyTorch 实现：Switch Transformer MoE Layer

class SwitchMoELayer(nn.Module):
    """Switch Transformer style MoE layer with Top-1 routing."""

    def __init__(self, hidden_dim: int, ffn_dim: int, num_experts: int,
                 capacity_factor: float = 1.25):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.router = nn.Linear(hidden_dim, num_experts, bias=False)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim, ffn_dim),
                nn.ReLU(),
                nn.Linear(ffn_dim, hidden_dim)
            ) for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor):
        batch_size, seq_len, hidden_dim = x.shape
        x_flat = x.view(-1, hidden_dim)  # (B*S, D)
        num_tokens = x_flat.shape[0]

        # Router: Top-1 selection
        router_logits = self.router(x_flat)  # (B*S, E)
        router_probs = F.softmax(router_logits, dim=-1)
        expert_indices = router_probs.argmax(dim=-1)  # (B*S,)
        expert_gates = router_probs.gather(1, expert_indices.unsqueeze(-1)).squeeze(-1)

        # Capacity: max tokens per expert
        capacity = int(self.capacity_factor * num_tokens / self.num_experts)

        # Dispatch tokens to experts
        output = torch.zeros_like(x_flat)
        for i in range(self.num_experts):
            mask = (expert_indices == i)
            if mask.sum() == 0:
                continue
            selected = x_flat[mask][:capacity]  # enforce capacity
            expert_out = self.experts[i](selected)
            gates = expert_gates[mask][:capacity].unsqueeze(-1)
            output[mask][:capacity] = expert_out * gates

        return output.view(batch_size, seq_len, hidden_dim)

Mixtral 8x7B 架构详解

设计理念

Mistral AI 于 2023 年 12 月发布的 Mixtral 8x7B，以 Mistral 7B 架构为基础，把每个 Transformer 层的 FFN 替换成由 8 个专家组成的 MoE 层。它使用Top-2 路由，每个 token 激活 2 个专家。

关键数据

总参数量：46.7B（8 个专家 x 约 5.6B FFN + 共享注意力参数）
激活参数量：约 13B（每个 token 2 个专家的 FFN + 共享参数）
专家函数：SwiGLU FFN
注意力：Grouped Query Attention（GQA）
上下文长度：32K token
应用 Sliding Window Attention

Top-2 路由公式

Mixtral 的 MoE 层输出计算如下。

y = \sum_{i \in \text{Top2}(g(x))} g(x)_i \cdot \text{SwiGLU}_i(x)

门控函数 $g(x)$ 对输入 $x$ 计算 Softmax 概率分布，并选出前 2 名专家。选中的两个专家的门控权重会重新归一化（renormalization），使其和为 1。

SwiGLU 专家网络

每个专家都是使用 SwiGLU 激活函数的 FFN。

\text{SwiGLU}(x) = (\text{Swish}(xW_1) \odot xV) W_2

PyTorch 实现：Mixtral MoE Block

class MixtralMoEBlock(nn.Module):
    """Mixtral-style MoE block with Top-2 SwiGLU experts."""

    def __init__(self, hidden_dim: int, ffn_dim: int, num_experts: int = 8):
        super().__init__()
        self.num_experts = num_experts
        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
        self.experts = nn.ModuleList([
            SwiGLUExpert(hidden_dim, ffn_dim) for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor):
        # x: (batch, seq_len, hidden_dim)
        gate_logits = self.gate(x)  # (batch, seq, num_experts)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-2 selection
        top2_probs, top2_indices = gate_probs.topk(2, dim=-1)
        # Renormalize gates to sum to 1
        top2_probs = top2_probs / top2_probs.sum(dim=-1, keepdim=True)

        # Compute expert outputs and combine
        batch, seq, dim = x.shape
        output = torch.zeros_like(x)
        for k in range(2):
            expert_idx = top2_indices[:, :, k]  # (batch, seq)
            gate_val = top2_probs[:, :, k].unsqueeze(-1)  # (batch, seq, 1)
            for i in range(self.num_experts):
                mask = (expert_idx == i)
                if mask.any():
                    expert_input = x[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += gate_val[mask].squeeze(-1).unsqueeze(-1) * expert_output

        return output


class SwiGLUExpert(nn.Module):
    """SwiGLU Feed-Forward Network used as expert."""

    def __init__(self, hidden_dim: int, ffn_dim: int):
        super().__init__()
        self.w1 = nn.Linear(hidden_dim, ffn_dim, bias=False)
        self.v = nn.Linear(hidden_dim, ffn_dim, bias=False)
        self.w2 = nn.Linear(ffn_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor):
        return self.w2(F.silu(self.w1(x)) * self.v(x))

DeepSeek-V2/V3 的创新：DeepSeekMoE

Fine-Grained Expert 划分

DeepSeek-V2（2024）采取了与既有 MoE 根本不同的方法。核心思路是Fine-Grained Expert Segmentation，把专家切分成更小、数量更多的单元。

把原本的 $N$ 个专家增加到 $mN$ 个，同时把每个专家的隐藏维度缩小到 $1/m$ 。与此同时，激活的专家数量也从 $K$ 按比例增加到 $mK$ ，在保持每个 token 总运算量不变的前提下，实现更细粒度的专家组合。

DeepSeek-V3 架构

DeepSeek-V3（2024 年 12 月）具备以下核心配置。

总参数量：671B
激活参数量：37B（每个 token）
路由专家：256 个（每层）
共享专家：1 个（每层，始终激活）
激活路由专家数：8 个（每个 token）
注意力：Multi-head Latent Attention（MLA）

无辅助损失的负载均衡

DeepSeek-V3 最具创新性的贡献之一，是无辅助损失的负载均衡策略。既有的 MoE 模型为了实现负载均衡会使用辅助损失（auxiliary loss），但这个辅助损失的系数很难恰当设定，取值过大又会损害模型性能。

DeepSeek-V3 转而为每个专家添加一个偏置项 $b_i$ ，仅用于路由决策。

i^* = \text{TopK}(s(x)_i + b_i)

g(x)_i = \frac{s(x)_i}{\sum_{j \in \text{TopK}} s(x)_j}

偏置项 $b_i$ 只影响路由决策，不参与实际门控权重的计算。过载专家的 $b_i$ 会被调低，欠载专家的 $b_i$ 会被调高，以此在不污染损失函数的前提下实现负载均衡。

Device-Limited 路由

为了限制通信成本，DeepSeek-V3 规定每个 token 最多只能发往 $M$ 个节点。方法是根据分布在各节点上的专家的亲和度分数选出前 $M$ 个节点，再只在这些节点内部的专家之间执行 Top-K 路由。

路由策略对比

Top-1 Routing（Switch Transformer）

每个 token 只精确激活 1 个专家。通信成本最低、实现最简单，但依赖单一专家，表现力可能受限。

Top-2 Routing（Mixtral、GShard）

每个 token 激活 2 个专家并加权求和。相比 Top-1 能获得更丰富的表达，但通信成本翻倍。

Expert Choice Routing（Zhou et al., 2022）

与传统方式相反，由专家选择 token。由于每个专家选择固定数量的 token，负载均衡自动得到保证。不过，同一个 token 可能被 0 个或多个专家选中，具有非确定性。

Soft MoE（Puigcerver et al., 2023）

不采用离散路由，而是把所有 token 的加权组合传给每个专家。完全可微、不会丢弃 token，但由于每个专家都会处理全部 token 的信息，并非真正意义上的稀疏。

PyTorch 实现：Expert Choice Routing

class ExpertChoiceRouter(nn.Module):
    """Expert Choice Routing: experts select tokens."""

    def __init__(self, hidden_dim: int, num_experts: int, capacity_factor: float = 1.0):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.router = nn.Linear(hidden_dim, num_experts, bias=False)

    def forward(self, x: torch.Tensor):
        # x: (num_tokens, hidden_dim)
        num_tokens = x.shape[0]
        capacity = int(self.capacity_factor * num_tokens / self.num_experts)

        # Compute affinity scores
        scores = self.router(x)  # (num_tokens, num_experts)
        scores = F.softmax(scores, dim=0)  # softmax over tokens (not experts)

        # Each expert selects top-capacity tokens
        # Transpose: (num_experts, num_tokens)
        expert_scores = scores.t()

        # Top-capacity selection per expert
        top_scores, top_indices = expert_scores.topk(capacity, dim=-1)
        # top_scores: (num_experts, capacity)
        # top_indices: (num_experts, capacity)

        return top_scores, top_indices

训练稳定性技巧

Load Balancing Loss

Switch Transformer 提出的负载均衡损失定义如下。

\mathcal{L}_{balance} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i

其中 $N$ 是专家数量， $f_i$ 是路由到专家 $i$ 的 token 比例， $P_i$ 是路由器分配给专家 $i$ 的概率的平均值。系数 $\alpha$ 通常设置在 0.01~0.1 之间。

在理想的均匀分布下， $f_i = P_i = 1/N$ ，损失即为 $\alpha$ ；不均衡程度越大，损失越高。

Router Z-Loss

ST-MoE（2022）提出的 Router Z-Loss，通过限制路由器 logit 的大小来提升训练稳定性。

\mathcal{L}_{z} = \frac{1}{B} \sum_{x \in B} \left( \log \sum_{i=1}^{N} e^{z_i(x)} \right)^2

其中 $z_i(x)$ 是路由器的 logit。这一损失能防止 logit 过度增大，从而缓解路由决策的不稳定性与收敛问题。

PyTorch 实现：Load Balancing + Z-Loss

def compute_moe_auxiliary_losses(
    router_logits: torch.Tensor,
    expert_indices: torch.Tensor,
    num_experts: int,
    alpha_balance: float = 0.01,
    alpha_z: float = 0.001
):
    """Compute load balancing loss and router z-loss.

    Args:
        router_logits: Raw router logits (batch*seq, num_experts)
        expert_indices: Selected expert indices (batch*seq, top_k)
        num_experts: Total number of experts
        alpha_balance: Weight for load balancing loss
        alpha_z: Weight for router z-loss
    """
    num_tokens = router_logits.shape[0]
    router_probs = F.softmax(router_logits, dim=-1)

    # --- Load Balancing Loss ---
    # f_i: fraction of tokens routed to expert i
    expert_mask = F.one_hot(expert_indices, num_experts).float()
    if expert_mask.dim() == 3:
        expert_mask = expert_mask.sum(dim=1)  # sum over top_k
    expert_mask = (expert_mask > 0).float()
    f = expert_mask.mean(dim=0)  # (num_experts,)

    # P_i: mean router probability for expert i
    P = router_probs.mean(dim=0)  # (num_experts,)

    balance_loss = alpha_balance * num_experts * (f * P).sum()

    # --- Router Z-Loss ---
    log_z = torch.logsumexp(router_logits, dim=-1)  # (num_tokens,)
    z_loss = alpha_z * (log_z ** 2).mean()

    return balance_loss + z_loss

推理优化

Expert Offloading

MoE 模型由于总参数量庞大，可能难以把全部专家都放进 GPU 显存。Expert Offloading 是一种把当前未激活的专家保存在 CPU 内存或磁盘上、仅在需要时才加载到 GPU 的技巧。

主要技巧如下。

LRU Cache：把最近使用过的专家缓存在 GPU 上
Predictive Prefetch：异步预先加载下一层将会用到的专家
Speculative Decoding + Offloading：与投机解码结合，隐藏卸载带来的延迟

量化（Quantization）

MoE 模型的量化与 Dense 模型类似，但由于各专家的权重分布可能不同，需要额外考虑。

GPTQ/AWQ：可以为每个专家独立设置量化配置
Mixed Precision：常用专家采用高精度，罕用专家采用低精度
MiLo（2025）：为极度量化的 MoE 添加 Low-Rank 补偿器以恢复精度

Expert Parallelism

在 MoE 推理中，Expert Parallelism 是把每个专家部署到不同 GPU 上并行处理的策略。通过 All-to-All 通信把 token 发送到对应专家所在的 GPU，处理完成后再收集回来。

主要 MoE 模型对比

模型	发布	总参数量	激活参数量	专家数	路由	专家类型	特点
Sparsely-Gated MoE	2017	137B	-	4096	Top-K	MLP	最初的大规模 Sparse MoE
Switch Transformer	2021	1.6T	-	2048	Top-1	FFN	简化路由，基于 T5
GLaM	2022	1.2T	97B	64	Top-2	FFN	相比 GPT-3 仅用 1/3 能耗
ST-MoE	2022	269B	-	32	Top-2	FFN	Z-Loss，注重稳定性
Expert Choice	2022	-	-	-	Expert Choice	FFN	由专家选择 token
Mixtral 8x7B	2023	46.7B	13B	8	Top-2	SwiGLU	开源，GQA
DeepSeek-V2	2024	236B	21B	160	Top-6	Fine-Grained	Auxiliary-Loss-Free
DeepSeek-V3	2024	671B	37B	256+1	Top-8	Fine-Grained	MLA + 共享专家
Llama 4 Scout	2025	109B	17B	16	Top-1	-	Meta 首个 MoE

综合实现示例：自定义 MoE Transformer Block

以下是把注意力层与 MoE FFN 结合在一起的完整 Transformer block 实现。

class MoETransformerBlock(nn.Module):
    """Complete Transformer block with MoE FFN layer."""

    def __init__(
        self,
        hidden_dim: int = 768,
        num_heads: int = 12,
        ffn_dim: int = 3072,
        num_experts: int = 8,
        top_k: int = 2,
        capacity_factor: float = 1.25,
        dropout: float = 0.1
    ):
        super().__init__()
        # Multi-Head Attention
        self.attn_norm = nn.LayerNorm(hidden_dim)
        self.attention = nn.MultiheadAttention(
            hidden_dim, num_heads, dropout=dropout, batch_first=True
        )

        # MoE FFN
        self.ffn_norm = nn.LayerNorm(hidden_dim)
        self.router = nn.Linear(hidden_dim, num_experts, bias=False)
        self.experts = nn.ModuleList([
            SwiGLUExpert(hidden_dim, ffn_dim)
            for _ in range(num_experts)
        ])
        self.top_k = top_k
        self.num_experts = num_experts
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask=None):
        # Pre-norm Attention
        residual = x
        x_norm = self.attn_norm(x)
        attn_out, _ = self.attention(x_norm, x_norm, x_norm, attn_mask=mask)
        x = residual + self.dropout(attn_out)

        # Pre-norm MoE FFN
        residual = x
        x_norm = self.ffn_norm(x)
        moe_out, aux_loss = self._moe_forward(x_norm)
        x = residual + self.dropout(moe_out)

        return x, aux_loss

    def _moe_forward(self, x: torch.Tensor):
        B, S, D = x.shape
        x_flat = x.view(-1, D)

        # Router
        logits = self.router(x_flat)
        probs = F.softmax(logits, dim=-1)
        top_k_probs, top_k_idx = probs.topk(self.top_k, dim=-1)
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # Dispatch and combine
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            for i in range(self.num_experts):
                mask = (top_k_idx[:, k] == i)
                if mask.any():
                    expert_out = self.experts[i](x_flat[mask])
                    output[mask] += top_k_probs[mask, k].unsqueeze(-1) * expert_out

        # Auxiliary loss
        aux_loss = compute_moe_auxiliary_losses(
            logits, top_k_idx, self.num_experts
        )

        return output.view(B, S, D), aux_loss

结论与未来展望

MoE 架构已经确立为能够同时实现「模型容量扩展」与「计算效率」这两个目标的最实用方法。从 Switch Transformer 的 Top-1 简化出发，Mixtral 8x7B 把 MoE 普及到开源生态，DeepSeek-V3 又以 Fine-Grained Expert 和 Auxiliary-Loss-Free 策略树立了新的标准。

未来的研究方向包括：

动态专家激活：根据输入难度调整激活专家数量的自适应路由
训练-推理一致性：让训练时的路由模式在推理时也能保持一致的技巧
专家特化分析：关于各个专家分别专精于哪些知识/功能的可解释性研究
面向边缘设备的 MoE：面向移动端与边缘环境的轻量化 MoE 设计

参考资料

Shazeer, N., et al. "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer." ICLR 2017.
Fedus, W., et al. "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity." JMLR 2022.
Jiang, A.Q., et al. "Mixtral of Experts." arXiv:2401.04088, 2024.
DeepSeek-AI. "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model." arXiv:2405.04434, 2024.
DeepSeek-AI. "DeepSeek-V3 Technical Report." arXiv:2412.19437, 2024.
Zoph, B., et al. "ST-MoE: Designing Stable and Transferable Sparse Expert Models." arXiv:2202.08906, 2022.
Zhou, Y., et al. "Mixture-of-Experts with Expert Choice Routing." NeurIPS 2022.
Puigcerver, J., et al. "From Sparse to Soft Mixtures of Experts." ICLR 2024.
Du, N., et al. "GLaM: Efficient Scaling of Language Models with Mixture-of-Experts." ICML 2022.
Jacobs, R.A., et al. "Adaptive Mixtures of Local Experts." Neural Computation, 1991.