对 Mixture of Experts (MoE) 架构进行了系统的学习和梳理。MoE 作为当前大语言模型突破算力瓶颈的核心技术,其背后的稀疏激活和条件计算思想非常优雅。在此做一份核心要点的学习笔记。
1. 为什么需要 MoE?
模块化模型和集成学习通过组合多个小型网络来形成一个更强的预测器,其核心思想可以表示为:
在现代大模型中,MoE 架构由多个专家网络(Expert networks)和一个门控/路由网络(Gating/Routing network)组成。MoE架构带来三大核心优势:
节省计算量:
标准模型计算量:
- :模型维度
- :FFN隐藏层维度
- :序列长度
MoE模型计算量:
更好的扩展性:使得MoE能在相同计算预算下训练更大规模的模型
专家专业化:不同专家会自动学习处理不同类型的数据
2. MoE的架构和变体
标准的 MoE 层通常会替换 Transformer Block 中的 Dense FFN 层。它包含:
- 专家网络:处理特定 Token
- 门控网络 (Router):通过 计算 Token 分配的概率分布
- 稀疏激活 (Top-k):只激活得分最高的 k 个专家,最终输出为被选中专家计算结果的加权和:
DeepSeek-MoE的三大创新:
细粒度专家 (Fine-grained Expert):将大专家拆分为更小的细粒度专家,总参数量不变,但提升了专家的特化程度,并缓解了路由歧义
共享专家 (Shared Expert):将通用知识解耦到总是被激活的共享专家中,减少冗余
专家选择范式(Expert Choice):颠覆了传统的 Token 选择 Expert(按行 Top-k),改为 Expert 选择 Token(按列 Top-k),天然保证了负载均衡
3. 专家并行与负载均衡
为了尽可能最大化利用内存,MoE 引入专家并行 (Expert Parallelism, EP) ,但这也带来了巨大的通信开销,当前主流 MoE 大模型的训练中,平均约有 47% 的时间花费在 EP 通信上。
此外,由于反馈循环中优化越好的专家越容易被 Router 选中,即“富者更富”,少数专家会变得过载,而其他专家处于饥饿状态,导致 OOM 风险和参数利用率低下,即负载不均衡 (Load Imbalance)。
解决负载不均衡的主流方案包括:
Auxiliary Loss:引入额外的正则化损失函数 ,惩罚偏离均匀分布的路由结果
Auxiliary-Loss-Free:由于显式的辅助损失可能与主任务损失冲突,该方法通过为每个专家设置可学习的偏置项来主动调整路由概率
Z-Loss:惩罚 Router logits 中过大的绝对值,防止梯度消失,鼓励模型探索 。
4. 路由策略与训练稳定性
由于 Top-k 操作具有不可导和不连续的特性,MoE 在训练早期极易出现 Gate-Expert Gradient Mismatch 的鸡生蛋问题。为了实现端到端训练和探索稳定性,研究者提出了多种可微路由和动态分配策略:
使用 Gumbel-Softmax 注入噪声 ,结合 Straight-Through Estimator (STE) 允许梯度回传
Top-p 路由与 ReLU 路由:不再固定 k 值,而是根据 Token 的难度动态决定激活专家的数量
引入 Auxiliary Experts:在路由中加入 Zero 或 Copy 这种计算量极小的专家单元,在增加适度成本的情况下提升鲁棒性