Mixture of Experts

Why sparse experts?

Inside a Transformer, attention mixes information between tokens. The feed-forward network (FFN) then transforms each token’s hidden vector independently. MoE usually replaces that FFN—not the whole Transformer—with a bank of FFNs with separate weights. An “expert” is a learned transformation, not a hand-written tool or a pre-assigned subject specialist.

Let E be the number of equal-sized routed experts and k the number selected per token. Stored expert parameters scale with E; active expert FFN work scales with k. Router computation, attention, shared experts and communication add further costs. Increasing E at fixed k can add capacity without a proportional increase in active FFN FLOPs. It still increases memory requirements.

For training estimates such as 6ND, use active parameters in N and count every executed pass. Use total parameters for stored weights and much of the optimizer-state budget. Routing, long-context attention, imbalance and recomputation require additional accounting.

The sum is the defining mixture: a token’s output can combine several transformations rather than pass through a single fixed FFN. To understand what that mixture learns, first separate the router’s choices from the experts’ computation.

Follow one token: score → select → execute → combine

Let x be the hidden state arriving at the MoE block. A small learned router projects it to E logits: raw scores before activation or normalization. Softmax or sigmoid converts them to affinity scores. Top-k selects the expert indices, and gate weights determine each selected expert’s contribution. Computing a score does not require executing that expert’s FFN.

In the illustration, E1 and E2 transform the same input vector independently. The values +2 and −1 show one output coordinate. The real weighted sum operates across all d coordinates. The five stages are representation, scores, selection, expert execution and combination.

Two score parameterizations—not two different dispatchers

Hard top-k chooses indices from either score parameterization. The selected index does not provide an ordinary gradient to every unselected expert. Executed experts and differentiable gate weights still receive task gradients. An auxiliary objective can add another router training signal.

Also inspect selected gate weights. Balanced assignment counts can hide experts with almost zero contribution to the mixture. Counts measure which FFNs execute, not how much their outputs affect the result.

For top-1, renormalizing the only selected score to one removes its weighting gradient. Keep Switch’s selected probability weight; do not substitute the normalized top-2 recipe.

One token is straightforward. The difficulty appears when many tokens independently prefer the same expert.

A batch turns routing into an allocation problem

Consider the eight-token synthetic batch used below. Top-1 routing initially sends six tokens to E1, two to E2, and none to E3 or E4. Each local decision can look reasonable from its scores, yet the combined assignment is poor for utilization.

Imbalance affects both execution and learning. Expert workers wait for the busiest worker, so idle workers cannot reduce its processing time. Selected experts receive task updates. Unused experts receive no task gradient through their FFNs on this batch. Persistent imbalance can reinforce a small subset and leave other stored experts effectively untrained. Those unused weights still consume memory.

Balancing should prevent starvation and excessive load without removing useful preferences. A flat histogram measures assignments, not expert quality or output contribution. Balancing scope determines which population that histogram represents.

Before changing how the router learns, there is an immediate systems question: what happens when an expert receives more assignments than its buffers can hold?

Capacity: where an overloaded route becomes a dropped assignment

A fixed-capacity dispatcher allocates a bounded number of token slots per expert. The average assignment count is kT/E for T valid tokens, so a capacity factor γ adds slack around that average.

For our top-1 example, T = 8, E = 4 and γ = 1 give C = 2. E1 receives six assignments but has two slots: four overflow. With dropping, their E1 contribution is omitted; that does not necessarily delete those tokens from the sequence, since the surrounding Transformer can retain its residual path.

Increasing γ adds room but can enlarge padded buffers, memory use and transfers. A dropless dispatcher instead processes all assignments with variable expert batch sizes. It avoids missing expert contributions but must still execute the busiest expert. Dropless routing and balanced routing are different properties.

Capacity handles the consequence of imbalance. The next methods try to change the assignments that cause it.

Auxiliary balancing: teach the router to spread work

The first approach adds a load-balancing term to the task loss. The router must then select useful experts for prediction while avoiding repeated concentration on a small subset.

With top-1 dispatch, let $f_i$ be the fraction of valid tokens sent to expert i and $p_i$ the batch-average router probability for that expert.

Treat discrete counts as fixed for differentiation. The probability term supplies the router gradient. Too little balancing leaves overloaded or unused experts. Too much can worsen the language-model objective. Equal counts do not prove that the model allocates computation optimally. Switch ↗

Why multiply a hard load by a soft probability?

