MoE:用稀疏激活突破密集模型的效率瓶颈

混合专家模型(Mixture of Experts)通过稀疏激活机制,让模型总参数量巨大但每次推理只激活一小部分。DeepSeek-V3的671B参数中每次只激活37B,Llama 4的500B+参数只激活30B——这就是MoE的效率魔法。

MoE的数学基础

标准MoE层

给定输入 $x$,MoE层的计算过程:

$$y = \sum_{i=1}^{N} g_i(x) \cdot f_i(x)$$

其中 $f_i$ 是第 $i$ 个专家(一个FFN),$g_i$ 是路由函数:

$$g_i(x) = \text{softmax}(W_g \cdot x)_i$$

Top-K路由只保留 $K$ 个最大的 $g_i$,其余置零:

$$g_i(x) = \begin{cases} \frac{\exp(W_g \cdot x)i}{\sum{j \in \text{TopK}} \exp(W_g \cdot x)_j} & i \in \text{TopK} \ 0 & \text{otherwise} \end{cases}$$

专家容量因子

为防止所有token路由到同一专家(路由崩塌),每个专家有一个容量上限:

$$\text{capacity} = \frac{T \cdot K}{N} \times \text{capacity_factor}$$

其中 $T$ 是token数,$K$ 是Top-K,$N$ 是专家数。超出容量的token被丢弃或传递到下一层。

路由机制设计

标准Top-K路由

最简单的路由方式。每个token独立选择Top-K个专家:

class TopKRouter(nn.Module):
    def __init__(self, dim, n_experts, top_k=2):
        self.gate = nn.Linear(dim, n_experts)
        self.top_k = top_k
    
    def forward(self, x):
        # x: [batch * seq_len, dim]
        logits = self.gate(x)
        scores = F.softmax(logits, dim=-1)
        
        # Top-K选择
        topk_scores, topk_indices = scores.topk(self.top_k, dim=-1)
        topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True)
        
        return topk_scores, topk_indices

专家选择路由

DeepSeek-V3使用了一种"专家选择"的反向路由:每个专家选择Top-K个token,而非token选择专家。这平衡了专家负载:

class ExpertChoiceRouter(nn.Module):
    def forward(self, x):
        logits = self.gate(x)  # [T, N]
        # 转置:从token角度变为专家角度
        logits_T = logits.T  # [N, T]
        
        # 每个专家选择Top-K个token
        capacity = T * K // N
        scores, indices = logits_T.topk(capacity, dim=-1)
        return scores, indices

共享专家机制

DeepSeek-V2引入了共享专家:一组始终激活的专家处理通用信息,路由专家处理特定信息:

class MoELayer(nn.Module):
    def __init__(self, dim, n_shared=2, n_routed=256, top_k=8):
        self.shared_experts = nn.ModuleList([FFN(dim) for _ in range(n_shared)])
        self.routed_experts = nn.ModuleList([FFN(dim) for _ in range(n_routed)])
        self.router = TopKRouter(dim, n_routed, top_k)
    
    def forward(self, x):
        # 共享专家:始终激活
        shared_output = sum(expert(x) for expert in self.shared_experts)
        
        # 路由专家:稀疏激活
        scores, indices = self.router(x)
        routed_output = torch.zeros_like(x)
        for i, expert in enumerate(self.routed_experts):
            mask = (indices == i).any(dim=-1)
            if mask.any():
                routed_output[mask] += expert(x[mask]) * scores[mask]
        
        return shared_output + routed_output

共享专家的优势:减少了路由专家的学习压力,通用知识不需要被冗余存储在多个专家中。

训练挑战与解决方案

负载均衡

路由崩塌是MoE训练中最常见的问题:大部分token被路由到少数专家,其余专家得不到训练。

辅助损失法:添加一个正则化项鼓励均衡分布:

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

其中 $f_i$ 是路由到专家 $i$ 的token比例,$P_i$ 是路由到专家 $i$ 的平均概率。

噪声注入:在路由前添加噪声增加探索:

logits = self.gate(x) + torch.randn_like(logits) * noise_scale

通信开销

分布式训练中,All-to-All通信是MoE的主要瓶颈。当专家分布在不同GPU上时,token需要跨GPU传输:

GPU0: token A,B,C → 专家在GPU1,2,3 → All-to-All → 计算完毕 → All-to-All回来

优化策略:

  1. 专家并行:将不同专家放在不同GPU上
  2. token重排序:按目标专家分组,减少通信量
  3. 通信重叠:将通信与计算重叠执行

知识遗忘与专家特化

MoE训练中常见的问题是专家没有发展出明确的特化方向。通过分析训练后专家的路由模式,可以判断专家是否真正"专业化":

def analyze_expert_specialization(model, dataloader):
    expert_domains = defaultdict(list)
    for batch in dataloader:
        _, indices = model.router(batch)
        for token, idx in zip(batch, indices):
            for expert_id in idx:
                expert_domains[expert_id].append(token.domain_label)
    
    # 计算每个专家的领域分布熵
    for expert_id, domains in expert_domains.items():
        entropy = compute_entropy(Counter(domains))
        print(f"专家 {expert_id}: 熵={entropy:.2f}")
        # 低熵 = 高特化度

推理优化

专家缓存

MoE推理时只需要激活的专家,可以利用这一特点优化显存:

class CachedMoE:
    def __init__(self, experts, cache_size=8):
        self.experts = experts  # 所有专家存储在CPU/SSD
        self.cache = {}  # GPU上的专家缓存
        self.cache_size = cache_size
    
    def forward(self, x, expert_indices):
        outputs = []
        for i, idx in enumerate(expert_indices):
            if idx not in self.cache:
                self._load_to_cache(idx)
            expert = self.cache[idx]
            outputs.append(expert(x[i:i+1]))
        return torch.cat(outputs)
    
    def _load_to_cache(self, idx):
        if len(self.cache) >= self.cache_size:
            self.cache.pop(next(iter(self.cache)))  # LRU
        self.cache[idx] = self.experts[idx].to('cuda')

专家合并

对于高度相似的专家,可以在推理时合并以减少计算:

def merge_similar_experts(experts, threshold=0.95):
    merged = []
    for expert in experts:
        merged_flag = False
        for m in merged:
            similarity = compute_weight_similarity(expert, m)
            if similarity > threshold:
                # 加权合并
                m.weight = (m.weight + expert.weight) / 2
                merged_flag = True
                break
        if not merged_flag:
            merged.append(expert)
    return merged

MoE vs Dense模型对比

以相近推理成本对比:

维度Dense 14BMoE 50B/14B激活
推理速度基准相近
总知识量14B容量接近50B容量
训练成本基准约1.5倍
显存占用28GB(FP16)100GB(FP16)
MMLU72%79%

MoE在相近推理速度下提供了显著更好的模型质量,代价是更大的显存需求。

结语

MoE已成为超大模型的主流架构选择。DeepSeek、Llama、Qwen的最新旗舰模型都采用了MoE。随着路由算法的改进和分布式训练框架的优化,MoE的工程门槛正在降低。未来,从dense到MoE的迁移将像从RNN到Transformer的迁移一样成为行业共识。