Mixture-of-Experts (MoE) architecture: Visual walkthrough

Step 1 — A token is just an arrow of numbers

WHAT. The thing flowing through the network is a token embedding, written $x$ . It is a vector: a fixed-length list of $d$ numbers. Think " $d = 4$ numbers" for pictures; real models use $d$ in the thousands.

WHY. Everything downstream — scoring, routing, blending — is arithmetic on these numbers. Before we can route a token we must agree on what a token is: a point in space, equivalently an arrow from the origin to that point.

PICTURE. The blue arrow is our token $x$ . Its two shown coordinates $x_{1}, x_{2}$ are just the numbers that place its tip. (Real tokens live in $d$ dimensions; we draw 2 so your eyes can follow.)

Figure — Mixture-of-Experts (MoE) architecture

Step 2 — Each expert has a "preference arrow"; the dot product measures how well $x$ matches it

WHAT. We have $N$ experts. The router stores one preference vector per expert, stacked as rows of a matrix $W_{g}$ of shape $N \times d$ ( $N$ rows, $d$ columns). Row $i$ , call it $w_{i}$ , is expert $i$ 's favourite direction. We score the match between token $x$ and expert $i$ with the dot product.

WHY THIS TOOL — why a dot product and not, say, a difference? We want one number saying "how aligned is $x$ with what expert $i$ likes." The dot product $w_{i} \cdot x = \sum_{j} w_{ij} x_{j}$ is large and positive when the two arrows point the same way, zero when perpendicular, negative when opposed. That is exactly "alignment," and it is a single number — perfect for scoring.

PICTURE. Two red expert-preference arrows. The token (blue) is nearly parallel to expert $A$ (big projection → big score) and nearly perpendicular to expert $B$ (tiny projection → small score). The dashed line is $x$ 's shadow onto $w_{A}$ — its length is the dot product.

Step 3 — Turn the raw scores into probabilities with softmax

WHAT. After Step 2 we have $N$ raw scores $h = [h_{1}, \dots, h_{N}]$ . Some may be negative, none sum to anything nice. We convert them into probabilities $p_{i}$ that are all positive and sum to $1$ , using softmax.

WHY THIS TOOL — why exponentiate, why not just divide by the sum? Three demands: (1) every weight must be positive (a negative blend weight is meaningless), (2) they must sum to 1 (so the output is a proper average), (3) the map must be smooth so the router can learn by gradients. Plain division fails demand (1) with negative scores. The exponential $e^{h_{i}}$ is always positive and grows monotonically — bigger score, bigger share — and it is smooth everywhere.

PICTURE. Left bars: the raw logits (some negative). Middle: after $e^{(\cdot)}$ everything is positive and the tall ones stretch taller. Right: after dividing by the total, the bars become a probability distribution summing to $1$ .

Step 4 — Keep only the Top- $k$ experts (this is where the savings live)

WHAT. We do not run all $N$ experts. We pick the $k$ experts with the largest $p_{i}$ (usually $k = 1$ or $k = 2$ ). Call that set of chosen indices $T$ . Everyone outside $T$ is set to gate $0$ .

WHY. An expert with gate $0$ contributes $0 \times E_{i} (x)$ — so we never compute $E_{i} (x)$ at all. That is the whole reason MoE is cheap: compute scales with $k$ , not $N$ (see Sparse vs Dense Models and Conditional Computation). You store $N$ experts (knowledge) but fire only $k$ (cost).

PICTURE. The four probability bars from Step 3. A yellow "keep line" sits under the top-2; the top-2 bars glow, the bottom-2 are crossed out (gate $0$ , greyed). Only the glowing experts will do arithmetic.

Step 5 — Re-normalize the survivors so their weights sum to 1

WHAT. After deleting the losers, the surviving probabilities no longer sum to $1$ (we threw some mass away). We divide each survivor by the survivors' total, giving the final gates $g_{i}$ .

WHY. The layer output will be a weighted average of expert outputs. For it to be a proper average — a convex blend, no accidental scaling up or down — the weights must sum to $1$ . Re-normalizing over just $T$ restores that.

