Mixture of Experts (MoE) in LLMs: A Beginner's Guide
The Scaling Problem MoE Solves
Larger language models can learn more patterns because they have more parameters, but a standard dense Transformer uses the same full set of layer weights for every token. Adding parameters therefore increases both model capacity and the computation required to process each token. At large scale, this coupling becomes expensive during training and inference.
What role does the FFN play in a Transformer?
A Transformer block divides its work between attention and a feed-forward network. Attention mixes information across token positions, allowing a token to gather relevant context from the sequence. The FFN then transforms the features inside each token representation independently. It applies the same learned function to every position, but each position receives a different context-aware vector from attention.
Self-attention
Decides which other token positions are relevant and combines information from them.
Feed-forward network
Expands each token vector, applies a nonlinear or gated transformation, and projects it back to the model width.
The FFN performs feature mixing and nonlinear computation that attention alone cannot provide. Its intermediate dimension is usually wider than the model dimension, so its projection matrices contain a substantial share of the parameters and arithmetic in each block. In a dense Transformer, every token uses the same FFN weights. Increasing that FFN's width increases model capacity, but every token must execute the larger matrix multiplications.
Context-aware token→Expand features→Nonlinearity or gate→Project back
Scaling approach
Stored capacity
Work per token
Wider dense FFN
Increases
Increases with the wider FFN
Sparse expert FFNs
Increases with more experts
Uses only the selected experts
Mixture of Experts, or MoE, addresses this problem with conditional computation. Instead of one feed-forward network, an MoE layer contains several independent feed-forward networks called experts. A learned router examines each token representation and selects only a small number of those experts to process it.
This separates two quantities that are identical in a dense layer: total parameters, which determine how much expert capacity the model stores, and active parameters, which determine how much of that expert capacity one token uses. The model can increase the first quantity without increasing the second at the same rate.
The design objective
Increase model capacity without making every token execute the full expanded model. The cost is additional routing logic, expert-weight storage, load-balancing constraints, and communication between devices.
1. Start With a Dense Feed-Forward Network
A Transformer block has two main sublayers. Self-attention moves information between token positions. A feed-forward network, or FFN, then transforms each token position independently. In a dense Transformer, every token passes through the same FFN weights.
{FFN}(x) = W_{2}\,\phi(W_{1}x + b_{1}) + b_{2}
The first projection expands the hidden vector, the activation adds nonlinearity, and the second projection returns it to the model width.
Many current LLMs use a gated variant such as SwiGLU, which has three large matrices instead of two. These FFN matrices account for a substantial share of a Transformer's parameters and arithmetic. That makes the FFN a natural place to add conditional computation.
| Sublayer | Mixes across tokens? | Typical role |
|---|---|---|
| Self-attention | Yes | Gather context from other token positions |
| Feed-forward network | No | Transform each token representation with shared weights |
2. Where MoE Fits Inside a Transformer
A decoder-only Transformer is built by repeating the same basic unit many times. That unit is called a Transformer block or Transformer layer. Each block receives one vector for every token, updates those vectors with attention and an FFN, and passes the result to the next block.
The data flow through one block
1. Normalize
Normalize the incoming token vectors so the next sublayer receives values on a stable scale.
2. Self-attention
Let every token gather relevant information from earlier token positions.
3. First residual addition
Add the attention output back to the block input so the original representation has a direct path forward.
4. Normalize again
Prepare the context-aware token vectors for the feed-forward sublayer.
5. Dense FFN or MoE
Transform each token independently. This is the only step that changes when a dense block becomes an MoE block.
6. Second residual addition
Add the FFN or MoE output back to the residual stream, producing the block output.
In the common pre-normalization design, this flow can be written compactly. Let x be the input to the block, h the representation after attention, and y the final block output:
h = x + {Attention}({Norm}(x))
y = h + {FFN}({Norm}(h))
What changes in an MoE block?
The attention sublayer, both residual connections, and normalization layers remain in place. The dense FFN in step 5 is replaced by an MoE sublayer containing a router and several expert FFNs:
h = x + {Attention}({Norm}(x))
y = h + {MoE}({Norm}(h))
| Part of the block | Dense block | MoE block |
|---|---|---|
| Self-attention | Runs for every token | Runs for every token |
| Normalization | Shared | Shared |
| Residual paths | Two residual additions | The same two residual additions |
| Token-wise transformation | One shared FFN | Router plus selected expert FFNs |
| Expert activation | Not applicable | Only top-k experts run for each token |
Model designers decide which blocks use this replacement. A model may use MoE in every block, alternate dense and MoE blocks, or begin with several dense blocks before switching to MoE blocks. Regardless of the schedule, an expert is only an FFN inside one block. It is not a complete Transformer block and does not contain its own attention layer.
Where Does MoE Fit?
MoE replaces one specific component in the Transformer architecture
Showing: Standard Transformer (click to see MoE)
Output
↑
Repeat × N layers
Add & Normalize
↑
Feed-Forward Network (FFN)
Standard
Same network processes every token the same way
FFN(x) = W₂ · ReLU(W₁ · x)
↑
Add & Normalize
↑
Self-Attention
Tokens look at each other to understand context
↑
Input Tokens
"The cat sat on the mat"
The Standard Approach
In a standard Transformer, every token passes through the same FFN with the same weights. This means all parameters are used for every token, which gets expensive as models grow.
100%
Parameters used per token
Standard FFN
~10-25%
Parameters used per token
MoE (sparse)
Sparse MoE therefore does not make attention sparse. Long-context attention, the KV cache, and attention communication remain separate costs. MoE changes which token-wise FFN parameters are active inside the blocks that use it.
3. Router, Experts, and Combiner
A basic sparse MoE layer has three conceptual pieces. Suppose a batch contains T token states, each with width d and the layer has E experts.
1
Router
Reads each token state and produces E expert scores. The router is small compared with the expert bank.
2
Dispatch
Keeps the top-k choices, groups tokens by expert, and sends each group to the device that owns that expert.
3
Experts
Run independent FFNs over their assigned token groups. Experts usually share the same shape but have different weights.
4
Combine
Restores the original token order and adds the selected expert outputs using their routing weights.
Expert does not mean a separate complete LLM
Each expert is usually one FFN inside one Transformer layer. Attention, embeddings, normalization, the language-model head, and often some dense FFNs are shared by every token.
4. The Routing Math
For one token state x \in \mathbb{R}^{d} a linear router produces one logit per expert:
r = W_{\text{router}}x + b,\qquad r \in \mathbb{R}^{E}
p_{i} = \frac{\exp(r_{i})}{\sum\limits_{j = 1}^{E}\exp(r_{j})}
The model keeps only the k experts with the largest scores. Let S(x) be that selected set. A common implementation renormalizes the selected probabilities and forms a weighted sum:
{\overset{\sim}{p}}_{i} = \frac{p_{i}}{\sum\limits_{j \in S(x)}p_{j}},\qquad i \in S(x)
y = \sum\limits_{i \in S(x)}{\overset{\sim}{p}}_{i}\, E_{i}(x)
With top-1 routing, one expert produces the output. With top-2 routing, two experts run and their results are blended. Some architectures use sigmoid routing rather than a single softmax across all experts, and details such as score normalization, expert grouping, routing bias, and shared experts vary by model.
E
Number of experts available in the layer
k
Number of routed experts active for one token
pᵢ
Router weight assigned to expert i
Interactive: Top-k Expert Routing
Select different tokens, switch between top-1 and top-2 routing, and reveal the raw router logits. Notice that the router makes a new decision for every token at every MoE layer.
MoE Gating: Token → Expert Routing
Each token computes p(expert | token) via softmax over gate logits
Select Token
"The"
"model"
"learns"
"math"
Top-K:
1
2
3
4
Show Logits
Token Embedding
"The"
0.8
0.2
-0.5
0.1
softmax
p(expert | token)
Expert 1
39.9%✓
Expert 2
14.1%
Expert 3
21.3%
Expert 4
24.7%✓
Selected Experts (Top-2) with Normalized Weights
Expert 1
weight: 61.8%
Expert 4
weight: 38.2%
Output = 0.62 × Expert1(x) + 0.38 × Expert4(x)
Math Connection
This is conditional probability from your Probability module! The gating network computes p(expert | token) using softmax over learned gate weights, exactly like computing class probabilities in classification.
5. One MoE Forward Pass, End to End
An MoE forward pass has one mathematical path and one systems path. Mathematically, the router selects expert functions and combines their outputs. Operationally, the runtime reorganizes token states into efficient expert batches, moves them when experts are sharded, and restores the original token order before the residual addition.
Inside one sparse MoE layer
Follow four token states through top-2 routing. Select a stage to inspect the tensors and system operation at that point.
One device
Sharded experts
1Normalize
stable inputs
2Route
scores + top-2
3Dispatch
group by expert
4Experts
batched FFNs
5Combine
weighted outputs
6Residual
block output
Stage 1 · normalization
Prepare the post-attention token states
The block normalizes each token vector independently before routing. Sequence order and token identity do not change.
t1Thet2modelt3routest4tokens
RMSNorm
[T, d_model] → [T, d_model]
t1Thet2modelt3routest4tokens
Available experts4
Active per tokentop-2
Expert placement2 devices
The visualization uses four experts and top-2 routing, so each token creates two expert assignments. Capacity checks happen before dispatch. Inter-device exchange is required only when the selected experts are stored on different devices.
6. What an Expert Actually Computes
Each expert is typically an ordinary FFN with its own weights. If Expert 3 and Expert 7 receive the same vector, they can produce different outputs because training has changed their matrices in different ways.
Experts as Linear Transformations
Each expert applies a different matrix y = W_(e)x to the input
Expert 1
Expert 2
Expert 3
Expert 4
Compare All
Stretches along x, slight rotation
Transformation Matrix W₁
1.5
0.3
0.2
0.8
×
x₁x₂
=
y₁y₂
Math Connection
This is matrix multiplication from your Linear Algebra module! Each expert is just a different W matrix that transforms the input vector into a different subspace. The gating network decides which transformation to apply per token.
The two-dimensional transforms above are a visual analogy. A real expert operates in thousands of dimensions, uses nonlinear activations, and may contain a gated FFN such as SwiGLU. The important point is that the router chooses a parameterized transformation, not a stored answer or a database entry.
7. Tokens Route Independently
Routing happens for token representations, not whole prompts. The tokens in one sentence can visit different experts, and the same token string can route differently when its context changes its hidden state. The decision is also repeated at every MoE layer, so a token can take a different route deeper in the model.
Token-by-Token Expert Routing
Each token in a sequence gets routed to its own expert, independently and in parallel
Disambiguation
Code & Math
Mixed Content
Input Sentence
on
the
river
bank
E1: Nouns/Entities
E2: Verbs/Prepositions
E3: Articles/Grammar
E4: Code/Technical
Route All Tokens
Per-Token Routing
Unlike attention (which mixes information across tokens), expert routing is token-independent. Each token gets its own expert assignment based solely on its embedding. This means the word "bank" might route to different experts depending on whether the context suggests finance or nature, but that context is already embedded in the token representation by the attention layer before it.
The expert labels are an illustration
Real experts do not reliably divide into neat human categories. Researchers may observe preferences for languages, token types, or domains, but the learned computation can remain distributed and difficult to name.
8. Total Parameters vs Active Parameters
The most useful MoE distinction is between parameters stored by the model and parameters used for one token. Suppose every MoE layer has E experts, each expert has P_{\text{expert}} parameters, and the router selects k of them.
Stored expert parameters per layer
P_{\text{stored}} \approx E\, P_{\text{expert}}
Active expert parameters per token
P_{\text{active}} \approx k\, P_{\text{expert}}
If a layer has eight equally sized experts and uses top-2 routing, a token activates two of the eight expert FFNs. That is one quarter of the expert bank, plus the shared attention, router, normalization, embedding, and output parameters elsewhere in the model.
| Resource | Scales with all experts? | Why |
|---|---|---|
| Weight storage | Yes | Every expert weight must exist in host or accelerator memory. |
| Expert arithmetic per token | Usually no | Only k selected experts run for that token. |
| Optimizer state during training | Yes | Every trainable expert needs optimizer statistics and gradients when active. |
| Communication | Depends | Remote expert choices cause token exchange across devices. |
| Latency | Depends | Sparse kernels, batch size, imbalance, and networking determine realized speed. |
MoE only has a fair compute comparison when expert width, top-k, layer placement, and shared dense work are specified. Two same-sized active experts can cost more arithmetic than one dense FFN. Model designers often adjust expert width or the number of MoE layers to meet a target compute budget.
9. How an MoE Model Learns
The language-model objective stays familiar. The model predicts the next token, computes cross-entropy, and backpropagates the loss. MoE changes the path that each token takes through selected FFN parameters.
\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{language}} + \alpha\mathcal{L}_{\text{balance}} + \beta\mathcal{L}_{\text{router}}
The auxiliary terms vary by architecture. They can encourage balanced routing, control large router logits, or stabilize numerical behavior.
Selected experts
Receive task gradients from the tokens they processed, so their FFN weights update on those examples.
Router
Receives gradients through selected routing weights and any routing-specific losses. The discrete top-k boundary remains piecewise and requires careful optimization.
This creates a feedback loop. Early routing decisions determine which experts receive examples. Those experts improve on their assigned traffic, which can make the router prefer them again. Useful specialization can emerge, but uncontrolled feedback can also cause expert collapse.
10. Expert Collapse and Load Imbalance
If the router sends most tokens to the same expert, several problems arrive together. The popular expert becomes a compute bottleneck, its capacity fills, other experts learn slowly, and the model wastes most of its stored parameters. This failure mode is often called expert collapse.
Load Balancing as Entropy Control
See how expert distribution affects training loss and model performance
Collapsed (No Balance)
Moderate Imbalance
Well Balanced
Expert Usage Distribution
35%
E1
28%
E2
22%
E3
15%
E4
Entropy1.94 / 2.00 bits
Training Loss Over Time
Simulate Training
3.0
1.5
0.0
Loss: 3.007
Step: 1/50
Training Steps
Final Loss Comparison
1.23
Collapsed
Capacity wasted
0.75
Moderate
Suboptimal
0.46
Balanced
Best performance
Math Connection
This is Entropy and KL Divergence from your Information Theory module! The load balancing loss KL(usage || uniform) pushes expert usage toward a high-entropy (uniform) distribution. Higher entropy = better balance = lower final loss.
Perfectly uniform routing is not always the goal. Real data is not uniform, and some experts may deserve more traffic. The engineering goal is to avoid severe hotspots while preserving enough freedom for the router to learn useful assignments.
11. Expert Capacity and Token Overflow
Accelerators work best with bounded tensor shapes. Training systems therefore often give each expert a fixed number of assignment slots for one batch. If there are T tokens, k choices per token, E experts, and capacity factor c a common capacity rule is:
C = \left\lceil c\,\frac{Tk}{E} \right\rceil
Expert Capacity Factor
Each expert can only handle a limited number of tokens per batch
Capacity Factor1.25
0.5 (strict)2.0 (relaxed)
Tokens per Batch32
1664
Expert Capacity = (Tokens / Experts) × Capacity Factor
= (32 / 4) × 1.25 = 10 tokens/expert
Token Distribution to Experts
+2 dropped
cap
Expert 1
10/12 tokens
cap
Expert 2
9/9 tokens
cap
Expert 3
6/6 tokens
cap
Expert 4
3/3 tokens
Processed
Overflow (dropped)
Capacity limit
28
Tokens Processed
2
Tokens Dropped
88%
Utilization
!
Token Overflow Detected
2 tokens were dropped because Expert 1 exceeded its capacity. Increase the capacity factor or improve load balancing to prevent this.
Why Capacity Factor Matters
The capacity factor is a trade-off: too low means tokens get dropped when experts are overloaded, too high wastes memory by reserving space that may not be used. Typical values are 1.0 to 1.25. The Switch Transformer paper recommends 1.25 for training stability.
Low capacity
Less padding and memory, but more assignments can overflow.
High capacity
Fewer overflows, but more empty slots waste compute and memory.
Overflow policy
A system may drop, reroute, queue, or dynamically allocate extra work.
12. Auxiliary Load-Balancing Loss
One widely used strategy adds a small auxiliary objective. Let f_{i} be the fraction of hard token assignments sent to expert i and let P_{i} be that expert's average router probability. A Switch-style form is:
\mathcal{L}_{\text{balance}} = E\sum\limits_{i = 1}^{E}f_{i}P_{i}
Concentrating both assignments and probability on a few experts raises this term. The coefficient must be small enough that language modeling remains the main objective. Some newer systems use constrained assignment, routing biases, or auxiliary-loss-free balancing because a global penalty can interfere with the main task.
Auxiliary Load Balancing Loss
Watch how the auxiliary loss pushes expert usage toward balance during training
Switch Transformer Auxiliary Loss
\mathcal{L}_{\text{aux}} = \alpha \cdot N \cdot \sum_{i = 1}^{N}f_{i} \cdot P_{i}
\alpha
Balance weight
0.010
N
Number of experts
4
f_{i}
Fraction to expert i
tokens / total
P_{i}
Avg router prob for i
softmax output
\alpha (Balance Weight)0.010
0.001 (weak)0.1 (strong)
Simulate Training
Reset
Step: 0 / 50
Expert Token Distribution
Expert 154.6% × 49.0% = 26.78%
Expert 225.0% × 28.3% = 7.09%
Expert 311.4% × 13.4% = 1.52%
Expert 48.5% × 7.6% = 0.65%
f_{i}
P_{i}
Ideal (uniform):25.0% each
Auxiliary Loss During Training
0.012
0.010
ideal
\mathcal{L}_{aux} 0.0144
Training Steps
Live Calculation
\mathcal{L}_{aux} = \alpha \cdot N \cdot \sum_{i = 1}^{N}f_{i} \cdot P_{i}
=0.010×4× (0.268 + 0.071 + 0.015 + 0.006)
=0.0144
Total Training Loss
\mathcal{L}_{total} = \mathcal{L}_{task} + \mathcal{L}_{aux}
= Cross-Entropy + 0.0144
Ideal Case (Balanced)
When f_{i} = P_{i} = \frac{1}{N} for all experts:
\mathcal{L}_{aux} = \alpha \cdot N \cdot N \cdot \frac{1}{N^{2}} = \alpha
Minimum possible loss = 0.010
Collapse Case (All to One)
When f_{1} = P_{1} = 1 (everything to Expert 1):
\mathcal{L}_{aux} = \alpha \cdot N \cdot 1 = \alpha N
Maximum loss = 0.040 (4× higher)
Why This Works
The product f_{i} \cdot P_{i} is key: it penalizes when the router both assigns high probability to an expert AND actually routes many tokens there. This breaks the "rich get richer" problem. The factor of N ensures the penalty scales appropriately with model size.
13. Sparse Gradients and Uneven Learning
An expert receives ordinary backpropagation through the tokens it processed. An unselected expert has no task-dependent computation for that token, so it receives no task gradient from that path. Across a sufficiently large and balanced batch, every expert should still receive useful work.
Sparse Gradients in MoE
Only selected experts receive gradients: ∂L/∂W_(e) = 0 for non-selected experts
Select Experts (Top-K)
Compare Dense
Expert 1
Expert 2
Expert 3
Expert 4
x
Input
Expert 1
12M params
Expert 2
12M params
Expert 3
12M params
Expert 4
12M params
ŷ
Output
Forward
→
Backward
→
Update
Run Forward + Backward Pass
Sparse MoE (Top-2)
24M
params updated per step
Dense Model
48M
params updated per step
Math Connection
This is the Chain Rule from your Calculus module! Gradients only flow through the computational path that was actually used. Non-selected experts have ∂L/∂W = 0 because they never contributed to the output. This is structural gradient sparsity.
Sparse activation reduces expert arithmetic, but training must still maintain all expert weights, gradients, and optimizer state. Checkpointing and optimizer memory therefore follow total parameters much more closely than active parameters.
14. Shared Experts and Routed Experts
Some architectures include one or more shared experts that run for every token alongside sparsely selected routed experts. The shared path can learn broadly useful computation while routed experts spend more capacity on conditional behavior.
Shared vs Routed Experts
Shared experts process all tokens, routed experts specialize
Routed Only (Mixtral)
Shared + Routed (DeepSeek/Llama 4)
Select Token
The
model
learns
code
x = "The"
Shared Expert
Processes ALL tokens
E_(shared)(x)
Universal knowledge
Routed Expert: Grammar
Selected for "The"
E₃(x)
Articles and syntax
+
y = E_(shared)(x) + g·E₃(x)
Animate Forward Pass
Why Shared Experts?
Without Shared Expert
- • Each routed expert must learn common patterns
- • Redundant knowledge across experts
- • Experts can't fully specialize
With Shared Expert
- • Shared expert handles universal patterns
- • Routed experts focus on unique knowledge
- • Better parameter efficiency
The Hiring Manager Analogy
Think of it like a company: instead of hiring specialists who each need to know basic communication and computer skills, you have a shared "operations" team that handles common needs for everyone. Your specialists can then truly specialize in their domain without duplicating basic infrastructure.
Shared experts add guaranteed computation, so they reduce sparsity. They can still be valuable when the architecture benefits from separating common transformations from routed capacity. As with routed experts, human-readable labels in a diagram are explanatory examples rather than guaranteed learned roles.
15. Expert Parallelism Across GPUs
Large expert banks do not fit efficiently on one accelerator. Expert parallelism places different experts on different devices. After routing, tokens travel to the devices that own their selected experts, then expert outputs travel back.
Local tokens
→
All-to-all dispatch
→
Grouped expert FFNs
→
All-to-all return
Expert parallelism
Partitions experts so each device owns a subset of the expert bank.
Data parallelism
Replicates a model partition while different replicas process different examples.
Tensor parallelism
Splits large matrices within one layer across devices.
Pipeline parallelism
Places different ranges of Transformer layers on different device groups.
Real training jobs combine several forms of parallelism. The main MoE-specific risk is that useful FLOPs wait behind network transfers or one overloaded expert. Fast interconnects, local expert placement, balanced batches, token permutation kernels, and grouped GEMMs are therefore part of the architecture in practice.
16. What Changes During MoE Inference
At inference time, the router still makes a decision for every token at every MoE layer. There is no need for backpropagation or optimizer state, but the full expert bank must remain available and decode batches can be small or irregular.
| Concern | Training | Inference |
|---|---|---|
| Batch shape | Large token batches help fill experts | Autoregressive decode may provide few tokens per step |
| Capacity overflow | Fixed capacity and dropping may be tolerated by the training recipe | Dropping user tokens can damage output, so dynamic handling is preferable |
| Memory | Weights, gradients, activations, and optimizer state | Weights and runtime caches, with all experts still stored |
| Communication | Forward and backward all-to-all | Forward all-to-all on every routed layer |
| Kernel efficiency | Large expert batches can use hardware well | Small or skewed batches can underutilize expert matrices |
Active parameters are not a latency promise
Two models with the same active parameter count can have different latency. Expert width, quantization, batch size, device count, interconnect, routing skew, attention cost, and implementation quality all affect the result.
17. How Real MoE Architectures Differ
MoE names a family of designs rather than one fixed architecture. The examples below show how routing choices changed across influential systems. Parameter figures use each paper's own reporting convention, so they should not be treated as perfectly standardized comparisons.
| Work | Routing design | Why it matters |
|---|---|---|
| Sparsely-Gated MoE | Sparse learned gating over large expert banks | Established conditional computation at very large scale for language tasks. |
| Switch Transformer | Top-1 routing | Simplified routing and demonstrated trillion-parameter sparse Transformers. |
| ST-MoE | Sparse experts with router stability techniques | Studied stable training and transfer behavior, including router z-loss. |
| Mixtral 8x7B | Eight FFN experts, top-2 per token | Reported 47B total and 13B active parameters in an openly released decoder-only model. |
| DeepSeek-V3 | Fine-grained routed experts, shared experts, and routing bias | Reported 671B total and 37B active parameters with auxiliary-loss-free load balancing. |
The progression is not a simple march toward more experts. Top-k, expert size, shared capacity, balancing method, placement, and communication topology are co-designed. A routing method that looks elegant in an equation may perform poorly when it creates tiny matrix multiplications or excessive network traffic.
18. Common MoE Misconceptions
01
“An 8x7B model has 56B parameters.”
The multiplication is only a name-level shortcut. Shared attention, embeddings, norms, and other weights are not duplicated eight times. Use the model report's exact total and active counts.
02
“Only active parameters need memory.”
All expert weights must be stored or fetched. Sparsity reduces expert computation per token, not the full storage requirement.
03
“Each expert owns one subject.”
Some routing preferences can emerge, but experts are learned functions and may not map cleanly to human topics.
04
“MoE makes the entire Transformer sparse.”
The typical sparse component is the FFN sublayer. Attention and other shared layers still run for every token.
05
“Top-1 is always faster than top-2.”
Top-1 runs fewer experts, but realized latency also depends on utilization, routing balance, kernel shapes, and communication.
06
“More experts always improve quality.”
Extra capacity helps only when routing, data, optimization, and systems efficiency let those experts learn and serve useful functions.
07
“Inference cost equals active parameter count.”
Active parameters are a useful capacity measure, but FLOPs, memory bandwidth, weight precision, attention, and inter-device traffic determine cost.
19. A Minimal MoE Implementation
This framework-neutral pseudocode captures the forward pass. Production code replaces the Python loop with token permutation, capacity handling, grouped matrix multiplication, and collective communication kernels.
def sparse_moe(x, router, experts, top_k=2):
# x: [tokens, d_model]
router_logits = router(x) # [tokens, num_experts]
router_probs = softmax(router_logits, -1)
weights, expert_ids = topk(router_probs, top_k)
weights = weights / weights.sum(-1, keepdim=True)
output = zeros_like(x)
for expert_id, expert in enumerate(experts):
token_ids, slots = where(expert_ids == expert_id)
if len(token_ids) == 0:
continue
expert_input = x[token_ids]
expert_output = expert(expert_input)
output[token_ids] += weights[token_ids, slots, None] * expert_output
return output
Correctness checks
- Selected weights sum to one for each token when renormalization is intended.
- Dispatch and combine preserve the original token order.
- Every selected assignment contributes exactly once.
- Padding or overflow slots cannot leak into valid outputs.
Training checks
- Track tokens per expert and probability mass per expert.
- Measure dropped or rerouted assignment rate.
- Monitor router logits, entropy, and auxiliary losses.
- Verify that every expert receives gradients over time.
A useful test starts with one expert and top-1 routing. That result should match a dense FFN with the same weights. Next, use two identical experts and verify that weighted combination still matches. Only then introduce distinct experts, capacity limits, and distributed dispatch.
Primary Sources and Further Reading
The architecture descriptions and reported model figures in this guide are grounded in the original papers and technical reports below.
- Adaptive Mixtures of Local Experts, Jacobs et al. (1991)
- Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, Shazeer et al. (2017)
- GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, Lepikhin et al. (2020)
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, Fedus et al. (2021)
- ST-MoE: Designing Stable and Transferable Sparse Expert Models, Zoph et al. (2022)
- Mixtral of Experts, Jiang et al. (2024)
- DeepSeek-V3 Technical Report, DeepSeek-AI (2024)