Mixture-of-Experts: Sparse Computation for Larger Models
A dense neural network has one knob — make it bigger and you pay more compute on every input. A mixture-of-experts network has two knobs. You can grow the total parameter count almost arbitrarily, while holding the FLOPs spent per input nearly constant. Mixtral 8x7B stores 47B parameters but spends ~13B of compute per token. DeepSeek-V3 stores 671B parameters but spends ~37B per token. GPT-4 is widely rumored to do the same. This decoupling — params and FLOPs as two independent dials — is the single most important architectural primitive of the 2020s, and it is not a transformer trick. It is a layer-replacement trick that works on vision transformers, recsys MLPs, and dense feed-forward blocks of any kind. This lesson treats MoE as the general primitive it is.
Learning Objectives
After this lesson, you will be able to:
Explain the dense-vs-sparse compute split — why dense models scale linearly in compute with parameters, while MoE scales sublinearly by activating only k of E experts per input
Walk through the routing math — gating logits, top-k selection, weighted expert combination — and the auxiliary load-balancing loss that prevents router collapse
Diagnose the load-balancing problem, expert-capacity dropping, and the 2024-frontier auxiliary-loss-free routing scheme from DeepSeek-V3
Pick MoE vs dense for non-language workloads — vision transformers, recsys, dense FFNs — and recognize the inference-time costs (memory, all-to-all latency) that make MoE a bad fit for single-GPU serving
Don't worry if "mixture of experts" sounds frontier-only. The math is two equations — gating with top-k, then weighted sum — and the hard parts are infrastructure, not modeling. Once you see the routing as a sparse softmax that picks a few experts, the rest is plumbing.
Every architecture before sparsely-gated MoE assumed a fixed compute budget per input meant a fixed parameter budget per input. MoE breaks that link without changing anything else about the model — same backprop, same optimizer, same loss, same data pipeline.
An MoE block is a drop-in replacement for any dense feed-forward layer. The interface is identical — input shape, output shape, gradient flow — but the per-input compute is much smaller than the per-input parameter count would suggest.
Top-1 (Switch Transformer). Each input routes to exactly one expert. FLOPs per input = single-expert FLOPs. Maximally sparse, maximally cheap. The downside: routing decisions are binary, and a single bad route is catastrophic for that input. Switch used a "noisy top-1" with stochastic routing during training to mitigate.
Top-2 (GShard, Mixtral). Each input routes to two experts and their outputs blend. FLOPs per input = 2x single-expert FLOPs. The redundancy hedges against bad routing decisions: even if the top expert is wrong, the second-best usually catches the slack. This is the modern default for open-source MoE.
Top-k for k ≥ 3 is rare. Diminishing returns on quality, growing communication and capacity overhead.
Send tokens through the router and watch the gate light up only the top-k experts while the rest sit dark — sparse activation in motion.
Loading visualization...
What Do You Think?
Mixtral 8x7B uses top-2 routing across 8 experts plus shared attention. Total parameter count is 47B, but only ~13B participate in any single forward pass. What does '8x7B' mean — is each expert exactly 7B parameters?
Without intervention, MoE training collapses into a Matthew-effect loop:
Early in training, one or two experts are randomly slightly better at predicting common patterns.
The router learns this and routes more inputs to those experts.
Those experts get more gradient updates and improve faster than the underutilized ones.
Within a few thousand steps, one expert hogs 80–95% of tokens. The rest sit idle, get no learning signal, and atrophy.
The system has degenerated into approximately-dense behavior with a few wasted experts hanging off the side. All the parameter-efficiency wins of MoE are gone.
What Do You Think?
You train an MoE model with E=8 experts and top-2 routing, but you forget to include any load-balancing mechanism. After 5,000 training steps, what does the token-to-expert distribution look like?
The classical Shazeer 2017 fix is an auxiliary loss term that penalizes uneven distributions. Define f_i as the fraction of inputs routed to expert i, and P_i as the mean routing probability assigned to expert i across the batch. The aux loss is:
Laux=α⋅E⋅i=1∑Efi⋅Pi
This is added to the main task loss, and the router learns to balance load while still picking experts that minimize task error. It works, but it has well-known drawbacks: the aux gradient conflicts with the main gradient (the router gets pulled in two directions), and the alpha hyperparameter is finicky — too low and load is unbalanced, too high and routing degenerates to uniform random.
There's a subtlety the routing equation hides. Real MoE training happens in batches, and each expert in the batch lives on some GPU with finite memory and a fixed compute budget. If the router's load happens to spike — say 60% of tokens in this batch want expert 3 — then expert 3 cannot physically process all of them in one step without exceeding its memory or its compute budget. So MoE implementations introduce expert capacity.
Expert capacity is the maximum number of tokens an expert can process in a single batch:
A typical capacity_factor is 1.0 (each expert handles its fair share if load is uniform) up to 1.25 (each can handle 25% over fair share to absorb routing imbalances). If load exceeds capacity, the overflow tokens are dropped — they bypass the MoE block entirely, often via a residual connection that lets them flow unchanged to the next layer.
Dropping tokens sounds catastrophic, but in practice it's manageable: with a good load balancer, the drop rate is 1–5% per layer, and the residual flow gives those tokens at least the previous layer's representation. Modern implementations (Mixtral, DeepSeek-V3) use various overflow tricks: route the overflow to the second-choice expert, or to a shared expert, or rely on dynamic capacity scaling.
Quick check
Why does DeepSeek-V3's auxiliary-loss-free routing avoid the gradient-conflict problem that the classical aux loss has?
Single-GPU MoE training is rarely interesting. The whole point of MoE is to have far more parameters than fit on one GPU. The standard pattern is expert parallelism: each GPU hosts a subset of the experts, and tokens are shuffled across GPUs to reach the experts they were routed to.
Concretely, with E=8 experts and 8 GPUs in the expert-parallel group:
Each GPU starts with a batch of tokens — its own data-parallel slice.
The router runs locally on each GPU and decides which expert each token wants.
All-to-all communication: each GPU sends its tokens-for-expert-i to the GPU that hosts expert i. Every GPU sends to every other GPU.
Each expert runs locally on its share of tokens.
All-to-all communication (reverse): results flow back to the originating GPUs.
The MoE block output is now correctly distributed; training continues.
The two all-to-all operations are the dominant cost of MoE training, and they only work well if the GPUs share high-bandwidth interconnect — NVLink within a node, InfiniBand or NVL-NIC between nodes. Over slow ethernet, MoE training is borderline unworkable because every layer pays an all-to-all latency tax.
Quick check
Why does expert parallelism require all-to-all communication rather than the all-reduce that data-parallel training uses?
Below is a tiny end-to-end MoE block in numpy: 4 experts, each a single linear layer, with top-2 routing trained on a 4-cluster synthetic dataset. With the seed and learning rates as shipped, each of the four experts ends up owning a different cluster after 800 steps — the routing table comes out as a clean permutation (cluster 0 → E3, cluster 1 → E1, cluster 2 → E0, cluster 3 → E2), and the block classifies all 800 points correctly.
That outcome is not automatic, and the reason is the section you just read. The router here trains at lr_router = 0.01 while the experts train at lr = 0.05. Raise lr_router to 0.05 and rerun: the router's weights blow up to roughly 200 in magnitude, the top-2 softmax saturates to hard 0/1, and two of the four experts stop receiving any traffic at all — clusters 0 and 3 both pile onto E2, clusters 1 and 2 both onto E0. Accuracy stays at 1.000, which is exactly what makes the failure sneaky: the loss is happy while half the parameters sit idle. That is the Matthew-effect collapse from the quiz above, reproduced in twenty lines. Expert specialization is also seed-dependent — change default_rng(7) and you may get three experts sharing four clusters even at the lower rate, which is precisely why real MoEs need an explicit balance mechanism rather than a well-chosen learning rate.
Loading visualization...
The diagnostic table at the end is the payoff. Each row is a cluster; each column is an expert; each cell is the average routing weight. Read that matrix rather than assuming what it will say. If each cluster concentrates on its own expert, that is specialization learned end-to-end from a single cross-entropy signal. If instead some experts receive almost no traffic while others absorb several clusters, you are looking at router collapse — the failure mode described above, reproduced live.
This is the same primitive that powers Mixtral and DeepSeek-V3, just at 10,000x the scale and with infrastructure (expert parallelism, all-to-all, capacity dropping) bolted on. The math underneath is exactly what you see in the playground.
What Do You Think?
DeepSeek-V3 reports 256 routed experts plus 1 shared expert per MoE layer, top-8 routing among the routed experts. Suppose each expert FFN is sized comparably to a ~7B-dense-model FFN (~1.4B params for the FFN portion). Roughly how many total FFN parameters does this give the model, and how many are active per token? Report total in trillions and active in billions.
Most MoE coverage in the wild focuses on language models, but MoE is a layer-level primitive, not a language-model trick. Any architecture with dense feed-forward blocks can swap them for expert banks.
V-MoE (Riquelme et al., NeurIPS 2021) was the first big sparse vision transformer. The authors swapped the FFN of every other ViT block with an MoE block (E=32, top-k=2), scaled to a 14.7B-parameter V-MoE-H model (vs the dense ViT-H at 632M), and matched or exceeded dense performance on ImageNet at much lower per-image FLOPs. Key finding: vision experts specialize on visual concepts (textures, object parts, scene types) in ways that mirror how language experts specialize on syntactic categories.
GLaM-Vision and follow-ups continued this line, applying fine-grained MoE to image classification and detection. DeepSeek-VL and Qwen2-VL use MoE in their language-model backbone but route on multimodal tokens (image patches and text), with expert specialization that respects modality boundaries.
Recsys models are some of the largest deployed neural networks in production — Meta's DLRM-style models routinely cross 1T parameters, almost entirely in embedding tables. Their dense MLP layers (used for feature interaction and ranking) are prime MoE targets. Meta's published work on GShard-style routing for ranking models and ByteDance's TikTok recommender internals both use MoE in the ranking MLP stack — different experts handle different user-content interaction patterns (e.g., short-form vs long-form, new vs returning users).
The trade-off in recsys is sharper than in language: latency budgets are tight (often single-digit milliseconds for a ranking call), so MoE has to pay for itself in quality lift while staying under the latency ceiling. The wins are real but more constrained than in language.
The general pattern is "any dense block, sparsified": dense classifier heads, dense projection layers in multimodal fusion, dense bottleneck MLPs in autoencoders. The infrastructure overhead (routing, capacity, all-to-all) only pays for itself above a certain scale — typically when total parameters exceed 1–10B — but the primitive itself is architecture-agnostic.
Single-GPU inference. All E experts must fit in GPU memory because routing decisions are made at runtime — you cannot predict in advance which experts you'll need. So memory is dominated by total params, not active params. A 47B-param Mixtral takes the same VRAM as a dense 47B model; the "13B active" only saves FLOPs, not memory. For single-GPU serving of small models, dense is simpler and roughly as fast.
Strict low-latency requirements. All-to-all communication for expert parallelism adds latency that scales with cluster size. For ms-level latency budgets (recsys ranking, real-time speech), the comm overhead can dominate. Some systems pin all experts to one GPU to dodge the all-to-all, but then you've lost the parameter-scaling win.
Small models. MoE only pays off above 1–10B total parameters. At smaller scales, the routing overhead, capacity dropping, and load-balancing headache outweigh the FLOPs savings. A dense 1B model is much simpler and often faster than a 4B-param MoE with similar quality.
Limited training data. MoE's parameter count grows fast, and the experts need data to specialize on. If you only have a few hundred million tokens of training data, you can't fit a 47B-param sparse model — the experts won't get enough samples each to differentiate.
A condensed timeline of MoE models you should know:
Switch Transformer (Google Brain, 2021) — first big transformer-MoE LLM. E up to 2048, k=1 (top-1 routing). Trained on 750B tokens. Demonstrated that simpler routing (k=1) plus aggressive capacity factors can train stably at trillion-parameter scale.
GShard (Google, 2020) — the infrastructure paper that made top-2 routing on TPU pods practical. Introduced expert parallelism and the all-to-all-based MoE training pattern that everyone still uses.
ST-MoE (Google, 2022) — stable training of sparse MoE. Introduced the "router z-loss" trick for numerical stability, and characterized when MoE training diverges (large variance in routing logits causes softmax overflow during fp16 training).
Mixtral 8x7B (Mistral AI, December 2023) — the first widely-adopted open-source MoE LLM. 47B total params, ~13B active per token via top-2 routing among 8 experts. Single-handedly made MoE a mainstream architecture choice.
DeepSeek-V2 (May 2024) — 236B total, 21B active. Introduced DeepSeekMoE: fine-grained experts (more, smaller) + shared experts (always active for general knowledge) + Multi-head Latent Attention. Showed the recipe could scale beyond Mixtral's coarse design.
DeepSeek-V3 (December 2024) — 671B total, 37B active. Added the auxiliary-loss-free routing scheme. Demonstrated GPT-4-class performance at a fraction of the training cost; reset everyone's assumptions about MoE economics.
DBRX (Databricks, March 2024) — 132B total, 36B active. 16 experts, top-4 routing. Notable for its fine-grained design and strong open-source code performance.
Qwen2-MoE (Alibaba, 2024) — 57B total, 14B active. Demonstrated MoE quality at relatively small total-parameter sizes.
Falcon Mamba MoE (TII, 2024) — combined Mamba state-space layers with MoE FFN. Showed the primitive generalizes outside attention-only transformers.
V-MoE (Google, 2021) and GLaM-Vision — the vision branch. Established that MoE specialization happens on visual concepts the same way it does on linguistic ones.
The lineage you should remember: Jordan & Jacobs 1991 → Shazeer 2017 → Switch & GShard 2021 → Mixtral 2023 → DeepSeek-V3 2024. Each step changed one major thing: from theory to practice, from LSTM to transformer, from dense routing to top-k, from research lab to open source, from auxiliary loss to bias-controller routing.
#Architect's View: Picking MoE for a Real Workload
A practical decision tree for "should I use MoE for my next model":
Is your model going to be larger than ~10B total parameters and served on multi-GPU infrastructure? If no, dense is probably better. If yes, continue.
Do you have enough training data to drive E experts to specialize? Rule of thumb: at least 100B tokens for a small (8-expert) MoE, more for fine-grained designs. If no, dense is probably better.
Is your inference workload throughput-bound (batches of inputs amortize over the same experts) or latency-bound (single-input low-latency calls)? MoE wins big on throughput-bound serving. For latency-bound serving on a single GPU, dense often wins because the per-token FLOPs savings are eaten by memory-bound inference (all experts live in VRAM regardless).
Is your interconnect high-bandwidth (NVLink within node, InfiniBand between nodes)? If yes, expert parallelism is practical. If no (ethernet-only), MoE training will be agonizingly slow.
Are you willing to absorb the operational complexity? MoE adds load monitoring, capacity tuning, expert-utilization dashboards, and a new family of failure modes (router collapse, capacity drops, hot experts). Production MoE is non-trivial; if you have a small team, dense is operationally simpler.
If all five are "yes," MoE is likely the right call. For most teams below the frontier-lab scale, the honest answer is "dense first, MoE when you've hit the wall on dense."
MoE is a primitive, not a transformer trick. Any dense feed-forward block — vision transformer FFN, recsys MLP, classifier head — can be swapped for an expert bank plus a router. The decoupling of total parameters from per-input FLOPs is the whole point and it generalizes across architectures.
Two equations carry the math. Gating with top-k masking, then weighted sum over the surviving experts. Everything else (capacity, load balance, all-to-all) is infrastructure.
Router collapse is the default failure mode and must be addressed explicitly. The classical fix is an auxiliary load-balancing loss; the 2024-frontier fix is DeepSeek-V3's bias-controller scheme that avoids gradient conflicts entirely.
Top-1 is cheap and stable at huge E; top-2 is the practical default; top-k≥3 is rarely worth it. Modern fine-grained designs (DeepSeek-V3) push k up to 8 but with much smaller experts to keep per-token compute flat.
MoE saves inference compute, not training compute. Training a 47B-MoE costs roughly what training an equivalent-quality dense model would — the win is at inference time, especially in throughput-bound workloads.
Communication is the dominant cost. Expert parallelism needs all-to-all, which needs high-bandwidth interconnect (NVLink + InfiniBand). On slow networks, MoE training stalls.
Memory is the silent inference cost. All E experts live in GPU memory because you can't predict routing in advance; "active params" describes FLOPs, not memory footprint.
A team trains a dense 70B-parameter model and notices inference costs are crushing their unit economics. They switch to a 70B-active / 280B-total MoE with top-2 routing among 8 experts. Quality is comparable. What changes in their cost structure?
MoE is the most important architectural primitive of the 2020s precisely because it is so general. Anywhere a dense feed-forward block lives, an expert bank can take its place. Once you see the routing as a sparse softmax over a learned weight matrix, the rest — capacity, load balance, all-to-all — is engineering details that the field has slowly worked out over almost a decade. Build the primitive once at small scale; you've understood every modern frontier model.