The hard fraction $f_i$ records actual assignments. The soft average $p_i$ is differentiable, so the objective can change future scores. Treat dispatch counts as fixed for this gradient. An expert with larger $f_i$ then has a larger coefficient on $p_i$. Moving probability away from heavily loaded experts can reduce the auxiliary term.

For a concrete snapshot, take $f=[0.75,0.25,0,0]$ and average probabilities $p=[0.50,0.25,0.15,0.10]$. The normalized auxiliary term is:

If f remains fixed for this local comparison, shifting 0.10 probability mass from E1 to E3 lowers the normalized term by $4\cdot0.75\cdot0.10=0.30$. In training, later assignments also change; this is a gradient intuition, not a complete optimization trajectory.

The coefficient α controls the tradeoff. Too small a value leaves imbalance; too large a value can worsen the language-model loss. Can a separate controller balance the batch without adding this competing gradient to representation learning?

Loss-free balancing: separate preferences from load control

The learned affinity describes an expert’s suitability for the current representation. A correction bias reflects recent load. The controller combines them for selection but retains original affinities for output mixing. This changes assignments without treating load correction as a learned task preference.

Watch a routing boundary move

When the E3 bias crosses the selection boundary, E3 replaces E2. Its mixture weight still uses the original 0.30 affinity, not the biased selection score. The controller changes which FFNs execute, not the affinities that weight their outputs.

Add a correction bᵢ only to the score used for expert selection. This controller replaces the main batch-balancing gradient. It increases bias for underloaded experts and decreases it for overloaded experts.

Reduce counts across workers before updating global biases. Save controller state in checkpoints for deterministic resume. Compare update speed with routing churn, not just mean load. “Loss-free” describes the main batch controller. DeepSeek-V3 still uses a small sequence-level auxiliary term. V3 §2.1.2 ↗

The single-token slider shows what a bias changes. The batch controller shows where that bias comes from. Real training changes affinities and representations between steps. This demonstration holds them fixed to isolate the load controller.

A sign controller uses only “too high” or “too low.” It ignores how far scores lie from a selection boundary. Quantile balancing tries to use that missing geometry.

Quantile balancing: estimate the boundary you need

If too many tokens exceed an expert’s selection threshold, raising that threshold reduces its load. A fixed-step controller approaches the cutoff gradually. A quantile estimates the threshold that leaves the target fraction of margins above it.

Adding a bias moves a margin relative to zero: a threshold τ corresponds to bias −τ in this simplified view. The quantile converts a target count into a score-space correction rather than requiring a tuned fixed increment.

Kimi K3 uses a quantile controller to infer a threshold from the current score distribution. For each token, take the (k+1)-th cutoff from scores plus the current biases. For each expert, subtract that biased cutoff from the raw score. Estimate the threshold that targets mk/E selections for m valid tokens.

The K3 report describes all-reduced per-expert histograms with 1,000 bins instead of transmitting every margin. It reports communication below 1% of raw-margin exchange. Bin width limits quantile resolution. The controller mean-centers biases and freezes them for inference; mixing still uses original affinities. These are K3 settings, not general MoE defaults.

Kimi K3 report ↗ · K3 training recipe →

Histograms make this calculation scalable. Workers summarize margin distributions, combine counts and estimate the threshold. The controller changes selection but still uses original affinities for output mixing.

Balanced across which tokens?

Batch balance can hide different local distributions. Two sequences can each use only one expert while their combined histogram is flat.

V3’s small-model ablations found similar validation loss for batch-wise auxiliary balancing and loss-free batch balancing. Stronger sequence-wise balancing gave worse results. Thus, comparisons must hold balancing scope fixed: changing the population can explain part of the apparent method advantage. V3 §4.5.3 ↗

Uniform counts help detect collapse; they do not establish optimal compute allocation. Difficult code, reasoning or rare-language tokens may need different computation from easy or duplicated tokens. A better objective would allocate computation by marginal loss reduction. Under inference domain shifts, previously balanced experts can become overloaded. Placement and replication address serving load, not model quality.

Text tokens can also compensate for concentrated image-token assignments in an aggregate histogram. DeepSeek-V4.1-Flash keeps separate load measurements and correction biases for text and image tokens. Each modality adjusts selection from its own load; mixing still uses original scores.

Its recipe uses a modality-specific bias update of 0.001 and a sequence-level balance coefficient of 0.0001. These are model-specific settings, not recommended defaults. Per-sequence, per-domain and per-modality measurements answer different questions. Strong constraints can also suppress useful specialization.

