Gating Dropout: Communication-efficient Regularization for Sparsely Activated Transformers
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.