最近需要做MOE相关的工作,简单分析记录下Hinton团队2017ICLR上发表的论文
论文题目:《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》
论文地址:https://arxiv.org/abs/1701.06538
背景介绍:
随着深度学习的发展,数据规模和模型容量已经是深度学习的关键因素。传统的深度学习模型中,对于每一个样本输入,完整的模型都需要被激活。随着数据和模型的扩大,这样的开销是二次的。因此,引入条件计算的概念,即通过动态激活部分神经网络,从而在增加模型参数量的情况下,而不增加计算量。
但是条件计算提出以来,面临着几个困境:
1.现代的计算机硬件,尤其是GPU,精通计算操作而不擅长分支。因此先前很多的工作每个门控制大chunk网络以减少分支。
2.条件计算会减少batch size。
3.网络带宽是瓶颈。
4.需要设计特定的损失函数,Bengio先前的工作就设计了三种loss,这些loss的设计会影响模型的效果和负载均衡。
5.现有的相关工作都是在小数据集上的小模型实验。
方法:
提出稀疏门控的多专家混合网络,通过一个可学习的门控网络来稀疏地选择专家。
相关工作:
介绍了先前机器学习&深度学习领域的一些专家网络的工作,但是这些网络都是偏上层的专家结合,每个专家网络其实是一个完整的模型。而该论文的工作实际上是将MOE做成一个通用的网络模块,同时作为一种实际的方式来控制模型的容量。
模型结构:
如上公式,模型输出就是通过门控网络G来赋予不同专家E的输出不同权重。文中提到,如果专家的数量过多,也可以构建多层的MOE结构。
这里文中有一段对于MOE结构的理解。在MOE中,专家网络是前馈神经网络类似于参数化的权重矩阵,而对于多个专家激活的情况就对应于一种block-wise的dropout结构。
门控网络:
简单的softmax网络:
有噪的topk网络:
·通过topk选取前k个专家,其余专家的系数为0。稀疏门控的设计能够节省计算量。
·通过加入噪声(其中噪声权重矩阵是可学习的)来控制负载均衡。
解决问题:
batch size的问题:简单来说,假设batch size为b,从n个专家中选取topK个。由于稀疏激活的原因,每个专家接收的样本数是k*b/n << b,因此会造成batch size减小。
同步的数据和模型并行:控制每一张卡上只有一个专家网络,通过门控网络实现数据分发到哪个专家网络。该方法控制每张卡上内存和通信消耗几乎一致,因此只需要增加卡的数量d,就可以等比例的增大batch size。
卷积性:网络的每一个时间步的MOE都相同,如果把LSTM的多时间步拆开,相当于形成一个很大的batch输入,因此也可以增大batch size。
其他一些优化显存的方法增大batch size。
网络带宽问题:通过增大隐层维度或者隐层的数量来提高计算效率。
负载均衡问题:具体而言,模型训练时会趋向于总激活某几个专家,这是模型的“自我强化”导致的不平衡性。先前的一些工作会加入一些硬限制和软限制。
MOE中通过增加两种loss设计,importance loss&load loss。前者定义重要性,表示某个专家训练的样本数量,从而鼓励所有专家都参与训练。后者是为了解决有些专家接收少量大权重的样本,有些专家接收大量小权重的样本的情况。
实验:
实验部分展示了在语言建模以及翻译等任务上MOE的表现。基本来说,MOE在更小计算量的情况下,模型拥有更高的参数量(最高为1370亿参数量),并且在test perplexity上更低。
结论:
MOE确实提供了一个很好的思路来提升模型容量。在理想情况下,只需要增加专家的数量(设备数量)就可以增大模型参数量。但是实际训练操作比较困难,同时文中的专家网络仅是简单的前馈神经网络,其他网络结构还需要近一步探索。
之后会记录一系列的MOE相关论文,希望能够获得启发。