Mixture of Experts at Scale (2024) › The Solution

MoE architectures addressed these limitations by introducing sparsity into the model architecture, allowing only a subset of parameters to be active for each input. The key innovation was the use of a gating network that learned to route each input to the most relevant subset of expert networks. Instead of using all model parameters for every computation, MoE models would activate only 1-2 experts out of many, typically 8-128 experts depending on the model size. This sparse activation pattern meant that while the total number of parameters in the model could be very large, the computational cost of each forward pass was much smaller, as only the active experts needed to be computed.

The architecture of MoE models consists of multiple expert networks, each a complete neural network capable of processing inputs, along with a gating or routing mechanism that determines which experts should handle each input. During training, the gating network learns to identify patterns in inputs and route them to the most appropriate experts. For example, an expert might specialize in scientific terminology, another in code understanding, and another in conversational language. The gating network would route scientific text to the first expert, code snippets to the second, and casual conversation to the third.

The scaling of MoE models in 2024 involved several key technical innovations that made these architectures practical at large scale. First, researchers developed more sophisticated routing algorithms that could learn to assign inputs to experts more effectively. These routing mechanisms needed to balance two competing objectives: routing inputs to the best-suited experts for quality, while ensuring that all experts received roughly equal amounts of work for training stability. Earlier MoE models had often collapsed to using only a few experts, reducing the benefits of the architecture.

Second, researchers developed better load balancing mechanisms to ensure that all experts were used roughly equally during training. This load balancing was critical because if one expert received all the work while others were idle, the model would effectively reduce to a smaller dense model, losing the efficiency benefits. Techniques like auxiliary load balancing losses and top-k gating with capacity constraints helped ensure that work was distributed across experts while still maintaining quality routing decisions.

Third, researchers developed more efficient implementations that could handle the dynamic routing and sparse computation patterns required by MoE architectures. Traditional neural network implementations were optimized for dense computation patterns where the same operations were applied to all inputs. MoE models required infrastructure that could dynamically route inputs, manage expert loading and unloading from memory, and efficiently handle the sparse activation patterns. These implementation improvements were essential for making MoE models practical at scale.

The routing mechanism itself works by computing a probability distribution over experts for each input token. The gating network takes the input and produces scores for each expert, which are then used to select the top-k experts (typically k=1 or k=2) that will process that input. The selected experts process the input in parallel, and their outputs are combined based on the routing scores. This process happens for every token or group of tokens, creating a dynamic routing pattern that adapts to the input content.