MoE 混合专家架构深度解析
MoE(Mixture of Experts)通过稀疏激活实现模型容量的扩展而不增加推理计算量,是大模型 scaling 的重要方向。
核心机制
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 采用无辅助损失的负载均衡策略,通过偏置项动态调整专家选择,避免了辅助损失对主任务的干扰。