DL0109 Mixture-of-Experts MoE Routing

A sparse mixture-of-experts layer replaces one feed-forward block with parallel feed-forward blocks (the experts) plus a tiny linear router, and sends each token to only of them. This is conditional computation: parameter count grows with while the FLOPs per token grow only with , so Mixtral 8x7B holds 46.7B parameters but activates about 12.9B per token, and DeepSeek-V3 holds 671B while activating 37B. The interesting engineering is not the experts, which are ordinary MLPs, but the routing strategy, because a learned router is free to collapse onto a few favorite experts and then most of the capacity you paid for is never trained. The dominant strategy is token-choice top-k: every token picks its own experts, each expert has a fixed buffer set by a capacity factor, and tokens that arrive at a full expert are dropped through the residual connection. The main alternatives invert or remove that choice: expert-choice routing lets each expert select its top tokens (perfect load balance, but some tokens get no expert), hash layers assign tokens deterministically with no learned router at all, and BASE layers solve a global linear assignment per batch. The trade-off axis is always the same, namely how much balance you buy and what you pay for it in dropped tokens, extra loss terms, batch dependence, and all-to-all communication.

(1) Parameters Decouple From FLOPs: with a token touches 25% of the expert weights, so quality scales with total parameters while latency scales with active parameters.
(2) The Router Is One Matrix: a single projection followed by top-k selection and a renormalized softmax over the selected logits, which is why gradients reach the router only through the chosen experts.
(3) Load Balancing Is The Central Problem: Switch-style training adds an auxiliary balance loss (typically ) and a router z-loss to keep logits small and the assignment spread out.
(4) Capacity Factor Controls Token Dropping: a buffer of slots per expert means near 1.0 is cheap but discards tokens whenever the router skews.
(5) Routing Families Differ In Who Chooses: token-choice, expert-choice, deterministic hashing, and global assignment sit on a spectrum from fully learned and imbalanced to fully balanced and inflexible.
(6) Serving Cost Is Memory, Not Compute: every expert must be resident in HBM even though each token uses of them, and expert parallelism adds two all-to-all collectives per MoE layer.

Figure 1: The two families differ in the direction of selection. Token-choice top-k guarantees every token gets experts but not that experts get equal load, so overflow is dropped at capacity; expert-choice guarantees equal expert load but not that every token is served.

Token-choice top-k stays the production default for autoregressive language models for one structural reason: it is a per-token function, so the routing decision for token does not depend on any other token in the batch. Expert-choice and BASE layers both rank or match tokens against each other, which makes the forward pass batch-dependent, breaks the causal guarantee during teacher-forced training, and cannot be reproduced at decode time when the batch is a single token. Hash layers show how much of MoE’s benefit comes from capacity rather than from clever routing, since a fixed hash of the token ID is perfectly balanced by construction and still recovers a large part of the gain, but it can never learn semantic specialization. In practice the fix for token-choice imbalance is not to abandon it: DeepSeek-V3 keeps token-choice top-8 over 256 fine-grained experts, adds one always-on shared expert to absorb common knowledge, restricts each token to experts on at most 4 nodes to bound communication, and replaces the auxiliary loss with a per-expert bias that is nudged up or down to equalize load.

Mathematical Formulation:

Where:

• is the token hidden state and the MoE layer output, written with the residual path that a dropped token falls back to.
• is the router, its logits, and the -th expert MLP; is the expert count and the number kept.
• is the selected index set and the gate weight, renormalized over the selected logits only so the weights sum to 1.
• is the fraction of tokens in the batch dispatched to expert (piecewise constant, no gradient) and the mean router probability for expert (differentiable), so pushes probability mass away from overloaded experts.
• is the balance-loss coefficient, commonly ; the product equals 1 under a perfectly uniform assignment and grows toward under total collapse.
• is the tokens per device, the capacity factor (usually in training, larger at eval), and the per-expert buffer; any token beyond is dropped.

Figure 2: Token dropping is a joint function of the capacity factor and the router’s skew. Under a balanced router, drops nothing; under a collapsed router, even doubling capacity leaves a quarter of the assignments discarded, which is why the balance loss matters more than the buffer size.

Figure 3: The full MoE layer in one picture. A single linear router produces logits over all experts, the top-k selection picks only of them, and the selected expert outputs are combined with renormalized gate weights and added back through the residual connection. Non-selected experts stay resident in memory but consume zero FLOPs for this token, which is why total parameters scale with while active compute scales with .