Mixture of Experts (MoE): Architecture, Routing, and Scaling

Introduction

Mixture of Experts (MoE) is a neural network architecture that dramatically increases model capacity without proportionally increasing computational cost. Instead of applying the same feed-forward network to every token, MoE replaces each FFN with multiple specialised “expert” networks and a learned routing mechanism that activates only a subset of experts per token. The result is a model that has billions more parameters than its dense counterpart but uses a fraction of them for any single forward pass.

The importance of MoE lies in its ability to decouple model capacity from compute cost. A dense 1-trillion-parameter model would be computationally prohibitive for both training and inference. An MoE model with 1 trillion total parameters but only 100 billion activated parameters per token achieves comparable representational capacity at roughly 10% of the compute cost. This scaling efficiency has made MoE the default architecture for frontier AI systems: Mixtral 8x7B (Mistral, 2024), GPT-4 (OpenAI, 2023), DeepSeek-V2 and V3 (DeepSeek, 2024-2025), Qwen2.5-MoE (Alibaba, 2025), and Switch Transformer (Google, 2022) all use variants of the MoE architecture.

For a broader understanding of how MoE fits into the modern Transformer landscape, see our complete guide to Transformer architecture.

Core Architecture

In a standard Transformer, each layer contains a multi-head attention mechanism followed by a position-wise feed-forward network. The FFN consists of two linear projections with a non-linear activation: FFN(x) = W_2 * GELU(W_1 * x + b_1) + b_2. This dense FFN uses the same set of parameters for every input token, regardless of whether the token requires the full representational capacity.

MoE replaces this single FFN with a set of E expert networks — each an independent FFN with its own parameters — and a gating network (router) that learns to assign each token to a subset of experts. For each input token, the router produces a probability distribution over experts, selects the top-k experts by probability, and computes the output as a weighted combination of the selected experts' outputs:

y = sum_{i in top-k} G(x)_i * E_i(x)

where G(x) is the router output (a probability distribution over E experts) and E_i(x) is the output of expert i. The router is typically a simple linear layer followed by softmax: G(x) = softmax(W_r * x), where W_r is a learned weight matrix of shape d_model x E.

The router architecture is remarkably simple compared to its impact. A single linear projection determines which experts activate. Despite its simplicity, the router learns meaningful specialisation: experts naturally cluster by data domain, syntactic pattern, or reasoning type during training. In DeepSeek-V2, with 256 experts, post-hoc analysis reveals that approximately 30% of experts specialise in mathematics and code, 25% in factual knowledge, 20% in linguistic patterns, and the remaining 25% handle general-purpose representation [1].

import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseMoE(nn.Module):
    def __init__(self, d_model, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, 4 * d_model),
                nn.GELU(),
                nn.Linear(4 * d_model, d_model),
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)
        router_logits = self.router(x_flat)
        router_weights = F.softmax(router_logits, dim=-1)
        top_k_weights, top_k_indices = torch.topk(router_weights, self.top_k, dim=-1)
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
        outputs = torch.zeros_like(x_flat)
        for i in range(self.num_experts):
            mask = (top_k_indices == i).any(dim=-1)
            if mask.any():
                expert_out = self.experts[i](x_flat[mask])
                weight = top_k_weights[mask][top_k_indices[mask] == i].unsqueeze(-1)
                outputs[mask] += weight * expert_out
        return outputs.view(batch_size, seq_len, d_model)

Routing Strategies

The routing mechanism determines which experts process each token and how their outputs are combined. The choice of routing strategy has a direct impact on model quality, computational efficiency, and training stability.

Top-k Routing

The most common strategy, introduced by Shazeer et al. in the original MoE paper (2017), selects the k experts with the highest router probabilities for each token. Top-1 routing (k=1) is the most computationally efficient — each token activates exactly one expert — but forces the router to make a discrete decision that may not capture the full complexity of the input. Top-2 routing (k=2), used by Mixtral 8x7B, provides a balance between efficiency and representational quality: each token activates two experts, and the computational cost is approximately 2/E of the dense equivalent [2].

