Switch Transformer：训练万亿级参数的语言模型.
- paper：Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
- arXiv：link
在NLP领域，模型的大小(用模型的参数量衡量)通常与模型的性能成正比。GPT-3模型具有$1750$亿参数($175$ billion)，使用稀疏注意力(sparse attention)运算，这类运算很难发挥GPU等硬件的性能。作者提出了Switch Transformer，采用稀疏路由(sparse routing)设计网络结构，但计算上仍然是密集的，充分使用GPU等硬件的性能，可以训练出具有$1.6$万亿参数($1.6$ trillion)的语言模型。
Switch Transformer是按照MoE结构范式设计的，采用特殊的分布式训练设置和一些训练技巧，分别介绍如下。
Mixture of Expert (MoE)
混合专家系统(Mixture of Expert, MoE)是一种神经网络的设计范式，于上世纪$90$年代被提出。MoE是指根据数据的不同产生方式分别训练不同子模型，每个模型被称为一个专家(expert)，对于每次数据输入使用一个门控模块决定使用哪个专家，模型的实际输出为各个专家输出的加权组合。每个专家通常具有不同的网络结构。
在深度学习中，模型对于输入通常重复使用相同的模型参数。将MoE的思想引入，设计一个具有稀疏激活(sparsely-activated)特征的模型。尽管模型本身具有较多的参数量，但对于每次输入只激活其中的一部分，并没有额外的计算量。这个过程是通过动态路由实现的。
Partitioning Strategy
在深度学习中，可能会出现由于模型过大或数据集过大，无法在单张GPU或单个core上训练的情况。选择合适的分布式训练方法能够有效加快训练过程，减少训练难度。模型参数和数据集的分布式训练在MoE范式下有以下几种形式：
- 数据并行(data parallelism)：模型参数在所有core上共享；数据被划分到每个core上；
- 模型并行(model parallelism)：模型参数被划分到每个core上；数据在所有core上共享；
- 数据和模型并行(data and model parallelism)：模型参数被划分到每组core上，而在不同组之间共享；数据被划分到每组core上，而在不同组之间共享；
- 专家和数据并行(expert and data parallelism)：模型参数在每个core上都不同，对应不同的专家；数据被划分到每个core上；
- 专家,模型和数据并行(expert, model and data parallelism)：模型参数被划分到每组core上，而在不同组之间不同，对应不同的专家；数据被划分到每组core上，而在不同组之间共享。
Switch Transformer
将MoE引入Transformer的过程如下。
Transformer的主体部分是由多头自注意力层MHA和前向传播层FFN堆叠组合而成。MHA实现不同token之间的交互，FFN是对每个token进行非线性变换，其输出作为下一层的输入，可以看作其实现了不同层之间的交互。
由于FFN通常具有更多参数，因此将其作为专家。通过设置多个FFN，对于每个输入token计算一个线性得分，将其作为路由选择合适的FFN。对于模型整体，由于引入多个FFN增加的模型的参数量，但每次计算中只使用了其中一个FFN，并没有增加计算量。
作者还提出了一些使训练更稳定的技巧：
- 精度转换。在float16精度下训练使得模型不稳定。作者提出，在计算路由时将数据扩张为float32精度，计算结束后再恢复成float16精度，能够提高模型训练的稳定性。
- 更小的参数初始化。作者提出，将模型参数初始化为原来的$1/10$，能够提高模型训练的稳定性。
- 正则化。作者提出，采用expert dropout，即在专家层采用更大的dropout rate($0.4$,其余层设置为$0.1$)，能够提高模型训练的稳定性。
