How Mixture of Experts Works: Sparse Gating, Token Routing, and Load Balancing

Large language models face a fundamental scaling trade-off: as parameter counts grow to absorb more world knowledge, the floating-point operations (FLOPs) required per token increase proportionally. In a standard dense Transformer, every single parameter participates in processing every incoming token.

Mixture of Experts (MoE) decouples total model capacity from per-token compute cost. By replacing dense feed-forward layers with a dynamic pool of specialized sub-networks ("experts") and routing each token to only a small subset of them, MoE architectures allow models to scale parameter counts by orders of magnitude while keeping inference latency and compute budgets manageable.

Here is how sparse gating, token routing, load balancing, and fine-grained expert architectures work under the hood.

The Dense Scaling Bottleneck

In a standard decoder-only Transformer, each layer consists of two core components:

Multi-head self-attention, which models interactions across tokens in the context window.
A feed-forward network (FFN), typically a two-layer multi-layer perceptron (MLP) with non-linear activation (such as SwiGLU or GELU).

FFN layers account for roughly two-thirds of a model's total non-embedding parameters. Research into Transformer interpretability demonstrates that FFN layers function primarily as associative key-value memories, storing factual knowledge and domain-specific patterns.

In dense models like Llama 3 or standard GPT-style architectures, computing a 70-billion-parameter model requires executing 70 billion parameters worth of matrix multiplications for every token generated. Doubling parameter capacity to store broader knowledge doubles the matrix arithmetic per token.

MoE addresses this by introducing conditional computation: activating only the parameters relevant to the immediate token.

The MoE Architecture: Sparse Expert Layers

In an MoE Transformer, the self-attention mechanism remains dense and shared across all tokens to maintain full contextual visibility across the sequence. However, the monolithic FFN block in each layer is replaced by:

A set of $N$ parallel expert networks: $\{E_1, E_2, \dots, E_N\}$ , where each expert is an independent FFN with identical structure but distinct weights.
A gating network (router): $G(x)$ , parameterized by a lightweight linear projection layer $W_g$ .

Mixture of Experts Architecture and Token Routing

When a token representation $x$ arrives at the MoE layer, the router evaluates $x$ across all $N$ experts and produces a routing distribution. The layer output $y$ is computed as the weighted sum of outputs from the top- $k$ selected experts:

$y = \sum_{i \in \text{TopK}(G(x), k)} G(x)_i \cdot E_i(x)$

Where $k \ll N$ . For instance, in Mixtral 8x7B (Jiang et al., 2024), $N=8$ and $k=2$ . Although the model contains 46.7 billion total parameters, each token only activates approximately 12.9 billion parameters during the forward pass.

Gating Mechanisms and Routing Strategies

The router determines which experts receive each token. The evolution of routing algorithms reflects a balance between routing granularity and training stability:

1. Noisy Top-K Gating

Introduced by Shazeer et al. (2017), this approach computes routing logits $H(x)$ with injected Gaussian noise to encourage exploration:

$H(x)_i = (x \cdot W_g)_i + \epsilon \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i)$

where $\epsilon \sim \mathcal{N}(0, 1)$ . The router selects the top- $k$ logits and applies a softmax over those selected indices:

$G(x)_i = \frac{\exp(H(x)_i)}{\sum_{j \in \text{TopK}(H(x), k)} \exp(H(x)_j)}$

2. Switch Routing (Top-1)

The Switch Transformer (Fedus et al., 2021) simplified routing by selecting only the single highest-scoring expert ( $k=1$ ):

$G(x)_i = \text{Softmax}(x \cdot W_g)_i \quad \text{for } i = \text{argmax}(x \cdot W_g)$

Top-1 routing maximizes computational efficiency and reduces communication overhead in distributed clusters, but it can lead to unstable training dynamics and lower sample efficiency compared to top-2 routing.

3. Top-2 Gating with Normalization

Modern open-weight models like Mixtral 8x7B utilize top-2 routing without noise injection. The router computes standard linear logits, selects the top two indices, and renormalizes the softmax over those two scores so their weights sum to 1. This ensures gradient flow to multiple experts per token while maintaining a predictable compute budget.

The Routing Collapse Trap and Load Balancing

A major challenge in training sparse MoE models is routing collapse. If one expert initially performs slightly better on common token patterns, the router assigns more tokens to it. That expert receives more gradient updates and becomes even more competent, while underutilized experts starve. Eventually, the router routes all tokens to the same one or two experts, collapsing the model into an inefficient dense network.