The choice of k involves a trade-off. Larger k improves model quality by aggregating more expert perspectives but increases computational cost linearly. Empirical studies on models up to 1 trillion parameters show that k=2 captures roughly 95% of the quality improvement over k=1, while k=3 adds only marginal gains at 50% more compute. DeepSeek-V3 uses k=6 with 256 experts, achieving a sparsity ratio (k/E) of approximately 2.3%, meaning only 6 of 256 experts activate per token.

Noisy Top-k Routing

Standard top-k routing suffers from a fundamental problem: experts that are rarely selected receive no gradient updates, creating a feedback loop where the router learns to ignore them because they have not been trained sufficiently. Noisy top-k routing addresses this by adding learnable Gaussian noise to the router logits before the softmax [3]:

G(x) = softmax(W_r * x + epsilon * softplus(W_noise * x))

where epsilon is a standard normal random variable and W_noise is a learned noise parameterisation. The noise encourages exploration: experts that are normally ranked just below the top-k threshold occasionally receive routed tokens due to the random perturbation, receiving gradient updates that prevent the expert from atrophying entirely.

Expert Capacity and Capacity Factor

In distributed training, each expert processes a variable number of tokens depending on router decisions. If routing is imbalanced — for example, if one expert receives 50% of the tokens while another receives 2% — the hardware utilisation becomes uneven, creating stragglers that slow down the entire training step.

Expert capacity caps the maximum number of tokens an expert can process in a single training step. It is defined as: capacity = capacity_factor * (total_tokens / num_experts). The capacity factor controls how much imbalance the system tolerates. A capacity factor of 1.0 means each expert processes exactly the expected number of tokens, enforcing perfect balance at the cost of dropping tokens when an expert receives more than capacity. A capacity factor of 1.25 allows 25% overload, reducing dropped tokens at the cost of wasted compute on under-utilised experts [4].

When an expert exceeds its capacity, excess tokens are either dropped (with the token passing through a residual connection or being processed by a fallback expert) or dispatched to the next training step. Token dropping introduces approximation error but is tolerable at moderate rates (1-3% of tokens). Mixtral 8x7B uses a capacity factor of approximately 1.25 with token dropping.

Auxiliary Load Balancing Loss

To encourage balanced expert utilisation during training, MoE models add an auxiliary loss that penalises imbalanced routing. The standard formulation is the importance loss, which measures the coefficient of variation of expert utilisation:

L_aux = alpha * E * sum_i (f_i * P_i)

where f_i is the fraction of tokens routed to expert i, P_i is the average router probability for expert i, and alpha is a scaling hyperparameter (typically 0.01). The product f_i * P_i is maximised when both the fraction of tokens assigned to expert i and the router's confidence in those assignments are uniformly distributed across experts. This loss is differentiable and provides a strong gradient signal to the router to distribute tokens evenly [3].

Load Balancing

Load balancing is the single most important engineering challenge in MoE. Without explicit balancing mechanisms, the router converges to a degenerate solution where a small subset of experts receive most of the tokens while the majority atrophy. This undermines the entire premise of MoE — if only 10% of experts are used, the effective model capacity is equivalent to a much smaller dense model [5].

Importance Loss and Z-Loss

The importance loss described above is the most widely adopted auxiliary balancing objective. However, it has a subtle limitation: it balances the router's output probabilities but not necessarily the actual token-to-expert assignments, because the router can assign high probability to an expert without routing many tokens there (if other experts have even higher probabilities). To address this, DeepSeek-V2 introduced a second auxiliary objective called z-loss, which regularises the logits entering the router softmax:

L_z = beta * (1/N) * sum_i (log(sum_j exp(z_ij)))^2

