🧩 Mixture of Experts

What is Mixture of Experts (MoE)?

A Mixture of Experts model replaces a dense feed-forward (FFN) layer with N parallel expert FFNs and a lightweight router (or gating network) that selects a small subset of experts per token. Only the chosen experts are activated, so the model's total parameter count can be vastly larger than the parameters actually used for any single token. This decouples model capacity from per-token FLOPs, the central appeal of MoE.

A typical MoE block: y = Σ_i G(x)_i · Expert_i(x), where G(x) is the router output and only top-k entries are non-zero.

The Gating Mechanism (Router)

The router is usually a single linear layer W_r that produces a logit per expert: logits = x · W_r. A softmax (or sigmoid) turns these into routing weights. Most production systems use top-k routing (k is typically 1 or 2): each token is dispatched to the k experts with the highest scores, and their outputs are combined using the normalized routing weights.

• Switch Transformer (k=1): Maximally sparse, one expert per token, minimal compute.
• GShard / Mixtral (k=2): Slightly more compute, better quality, smoother gradients to the router.
• Expert-Choice routing: Inverts the perspective so experts pick the top tokens they want, guaranteeing perfect load balance.

Historical Context

• Shazeer et al. (2017), "Outrageously Large Neural Networks": The original sparsely-gated MoE applied to LSTMs, introducing top-k gating and noisy gating for exploration.
• GShard (2020) and Switch Transformer (2021) brought MoE to Transformers at scale (Switch reached 1.6T parameters with constant per-token cost).
• GLaM (2021) demonstrated MoE could match dense LLM quality at a third of the training energy.

Modern MoE Models (as of May 2026)

• Mixtral 8x7B (Mistral, Dec 2023): 8 experts, top-2 routing, ~47B total / ~13B active parameters per token. Matched Llama 2 70B on most benchmarks at ~5× the inference speed.
• Mixtral 8x22B (Apr 2024): Same architecture, larger experts, ~141B total / ~39B active.
• DeepSeek-V2 (May 2024): 236B total / 21B active, introduced fine-grained experts + shared experts.
• DeepSeek-V3 (Dec 2024): 671B total / 37B active, 256 routed experts + 1 shared expert, top-8 routing, auxiliary-loss-free load balancing. Trained for ~$5.6M, demonstrating MoE's training efficiency at frontier scale.
• Grok-1 (xAI, Mar 2024): 314B total, 8 experts, top-2.
• DBRX (Databricks, Mar 2024): 132B total / 36B active, 16 experts, top-4 routing.

Load Balancing

Without intervention, the router collapses: a few "winner" experts get all the tokens, the others receive nothing and never learn. Standard fixes:

• Auxiliary load-balancing loss (Switch Transformer): adds α · N · Σ_i f_i · P_i to the training loss, where f_i is the fraction of tokens sent to expert i and P_i is the mean router probability for it. This pushes both quantities toward uniform.
• Expert capacity: Each expert can process at most capacity = (tokens / N) · capacity_factor tokens per batch. Tokens beyond the cap are dropped (skip the FFN via the residual connection) or rerouted.
• Auxiliary-loss-free balancing (DeepSeek-V3): Adds a learned per-expert bias to the routing logits and adjusts it online so utilization stays balanced, avoiding the gradient interference an auxiliary loss can cause.

The MoE routing playground shows this live: with no balancing pressure the router collapses onto a couple of experts while the rest go dead and over-capacity tokens are dropped — turn the pressure up and watch the load spread back toward uniform.

Training Challenges

• Expert collapse: Without load balancing, the router degenerates to using a handful of experts.
• Routing instability: Discrete top-k decisions create non-differentiable transitions; small parameter changes can flip a token's assigned expert and cause loss spikes. Mixing Z-loss (regularizing router logit magnitudes) and BF16 routers (instead of FP16) helps.
• All-to-all communication overhead: In expert parallelism, every device must shuffle tokens to wherever their assigned expert lives. This is bandwidth-heavy and is often the dominant cost in MoE training.
• Batch size sensitivity: MoE benefits from very large batch sizes so each expert sees enough tokens to learn.

Expert Parallelism vs Tensor Parallelism

In expert parallelism (EP), each expert is placed on a different device; tokens are routed via all-to-all collectives, processed locally, and shuffled back. EP scales naturally with expert count but is bottlenecked by network bandwidth.

Tensor parallelism (TP) shards a single dense matmul across devices using all-reduce. For MoE, TP can be combined within an expert (shard each expert across a few GPUs) while EP shards across experts. Production MoE training uses 3D or 4D parallelism: DP × TP × EP × PP.

Fine-Grained Experts (DeepSeek)

Instead of a few large experts, DeepSeek splits each expert into many smaller ones (e.g., 256 experts of d_ff/8 size instead of 8 experts of full size) while keeping the same total parameter count. With top-k raised to compensate, this increases the combinatorial diversity of expert subsets (C(256, 8) ≫ C(8, 2)), encouraging specialization. Shared experts (always-on experts that every token uses) capture general patterns so the routed experts can specialize.

MoE for Inference

The asymmetric resource profile of MoE matters in deployment:

• Memory: All experts must be loaded in VRAM (you don't know which ones a token will pick). Mixtral 8x7B needs the full ~94 GB in FP16, the same as a 47B dense model.
• Compute / latency: Only k experts run per token, so per-token FLOPs match a smaller dense model. Mixtral 8x7B's decoding throughput is comparable to a ~13B dense model.
• Serving strategies: Continuous batching helps load-balance experts across requests. Expert offloading (CPU/SSD) trades latency for VRAM. Speculative decoding with a small dense draft model pairs well with MoE.