Mixture of Experts · L2 Mechanics

The router, top-k gating, why load balance is a training-stability problem, and what expert parallelism does to a compute cluster.

Token embedding
Router
small linear layer → logits over experts
Top-k selection
pick k highest-scoring experts
Dispatch
send token to each chosen expert
Weighted sum
combine outputs by router score

One MoE feed-forward layer, per token. Attention layers are usually left dense.

Concretely, an MoE feed-forward layer replaces a single MLP with $E$ parallel MLPs ('experts') of the same shape, plus a router: a small linear layer producing $E$ logits per token. A softmax over those logits gives a probability-like weight per expert; top-k routing keeps only the $k$ largest (commonly $k = 1$ or $k = 2$ ) and zeroes the rest, so only $k$ experts do any compute for that token. The chosen experts' outputs are combined via a weighted sum using their (renormalized) router weights.