Mixture of Experts Routes Capacity

Mixture of experts buys more parameters at the same FLOPs for each token. It charges the gain back in routing discipline, load balance, and all-to-all traffic. For a fixed FLOP budget an MoE design is usually the correct choice. The forward pass is the easy part of it.

Replace each dense MLP with many smaller MLPs called experts. Add a router that picks a few experts for each token. The FLOPs for each step hold and the parameter count climbs.

More parameters, the same FLOPs for each token

The name misleads. Nobody hand-assigns a coding expert or an English expert. The experts sit inside the MLP block, and the router learns from data which few of them a token wakes. Attention stays dense.

Give each expert the size of the dense MLP and route each token to one expert. The FLOPs for each forward pass match the dense model. The parameter count rises with the number of copies. The gain lands only if those parameters help the model store more facts and patterns.

Routing must stay sparse to keep that arithmetic. Waking every expert for every token raises the FLOPs too far to pay.

Fedus and coauthors reported in 2022 that more experts give lower loss at fixed FLOPs. Downstream scores moved with it. OLMo-MoE from AI2 reports the same pattern in recent models. By 2025 the advantage over dense models at fixed training FLOPs is clear.

The tokenizer note set the frame: efficiency decides how much scale you can afford. MoE is that lever one level up, and it buys more parameters for each FLOP you already pay.

Upcycling is the cheap way in. Start with a trained dense model, copy each MLP into several experts, and add a router. Perturb the copies if you want, then keep training as an MoE. If training stays stable, you keep what the dense model learned for a small amount of extra training. MiniCPM and Qwen report strong results from this path.

The gains are largest where multi-node training and complex parallelism are already necessary. Below that size the extra complexity costs more than it returns. If you can carry the complexity, MoE is cost-effective.

The hard part is shipping every token to its expert and back

MoE looks simple on slides and messy in production. Top-k selection is discrete. Gradient descent prefers smooth choices, and you cannot backprop through a choice. Naive training goes unstable and leaves many experts unused.

The systems half is heavier. MoE adds expert parallelism on top of whatever split the cluster already runs. The cluster shards the experts across devices, so each device holds one or a few. After routing, every token travels to the devices that hold its chosen experts. The experts run, the outputs come back, and the layer combines them. That exchange is all-to-all, and it happens at every MoE layer of every step.

On a multi-node cluster the exchange can cost a few milliseconds for each layer. It is the main trade against the compute gains and the memory gains. An MoE layer spends fewer FLOPs for each token and more bytes on the interconnect. The ceiling follows the bytes.

Modern kernels and libraries fuse many small expert matmuls into larger sparse operations. Small experts still run efficiently that way. The repair works on the compute side. The traffic stays.

Token-choice top-k routing won because the router must stay cheap

Each token arrives at the layer as a hidden state x, the residual input to the MLP. The router holds one vector e_i for each expert, kept separate from the expert weights.

• Score. Each expert gets a score equal to the dot product of x and its vector e_i.
• Normalize. A softmax over the experts turns those scores into s_i.
• Select. Take the top k by s_i. The gates g_i equal s_i for the chosen experts and zero for the rest.
• Compute. Run the chosen MLPs on x and weight each output by its gate. Add the products and send the sum into the residual stream.

That is the entire router: a linear layer, a softmax, and a top-k. Router compute that grows eats the savings the layer exists to create. The learning signal is weak and indirect, so more router machinery does not pay. A small router is also easier to stabilize and debug.

k is a hyperparameter with a price. At k equal to one you wake a single expert. At k equal to two you buy redundancy and exploration, and the MLP compute roughly doubles.

Earlier work argued for two so the router can compare experts instead of committing early. A larger k costs more FLOPs and more communication. Fine-grained experts shrink each expert, so a larger k stays affordable.

Five other routing designs lost, each for a named reason.

• Expert choice. Each expert picks its top-k tokens. That balances the load and swaps the question to what fits this expert.
• Global assignment. A matching problem balances the load exactly. Solving it costs too much once the model is large enough to need it.
• Hash routing. Sending tokens by hash ignores the semantics and still beats dense models. Even a crude partition across many MLPs helps.
• RL routing. Researchers tried it early because routing is discrete. It costs more than the simple recipe and does no better, so they stopped.
• Stochastic routing. Noise on the router logits pushes tokens toward other experts. It helped less than the balancing heuristics, and most systems dropped it.

Most strong open MoE models route this way, including OLMo-MoE, DeepSeek, Qwen, Mixtral, and the Google MoE variants. The recipe under all of them is the same: token-choice top-k routing plus explicit balancing.

Without balancing, training collapses onto a few experts

Without a balancing term, most tokens route to a few experts. Those experts get good at every task and the rest never learn. OLMo-MoE ran the ablation. Two experts take half the tokens and the others stay unused. Validation loss rises, and you paid for parameters that do no work.

The goal is to spread the traffic so every expert learns and no device turns into a bottleneck. The Switch-style auxiliary loss builds two vectors over the experts. f(i) is the fraction of tokens that reach expert i after top-k. p(i) is the fraction of router probability mass on expert i before top-k.

