Scaling Neural Networks with Mixture of Experts

Summary & Methodology Analysis

Neural networks are traditionally limited in capacity by their number of parameters, and conditional computation schemes that activate parts of the network per example face severe algorithmic and performance challenges on modern computing hardware. Standard existing mechanisms train models using conventional deep learning architectures such as stacked Long Short-Term Memory layers and data-parallel training across distributed GPU clusters. To overcome these limits, the authors introduce the Sparsely-Gated Mixture-of-Experts Layer, consisting of multiple feed-forward expert networks and a trainable gating network. They implement Noisy Top-K Gating by adding tunable Gaussian noise to the gating network input and keeping only the top k values to enforce sparsity and select expert combinations per input position. They also combine data parallelism and model parallelism to solve the shrinking batch problem, wherein standard layers and gating networks run data-parallel while expert networks act as model-parallel shards receiving combined batches. Additionally, a soft load-balancing loss and an importance loss based on the coefficient of variation of gate values are added to prevent the gating network from always selecting the same few experts and to ensure balanced expert utilization.

The resulting models demonstrate substantial scaling capabilities. The 4-billion parameter mixture-of-experts model with 4096 experts achieved a 24% lower perplexity on the 1-Billion-Word Language Modeling Benchmark compared to computationally-matched baselines. On the 100 Billion Word Google News Corpus, the 68-billion-parameter mixture-of-experts model achieved a test perplexity 39% lower than the baseline model. The evaluation suite also includes the WMT’14 En -> Fr corpus, the WMT’14 En -> De corpus, the Google Production English to French dataset, GNMT, LSTM, and PBMT.

Despite the strong performance, limitations exist. Test perplexity degrades when using 131072 experts on the 100 billion word corpus, which is possibly a result of excessive sparsity. The paper does not specify other limitations or exact hardware costs outside of the detailed architectural configurations.