[Paper Review] Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini+5 more|arXiv (Cornell University)|Jan 23, 2017

Advanced Neural Network ApplicationsComputer Science268 citations

TL;DR

Introduces a Sparsely-Gated Mixture-of-Experts (MoE) layer with thousands of experts to massively scale model capacity (up to 137B parameters) while maintaining practical compute, demonstrated on language modeling and machine translation. Shows large capacity gains with modest drops in efficiency.

ABSTRACT

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

Motivation & Objective

• Motivate conditional computation to dramatically increase model capacity without a proportional rise in computation.
• Propose and implement a Sparsely-Gated Mixture-of-Experts layer with thousands of experts.
• Evaluate MoE-augmented architectures on language modeling and machine translation benchmarks.
• Address practical training challenges such as batching, bandwidth, and load balancing in large-scale MoE systems.

Proposed method

• Define an MoE layer with multiple expert networks and a trainable gating network that selects a sparse subset of experts for each input.
• Use either softmax or noisy top-k gating to produce sparse gating weights, enabling per-example expert selection.
• Train gating and expert networks jointly by back-propagation, with mechanisms to encourage load balancing and mitigate expert dominance.
• Handle performance challenges via mixing data and model parallelism to increase effective batch sizes for MoE computations.
• Leverage convolutional application of MoE between stacked LSTM layers to allow per-position gating decisions.
• Experiment with architectures embedding an MoE layer between LSTM layers, including very large MoE capacities (up to thousands of experts).

Experimental results

Research questions

• RQ1 How can conditional computation via a sparsely gated MoE layer scale neural network capacity while preserving computational efficiency?
• RQ2 What gating strategies (softmax vs. noisy top-k) and architectural placements yield best performance on language modeling and translation tasks?
• RQ3 Can MoE-based models achieve state-of-the-art results on large-scale language modeling and machine translation benchmarks with realistic compute budgets?
• RQ4 What practical training and deployment challenges arise (batching, bandwidth, load balancing) and how can they be mitigated?

Key findings

• MoE models with up to thousands of experts achieved significantly better results than state-of-the-art at lower computational cost on large language modeling benchmarks.
• A 1B-word language modeling setup showed up to 24% perplexity improvement with large MoE capacity.
• On a 100B-word Google News corpus, MoE models with up to 137B parameters continued to improve perplexity, with performance gains correlating to data size.
• In machine translation, MoE-augmented GNMT-like models achieved BLEU improvements over strong baselines across multiple language pairs, with reduced training time in some configurations.
• Multilingual translation experiments demonstrated substantial improvements over multilingual baselines, with gains in perplexity and BLEU across most language pairs.
• The paper demonstrates the feasibility of training trillion-parameter models by scaling hardware and using MoE-based conditional computation.