PaperlayerView exact revision on arXiv

Source-linked AI summary

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean

arXiv:1701.06538v1cs.LGcs.CLcs.NEstat.ML

TL;DR

Neural networks need more parameters to absorb large datasets, but activating entire models makes scaling costly and conditional computation has practical challenges. The paper introduces a sparsely-gated Mixture-of-Experts layer that substantially increases capacity while preserving efficiency and improves language-modeling and translation benchmarks.

  • Problem

    Increasing neural-network capacity improves prediction but can cause roughly quadratic training-cost growth, while conditional computation faced significant algorithmic and performance challenges.

  • Method

    The paper introduces a sparsely-gated Mixture-of-Experts layer in which a trainable gating network selects a sparse combination of feed-forward expert networks for each input.

  • Results

    MoE models improved language-modeling and translation benchmarks, including 39% lower test perplexity than a computationally matched baseline on 100 billion words.

  • Takeaways & Limitations

    The work demonstrates major gains from conditional computation in deep networks and suggests it may extend to other domains with sufficiently large training sets.

  • Takeaways & Limitations

    Network bandwidth can bottleneck conditional computation because inter-device communication may constrain example-parameter interactions.

Abstract

from arXiv · show

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.

1 INTRODUCTION AND RELATED WORK

The paper addresses the challenge of increasing neural-network capacity without proportional computation by introducing a sparsely-gated Mixture-of-Experts layer. Applied to language modeling and translation, the approach achieves greater than 1000x capacity improvements with minor computational-efficiency losses.

  • Motivation: Scaling parameter counts improves prediction accuracy when training datasets are sufficiently large, but capacity is limited by computational costs.This pattern is reported across text, images, and audio.
  • Motivation: Conditional computation activates network parts per example to increase capacity without proportional computation, but prior work had not demonstrated massive practical gains.Existing schemes vary in gating sparsity, continuity, stochasticity, and determinism.
  • Challenges: GPU branching inefficiency, reduced conditional batch sizes, network bandwidth bottlenecks, sparsity-loss requirements, and limited dataset scale impede practical conditional computation.GPU clusters can have computational power thousands of times greater than aggregate inter-device network bandwidth.
  • Contribution: 1000x improvements in model capacity are achieved with minor computational-efficiency losses and significantly better results on public language-modeling and translation datasets.The authors present this as addressing the identified challenges and realizing conditional computation’s promise.
  • Contribution: The Sparsely-Gated Mixture-of-Experts layer combines feed-forward experts with a trainable gating network that selects a sparse expert combination for each input.All components are trained jointly by back-propagation.
  • Related work and novelty: Unlike stacked MoEs with separate gating decisions, the convolutional MoE makes potentially different sparse gating decisions at each text position and demonstrates massive capacity scaling.The paper focuses on language modeling and machine translation, inserting the MoE between stacked LSTM layers.

2 THE STRUCTURE OF THE MIXTURE-OF-EXPERTS LAYER

The MoE layer combines a sparse gating network with multiple expert networks, activating only a small subset of experts for each input. Noisy Top-K gating enforces this sparsity, while hierarchical MoE can reduce branching when the expert count is very large.

  • MoE layer structure: An MoE layer contains n expert networks with separate parameters and a gating network whose output is a sparse n-dimensional vector.Experts accept identically sized inputs and produce identically sized outputs in the paper’s initial investigations.
  • Computational sparsity: Sparsity avoids computing experts whose gate value is zero, allowing up to thousands of experts while evaluating only a handful per example.A two-level hierarchical MoE can further reduce the branching factor by having a primary gate select secondary mixture-of-experts.
  • Gating mechanism: Noisy Top-K gating adds tunable Gaussian noise before Softmax, retains only the top k values, and sets all remaining gates to zero.The zero gates result from setting non-top-k values to −∞ before Softmax, and the sparsity is intended to save computation.
  • Gating-network training: The gating network is trained by back-propagation with the rest of the model, with nonzero derivatives for the top k experts when k > 1.Gradients also back-propagate through the gating network to its inputs.

3 ADDRESSING PERFORMANCE CHALLENGES