where z_ij is the router logit for expert j on token i, and beta is a scaling hyperparameter. Z-loss penalises large logit magnitudes that would produce overly confident routing decisions, encouraging smoother expert assignments and reducing the variance in expert utilisation. DeepSeek-V2 found that combining importance loss with z-loss at alpha=0.01 and beta=0.001 achieved near-perfect load balance across 256 experts [1].

Batch Priority Routing

Batch Priority Routing (BPR), introduced by Google in 2024, takes a different approach: instead of adding auxiliary losses, it modifies the routing decision itself. In BPR, tokens within a batch are processed in order of router confidence. The highest-confidence tokens are routed first, ensuring that experts receive the tokens they are best suited for. Lower-confidence tokens are routed to whatever capacity remains. BPR eliminates the need for auxiliary balancing losses entirely, achieving perfect load balance by construction while maintaining model quality [6].

Expert Choice Routing

Expert Choice Routing (ECR) inverts the standard routing paradigm. Instead of each token choosing its top-k experts, each expert chooses its top-k tokens. This guarantees that every expert processes exactly k tokens, achieving perfect load balance by definition. The tokens that are not selected by any expert receive a zero contribution from the MoE layer for that token position (or are processed by a shared expert). ECR eliminates the need for capacity factors, token dropping, and auxiliary losses. However, it introduces the risk that some tokens receive no expert processing at all, which can degrade quality for long-tail inputs [7].

Training Stability

Training MoE models at scale introduces unique stability challenges that do not exist in dense model training. The discrete routing decisions, the interaction between the router and experts, and the auxiliary loss landscape all contribute to training dynamics that require careful management.

Router Collapse

Router collapse occurs when the router converges to a solution that routes all tokens to the same expert, making the MoE layer effectively dense. This is the most common training failure mode for MoE models. Router collapse typically emerges within the first 1,000 training steps and is characterised by the router logits for one expert dominating all others. The primary mitigation is a sufficiently high auxiliary loss coefficient (alpha between 0.01 and 0.1) combined with large initial router weights that break symmetry [3].

Gradient Noise

The sparse routing mechanism introduces gradient noise because each token activates only a subset of experts. An expert receives gradient updates only from the tokens routed to it, which means expert gradients are a sparse subset of the total gradient signal. This increases gradient variance, especially for experts that process few tokens. The noise is proportional to 1/sqrt(expert_capacity): experts that process 128 tokens per step have roughly 4x the gradient variance of experts that process 2,048 tokens. Large batch sizes (4,096+ sequences) are essential for stable MoE training to compensate for this variance [8].

Initialization Strategies

Router initialisation is critical for training stability. If the router weights are too small, the initial routing distribution is nearly uniform, and all experts receive roughly equal tokens by chance. If they are too large, the router collapses to a near-deterministic assignment that prevents early exploration. The standard practice is to initialise the router weight matrix with small values (standard deviation 0.02-0.05) and use a small positive bias for the auxiliary loss — typically a bias of 1.0 on the importance loss coefficient, which decays linearly over the first 10,000 steps. Switch Transformer and Mixtral both use this warm-up strategy [4].

Scaling Learning Rate for MoE

MoE models generally require lower learning rates than equivalent dense models. The interaction between the router and experts creates a co-adaptation dynamic where the router and experts must learn compatible specialisation. If the learning rate is too high, the router changes its assignments faster than the experts can adapt, leading to oscillations. DeepSeek-V3 found that an optimal learning rate for MoE is roughly 0.5x to 0.7x the optimal learning rate for an equivalent dense model at the same total parameter count. The router learning rate is often further reduced by a factor of 0.1 to decouple router convergence from expert training [1].

Memory and Compute Trade-offs

MoE's primary advantage is the decoupling of total parameter count from per-token compute cost. However, this decoupling is not free — it introduces memory overhead and communication costs that must be managed carefully.

Activated vs total parameters. In a dense model, every parameter is used for every forward pass. In an MoE model, only a fraction of experts activate per token. Mixtral 8x7B has 47 billion total parameters but activates only 13 billion per token (2 of 8 experts), giving an effective parameter-to-compute ratio of approximately 3.6:1. DeepSeek-V2 has 236 billion total parameters but activates only 21 billion per token (6 of 256 experts plus shared expert), achieving a ratio of 11:1. The higher the ratio, the more model capacity per unit of compute.

