Mixture of Experts: sparse capacity.

Replace one dense MLP with conditional computation.

A decoder block normally applies the same feed-forward network to every token. An MoE block keeps the attention and residual structure, but replaces selected feed-forward layers with a bank of N experts. Each expert is usually an independent MLP with the same input and output width.

Sparsity changes the systems contract. The model must store every expert even though a token uses only a few. During distributed training, tokens may cross devices to reach their assigned experts. A sparse FLOP count can therefore hide memory and communication costs.

The router is a learned traffic controller.

For token representation x ∈ ℝᵈ, a linear router produces one logit per expert. A softmax converts these logits to routing weights. The layer selects the largest k weights, evaluates only those experts, and mixes their outputs.

scores = softmax(router(x), dim=-1)      # [tokens, experts]
weight, expert_id = topk(scores, k=2)
expert_in = dispatch(x, expert_id)
expert_out = experts(expert_in)
y = combine(expert_out, weight)

Top-1 routing reduces expert compute and simplifies combining; top-2 offers a second path and often smoother learning at a higher cost. Router noise, temperature, and score precision all influence early assignment. Router logits should usually remain in a stable numerical format even when expert MLPs use lower precision.

Every expert needs a token budget.

Hardware prefers bounded, rectangular batches. If a batch contains T routed token assignments and N experts, an even allocation gives T/N assignments per expert. A capacity factor c ≥ 1 reserves headroom.

A small capacity factor saves memory but increases dropped assignments when routing is uneven. A large factor wastes padded slots and makes step time less predictable. Capacity should be tuned from measured expert histograms, not chosen once and forgotten.

Specialization without balance becomes collapse.

If one expert receives the most probability early, it trains on more tokens, improves faster, and becomes even more attractive to the router. This positive feedback can leave other experts undertrained. An auxiliary objective counteracts that loop by penalizing uneven routing probability and uneven token counts.

Perfect uniformity is not the goal. Useful specialization can follow language, syntax, domain, or token type. The target is healthy utilization: no dead experts, no persistent hotspots, and no quality dependence on a single path. Track routing by corpus source as well as globally; an aggregate histogram can hide one domain monopolizing an expert.

Who chooses whom changes the optimization problem.

In token-choice routing, every token ranks the experts. Top-1 routing, popularized by Switch Transformers, sends each token down a single expert path and keeps active expert compute close to a dense MLP. Top-2 routing mixes two expert outputs, providing a second path at roughly twice the expert work. Larger k makes the layer less sparse and increases dispatch traffic.

Architectures also differ in where sparse layers appear. Replacing every MLP maximizes parameter capacity but creates frequent communication. Interleaving dense and sparse blocks reduces routing points. Shared experts provide a stable common path, while routed experts learn conditional transformations. These choices should be reported alongside total parameters and active parameters; the label “MoE” alone is not a complete architecture description.

Expert parallelism turns routing into all-to-all traffic.

When experts are sharded across devices, each worker groups local tokens by destination, exchanges token buffers, runs the resident experts, and returns outputs to the original token order. The exchange can dominate step time when token batches are small, experts are poorly placed, or routes are imbalanced.

Measure active parameters, total resident parameters, bytes exchanged, expert compute utilization, and padding separately. “Only two experts are active” does not imply a cheap system if the full expert bank must stay in memory or if every layer produces a network exchange.

Parallelism has multiple axes. Data parallelism replicates the model, tensor parallelism splits matrix operations, pipeline parallelism assigns layer ranges, and expert parallelism distributes experts. Their communication schedules can collide. A useful deployment plan places experts with topology in mind and profiles real message sizes instead of trusting ideal FLOP counts.

Sparse compute introduces dense operational costs.

The complete expert bank still occupies parameter memory, checkpoint storage, and optimizer state during training. A route that looks cheap in arithmetic can stall on all-to-all exchanges. Uneven token counts create stragglers because the step finishes at the pace of the busiest expert. Capacity overflow can silently drop or reroute assignments, changing the function being trained.

Expert specialization can also be fragile across domains: a route learned from the training mixture may become overloaded when production traffic shifts. The remedy is observability across layers, experts, domains, and time—not one global utilization average.