PICTURE. Left: the two survivors, e.g. $p = 0.63$ and $p = 0.23$ , summing to only $0.86$ (a gap). Right: after dividing by $0.86$ they become $0.73$ and $0.27$ , refilling the bar to a full $1.0$ .

Step 6 — Blend the chosen experts' outputs into the answer

WHAT. Now, and only now, we actually run each chosen expert $E_{i}$ on the token: $E_{i} (x)$ . Each $E_{i}$ is an ordinary feed-forward network. We combine their outputs, each scaled by its gate.

WHY. The gate $g_{i}$ says "how much to trust expert $i$ for this token." A gate-weighted sum lets the two specialists co-author the answer, weighted by relevance. The result $y$ is the MoE layer's output — same shape as $x$ , so it slots into the network like any other layer.

PICTURE. Two glowing experts each emit an output arrow. Expert $A$ 's arrow is scaled by $g_{A} = 0.73$ (long), expert $B$ 's by $g_{B} = 0.27$ (short); tip-to-tail addition gives the final green output arrow $y$ .

Step 7 — Edge case: the "rich get richer" collapse, and the balancing nudge

WHAT. Steps 1–6 describe one token. Over a whole batch, a nasty degenerate case appears: the router keeps sending most tokens to a handful of experts, and the rest go dead. This is routing collapse.

WHY IT HAPPENS. A slightly-better expert gets more tokens → trains faster → gets even better → gets even more tokens. A runaway feedback loop. The result wastes almost all your stored parameters — the entire MoE bet is lost. We counter it with an auxiliary load-balancing loss (see Load Balancing and Auxiliary Losses).

PICTURE. Left board: "no aux loss" — one giant bar (an overloaded expert), the rest flat on the floor (dead). Right board: "with aux loss" — bars pushed toward the even dashed line $1/ N$ .

Step 8 — Edge case: capacity overflow and token dropping

WHAT. Even with balancing, each expert has a hard seat limit per batch, the capacity $C$ . Tokens beyond $C$ for an over-subscribed expert are dropped — they skip the experts and pass straight through the residual connection unchanged.

WHY. Real hardware allocates fixed-size buffers per expert (see Model Parallelism & Expert Parallelism). If more tokens arrive than seats exist, the extras cannot fit. The capacity factor tunes how many spare seats you buy: bigger = fewer drops but more memory/compute.

PICTURE. An expert "bus" with $C$ seats. Tokens board up to the line; the overflow tokens (pink) fall off and take the residual "bypass road" around the experts.

The one-picture summary

One token flows left-to-right: arrow → dot products → softmax → keep top- $k$ → re-normalize → weighted blend → $y$ , with the two guardrails (balance, capacity) hanging underneath.

Recall Feynman: the whole walkthrough in plain words

A token arrives as an arrow of numbers. Each expert keeps a "favourite direction," and we measure how well the token points that way with a dot product — one number per expert (Step 2). Those numbers are messy, so softmax squashes them into clean positive percentages that add up to $100%$ (Step 3). We keep only the two biggest and throw the rest to zero, because a zeroed expert does no work — that's the whole reason MoE is cheap (Step 4). The two survivors' percentages no longer add to $100%$ , so we rescale them (Step 5), then run just those two experts and mix their answers by those percentages to get the output (Step 6). Two safety rails watch the batch: a balancing nudge stops one expert from hogging everyone (Step 7), and a seat limit drops overflow tokens onto a bypass road when an expert is full (Step 8).

Recall Quick self-check

Why does compute scale with $k$ and not $N$ ? ::: Because non-chosen experts get gate $0$ , so $E_{i} (x)$ is never computed for them; only the $k$ chosen experts run. Why softmax and not plain division of the logits? ::: Logits can be negative; softmax's exponential makes every weight positive, keeps them summing to 1, and stays smooth so the router can learn by gradients. What breaks if we skip the load-balancing loss? ::: Routing collapses onto a few experts (rich-get-richer), leaving most parameters dead and wasting the model's capacity.