论文阅读笔记——Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsit

Switch Transformers 论文
Switch Transformer 在 Mixtral of Experts 的基础上，简化路由机制，仅选择一个专家，采用动态分配 token，结合数据并行+模型并行+专家并行，预训练速度相比 T5 提升 4-7 倍，支持万亿参数规模。

• 专家: 分布在不同的 device 上, 每个专家是一个 FFN 网络, 有独立的权重；
• 专家容量: 每个专家处理的 batch 大小, 公式为 $\frac{\text{tokens per batch}}{\text{num experts}}\cdot \text{capacity factor}$
• 容量系数: 计算专家容量时乘上系数, 可以为专家多分配一些 buffer 来改善 token 溢出的情况
  如果专家容量系数设置过小,，太多的 token 被路由到一个专家上,，会造成溢出的 token 走了残差分支直接传给下一层； 如果专家容量系数设置过大,，会造成内存和计算的浪费
  $$\text{expert capacity}=\left(\frac{\text{tokens per batch}}{\text{number of experts}}\right)\times \text{capacity factor}.$$
  对于 Switch Transformer，辅助损失会被加到总损失上，给定 i 从 1 到 N 的 N 个专家，一个 batch $\mathcal{B}$ 中有 T 个 token，损失计算如下：
  l o s s = α ⋅ N ⋅ ∑ i = 1 N f i ⋅ P i f i = 1 T ∑ x ∈ B 1 { a r g m a x p ( x ) = i } P i = 1 T ∑ x ∈ B p i ( x ) \begin{aligned} loss=\alpha · N · \sum_{i=1}^Nf_i·P_i \\f_i =\frac{1}{T}\sum_{x\in\mathcal{B}}1\{argmax\quad p(x)=i\} \\P_i=\frac{1}{T}\sum_{x\in\mathcal{B}}p_i(x) \end{aligned} loss=α⋅N⋅