The section addresses MoE performance bottlenecks caused by shrinking per-expert batches, limited activation memory, and distributed communication. It proposes enlarging effective batches through synchronized parallelism and convolutionality while scaling experts with cluster size.

  • Batch-size challenges: Large batch sizes amortize parameter-load and update overhead, but routing k of n experts across b examples gives each expert only about kb/n examples.As n grows, this shrinking batch makes naive MoE implementations increasingly inefficient.
  • Mixing Data Parallelism and Model Parallelism: Synchronous mixing of data and model parallelism combines batches across d devices, increasing each expert’s batch size by a factor of d.Each device processes a batch of size b, while the distributed MoE combines the batches for expert computation.
  • Mixing Data Parallelism and Model Parallelism: Increasing the number of devices proportionally increases experts and parameters while keeping per-device memory, bandwidth requirements, and step times constant.The total batch size increases so the batch size per expert remains constant.
  • Taking Advantage of Convolutionality: Applying the same MoE across all time steps after the previous layer finishes increases its input batch by the number of unrolled time steps.This convolutional arrangement is unavailable to recurrent MoEs whose timestep inputs depend on preceding MoE outputs.

4 BALANCING EXPERT UTILIZATION

The gating network can collapse onto a few experts because specialization reinforces selection imbalance. The method addresses this with losses balancing both aggregate gate importance and the number of assigned examples.

  • Problem: Gating can favor the same few experts, creating a self-reinforcing imbalance because favored experts train faster and are selected more often.This behavior resembles a reported local-minimum problem addressed elsewhere with a hard constraint.
  • Importance balancing: The importance loss penalizes variation in experts’ batchwise summed gate values, encouraging equal expert importance.Limportance equals the squared coefficient of variation of importance values multiplied by wimportance.
  • Additional constraints: Per-example sparsity and gate-value diversity require no additional losses because sparsity follows fixed k and specialization naturally diversifies gate values.The authors describe this diversification as a virtuous cycle.
  • Load balancing: Equal importance does not guarantee equal example counts, so Lload balances loads to prevent memory and performance problems on distributed hardware.Experts may otherwise receive few high-weight examples or many low-weight examples.

5 EXPERIMENTS

Experiments show that sparsely gated MoE layers substantially improve language-modeling quality at matched computation while scaling capacity to tens of billions of parameters. They also improve machine-translation performance, including multilingual translation, with modest computational efficiency costs.

  • Language modeling: 24% lower test perplexity was achieved by the 4096-expert model than the computationally matched baseline at roughly 8 million ops/timestep.The comparison used flat MoEs with 4, 32, and 256 experts and hierarchical MoEs with 256, 1024, and 4096 experts, with about 1 million parameters per expert.
  • Language modeling: 6% of the computation was sufficient for the fastest high-capacity MoE model to beat the best published result when training epochs were controlled.The high-capacity models contained about 4 billion parameters and varied their computational budgets.
  • Machine translation: 40.56 and 26.03 BLEU were achieved on WMT’14 En→Fr and En→De, respectively, improving strong baselines by 1.34 and 1.12 BLEU without RL refinement.On the Google Production dataset, the model achieved a 1.01 higher test BLEU score.
  • Multilingual translation: 19% lower dev-set perplexity was achieved by the multilingual MoE model than the multilingual GNMT model, with gains on 11 of 12 language pairs.The improvement reached as much as 5.84 BLEU points, and the MoE model beat monolingual GNMT models on 8 of 12 language pairs.

6 CONCLUSION

The work demonstrates major gains from conditional computation in deep networks by addressing its design and performance challenges through combined algorithmic and engineering solutions.

  • 6 CONCLUSION: The work is the first to demonstrate major wins from conditional computation in deep networks.The authors identify and address conditional computing’s design considerations and challenges through algorithmic and engineering solutions.
  • 6 CONCLUSION: Conditional computation may benefit other domains when sufficiently large training sets are available.The paper’s experiments focused on text, while the conclusion points to broader applications under this data condition.

APPENDICES · A LOAD-BALANCING LOSS · B HIERACHICAL MIXTURE OF EXPERTS

