MoE架构深度解析:混合专家模型的训练与推理优化
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选择专家。这平衡了专家负载: ...