To prevent collapse, MoE architectures introduce auxiliary objectives and routing constraints:

Auxiliary Load Balancing Loss

Switch Transformers and GShard enforce uniform expert utilization across a batch of tokens using an auxiliary loss:

$\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^N f_i \cdot P_i$

$f_i$ is the fraction of tokens in the batch dispatched to expert $i$ :

$f_i = \frac{1}{T} \sum_{t=1}^T \mathbb{I}(\text{token } t \text{ routes to expert } i)$

$P_i$ is the average routing probability assigned to expert $i$ across all tokens in the batch:

$P_i = \frac{1}{T} \sum_{t=1}^T G(x_t)_i$

$\alpha$ is a scaling hyperparameter (typically $0.01$ to $0.1$ ).

The product $f_i \cdot P_i$ is minimized when both token counts ( $f_i$ ) and router probabilities ( $P_i$ ) are uniformly distributed across all $N$ experts ( $1/N$ ).

Router z-loss

In ST-MoE (Zoph et al., 2022), researchers observed that large routing logits cause numerical instability in 16-bit floating-point (fp16/bf16) training. ST-MoE introduced a router z-loss:

$\mathcal{L}_z = \frac{c_z}{B} \sum_{b=1}^B \left( \log \sum_{i=1}^N \exp(h_{b,i}) \right)^2$

This penalizes extreme router logits, stabilizing training without degrading model quality.

Auxiliary-Loss-Free Balancing

While auxiliary losses enforce balance, they force a trade-off: the model must optimize for load distribution at the expense of pure language modeling perplexity.

To eliminate this conflict, DeepSeek-V3 (DeepSeek-AI, 2024) introduced auxiliary-loss-free load balancing. Instead of adding a loss term, the router maintains dynamic bias terms $b_i$ for each expert. When computing routing decisions, the bias is added to the routing score:

$s_i = \text{Softmax}(x \cdot W_g)_i + b_i$

If an expert is overloaded, its bias $b_i$ is decreased; if underutilized, $b_i$ is increased. The gating weights applied to the expert outputs still use the unbiased probabilities $\text{Softmax}(x \cdot W_g)_i$ . This ensures balanced hardware utilization across GPUs without penalizing the core training objective.

Fine-Grained and Shared Experts

Early MoE models used a small number of large experts (e.g., 8 or 16 experts per layer). Recent architectures have shifted toward finer granularity:

DeepSeekMoE Architecture

DeepSeekMoE (Dai et al., 2024) demonstrated two architectural principles:

Fine-Grained Expert Segmentation: Rather than using 8 large FFNs with top-2 routing, DeepSeekMoE splits each FFN into smaller modules (e.g., 64 or 256 fine-grained experts) and activates a larger count (e.g., top-8 or top-32). This expands the combinatorics of possible expert combinations, enabling specialized knowledge isolation.
Dedicated Shared Experts: Certain patterns (such as punctuation, conjunctions, and general syntactic framing) are universal across all tokens. In traditional MoE, redundant copies of this basic knowledge appear in multiple experts. DeepSeekMoE isolates 1 or 2 experts as permanently active shared experts that process every token, allowing routed experts to focus exclusively on specialized domains.

Hardware and Serving Trade-offs

While MoE provides a favorable FLOPs-to-quality ratio, it introduces specific systems-level engineering challenges:

Memory Capacity vs. Memory Bandwidth: An MoE model requires VRAM to hold all $N$ experts simultaneously. A 671B MoE model (such as DeepSeek-V3) requires the same memory footprint as a 671B dense model (~1.3 TB in 16-bit, or ~340 GB in 4-bit quantization), even though each token only executes ~37B parameters.
Expert Parallelism and All-to-All Communication: In distributed multi-GPU environments, experts are distributed across different devices. Routing tokens to their assigned experts requires an all-to-all communication primitive across GPUs. High-bandwidth interconnects (like NVLink and InfiniBand) are necessary to prevent communication bottlenecks from dominating compute time.
Dynamic Batching and Kernel Optimization: Because token routing varies per input, individual experts receive variable batch sizes. Modern inference engines use specialized ragged batching kernels (such as MegaBlocks, DeepGEMM, and vLLM custom MoE kernels) to avoid padding overhead and maximize tensor core saturation.

Summary

Mixture of Experts decouples parameter capacity from operational compute cost. By replacing dense feed-forward layers with routed expert sub-networks, MoE enables large language models to scale world knowledge while maintaining high inference throughput.