The loss is the dot product of f and p across the experts. An expert that already takes many tokens gets its router probability pushed down. Variants apply the loss for each expert and for each device.

DeepSeek v3 balances without an explicit auxiliary loss. It keeps a bias b_i for each expert and counts the tokens each expert took after every batch. An underused expert gets a higher b_i and an overused one gets a lower b_i.

Routing adds b_i to the score before the softmax or sigmoid. It stays out of the final gate weights. Tokens move toward the underused experts and the weights the layer applies do not change. A sequence-wise balancing loss sits on top. It corrects imbalance inside one sequence at inference rather than only across batches.

The router softmax goes numerically unstable, so systems compute it in float32. A z-loss on the router logits penalizes the log-softmax normalizer. Small logits mean fewer loss spikes.

The extra parameters overfit small fine-tuning sets. DeepSeek answers with very large supervised fine-tuning data, so the MoE choice reaches back into what you collect.

Fine-grained experts buy more than shared experts

The naive build copies the whole MLP into every expert, and the parameter count climbs fast. Fine-grained experts split the MLP expansion dimension instead. That dimension is often four times the hidden size. Cutting it into narrow pieces gives you many more experts while the compute for each token stays low.

Shared experts run for every token and hold the patterns that are useful everywhere. That is also their price. Their compute lands on the bill for every token. A fine-grained expert costs you only when the router picks it.

DeepSeek made both moves popular. Ablations in DeepSeek and OLMo-MoE show that more fine-grained experts improve the loss and the benchmarks. Shared experts help in some settings and are not necessary in all of them.

Early Google systems ran 8 to 16 experts in each layer with 1 or 2 active. GShard, Switch, and ST-MoE all sat in that band. DeepSeek and the labs that followed run dozens of fine-grained experts. Most of them keep at least one shared expert and several active for each token.

Spend the fine-grained split first. Add a shared expert when a measurement asks for one.

Token dropping makes inference non-deterministic at temperature zero

Experts and devices carry a capacity limit for each batch. When an expert receives more tokens than it can hold, some systems drop the overflow. Those tokens skip that expert, and the model does not use it for them.

The drops depend on what else is in the batch. Another request landing in the same batch changes which tokens the system drops. The same prompt at temperature zero then returns a different answer. A test you ran alone does not predict what a user gets inside a full batch.

I do not promise a reproducible answer from a system that drops tokens. Reproducibility turns into a property of how you batch. Measure the drop rate on a production-shaped batch before you make that promise.

The DeepSeek line changed routing and balancing while the layer types held

DeepSeek MoE v1 held about 16B parameters with about 2.8B active. Each layer carried 64 fine-grained experts and 2 shared experts, with several active for each token. Routing was token-choice top-k, with the softmax before the top-k and a weighted sum of expert outputs. Balancing was an auxiliary loss for each expert and for each device.

v2 reached about 236B total and about 21B active on the same MoE core. It added device top-m routing. The router picks a small set of devices for each token, then picks experts inside them. That cuts the traffic. Communication balancing losses came with it.

v3 reached about 671B total and about 37B active. Routing normalizes the gates so the expert outputs sum to one. It uses sigmoid gating in one part and keeps device top-m. Balancing moved to the online bias for each expert plus the sequence-wise loss.

From about 16B to about 671B, the changes sit in routing, balancing, and systems tuning. No new layer type appears. Device top-m is the change that attacks the traffic directly, and it stayed from v2 onward. Read the two parameter counts together, because the total sets the memory and the active count sets the FLOPs.

MLA answers the KV cache with a smaller cache

The architecture note named the KV cache as the reason attention turns memory-bound during decode. MLA is the answer DeepSeek v3 gives to that bill. It cuts KV cache memory and the FLOPs do not rise.

Standard attention caches the full keys and the full values for each token. MLA compresses the hidden state h_t into a smaller latent c_t.

The layer caches c_t alone. When it needs keys and values, it up-projects them.

Written that way, MLA adds matmuls. Merge the up-projection into the projections next to it and the count stays about the same. RoPE makes the merge harder, because it rotates Q and K between the projections. The MLA design keeps RoPE clean by rotating only the parts it leaves uncompressed.

Multi-token prediction adds signal without a new layer

Standard training predicts the next token. Multi-token prediction uses the hidden state at each position to predict further ahead. An extra head or a small layer does the work. DeepSeek v3 adds one head for one step ahead, which gives each sequence more learning signal and can improve efficiency.

The Builder Test

Read four numbers on the same window: expert load, routing by domain, quality by domain, and communication overhead. Routing by domain is the share of each domain that lands on each expert. One of them alone hides the failure. Expert load can look flat while one domain rides two experts. Average quality can look fine while that same domain gets worse.

Then change one routing knob and read the four again. If the loss improves while the drop rate rises, the win came with a cost the loss curve cannot show. When communication overhead climbs and quality holds, you bought the same model for more money.

What Carries

Conditional capacity is real. You pay for it at the router. Instrument the router first, because it decides which parameters wake and how many bytes cross the interconnect.

Routing turns compute into traffic, so the next question is what the hardware charges to move bytes.