🧩 Mixture of Experts (MoE) — Explainer & Interactive Simulator

Mixture of Experts is a neural network architecture where, instead of every input passing through one big dense feed-forward block, a layer is split into many smaller expert sub-networks, and a lightweight gating (router) network decides — per token — which one or few experts should process it.

This is how models like Mixtral 8x7B, DeepSeek-MoE, Switch Transformer, and GShard scale to huge parameter counts while keeping inference compute roughly constant: a model might have hundreds of billions of total parameters, but only activate a few billion of them for any given token.

How it works

  1. Experts. Each MoE layer replaces one dense FFN with N independent FFN "experts" (E_0 … E_{N-1}), each with its own weights.
  2. Gating network. A small linear layer scores each token against every expert: logits = x @ W_gate, then softmax turns scores into a probability distribution over experts.
  3. Top-k routing. Only the top-k highest-scoring experts are actually run for that token (commonly k=1 for Switch Transformer, k=2 for Mixtral/ GShard). The token's output is the weighted sum of the chosen experts' outputs, weighted by their (renormalized) gate probabilities.
  4. Capacity limits. To keep training/inference batched and efficient, each expert is only allowed to process a fixed number of tokens per batch (capacity). Tokens that overflow a full expert are dropped for that layer (they skip via the residual connection instead).
  5. Load balancing. If left unconstrained, gating tends to collapse onto a few "favorite" experts. An auxiliary load-balancing loss penalizes uneven routing so all experts stay utilized and trained.

Why it matters

  • Sparse activation → you get the representational capacity of a huge dense model, but the compute cost per token of a much smaller one.
  • Specialization → different experts can implicitly specialize (e.g. by language, syntax pattern, or topic), though in practice this specialization is soft and not perfectly interpretable.
  • The catch → routing imbalance, dropped tokens, communication overhead (in distributed training experts often live on different devices), and training instability are the main engineering challenges.

🎛️ Try it yourself below

Tune the simulation parameters and watch: which expert each token gets routed to, how balanced the load is across experts, the gate-weight distribution, and the theoretical FLOPs savings vs. an equivalent dense layer.

8 512
2 64
2 16
1 8
0 3
0.5 3
1 8
0 9999

Notes on this simulation

  • Token embeddings are synthetic (sampled from a few Gaussian clusters) purely to make routing patterns visually interpretable — no real language model is involved.
  • The gating network is a single random linear layer, matching the structure (not the trained weights) of a real MoE router.
  • Capacity-based dropping and the auxiliary load-balancing loss follow the formulation in the Switch Transformer paper (Fedus et al., 2021) and GShard (Lepikhin et al., 2020).
  • This is meant as an intuition-building tool, not a benchmark of any real model's routing behavior.