Avoiding starvation makes capacity trainable. It still does not tell us how that capacity should be divided between experts.

Expert layout: finer pieces and a shared path

Balancing controls token assignments. Layout controls the transformations those tokens can combine. DeepSeekMoE splits routed FFNs into narrower experts and reserves an always-active shared path. These changes address different sources of duplicated or inflexible computation.

Fine-grained experts: keep the work, increase the combinations

Start with K selected FFNs out of N. Split each into m experts with approximately 1/m of the original intermediate width. Select mK out of mN. Expert parameters and active FFN width remain approximately equal, but each token can combine more pieces.

More combinations can support finer specialization, but do not guarantee it. Additional router scores, dispatch entries and smaller matrix multiplications can reduce hardware efficiency. Compare quality at matched active compute and measure execution cost.

Shared experts: do not rediscover common work everywhere

Several routed experts may learn the same common transformation. A shared FFN provides an always-active path for that work. Routed experts can then specialize in conditional differences. Shared experts and fine-grained routing can coexist.

DeepSeekMoE’s matched comparisons support finer segmentation and shared expert isolation in its tested models; differences between tested shared-to-routed ratios were relatively small. A shared path is a design choice, not a stability guarantee. DeepSeekMoE §§3–4 ↗

Latent experts: reduce the width of the routed path

Fine segmentation increases the number of expert choices. A latent expert path reduces the width at which those experts operate. Kimi K3 projects the hidden state down to width 3,584, half its main hidden width.

K3 has 896 routed experts and selects 16 per token: $\frac{16}{896}=\frac{1}{56}$. Its two shared experts remain at full width. These counts describe different paths; do not assume that the shared experts use the latent width.

The narrower path reduces expert cost, but extra projections can increase activation scale. K3 applies RMSNorm before the up-projection and uses bounded SiTU-GLU activations. Quantile balancing controls load separately.

MAI-Base-1 makes a different choice: select 8 of 512 experts, project down before dispatch, but compute routing from the original representation. It uses dropless routing and alternates dense FFNs with highly sparse MoE blocks. Its study reports comparable scaling and better wall-clock efficiency than more uniform MoE allocation on its stack. This is evidence for that layout, not a rule that every layer should alternate.

K3 Stable LatentMoE → · MAI architecture and scaling →

Request-level selection: when moving weights is the bottleneck

Token-level routing assumes that selected expert weights are readily available. That assumption fails when most weights reside in flash storage. Loading a different expert set for every token can exceed the device's storage bandwidth.

Apple's AFM 3 Core Advanced stores a 20B sparse model in NAND flash and activates 1B–4B parameters. A lightweight dense block selects experts for the request, loads them into DRAM, and periodically reselects them during generation. A large shared path remains active.

Apple calls this Instruction-Following Pruning (IFP). It creates a request-conditioned subnetwork that runs for a period, rather than a fresh token-level top-k decision at every layer. Apple also describes a predetermined active-parameter budget chosen for task difficulty. The benefit is fewer weight transfers; the tradeoff is coarser selection.

The existing note does not disclose a complete expert layout, reselection interval or training recipe. Apple's cloud PT-MoE lineage is a separate design; do not infer its architecture from the on-device IFP description. Apple IFP and disclosure limits →

The router selects a logical expert ID. On a cluster, that ID also determines where the token must travel.

From logical experts to expert parallelism

On a cluster, tokens must be grouped by the worker that owns each selected expert. The mathematical operation is unchanged. Dispatch transfers the inputs, and the combine operation restores token order. These operations connect the router’s choices to communication cost.

Efficient execution groups assigned tokens by expert, runs a matrix batch, then scatters weighted outputs back to their original positions. Expert parallelism distributes experts across workers. Equal expert counts can still produce overloaded nodes, uneven link traffic, small matrix batches or slow steps.

DeepSeekMoE distinguishes expert-level and device-level balancing. DeepSeek-V3 limits the nodes reachable by each token. These policies constrain communication as well as selection. Aggregate text/image counts can hide modality-specific imbalance. A balanced training mix also cannot guarantee throughput during a domain-specific inference burst.

Measure worst-case tokens per expert, achieved cross-node all-to-all bandwidth and contention with tensor/data-parallel traffic. Aggregate bisection bandwidth matters more than one advertised link speed. Place experts to match the physical topology. Network-assisted reductions do not automatically accelerate arbitrary all-to-all traffic.