Memory hierarchy.All expert parameters must fit in GPU memory even though only a subset are used per token. With 256 experts, total expert parameters dominate memory usage. DeepSeek-V2's 256 experts require approximately 450 GB for the FFN parameters alone at FP8, requiring 6 H100 GPUs (80 GB each) just to hold the weights. Expert parallelism shards experts across GPUs: each GPU hosts a subset of experts, and tokens are communicated between GPUs via all-to-all operations.

Comparison of dense and MoE models. Mixtral 8x7B delivers roughly 4.6x the throughput of a dense 70B model per GPU while matching its quality. DeepSeek-V2's higher parameter-to-compute ratio comes with higher minimum GPU requirements.

Inference in MoE Models

Deploying MoE models for production inference introduces challenges beyond those of dense models. The key difference is that expert parameters must be distributed across GPUs, and tokens must be routed to the correct GPU — a communication pattern fundamentally different from dense model parallelism.

Expert Parallelism

Expert parallelism shards the experts across GPUs rather than sharding individual layers. With 8 experts and 8 GPUs, each GPU hosts exactly one expert (plus the attention layers, which are replicated). When a token's router selects expert 3, that token must be sent to GPU 3, processed, and the result returned. This all-to-all communication pattern is the dominant cost in MoE inference — for models with many experts, communication latency can exceed computation time [9].

All-to-All Communication

Each MoE layer requires an all-to-all communication operation: every GPU sends its tokens to the GPUs hosting their selected experts and receives tokens routed to its local expert. The communication volume is proportional to the number of tokens and the hidden dimension. For Mixtral 8x7B with 8 GPUs and batch size 256, each GPU sends approximately 32 tokens per MoE layer, with a per-token data size of 4,096 FP16 values (8 KB). The total communication per layer is 256 tokens * 8 KB = 2 MB per GPU, taking roughly 5 microseconds on NVLink (900 GB/s bidirectional) compared to roughly 100 microseconds for the expert computation. Communication overhead is manageable for small numbers of experts but becomes significant at 64+ experts [2].

Batch Processing Inefficiencies

Batch processing in MoE suffers from a unique inefficiency: the all-to-all pattern distributes tokens unevenly across experts, and GPUs must wait for all experts to complete before proceeding to the next layer. If one expert receives twice as many tokens as another, the faster GPUs stall waiting for the straggler. The capacity factor partially addresses this by capping per-expert load, but the fundamental variance in expert utilisation means that MoE inference typically achieves 70-85% of theoretical peak GPU utilisation compared to 90-95% for dense models [10].

Speculative Decoding for MoE

Speculative decoding — using a small draft model to propose multiple tokens that the target model verifies in parallel — is particularly effective for MoE models. The draft model can be a dense 1B-parameter model that runs on a single GPU, while the target MoE model verifies tokens using its full capacity. Because verification requires only a single forward pass of the MoE model (including all-to-all communication), the cost of speculative decoding is roughly 1.5x a single MoE forward pass regardless of the draft length. For Mixtral 8x7B, speculative decoding with a 1B draft model achieves 2.3-2.8x throughput improvement, compared to approximately 2.0x for dense models of similar quality [11].

For more on inference optimisation strategies, see our guide to LLM inference cost optimisation.

Fine-Tuning MoE Models

Fine-tuning MoE models presents unique challenges compared to dense model fine-tuning. The expert specialisation learned during pre-training must be preserved while adapting to the target task.

Adapter-Based Tuning (LoRA per Expert)

