Mixture of Experts for LLMs

Mixture of Experts (MoE) is a conditional computation architecture in which a learned routing function selects a sparse subset of specialized sub-networks ("experts") to process each input token, enabling model capacity to scale super-linearly with parameter count while keeping per-token FLOPs approximately constant. It matters because it is the dominant technique by which frontier LLMs (GPT-4, Gemini, Mixtral, DeepSeek-V2) achieve trillion-parameter scale at tractable training and inference cost. The counterintuitive fact: despite each expert being trained on only a fraction of tokens, MoE models consistently outperform dense models of equal compute budget — more parameters with fewer FLOPs per token is strictly better than the naive tradeoff suggests.

💡 The Core Insight

The fundamental insight is that not every part of a network needs to activate for every input. In a dense transformer, all parameters process every token — this couples model capacity tightly to compute. MoE breaks this coupling: by routing each token to only $k$ of $E$ experts, capacity grows as $O(E)$ while per-token compute grows as $O(k/E)$ of the total — making the two quantities independently controllable. The "aha moment" is recognizing that language is heterogeneous — mathematical tokens, code, French prose, and logical connectives do not benefit from the same learned features — so routing different token types to different specialists is not just a computational trick but a semantically meaningful decomposition.

⚠️ Common Misconception

The most pervasive error is assuming that MoE's efficiency advantage is simply "we use fewer parameters per forward pass, so it's faster." This conflates parameter count with memory bandwidth, communication volume, and load balance, all of which are the actual bottlenecks in practice. The correct mental model: in distributed training, each expert lives on a different device, so routing induces all-to-all communication whose cost scales with the number of tokens dispatched across devices — not with expert parameter count. A router that sends all tokens to two popular experts (a degenerate but low-loss solution) collapses the MoE to a dense model plus communication overhead, strictly worse than a baseline. The auxiliary load-balancing loss exists precisely to fight this collapse, and ignoring it or under-weighting it is the single most common failure mode when implementing MoE from scratch.

How It Works

Routing and Sparse Dispatch

Consider a transformer with $L$ layers. At each MoE layer, the standard FFN sublayer is replaced by $E$ expert networks $\{f_1, \ldots, f_E\}$ , each with its own parameters $\theta_i$ . A lightweight router — typically a linear projection followed by softmax — assigns per-token routing weights. Let $\mathbf{x} \in \mathbb{R}^d$ be the hidden state of an input token. The router computes logits:

\quad \mathbf{h}(\mathbf{x}) = \mathbf{x} W_r, \quad W_r \in \mathbb{R}^{d \times E}

Explanation: The router projects the token representation into a logit vector of dimension equal to the number of experts, producing an unnormalized score for each expert. The top- $k$ routing function selects the $k$ experts with highest logit values and softmax-normalizes only over those $k$ scores:

\quad G(\mathbf{x}) = \mathrm{Softmax}\!\left(\mathrm{TopK}\!\left(\mathbf{h}(\mathbf{x}),\, k\right)\right)

Explanation: Only the $k$ selected experts receive nonzero gating weights; all others are masked to $-\infty$ before the softmax, making the computation exactly sparse. The MoE sublayer output is then the weighted sum of selected expert outputs:

\quad \mathrm{MoE}(\mathbf{x}) = \sum_{i \in \mathcal{T}_k(\mathbf{x})} G_i(\mathbf{x})\, f_i(\mathbf{x})

Explanation: Each selected expert $f_i$ processes the full token representation independently, and their outputs are blended by the normalized routing weights $G_i(\mathbf{x}) > 0$ . Here $\mathcal{T}_k(\mathbf{x})$ denotes the index set of the top- $k$ experts for token $\mathbf{x}$ , $G_i(\mathbf{x})$ is the $i$ -th component of the gating vector (zero for non-selected experts), and $f_i : \mathbb{R}^d \to \mathbb{R}^d$ is typically a two-layer FFN with SwiGLU or ReLU activation.

Load Balancing

Without regularization, the router collapses: a few experts accumulate most tokens while the rest receive near-zero gradient signal — the expert collapse or rich-get-richer phenomenon. The standard remedy (Fedus et al., 2022; Shazeer et al., 2017) is an auxiliary load-balancing loss added to the primary language modeling objective. Define $f_i$ as the fraction of tokens routed to expert $i$ in a batch $\mathcal{B}$ of $T$ tokens, and $p_i$ as the mean routing probability assigned to expert $i$ :

