[论文解读] Gating Dropout: Communication-efficient Regularization for Sparsely Activated Transformers

Rui Liu, Young Jin Kim+2 more|arXiv (Cornell University)|May 28, 2022

Topic Modeling被引用 11

一句话总结

本文提出Gating Dropout，一种针对稀疏激活Transformer模型（特别是Mixture-of-Experts, MoE模型）的通信高效正则化技术。通过在训练过程中随机跳过将token路由至远程专家的操作，该方法在降低昂贵的跨机器全对全通信成本的同时引入正则化效应，提升泛化能力，从而实现更快的收敛速度，并在多语言翻译任务中实现最高0.6 BLEU的性能提升。

ABSTRACT

Sparsely activated transformers, such as Mixture of Experts (MoE), have received great interest due to their outrageous scaling capability which enables dramatical increases in model size without significant increases in computational cost. To achieve this, MoE models replace the feedforward sub-layer with Mixture-of-Experts sub-layer in transformers and use a gating network to route each token to its assigned experts. Since the common practice for efficient training of such models requires distributing experts and tokens across different machines, this routing strategy often incurs huge cross-machine communication cost because tokens and their assigned experts likely reside in different machines. In this paper, we propose \emph{Gating Dropout}, which allows tokens to ignore the gating network and stay at their local machines, thus reducing the cross-machine communication. Similar to traditional dropout, we also show that Gating Dropout has a regularization effect during training, resulting in improved generalization performance. We validate the effectiveness of Gating Dropout on multilingual machine translation tasks. Our results demonstrate that Gating Dropout improves a state-of-the-art MoE model with faster wall-clock time convergence rates and better BLEU scores for a variety of model sizes and datasets.

研究动机与目标

• 为解决Mixture-of-Experts (MoE)模型在分布式训练中因频繁全对全操作将token路由至远程专家而产生的高跨机器通信成本问题。
• 在不牺牲模型性能的前提下降低通信开销，尤其适用于大规模分布式训练环境。
• 引入一种正则化机制，通过强制在训练期间更多使用本地专家，提升模型泛化能力。
• 提升基于MoE的多语言机器翻译模型的训练效率和收敛速度。

提出的方法

• Gating Dropout在训练过程中随机禁用路由机制，使token保留在本地机器上，而非被发送到指定的远程专家。
• 该方法对门控网络的路由决策应用dropout率p，其中以概率p，token不被路由，而是直接在本地处理。
• 通过随机跳过全对全通信操作，显著减少了跨机器数据重排，降低了通信成本。
• 正则化效应源于专家因无法获得最优路由而被迫学习更具泛化能力的特征，从而提升泛化性能。
• 该方法兼容k=1和k>1的路由策略，可应用于任何MoE-based Transformer架构。
• 该方法仅在训练阶段应用，推理阶段禁用，从而保持模型准确性。

实验结果

研究问题

• RQ1 随机跳过路由操作是否能在不降低模型性能的前提下减少分布式MoE训练中的通信成本？
• RQ2 通过Gating Dropout实现的路由随机化是否能提升MoE模型的泛化能力？
• RQ3 dropout率如何影响MoE训练中通信效率、收敛速度与模型准确率之间的权衡？
• RQ4 Gating Dropout是否能提升多语言翻译基准上的训练吞吐量和收敛速度？
• RQ5 Gating Dropout是否对低资源语言翻译任务具有特殊优势？

主要发现

• 与基线MoE模型相比，Gating Dropout将达到目标BLEU分数的训练时间最多减少58%。
• 该方法提升了泛化性能，在多语言翻译任务中最高实现0.6 BLEU的分数增益。
• 由于减少了全对全通信，吞吐量最高提升16%，尤其在较高dropout率下更为显著。
• 在Gate-Expert-Drop中，最优性能出现在dropout率为0.2时；在Gate-Drop中，最优性能出现在dropout率为0.3时，更高率会导致性能下降。
• 正则化效应在低资源语言翻译任务中尤为有益，带来至少0.3 BLEU点的增益。
• 该方法在多种模型尺寸和数据集上均有效，展现出在收敛性和效率方面的一致性提升。