Low-Rank Adaptation (LoRA) is the most practical fine-tuning approach for MoE models. Rather than updating all expert parameters, LoRA injects low-rank adapters into the expert FFN layers. The key insight is that each expert can have its own LoRA adapters, allowing the fine-tuned model to preserve the expert specialisation learned during pre-training. Mixtral 8x7B fine-tuned with per-expert LoRA (rank 32, alpha 64) achieves comparable quality to full fine-tuning while using only 0.3% of the parameter budget — roughly 150 million trainable parameters instead of 47 billion [12].

Full Fine-Tuning Challenges

Full fine-tuning of MoE models requires careful management of the router-expert co-adaptation. If the router is updated aggressively, it may reassign tokens away from the experts that were originally best suited for them, degrading model quality. The standard practice is to freeze the router during fine-tuning (zero learning rate for router parameters) and only update expert weights, or to use a separate, much lower learning rate for the router (typically 0.01x to 0.05x the expert learning rate). Full fine-tuning also requires substantially more GPU memory than dense models of equivalent activated parameter count because all expert weights must be loaded even when only a subset are updated per batch [13].

Catastrophic Forgetting in Expert Specialisation

Fine-tuning MoE models on a narrow domain risks catastrophic forgetting of the expert specialisation learned during pre-training. If the fine-tuning data is concentrated in a single domain (e.g., legal documents), the router may learn to route all tokens to the same 2-3 experts, causing the remaining experts to lose their specialised knowledge. This is a form of router collapse triggered by data distribution shift. Mitigations include: (1) mixing 10-30% of general pre-training data into the fine-tuning batch, (2) adding a stronger auxiliary loss during fine-tuning (alpha = 0.05), and (3) per-expert learning rate scaling where frequently selected experts have lower learning rates to prevent over-specialisation.

Sparse vs Dense Updates

An important distinction in MoE fine-tuning is whether to update all experts (dense updates) or only the experts selected by the router (sparse updates). Sparse updates are computationally efficient — if each token activates 2 of 8 experts, only 25% of expert parameters need gradient computation per batch. However, sparse updates create a training instability: experts that receive few tokens on the fine-tuning dataset do not get updated, widening the quality gap between frequently- and infrequently-selected experts. The recommended approach is dense updates during fine-tuning, where each expert processes its allocated subset of tokens in sequence, regardless of whether the router selected it for those tokens — this ensures uniform expert adaptation [12].

Production Deployments of MoE

The most significant production deployment of MoE to date is Mixtral 8x7B (Mistral, 2024), which demonstrated that MoE models can match or exceed the quality of much larger dense models while maintaining practical inference economics.

Mixtral 8x7B uses 8 experts per MoE layer with top-2 routing, 47B total parameters, and 13B activated parameters per token. On standard benchmarks, Mixtral 8x7B matches or exceeds Llama 2 70B (a dense 70B model) while using 5.4x fewer FLOPs per token. On MMLU, Mixtral achieves 70.6% vs Llama 2 70B's 68.9%. On GSM8K (math reasoning), Mixtral achieves 58.4% vs 56.8%. The inference speed advantage is substantial: Mixtral generates approximately 110 tokens per second per H100 GPU compared to approximately 24 tokens per second for dense 70B models — a 4.6x throughput improvement [2].

Output quality vs inference speed. The primary trade-off in MoE deployment is between output quality and inference throughput. For the same compute budget, MoE models offer higher quality than smaller dense models (Mixtral 8x7B vs Llama 2 13B) and higher throughput than larger dense models (Mixtral 8x7B vs Llama 2 70B). The decision of when to choose MoE over dense models depends on the deployment constraints:

• Choose MoE when: you need quality close to a dense XB model but cannot afford the latency or GPU cost; you have high throughput requirements and can amortise the fixed memory cost of all experts; your workload benefits from diverse specialist knowledge (code + math + creative writing).
• Choose dense when: you are memory-constrained (consumer GPUs); you need the absolute highest quality for a specific domain and can accept the cost; your deployment requires very low latency and cannot tolerate all-to-all communication overhead.

For a practical comparison of model options, see our LLM comparison guide and our analysis of open-source vs closed-source LLMs.