K3’s MoonEP uses dynamic redundant experts and online routing plans to distribute expert-parallel communication. Replication changes physical placement, not the logical expert identity.

Hardware topology & parallelism → · GPU collectives →

Expert choice depends on the current representation, not just the original token. This matters when blocks are reused: recurrent MoEs can choose different experts on different visits to the same block. Shared parameters do not make repeated computation free. MoE with looped models →

Reusing a layer does not require reusing its expert set

A looped Transformer revisits the same physical layers. Its hidden state changes between visits, so the same router can select a different expert set. Sparse FFNs can therefore reduce the expressivity constraint of reusing one dense FFN.

The existing looped-MoE study compares four architectures: dense and sparse models, each with either 16 unique layers or 8 layers run twice. Its sparse blocks select 2 of 8 experts and use auxiliary balancing plus router z-loss.

Across most layers, only 4%–14% of tokens keep identical expert sets between the two passes. Disjoint sets occur for 25%–53%; most other tokens retain one expert and change one. Layer 7 has higher consistency. The authors suggest that it stabilizes representations before the loop-boundary vocabulary projection; this is a hypothesis.

Looped-MoE study and recurrence context →

Why a balanced router can still train badly

Assignments can look balanced while logits drift, expert updates have the wrong scale, or a sensitive reduction accumulates numerical error. Balancing controls allocation; it does not solve every stability failure. Identify the failure before increasing a balancing coefficient.

Router z-loss: stabilize the score scale

Softmax depends on differences between logits. Adding one constant to all logits leaves expert probabilities unchanged, but changes the unnormalized exponentials. Z-loss regularizes the resulting log-partition value.

Z-loss regularizes log-partition magnitude, not expert counts or entropy. ST-MoE used coefficient 0.001 in its selected experiment and reported improved stability without quality loss. This is not a universal coefficient. Switch stabilized its tested BF16 runs by computing local router operations in FP32. ST-MoE found that selective precision alone was insufficient at its largest scale. Keep these interventions distinct. Do not automatically apply a softmax z-loss to independent sigmoid affinities; specify the regularized logits and test quality and stability. ST-MoE ↗

For example, logits [0, 0] and [10, 10] both give probabilities [0.5, 0.5]. Their squared log-partitions are approximately 0.48 and 114.34. Assignments remain the same, but z-loss differs substantially. It controls this unconstrained numerical scale rather than balancing expert counts.

When the failure starts outside the router

The Laguna report describes expert collapse after hundreds of billions of tokens. The cause was mismatched Muon learning-rate and effective-weight-decay scaling. A separate failure involved the LM-head input-gradient all-reduce under tensor parallelism. BF16 reduction became a dominant error source as logits drifted upward; FP32 stabilized that reduction. These observations do not imply that every collective needs FP32. Laguna report ↗

Identical padding activations can all select the same expert, especially when padding does not mix with surrounding tokens. Exclude padding from dispatch where possible. Always exclude it from valid-token load statistics and auxiliary losses. Keep early layers dense only if tests support the quality and compute tradeoff.

Initialization changes what the router sees

A large random score advantage can concentrate the first updates on a few experts. Check raw logit scale, top-k margins and selected gate weights from the start. Equal assignment counts can hide experts whose selected weights are almost zero.

MAI identifies another source of early imbalance. Near-uniform attention initially resembles causal mean pooling, which makes token representations less distinct. Its recipe initializes attention output RMSNorm gains to zero. Cross-token interactions then increase during training, rather than dominating the initial router inputs.

Dense early layers can also delay routing until representations have been transformed. Neither approach guarantees stability. Compare them under the same quality and compute conditions. MAI initialization →

New structure tokens can disturb a pretrained router too. The Laguna-derived note initializes a new token from the average of its existing subtoken embeddings. It then warms up the input embeddings and LM head while freezing the rest of the model. This lets the new embeddings adapt before full post-training. Special-token initialization →

Post-training: router drift and route replay are different problems

Pretraining balance does not guarantee balance during RL. The policy sees a different data distribution, and router updates can change expert assignments quickly. MiMo-V2.6 provides a concrete failure case.

Restoring only the original router recovered balance without reducing benchmark quality. Freezing the router then kept load statistics stable while task performance improved. This isolates router drift in that experiment; it does not prove that router adaptation is always harmful.

MiMo Pro selects 8 of 384 experts; Flash selects 8 of 256. Both use an initial dense FFN and no shared experts in later MoE blocks. Hidden matrices use Muown, while the router, embeddings and LM head use AdamW. The report also records a rank receiving more than 30× the mean micro-batch token load.

