Reading Notes: “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”
date
Mar 3, 2025
slug
moe
status
Published
tags
NLP
summary
type
Post
Motivation
对于一些复杂的任务，可以通过增大模型容量来取得更好的效果。但是增大模型容量同时也会增大计算量，给实际应用带来瓶颈。
本文提出了 Sparsely-Gated MOE，作为一种 conditional computation 结构，可以用来在增加模型容量的同时不会成比例地增加计算量。
Approach
Sparsely-MOE 层由多个专家网络（experts）和一个门控网络组成：
- 稀疏激活：每个输入只会激活少量专家
- 可训练的门控网络：门控网络学习为每个输入选择适合的专家；门控网络 = Softmax Gating + Noisy Top-k Gating
- 负载均衡：通过特殊的损失函数确保专家的利用率均衡
Addressing Shrinking Batch Problem
由于每个样本会被分配到适合的专家网络，假设一共有 n 个专家，每个样本会被分配到 k 个专家，批量大小为 b，那么每个专家平均得到的批量大小为 ，为了缓解这个问题，本文通过混合数据并行和模型并行：
- 做正常的数据并行，对于非 MOE 层每个设备保存一份权重
- 对于 MOE 层，只保留一份权重，不同专家保存在不同设备上
- 由上，如果设备数量是 d，那么每个专家处理的批量大小平均增加到
Results
- 在大规模机器翻译以及语言建模 benchmark 上进行了评估，取得显著比 SOTA baselines 更好的性能
- 在同一性能水平下，需要的计算量相比于 baseline 显著减少