Gating Dropout: Communication-efficient Regularization for Sparsely Activated Transformers
This addresses a communication bottleneck for distributed training of large-scale MoE models, offering an incremental improvement with regularization benefits.
The paper tackles the high cross-machine communication cost in sparsely activated transformers like Mixture of Experts (MoE) by proposing Gating Dropout, which reduces communication by allowing tokens to stay local, resulting in faster convergence and improved BLEU scores on multilingual machine translation tasks.
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.