MoE 核心思想
MoE 将 FFN 层替换为多个”专家”子网络,每个 token 通过路由机制只激活部分专家。
路由机制
Top-K 路由是最常见的策略——每个 token 选择得分最高的 K 个专家:
def router(x, experts, k=2):
scores = softmax(gate(x))
top_k_indices = topk(scores, k)
output = sum(scores[i] * experts[i](x) for i in top_k_indices)
return output
负载均衡
为了避免某些专家”过劳”而其他专家”闲置”,需要引入负载均衡损失。Mixtral 8x7B 是 MoE 架构的成功案例。
对 Transformer 基础不熟悉的读者可以先阅读 Transformer 架构详解。