\quad \mathcal{L}_{\mathrm{aux}} = \alpha \cdot E \cdot \sum_{i=1}^{E} f_i \cdot p_i, \quad f_i = \frac{1}{T}\sum_{x \in \mathcal{B}} \mathbf{1}\!\left[i \in \mathcal{T}_k(\mathbf{x})\right], \quad p_i = \frac{1}{T}\sum_{x \in \mathcal{B}} G_i(\mathbf{x})

Explanation: This loss is minimized when all experts receive equal token fractions ( $f_i = 1/E$ ) and equal mean probability ( $p_i = 1/E$ ), penalizing any concentration of routing mass on a subset of experts. The coefficient $\alpha$ (typically $10^{-2}$ to $10^{-1}$ ) controls the tradeoff between routing quality and balance. Note that $f_i$ is non-differentiable (it depends on the discrete top- $k$ selection), so the gradient flows only through $p_i$ — a deliberate trick that makes this a bilinear surrogate loss.

Expert Capacity and Token Dropping

In hardware-efficient implementations, each expert is given a capacity buffer $C = \lceil cT/E \rceil$ where $c > 1$ is the capacity factor. Tokens routed to an expert beyond its capacity are dropped (replaced by zeros or their residual stream value). This introduces a controlled approximation that prevents stragglers in parallel dispatch:

\quad \text{Effective throughput} = \min\!\left(T,\; E \cdot C\right) = \min\!\left(T,\; c \cdot T\right)

Explanation: With capacity factor $c$ , up to $cT$ total expert-token slots exist; setting $c = 1.25$ wastes 25% of capacity but ensures fewer than 25% of tokens are ever dropped under balanced routing.

researcher · Noam Shazeer | b. 1975 | Invented the modern sparse MoE gating mechanism for neural networks

Shazeer was a longtime Google Brain researcher who co-authored the original transformer paper before departing to co-found Character.AI. His 2017 paper "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer" introduced the top- $k$ noisy gating mechanism and demonstrated for the first time that sparse MoE could be integrated into sequence-to-sequence models at scale, achieving state-of-the-art on machine translation with 137 billion parameters. Remarkably, his 2017 contribution was largely ignored for four years before the field rediscovered it via the Switch Transformer in 2021.

researcher · William Fedus | b. 1990 | Introduced Switch Transformer, simplifying MoE to top-1 routing at trillion-parameter scale

Fedus is a research scientist at Google Brain whose 2021 Switch Transformer paper demonstrated that routing to a single expert ( $k=1$ ) is not only sufficient but superior to $k=2$ in many regimes, dramatically simplifying the load-balancing problem. The Switch Transformer was the first publicly described trillion-parameter model and showed 7× speedup over T5-XXL at equal compute, catalyzing the modern era of sparse LLM research. Fedus has noted in interviews that the initial reviewer consensus was skeptical that top-1 routing could be stable — a misconception his paper definitively refuted.

Figure

MoE Layer Architecture: Token Routing Through Sparse Experts

Figure

Expert Utilization Under Balanced vs. Collapsed Routing

Intuition

🧠 The Mental Model

Imagine a large hospital where every patient — whether presenting with a broken bone, a psychiatric episode, a cardiac event, or a skin rash — is seen by the exact same generalist physician every single time. That's a dense transformer: one FFN, all tokens. Now imagine the hospital is reorganized into specialized departments: orthopedics, cardiology, psychiatry, dermatology. A triage nurse (the router) looks at each patient and sends them to the appropriate specialist — only one or two departments per visit, not all sixteen. The specialists become extremely good at their narrow domain because they see concentrated, relevant training cases. You can add new departments without making any individual department busier — the hospital's capacity scales without scaling the cost per patient. The MoE router is that triage nurse: it makes a fast, low-cost binary decision that unlocks massively parallel specialization. The pathology of expert collapse is equivalent to the triage nurse always routing everyone to the same two departments because they happen to be slightly more popular in the training distribution — a bureaucratic habit that destroys the entire point of specialization.

