The MoE Router Scores Experts Per Token

A Mixture-of-Experts (MoE) layer replaces the single FFN with several smaller FFNs called experts. Not every expert runs for every token. A small network called the router decides which experts each token should use.

The router is a simple linear projection: it takes the token's hidden state and produces one score (a router logit) per expert. These logits measure how relevant each expert is for this particular token.

The logits are then passed through softmax (or sigmoid, depending on the architecture) to get gating weights. In the softmax variant (which this course focuses on), the weights sum to 1. Why not just pick the top-k logits directly? Because the gating weights serve double duty: they select which experts to run and they determine how much each expert's output contributes to the final result. The sum-to-1 constraint ensures the combined output is a proper weighted average — not an arbitrarily scaled sum. Other routing variants exist (sigmoid gating, expert choice), but the softmax top-k pattern is the one you will encounter in the case studies.

Each token may pick a different set of experts. This is what makes MoE "sparse": the model has many total parameters (all the experts), but each token only activates a small subset. A model with 64 experts and top-2 routing uses only about 3% of its expert parameters per token — so it can be very large in parameter count while keeping per-token compute manageable.

In an ideal MoE, different tokens go to different experts — the workload is spread evenly. But without explicit encouragement, training can collapse into a few "popular" experts that handle most tokens while others are rarely used. This is called expert collapse.

Why does this happen? The router is trained to maximize prediction quality. If expert 2 happens to become slightly better early in training, the router sends more tokens to it, which gives it more gradient updates, which makes it even better — a positive feedback loop. Meanwhile, unused experts stagnate.

Training addresses this with a load-balancing loss: an auxiliary objective that penalizes the model when the token distribution across experts is too uneven. This is a training concern (not an inference concern), but you should know about it because it explains why real MoE models work well instead of degenerating — and why you sometimes see a "balance_loss" term in training configs.

MoE lets you scale the model's total knowledge (parameter count) independently from its per-token cost (active compute). A dense 7B-parameter model activates all 7B parameters for every token. A MoE model with 64 experts of 1B each has 64B total parameters but only activates ~2B per token (with top-2 routing), plus shared non-expert layers. The per-token compute is similar to a small dense model, but the total parameter count — and the capacity to specialize — is much larger.

The catch: all 64B parameters must fit in memory, even though only 2B are computed per token. MoE models are memory-heavy and bandwidth-heavy (loading expert weights), even if they are not compute-heavy. This is why MoE models are often memory-bandwidth-bound during decode — a bottleneck you will see in detail in the inference module, amplified here by the large total parameter count.