Mixture of Experts in LLMs: How Sparse Models Scale

The Core Idea

A standard transformer feed-forward network (FFN) applies the same set of weights to every token. In a dense model with 70 billion parameters, every token processed at inference touches all 70 billion weights. This creates a tight coupling between model capacity (total parameters) and inference cost (FLOPs per token).

Mixture of Experts breaks that coupling. Instead of one large FFN, an MoE layer contains multiple smaller FFNs, each called an expert. A learned router network receives each token and selects a small subset of experts to process it. The router's output is a set of scores over all experts; the top-k scoring experts are activated, their outputs are weighted by the router scores, and the weighted sum becomes the token's representation for that layer.

The result is that a model can have many more total parameters than a dense model of equivalent quality while activating only a fraction of them per token. A model like Mixtral 8x7B has eight experts per MoE layer and activates two per token, giving it roughly 46 billion total parameters but a per-token compute cost closer to a 12 to 13 billion parameter dense model.

The Router

The router is a small linear layer that maps each token's representation to a score over all experts. In Switch Transformers (Fedus et al., 2022), routing uses k=1, meaning each token goes to a single expert. In most subsequent work, k=2 is preferred because it allows the model to blend expert outputs and avoids the hard commitment problem where routing errors cannot be corrected within the layer.

The router is trained jointly with the rest of the model through standard backpropagation. Because the top-k selection is non-differentiable, training uses a straight-through estimator or a load-balancing auxiliary loss to encourage the router to distribute tokens roughly evenly across all experts.

Load Balancing

Without intervention, routers tend to collapse: a few experts receive most of the tokens, the rest are rarely activated and therefore rarely updated, and the model effectively loses most of its experts. This is called expert collapse and it is one of the central training challenges for MoE models.

The standard solution is an auxiliary load-balancing loss added to the training objective. This loss penalizes uneven expert utilization, encouraging the router to spread tokens more evenly. The weight given to this auxiliary loss is a hyperparameter: too high and it overrides the model's natural tendency to route semantically similar tokens to the same expert; too low and collapse still occurs.

Expert collapse is not an edge case. It is the default behavior of an MoE router without explicit regularization, and recovering from it during training is expensive.

Memory and Communication Costs

The efficiency advantage of MoE is primarily in compute per token, not in memory. A model with eight experts stores all eight expert weight matrices in memory, even though only two are used per token. For a 46 billion parameter model like Mixtral 8x7B, this means storing roughly 46 billion parameters on the serving hardware, which requires more GPU memory than a dense model of equivalent inference cost.

In distributed settings, MoE introduces additional communication overhead. When a token is routed to an expert that does not reside on the same accelerator as the token's current representation, the activation must be transferred across the network. This expert-parallelism communication cost can dominate at large scale if expert assignment is not carefully aligned with the hardware topology. Efficient MoE serving requires routing-aware placement strategies that minimize cross-device communication.

Switch Transformers

Switch Transformers (Fedus et al., arXiv:2101.03961) were among the first large-scale demonstrations that MoE could outperform dense models at the same compute budget. The paper trained models up to 1.6 trillion total parameters with k=1 routing and showed that at a fixed training compute budget, the sparse model achieved lower perplexity than a dense model using the same FLOPs. The key finding was that increasing the number of experts improved quality approximately log-linearly, even though the per-token compute cost remained constant.

Switch Transformers also identified and addressed several training instabilities specific to MoE, including loss spikes caused by routing instability early in training. Selective precision (keeping router weights in float32 while using bfloat16 elsewhere) substantially improved training stability without significant compute overhead.

Mixtral and the Open MoE Era

Mixtral 8x7B (Jiang et al., arXiv:2401.04088) brought MoE architecture to the open-weight ecosystem in a form that could be run on consumer hardware. Its 8 expert, k=2 design activates 2 experts per token and per layer, giving it roughly 12 to 13 billion active parameters at inference while storing 46 billion total. On a range of standard benchmarks, it matched or exceeded the performance of dense models with substantially more active parameters.

Mixtral's architecture places MoE layers only in the feed-forward blocks, not in the attention layers. The attention mechanism remains dense across all tokens, and only the FFN computation is sparsified. This is the standard approach because attention already has linear scaling in sequence length and because the FFN blocks account for a larger fraction of total parameters in standard transformer architectures.

Routing and Semantic Specialization

One of the most interesting empirical findings from MoE research is that experts develop loose specialization. When researchers analyze which tokens are routed to which experts, they find that different experts tend to handle different types of content: one expert receives a disproportionate share of tokens from code; another handles mathematical notation; another handles certain languages or grammatical constructs. This specialization emerges from training, not from any explicit design choice in the routing mechanism.

The degree of specialization is partial, not absolute. Most tokens could be handled adequately by several experts, and the router's choice often reflects the local context rather than a fixed assignment by token type. But the fact that specialization emerges at all suggests that MoE architectures are not simply scaling parameters wastefully. Different experts are encoding different aspects of the training distribution.

Training Instabilities and Remedies

MoE models are harder to train than dense models of equivalent active parameter count. Several failure modes are well-documented. Expert dropout, where certain experts receive so few tokens during a training step that their gradients are negligible, can cause those experts to fall behind their peers and eventually receive even fewer tokens in a self-reinforcing cycle. Token dropping, where a router sends more tokens to an expert than its capacity buffer allows, can cause information loss. Router z-loss, introduced in ST-MoE (Zoph et al., arXiv:2202.08906), adds an auxiliary penalty on large router logits and substantially reduces training instability without requiring capacity constraints.

The sensitivity of MoE training to hyperparameters, particularly the learning rate schedule and the weight of auxiliary losses, means that optimal configurations found for dense models do not transfer directly to MoE models. Organizations training MoE models from scratch should expect to run additional hyperparameter searches specific to the routing configuration.

Implications for Enterprise Deployment

For organizations deploying commercial models, MoE architecture is relevant primarily because it changes the relationship between model capability and inference cost. A frontier model with hundreds of billions of total parameters may have an active parameter count closer to 20 to 30 billion, which determines inference latency and GPU memory requirements per request.

Understanding whether a model is dense or sparse changes how you should interpret benchmark comparisons. A 70 billion parameter dense model and a 46 billion total parameter MoE model with 13 billion active parameters may perform comparably on some benchmarks while differing significantly on others, and their serving costs differ substantially. Evaluating models on tasks representative of your actual workload is more informative than comparing total parameter counts.

The memory requirement difference is also important for private deployment. A sparse model with 46 billion total parameters requires more GPU memory than a dense model with 13 billion active parameters, even though their inference compute cost may be similar. Organizations evaluating on-premises or private cloud deployment should account for total parameter storage requirements, not just active compute costs.