The appendices define a differentiable load-balancing objective for encouraging equal expert utilization and describe a hierarchical MoE that reduces branching when the expert count is large. Experiments show that either balancing loss preserves model quality, while stronger load weighting reduces overload.

  • A LOAD-BALANCING LOSS: A smooth Load(X) estimator enables backpropagation through the otherwise discrete number of examples assigned to each expert.The estimator is used to encourage experts to receive roughly equal numbers of training examples.
  • A LOAD-BALANCING LOSS: The load loss is the squared coefficient of variation of the load vector multiplied by the hand-tuned factor wload.This objective penalizes unequal expert loads.
  • A LOAD-BALANCING LOSS: Initializing Wg and Wnoise to all zeros yields no signal and some noise, producing approximately equal initial expert loads and helping avoid out-of-memory errors.The soft constraints need time to correct imbalance after initialization.
  • A LOAD-BALANCING LOSS: After 10 epochs, all combinations containing at least one balancing loss had similar model quality, whereas using no loss was much worse.The experiments measured perplexity, coefficients of variation for Importance and Load, and the most-overloaded-to-average load ratio.
  • A LOAD-BALANCING LOSS: Higher wload values reduced the load on the most overloaded expert.This ratio matters for load balancing on distributed hardware.
  • B HIERACHICAL MIXTURE OF EXPERTS: A two-level hierarchical MoE reduces branching by having a primary gate select weighted groups, each containing a secondary mixture-of-experts with its own gate.For groups of a experts with b experts each, the primary and secondary gates jointly weight the selected expert networks.
  • B HIERACHICAL MIXTURE OF EXPERTS: Hierarchical load is defined using primary and secondary gating loads so the loss retains a gradient with respect to the primary gate.The simpler formulation LoadH(X)i,j = Loadi(Xi)j would not provide that gradient.

C 1 BILLION WORD LANGUAGE MODELING BENCHMARK - EXPERIMENTAL DETAILS · C.1 8-MILLION-OPERATIONS-PER-TIMESTEP MODELS · C.2 MORE EXPENSIVE MODELS

The benchmark models combine two 512-dimensional LSTM layers with a sparsely gated MoE layer whose experts expand capacity while maintaining a fixed computational setting. Larger MoE models with 34M and 143M operations per timestep achieved test perplexities of 31.3 and 28.0, with the larger model improving on the best published model by 18%.

  • C 1 BILLION WORD LANGUAGE MODELING BENCHMARK - EXPERIMENTAL DETAILS: The five-layer architecture comprises word embeddings, an LSTM, an MoE layer, a second LSTM, and a softmax layer, with 512 dimensions outside the softmax.Dropout is applied to every layer except the softmax.
  • C.1 8-MILLION-OPERATIONS-PER-TIMESTEP MODELS: Each MoE expert uses a 1024-unit ReLU hidden layer and a 512-unit output layer, totaling 1M parameters per expert.Models use ordinary MoE layers with 4, 32, or 256 experts, or hierarchical layers with 256, 1024, or 4096 experts.
  • C.1 8-MILLION-OPERATIONS-PER-TIMESTEP MODELS: The benchmark compares sparse MoE models against computationally matched nonsparse alternatives, including wide and deep single experts and additional LSTM configurations.MoE-4 uses all four experts and therefore does not employ sparsity.
  • C.1 8-MILLION-OPERATIONS-PER-TIMESTEP MODELS: Training used 16 K40 GPUs, batches totaling roughly 300,000 words, 10 epochs, 27,000 steps, and Adam; most models trained in 12–16 hours.MoE-4 required 18 hours because expert computation ran on only 4 of 16 GPUs.
  • C.1 8-MILLION-OPERATIONS-PER-TIMESTEP MODELS: Evaluation reports holdout-set test perplexity alongside computational budget, parameter count, DropProb, and computational efficiency.The procedure sums over all words, including the end-of-sentence symbol.
  • C.2 MORE EXPENSIVE MODELS: The additional MoE-34M and MoE-143M models use 34M and 143M operations per timestep, respectively, with 1024-dimensional embeddings and MoE inputs and outputs.Both retain an MoE layer between two LSTM layers; MoE-34M uses 1024-unit LSTMs, while MoE-143M uses larger LSTM layers.
  • C.2 MORE EXPENSIVE MODELS: 31.3 and 28.0 were the test perplexities of MoE-34M and MoE-143M, respectively, demonstrating that additional computation remains useful with a large MoE.After 10 epochs, the larger model had a lower test perplexity by 18% than the best published model, with similar computational budget and training time.

