Home / Fundamentals / Transformer Architecture

Transformer ArchitectureSeniorQuestion 8 of 9

How does a Mixture-of-Experts transformer work, and what are its training and serving challenges?

30-second answer

MoE replaces the dense MLP with many expert MLPs and a learned router that sends each token to its top-k experts, so total parameters grow while per-token compute tracks only the active experts. Training needs load balancing (auxiliary losses or bias adjustment) to avoid router collapse, capacity limits that may drop tokens, and router stabilization such as z-loss. Serving must keep all experts in memory, uses expert parallelism with all-to-all communication, and needs large batches to keep experts busy; it pays off at high throughput, less so for small latency-sensitive deployments.

Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.

0–5 min Clarify and frame

Framing: "MoE decouples parameters from per-token compute. I'll describe routing, the load-balancing problem with its loss, capacity and token dropping, stability, and the systems implications for training and serving."

5–15 min Core explanation

An MoE layer replaces the dense MLP with E expert MLPs and a router:

logits = x W_r                     (E scores per token)
gates  = softmax(logits)
choose top-k experts by gate value
y = Σ_{e ∈ top-k} g_e · Expert_e(x)      (gates often renormalized over the top-k)

With E = 64 experts and k = 2 (or many small experts with larger k, plus one or more always-on shared experts), a model can hold hundreds of billions of parameters while each token uses a small fraction. Training and inference FLOPs scale with active parameters; capacity scales with total.

15–25 min Deep dive

Load balancing. Without pressure, the router collapses onto a few experts (they get better, so they're chosen more). A common auxiliary loss:

L_bal = α · E · Σ_e f_e · P_e
f_e = fraction of tokens dispatched to expert e (in the batch)
P_e = mean router probability assigned to expert e

It's minimized when routing is uniform. Too large an α hurts quality; some recent models instead adjust per-expert routing biases dynamically (no auxiliary loss) to balance load.

Capacity factor. Each expert processes at most capacity = cf · (tokens · k / E) tokens per batch. Overflow tokens are dropped (they skip the MLP via the residual) or rerouted. A low factor wastes less compute but drops more tokens; dropless implementations use variable-size grouped matmuls.

Stability. Router logits can grow and destabilize training; router z-loss penalizes (log Σ e^{logits})². Routers are often computed in fp32.

Fine-grained and shared experts. Many smaller experts (with larger k) give more flexible combinations; shared experts capture common knowledge so routed experts specialize.

25–32 min Systems trade-offs

Aspect Dense MoE
FLOPs per token All parameters Active parameters only
Memory All parameters All parameters must be resident
Parallelism Data, tensor, pipeline + expert parallelism with all-to-all
Efficiency at small batch Predictable Tokens scatter; experts under-utilized
Fine-tuning Standard Must keep load balanced; can overfit

Expert parallelism: experts are spread across GPUs; each MoE layer does an all-to-all to send tokens to their experts and another to bring results back. Interconnect bandwidth and load imbalance (stragglers) are the bottlenecks; overlapping communication with computation is key.

Serving: MoE shines at high throughput, where large batches keep all experts busy. For low-volume or latency-critical deployments, the memory footprint can make a dense model with the same active size more attractive.

32–37 min Production example

A team compared a dense model with an MoE model of similar quality whose active parameters were about a third of the dense model's. The MoE cut training FLOPs substantially and served more tokens per second per GPU at high load, but needed more GPUs to hold the weights. They used it for their high-volume API and kept a dense model for on-premises customers with small deployments.

37–40 min Likely follow-ups

  • Do experts specialize by topic? Mostly not in human-interpretable ways; specialization tends to be by token type, syntax or language.
  • How does top-k affect cost and quality? Compute scales roughly linearly with k; quality improves with diminishing returns.
  • What happens to dropped tokens? They pass through the residual connection unchanged for that layer.

What interviewers listen for

  • Active vs. total parameters.
  • The load-balancing loss, capacity factor and dropping.
  • All-to-all communication and serving trade-offs.

Red flags

  • "MoE uses less memory."
  • Ignoring load balancing.