Historical Development

⏳ History

1991

Jacobs, Jordan, Nowlan & Hinton introduce the original "Mixture of Experts" as a modular ensemble — separate networks trained on subsets of input space via a gating network — predating deep learning entirely.

1994

Jordan & Jacobs extend MoE to "Hierarchical Mixture of Experts," establishing the theoretical connection between gating networks and the Expectation-Maximization algorithm for fitting mixture models.

2017

Shazeer et al. (Google Brain) publish "Outrageously Large Neural Networks" — the first sparse MoE integrated into deep sequence models (LSTMs), introducing noisy top-

k

gating and demonstrating 137B parameter translation models with 1000+ experts per layer.

2020

Lewis et al. (Facebook AI) publish "BASE Layers," framing routing as a linear assignment problem to guarantee perfect load balance — an alternative to the auxiliary loss approach that achieves balance by construction.

2021

Fedus, Zoph & Shazeer publish "Switch Transformer" at Google Brain — simplifying to top-1 routing, scaling to 1.6 trillion parameters, and achieving 7× pre-training speedup over T5-XXL; first to demonstrate trillion-parameter training stability.

2022

Du et al. publish "GLaM" (Generalist Language Model), a 1.2T parameter top-2 MoE that matches GPT-3 quality using one-third the energy at inference; Zoph et al. release "ST-MoE," providing systematic stability recipes (router z-loss) for very large sparse models.

2023

Mistral AI releases Mixtral 8×7B as an open-weight top-2 MoE with 46.7B total parameters and 12.9B active parameters per token — matching Llama 2 70B quality at less than half the inference compute, democratizing MoE research beyond proprietary labs.

2024

DeepSeek-V2 introduces "DeepSeekMoE" with fine-grained expert decomposition (splitting FFN hidden dimensions into many smaller experts) and shared expert isolation — achieving state-of-the-art open performance; concurrently, "Expert Choice Routing" (where experts select tokens rather than tokens selecting experts) gains traction as a load-balanced alternative.

2024

Snowflake Arctic, Databricks DBRX, and Alibaba's Qwen-MoE are all released as MoE architectures, solidifying MoE as the dominant paradigm for efficient frontier LLM training.

Why It Matters

🔬 Real-World Impact

GPT-4 (OpenAI, 2023)Widely reported (though not officially confirmed) to be an MoE architecture with approximately 8 experts and top-2 routing across 220B active parameters per token out of ~1.8T total — making it the most commercially deployed MoE system in history, serving billions of queries daily.

Gemini 1.5 Pro (Google DeepMind, 2024)A sparse MoE with a 1M-token context window; the combination of MoE efficiency and long-context attention is specifically enabled by MoE's ability to maintain expert specialization across extremely diverse long-range token distributions without proportionally increasing per-token cost.

Mixtral 8×7B (Mistral AI, 2023)Open-weight top-2 MoE with 8 experts per layer; benchmarks demonstrate it matches or exceeds Llama 2 70B on most tasks while requiring only ~18% of the inference FLOPs, making it the canonical existence proof that open MoE can compete with closed dense models.

DeepSeek-V2 (DeepSeek, 2024)Uses fine-grained MoE with 160 routed experts and 2 shared experts per layer, activating 21B of 236B total parameters; introduces "shared experts" (always-active experts for common knowledge) alongside routed experts for rare knowledge — a hybrid dense-sparse architecture.

Google's Switch Transformer (2021, internal)Used in production for Google Search and Translate workloads; the key finding was that at fixed compute budget, a sparse MoE always outperforms a dense model, which reshuffled Google's internal infrastructure investment decisively toward sparse architectures.

Meta's research infrastructure"Megablocks" (Gale et al., 2022) provides a sparse CUDA kernel for MoE that eliminates token dropping entirely by reformulating dispatch as block-sparse matrix multiplication — now integrated into multiple production training frameworks and enabling MoE without the capacity factor hyperparameter.

JetMoE (MIT/UC San Diego, 2024)A 2.2B-active-parameter open MoE trained for

\$0.1

M demonstrating that MoE efficiency gains extend to modest compute budgets, not just hyperscaler regimes — achieved comparable quality to Llama 2 13B while using half the inference FLOPs.