D 100 BILLION WORD GOOGLE NEWS CORPUS - EXPERIMENTAL DETAILS

The experiments vary MoE scale from 32 to 131,072 experts and train models once through approximately 100 billion words on multi-GPU clusters. A 68-billion-parameter MoE achieves 39% lower perplexity than the baseline, while the largest model’s low measured efficiency is attributed to an unscaled batch size.

  • Model Architecture: MoE layers range from 32 ordinary experts to 131072 hierarchical experts, with hierarchical first-level branching factors from 32 to 256.The evaluated hierarchical configurations use 256, 1024, 4096, 16384, 65536, and 131072 experts.
  • Training: Models use approximately 2.5 million-word batches and train once-through over about 100 billion words on clusters of 32, 64, or 128 Tesla K40 GPUs.The last two models use 64- and 128-GPU clusters to fit all parameters in memory.
  • Training: Memory optimizations fit up to 1 billion parameters per GPU by recomputing expert hidden-layer activations during backpropagation and reducing optimizer auxiliary storage.The optimizer modifications are described as a second memory-saving strategy for expert parameters.
  • Results: 39% lower perplexity is achieved by the 68-billion-parameter MoE than by the baseline after 100 billion training words.Perplexity is evaluated on a holdout dataset and reported in Table 8.
  • Results: 0.30 TFLOPS/GPU is the largest model’s measured computational efficiency, likely because its training batch size was not increased proportionally with GPU count.The comparison also includes a computationally matched baseline and an unpruned 5-gram model.

E MACHINE TRANSLATION - EXPERIMENTAL DETAILS · F STRICTLY BALANCED GATING · G ATTENTION FUNCTION

The machine-translation experiments vary MoE scale and architecture, showing improving test perplexity with more experts and specialized expert behavior. They also introduce strictly balanced batchwise gating and a computationally efficient attention variant with little quality difference.

  • E MACHINE TRANSLATION - EXPERIMENTAL DETAILS: MoE layers are inserted in the encoder and decoder of a reduced-depth GNMT-based architecture.The encoder and decoder use 3 and 2 LSTM layers, respectively, with MoE layers between encoder layers 2–3 and decoder layers 1–2.
  • E MACHINE TRANSLATION - EXPERIMENTAL DETAILS: The multilingual model uses non-hierarchical 512-expert MoEs with k = 2 and expert hidden layers of size 8192.It uses noisy-top-k gating and doubles computation in the MoE layers, increasing the entire model’s computational budget.
  • E MACHINE TRANSLATION - EXPERIMENTAL DETAILS: 2048 experts continued to improve the model’s test perplexity as training progressed.Experiments compared no MoE, 32-expert flat MoE, and 512- or 2048-expert hierarchical MoE models; each input used exactly 4 experts per MoE layer.
  • E MACHINE TRANSLATION - EXPERIMENTAL DETAILS: All models except the no-expert model used the same 85M ops/timestep computational budget.Figure 4 reports perplexity on WMT’14 En→Fr and Google Production En→Fr as a function of processed words; beginning-of-training differences arose from batch sizes.
  • E MACHINE TRANSLATION - EXPERIMENTAL DETAILS: Experts became highly specialized by syntax and/or semantics in the 2048-expert WMT’14 En→Fr encoder MoE.One expert was used when the indefinite article “a” introduced the direct object in a verb phrase indicating importance or leadership.
  • F STRICTLY BALANCED GATING: Batchwise masking keeps the top m values per expert across a training batch, forcing every expert to receive the same number of examples.Here m = k|X|/n, so each example is sent to an average of k experts.
  • F STRICTLY BALANCED GATING: Inference uses learned per-expert thresholds to approximate the batchwise mask when large batches are unavailable.An additional training loss minimizes disagreement between the batchwise mask and threshold mask.
  • G ATTENTION FUNCTION: The modified attention function factorizes source and target transformations, enabling simultaneous computation across multiple source and target time steps.It supports optimized matrix multiplications, and the authors found little quality difference from the GNMT attention function.
Loading 1701.06538v1…