Mixture-of-Experts at Scale (2021) › The Solution

Mixture of Experts architectures addressed these limitations by introducing sparsity into neural network computation. Instead of using all model parameters for every input, MoE models divide the model into multiple expert networks, each a complete neural network capable of processing inputs independently. A learned gating or routing mechanism then selects which experts should handle each input, typically activating only 1-2 experts out of many. This sparse activation pattern means that while the total number of parameters in the model can be very large, the computational cost of each forward pass is much smaller, as only the active experts need to be computed.

The key innovation that made MoE architectures viable at scale in 2021 was the integration of expert networks into transformer layers. Researchers at Google Brain designed expert networks as feed-forward layers within the transformer architecture, replacing the standard dense feed-forward layer with multiple expert feed-forward networks. Each expert was a complete two-layer feed-forward network, and a routing mechanism would select which experts to use for each token or group of tokens. This design maintained the transformer architecture's proven structure while introducing the efficiency benefits of sparse activation.

The routing mechanism is the core component that determines which experts process each input. In the MoE architectures introduced in 2021, routing typically worked at the token level, where each token position in the sequence could be routed to different experts. The routing network, also called the gating network, takes the token representation as input and produces a probability distribution over all available experts. The mechanism computes scores for each expert, applies a softmax to create probabilities, and then selects the top-k experts (typically k=1 or k=2) based on these probabilities.

The mathematical formulation of routing begins with computing expert scores. For a token representation $h$, the gating network computes scores for each expert $i$:

$$g_i(h) = W_i \cdot h + b_i$$

where $W_i$ and $b_i$ are learned parameters for expert $i$. These scores are then normalized using a softmax function to create a probability distribution:

$$P_i(h) = \frac{\exp(g_i(h))}{\sum_{j=1}^{N} \exp(g_j(h))}$$

where $N$ is the total number of experts. The routing mechanism then selects the top-k experts with the highest probabilities. The final output combines the selected experts' outputs weighted by their routing probabilities or as a simple weighted sum, depending on the specific MoE variant.

One of the critical challenges with MoE architectures is ensuring that work is distributed evenly across experts during training. If one expert receives most of the inputs while others remain idle, the model effectively reduces to a smaller dense model, losing the efficiency benefits. Early MoE models often suffered from this collapse, where the routing mechanism would learn to always route to the same few experts, defeating the purpose of having multiple experts.

The 2021 MoE architectures introduced several techniques to address load balancing. The GShard model used capacity constraints that limited how many tokens could be routed to each expert, forcing more even distribution. The Switch Transformer introduced an auxiliary load balancing loss that penalized uneven expert usage. This loss term measured the variance in expert usage across a batch of examples and added it to the main training loss, encouraging the routing mechanism to distribute work more evenly while still maintaining quality routing decisions.

The load balancing objective typically measures how evenly experts are used. One common approach computes the coefficient of variation of expert usage, which penalizes cases where some experts receive many more inputs than others. By including this term in the training loss, the model learns to route inputs to appropriate experts while maintaining roughly equal usage across all experts. Without this balance, routing could collapse onto a few experts and prevent MoE architectures from scaling effectively.

In the MoE architectures from 2021, each expert was typically implemented as a feed-forward neural network within the transformer architecture. A standard transformer layer contains a multi-head self-attention mechanism followed by a feed-forward network. In MoE variants, the feed-forward network was replaced with multiple expert feed-forward networks and a routing mechanism. Each expert was a complete two-layer feed-forward network with the same architecture, typically expanding the input dimension, applying a non-linear activation, and then projecting back to the original dimension.

The expert networks learned distinct specializations during training, even though they started with identical architectures. Through the routing mechanism, different experts would process different types of inputs, leading them to develop specialized knowledge or patterns. Researchers observed that experts would naturally specialize in different domains, languages, or reasoning patterns without explicit supervision. This emergent specialization was one of the most intriguing aspects of MoE architectures. This showed how the routing mechanism could create a form of modular intelligence within a single model.

The computational advantages of MoE architectures became clear when scaling model size. In a dense model with $P$ parameters, every forward pass requires computing with all $P$ parameters. In an MoE model with $N$ experts each of size $E$ parameters, the total parameter count is $N \times E$, but each forward pass only activates $k \times E$ parameters, where $k$ is the number of active experts (typically 1-2). This means that an MoE model could have significantly more total parameters than a dense model while maintaining similar computational cost per forward pass.

For example, a dense model with 175 billion parameters requires computing with all 175 billion parameters for every input. An MoE model with 8 experts each of size 20 billion parameters would have 160 billion total parameters, but only compute with 20-40 billion parameters per input (depending on whether 1 or 2 experts are active). This sparse activation pattern enabled researchers to scale models to sizes that would be impractical with dense architectures, while maintaining reasonable computational costs for training and inference.