Recompute the same execution path

Even a frozen router can select different indices when training and inference use different numerics. Then recomputed token probabilities may not match the execution path that generated the rollout.

MiMo records inference-time expert indices and reuses them during training through Rollout Routing Replay (R3). It also replays the exact top-p candidate set and applies MXFP4 quantize-dequantize after each update. Matching weights alone does not establish this consistency.

DeepSeek-V4.1-Flash retains expert-routing states when interrupted trajectories resume with a newer checkpoint. Its training concatenates the recorded routes from each segment. Such a trajectory spans policy versions; it is not a fresh rollout under one current policy.

K3 uses MXFP4 routed weights and MXFP8 expert activations, with higher precision for sensitive non-expert components. Quantization-aware training continues through SFT and RL to align rollout and training numerics.

MiMo collapse and R3 → · V4.1 interrupted trajectories → · K3 deployment →

A poorly updated expert can become less useful, changing future scores and assignments. Numerical errors can also destabilize the representations that feed the router. A utilization failure therefore does not identify its cause by itself.

Implementation & diagnostics

The implementation follows the same order: compute affinities, select experts, execute grouped FFNs, combine outputs, then update controller state. The snippet below isolates selection and mixture weights. The checklist covers the surrounding implementation requirements.

A minimal selected-affinity operation, not a full dispatcher:

# affinity, bias: [valid_tokens, experts], [experts]
# Keep affinity in the autograd graph; update bias separately.
idx = (affinity + bias).topk(k, dim=-1).indices
weight = affinity.gather(-1, idx)
weight = weight / weight.sum(-1, keepdim=True).clamp_min(1e-9)
# Group tokens by expert, run FFNs, scatter weighted outputs.
# Do not use this normalization unchanged for Switch top-1.

Specify token masks, capacity and drop policy, dispatch order, inverse scatter indices and score scaling. Define cross-rank count reductions, checkpointed biases and expert placement. A short selection snippet cannot verify those contracts.

1. Validity. Padding masks, overflows, dropped assignments, deterministic resume.
2. Distribution. Counts, selected gate mass, dead-expert rate; split by sequence/domain/modality.
3. Numerics. Logit scale, logsumexp, activations, sensitive gradients and reductions.
4. Controller. Bias range, update speed, routing churn and oscillation.
5. Optimizer. Effective router/expert update scales, learning-rate and decay conventions.
6. Useful performance. LM loss, downstream quality, active FLOPs, all-to-all time, tail step latency.

Start with a simple auxiliary-balancing baseline unless the scale and measurements justify a separate controller. Add controls for observed failures. Do not force uniformity at the expense of task loss. Measure specialization by domain-conditioned usage and output contribution. Compare fine-grained and shared designs at matched compute and measured serving efficiency.

How these ideas developed

These milestones summarize overlapping design branches. Sparse gating enables conditional execution. Switch simplifies dispatch, ST-MoE addresses numerics, and DeepSeekMoE changes expert layout. Loss-free controllers separate batch-load correction from task learning. Later methods do not replace every earlier mechanism.

These are selected milestones in modern sparse neural models, not the origin of all mixture models. Dates are first paper submissions. Layout, numerical stability, and balance are parallel design axes—not a single chain of replacements.

1. Jan 2017
   Sparsely-gated MoE

   Conditional computation makes large banks of FFNs usable through a sparse learned gate.

   Shazeer et al. ↗
2. Jan 2021
   Switch

   Top-1 routing simplifies dispatch. Capacity, auxiliary balance, and selective precision become practical training concerns.

   Switch Transformers ↗
3. Feb 2022
   ST-MoE

   Router z-loss addresses unstable log-partition scale; it is not another name for load balancing.

   ST-MoE ↗
4. Jan 2024
   DeepSeekMoE

   Fine-grained and always-active shared experts change how capacity is composed.

   DeepSeekMoE ↗
5. Aug 2024
   Auxiliary-loss-free balance

   Load feedback changes selection biases outside the main gradient objective.

   Wang et al. ↗
6. Dec 2024
   DeepSeek-V3

   Combines bias-controlled batch balance with a small sequence auxiliary term, sigmoid affinities, and node-limited dispatch.

   DeepSeek-V3 ↗

Sources & connected notes

The primary papers are linked beside their mechanisms. These connected topics explain the optimizer, scaling, and execution details that extend beyond MoE.