All Lessons

Sparse Expert Routing and Capacity Control

Sparse expert layers increase parameter capacity while activating only a small subset for each token. This lesson derives routing, capacity limits, load balancing, and distributed communication costs.

AI Narration Press play to listen
0  / 6 paragraphs
Click any paragraph to jump · Scroll freely without breaking narration

A sparse expert layer contains several feed-forward networks but evaluates only a small number for each token. This separates total parameter count from active computation. The benefit is conditional capacity; the challenge is learning a router that uses experts productively without overloading a few and starving the rest.

For token representation $x$, router logits are $z=W_rx$ and probabilities are $p=\operatorname{softmax}(z)$. If $S(x)$ contains the selected top $k$ experts, the layer output can be written $$y=\sum_{e\in S(x)}\tilde p_e E_e(x),$$ where $E_e$ is expert $e$ and $\tilde p_e$ denotes routing weights normalized according to the implementation.

With $N$ tokens and $E$ experts, a nominal per-expert capacity is approximately $$C=\left\lceil c\frac{Nk}{E}\right\rceil,$$ where $c\geq1$ is a capacity factor. Tokens beyond capacity may be dropped, rerouted, or passed through a fallback. Each choice changes both computation and the effective function being trained.

A router can collapse by sending most tokens to a small expert subset. Auxiliary balancing losses compare average routing probability and actual token fractions across experts. These losses improve utilization but can conflict with specialization if weighted too strongly. Balanced traffic is a systems requirement, not proof that experts learned distinct capabilities.

Sparse routing creates communication. In distributed training, tokens are dispatched to devices holding selected experts and then gathered back into original sequence order. All-to-all traffic, padding to capacity, and uneven token counts can dominate runtime. Parameter count alone therefore says little about throughput.

Log expert utilization, probability entropy, overflow rate, dropped-token rate, load-balancing loss, and per-expert gradient norms. Evaluate whether routing changes with input content and whether expert removal has localized effects. A model can report low training loss while silently wasting experts or relying on a small routing bottleneck.