Skip to content

MoE 混合专家架构深度解析

MoE(Mixture of Experts)通过稀疏激活实现模型容量的扩展而不增加推理计算量,是大模型 scaling 的重要方向。

MoE 架构

核心机制

MoE 层将传统 FFN 替换为多个专家 FFN,通过门控网络选择激活的专家:

python
class MoELayer(nn.Module):
    def __init__(self, num_experts, top_k=2):
        self.experts = nn.ModuleList([FFN() for _ in range(num_experts)])
        self.gate = nn.Linear(d_model, num_experts)
        self.top_k = top_k
    
    def forward(self, x):
        gate_scores = F.softmax(self.gate(x), dim=-1)
        topk_scores, topk_indices = gate_scores.topk(self.top_k)
        
        output = torch.zeros_like(x)
        for i, expert in enumerate(self.experts):
            mask = (topk_indices == i).any(dim=-1)
            if mask.any():
                output[mask] += expert(x[mask]) * topk_scores[mask]
        return output

负载均衡

MoE 的关键挑战是负载不均衡——门控网络倾向于选择少数专家。辅助损失是常用的均衡策略。

DeepSeek-V3 的创新

DeepSeek-V3 采用无辅助损失的负载均衡策略,通过偏置项动态调整专家选择,避免了辅助损失对主任务的干扰。

相关资源

最近更新