Mixture of experts

The mechanism

A mixture-of-experts layer contains several expert networks and a router. The router selects a small subset for each token, then the system combines their outputs. Total stored parameters can therefore grow faster than the arithmetic used for one token. This creates an infrastructure problem: expert weights need capacity, routing can be imbalanced, and tokens may travel between devices. The experts are numerical functions, not independently conscious specialists.

An MoE router assigns tokens to selected experts, which changes both active computation and where gradients flow. A load-balancing auxiliary loss encourages use of experts under a particular formula; it is distinct from a hard capacity limit or dropping policy. The Switch top-one convention uses discrete dispatch fractions and mean router probabilities. Its dispatch counts are not differentiated through argmax. The lifecycle model is dense and does not execute expert routing or this auxiliary loss.

Score expert choices for each token.
Dispatch tokens to selected experts, possibly across devices.
Combine outputs and return them to the continuing representation.
Route tokens under an explicit top-k and tie convention.
Count dispatch fractions and mean router probabilities separately.
Evaluate the chosen auxiliary formula and identify its differentiable path.
Apply expert capacity and account for accepted, dropped or bypassed tokens independently.

Science and assumptions

Stored capacity versus activated computation

\begin{aligned}y&=\sum_{i\in S(x)}g_i(x)E_i(x)\\N_{\mathrm{total}}&=N_{\mathrm{shared}}+E N_{\mathrm{expert}}\\N_{\mathrm{active}}&\approx N_{\mathrm{shared}}+kN_{\mathrm{expert}}\end{aligned}

The output sums selected expert outputs weighted by the router. Stored parameters include all experts; approximate active parameters include only the selected k experts plus shared parameters.

Symbols and units

x: Input token representation
model-coordinate vector
y: Combined expert output
model-coordinate vector
S(x): Selected expert IDs for x
set of k experts
gᵢ(x): Routing weight for selected expert i
dimensionless
Eᵢ(x): Output of expert function i
same shape and units as y
E: Total expert count (distinct from function Eᵢ)
count
k: Experts selected per token
count
N_shared: Shared parameters included in this bookkeeping
parameters
N_expert: Parameters per equally sized expert
parameters
N_total: Total stored parameter count
parameters
N_active: Approximate parameters activated for one token
parameters

A router selects a set S(x) containing k of E experts. The selected outputs receive routing weights $g_{i}$ . With equally sized experts, increasing E can enlarge stored capacity without multiplying the number of expert matrices evaluated for each token. Communication and routing still consume real time.

Assumptions This is a simplified top-k layer with equal expert sizes. Routing weights may be renormalized differently, capacity limits can drop or reroute tokens, and shared experts alter the count. Parameter activation is not a complete FLOP or latency model.

Routing losses, sparse gradients and expert collapse · worked relation

L_{aux}=\alpha N\sum_i f_iP_i,\quad f_i=\frac1T\sum_x\mathbf1[\arg\max_j p_j(x)=i],\quad P_i=\frac1T\sum_xp_i(x)

Calculate routing loads and a declared auxiliary loss separately from expert capacity.

Symbols and units

L_aux: Switch-style auxiliary objective
dimensionless
alpha: Auxiliary coefficient
dimensionless
N: Experts
experts
T: Routed tokens
tokens
f_i: Discrete dispatch fraction
dimensionless
P_i: Mean router probability
dimensionless
p_i(x): Router probability for token x
dimensionless

Six tokens dispatch to three experts with counts [3,2,1], so f=[1/2,1/3,1/6]. Mean router probabilities are P=[1/2,1/3,1/6]. With alpha=0.01 and N=3, auxiliary loss is 0.01×3×(1/4+1/9+1/36)=0.011667. Uniform f=P=[1/3,1/3,1/3] gives 0.01. A hard capacity of two tokens per expert still rejects or bypasses one token from expert one under a declared policy.

Assumptions Switch top-one auxiliary convention, before capacity handling. Discrete f is treated as nondifferentiable; P supplies a gradient path. Worked-only routing.

Engineering example

Synthetic layer has eight experts with 1 million parameters each, plus 2 million shared parameters, and selects two experts per token. Stored parameters total 10 million; the simplified selected count is 4 million. At two bytes per stored parameter, weights occupy 20 MB before metadata.

Distinctions and tradeoffs

An expert is a learned computation, not a separate humanlike mind.

Sparse activation saves arithmetic relative to activating every expert, but increases routing and load-balancing complexity.

Model design and network topology become tightly connected.

Specifications and interfaces

Active vs total: Different parameter counts
Only selected experts execute per token.
Communication: Often all-to-all
Depends on expert placement and implementation.

Interfaces

Router Input: Token vectors
Residual stream Output: Weighted expert outputs

Check your understanding

Why can an MoE model have many stored but fewer active parameters?

Only selected experts run for each token
Its memory is imaginary
All expert weights disappear after training

Answer and reasoning

Only selected experts run for each token

Routing selects a subset of the available expert networks.

Does an auxiliary loss alone enforce the hard capacity of two?

Only if all model parameters are dense.
No; dispatch capacity and overflow behavior need a separate rule.
Yes; auxiliary losses physically prevent every overflow.

Answer and reasoning

No; dispatch capacity and overflow behavior need a separate rule.

No; dispatch capacity and overflow behavior need a separate rule. Switch top-one auxiliary convention, before capacity handling. Discrete f is treated as nondifferentiable; P supplies a gradient path. Worked-only routing.

Historical context

2021 · Switch Transformer simplifies sparse routing

The Switch Transformer preprint investigated top-one expert routing and training techniques that made very large sparse Transformer models easier to train; it was published in JMLR in 2022.

Fedus et al. — Switch Transformers

Evidence and scope

This is the lesson’s equal-size top-k accounting example. Selected parameter count is not a complete FLOP, latency or quality model; shared experts, router work, imbalance, capacity limits and communication change actual execution.

Mode: worked example. Switch top-one auxiliary convention, before capacity handling. Discrete f is treated as nondifferentiable; P supplies a gradient path. Worked-only routing. Numerical examples are authored calculations, except explicitly identified automated core fixtures. Primary sources support the scoped mechanism; no human specialist review or learner outcome is implied.

Use the execution investigation for the shared tensor-to-memory-to-result thread; no current standalone calculator models this lesson’s full mechanism.