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 架构详解