Mixture-of-Experts

A standard transformer block applies the same dense feed-forward network (FFN) to every token: an $d \to 4d \to d$ MLP with a nonlinearity in between, using the same weights for all tokens in all sequences. An MoE layer replaces this single FFN with $N$ parallel FFNs, called experts ($E_1, \dots, E_N$), each with independent weights. A router (or gating network) $G$, typically a single learned linear layer followed by a softmax, computes a distribution over experts for each token's hidden representation $x$:

where $W_g \in \mathbb{R}^{d \times N}$. Instead of using this full distribution to form a weighted average over all $N$ experts (which would cost the same as a dense layer, just relabeled), the router selects only the top-$k$ experts by gate score, typically $k=1$ (Switch Transformer) or $k=2$ (GShard, Mixtral, GLaM). The token's output is the weighted sum of just those $k$ expert outputs, re-normalized over the selected subset:

This is conditional computation: the computational graph executed for a given token depends on the token's own content. The crucial efficiency claim is that if $N$ experts each have the same size as a dense FFN, total parameters scale as $O(N)$, but active compute per token stays $O(k)$ — independent of $N$. A model can therefore have, say, 8 or 64 or 256 experts per MoE layer while each token only ever touches 1 or 2 of them. The key terminology to fix: sparse activation (only a subset of parameters fire per input), expert capacity (the maximum number of tokens a single expert can process per batch, a hard systems constraint), load balancing (encouraging roughly-even utilization of experts so none is starved or overloaded), and token dropping (what happens when an expert's capacity is exceeded — excess tokens either skip that expert via a residual pass-through, or are computed with reduced batching efficiency). MoE layers are typically interleaved with ordinary dense transformer blocks — replacing, say, every other FFN — rather than replacing every layer, both to control the routing overhead and to preserve stable gradient flow through the attention sublayers.

An MoE transformer layer processes a batch of tokens through three logical stages: routing (decide which experts each token goes to), dispatch and compute (send tokens to their chosen experts and run the expert FFNs), and combine (weight and sum the expert outputs back into the token's residual stream). In a multi-device setting, experts are typically sharded across accelerators — expert $i$ lives entirely on device $i$ — so dispatch and combine require an all-to-all communication step that physically moves tokens (or their hidden states) to the device holding their assigned expert and back.

Step-by-step mechanism for one MoE layer: 1. For each token $x$ in the batch, compute router logits $x W_g \in \mathbb{R}^N$ and a softmax over the $N$ experts. 2. Select the top-$k$ expert indices by softmax score for that token. 3. Re-normalize the top-$k$ scores so they sum to 1 (these become the combine weights). 4. Enforce expert capacity: each expert may only accept up to $C$ tokens in this batch (a fixed buffer size, computed as a capacity factor times the average tokens-per-expert). If more tokens are routed to an expert than $C$ allows, the lowest-scoring overflow tokens are dropped for that expert (they either bypass the layer via the residual connection or fall back to being unprocessed by that expert). 5. Dispatch: physically gather each expert's assigned tokens into a per-expert batch (in a distributed setting, this is the "dispatch" all-to-all). 6. Compute: run each expert's FFN on its batch of tokens, in parallel across experts/devices. 7. Combine: scatter the expert outputs back to their original token positions, multiply by the combine weight from step 3, and sum contributions from a token's $k$ chosen experts. 8. Add an auxiliary load-balancing loss to the training objective (computed from the router's soft probabilities and the actual hard dispatch fractions), so gradient signal discourages routing collapse. ``` function MoE_Layer(X, experts[1..N], W_g, k, capacity C):

# X: batch of T token vectors, each in R^d logits = X @ W_g # T x N probs = softmax(logits, axis=-1) # T x N, router probabilities topk_idx = [] # for each token, list of k expert ids topk_weight = [] # matching gate weights (renormalized) for t in 1..T: idx = indices of the k largest entries of probs[t] w = probs[t][idx] w = w / sum(w) # renormalize over the top-k subset topk_idx[t] = idx topk_weight[t] = w

# --- capacity-constrained dispatch --- expert_buffer = [ [] for e in 1..N ] # tokens routed to each expert dropped = set() for t in 1..T: for (e, w) in zip(topk_idx[t], topk_weight[t]): if size(expert_buffer[e]) < C: expert_buffer[e].append( (t, w) ) else: dropped.add( (t, e) ) # overflow: this (token, expert) pair is skipped

# --- expert compute (parallel across e) --- Y = zeros_like(X) for e in 1..N: if expert_buffer[e] is empty: continue tokens_e = gather(X, [t for (t, w) in expert_buffer[e]]) out_e = experts[e].forward(tokens_e) # standard FFN: d -> 4d -> d for (t, w), o in zip(expert_buffer[e], out_e): Y[t] += w * o # weighted combine

# tokens with all top-k experts dropped keep their residual (identity) path for t in 1..T: if all((t, e) in dropped for e in topk_idx[t]): Y[t] = X[t] # graceful fallback return Y, probs, expert_buffer # probs used for balancing loss ```

The load-balancing auxiliary loss. Left unconstrained, the router tends to converge to a self-reinforcing loop: an expert that is (by chance, early in training) slightly better on some tokens receives more gradient updates on those tokens, becomes even better on them, and is routed to even more — a rich-get-richer collapse. To counter this, Switch Transformer and GShard add an auxiliary loss that penalizes uneven distribution of the fraction of tokens dispatched to each expert ($f_e$) multiplied by the average router probability mass assigned to it ($P_e$):

This loss is minimized when all experts receive an equal share ($f_e = P_e = 1/N$ for all $e$, giving $L_{\text{aux}} = N \cdot N \cdot (1/N)^2 = 1$, its theoretical minimum given the $\sum f_e = \sum P_e = 1$ constraints), and is added to the main cross-entropy loss with a small coefficient (commonly $\alpha \approx 0.01$). Using $f_e$ (a hard, non-differentiable count) multiplied against $P_e$ (a soft, differentiable probability) is the trick that lets gradients flow: $f_e$ acts as a fixed importance weight in the backward pass while $P_e$ carries the gradient.

Compute (FLOPs) per token. A dense FFN of hidden width $d$ and expansion factor $4$ costs $O(8d^2)$ multiply-adds per token (two matrix multiplies of size $d \times 4d$ and $4d \times d$). An MoE layer with $N$ experts of the same shape and top-$k$ routing costs the router's negligible $O(dN)$ (a single small linear layer) plus $k$ expert FFN evaluations, i.e. $O(k \cdot 8d^2)$. Since $k \ll N$ (commonly $k \in \{1,2\}$ while $N$ ranges from 8 to hundreds), compute per token is independent of $N$ and matches a dense model with roughly $k$ times the FFN width, not $N$ times. This is the central efficiency claim, and it is exact at the level of FLOP counting — it's a direct consequence of top-$k$ sparse selection, not an approximation.

Parameters. Total parameter count of the MoE layer is $O(N \cdot 8d^2)$ — every expert's weights are stored and must be addressable (e.g., in accelerator memory or across a cluster), even though only $k$ of them are touched per token. This is why MoE models are described by two separate numbers: total parameters (e.g., Mixtral 8x7B has ~47B total parameters) and active parameters per token (Mixtral activates ~13B per token, roughly $2$ of $8$ experts' worth plus shared attention layers). The capacity/compute decoupling is real, but it shifts the bottleneck to memory bandwidth and interconnect: all expert weights must reside somewhere accessible, and tokens must physically travel (via all-to-all communication in distributed training/inference) to whichever device holds their assigned expert.

Communication cost in distributed settings. If experts are sharded one-per-device across $N$ devices, each all-to-all dispatch/combine step moves $O(T \cdot d)$ data (each of $T$ tokens sends a $d$-dimensional vector to its expert's device), and the shape of this exchange is irregular — it depends on the token-to-expert routing decided at runtime, which is why fixed expert capacity $C$ is enforced: it upper-bounds the communication and compute per expert to a static, schedulable size, at the cost of dropping overflow tokens. Typically $C$ is set as a capacity factor $\gamma$ times the ideal even split, $C = \gamma \cdot \frac{k T}{N}$, with $\gamma$ slightly above $1$ (e.g., $1.25$) to absorb natural imbalance while bounding worst-case cost.

Correctness / why the routing signal is trainable at all. The top-$k$ selection is a discontinuous, non-differentiable operation — a naive view suggests gradients can't flow to the router. The resolution is that the router does receive gradient signal, but only through the combine weights of the tokens it actually routes: $\partial y / \partial (\text{gate score of a selected expert})$ is well-defined via the softmax renormalization in step 3 of the algorithm, so the router learns to increase the weight (and thus future likelihood of selecting) experts that reduce loss, and decrease it for experts that don't, exactly analogous to how attention weights are learned. What the router does not get is gradient information about experts it didn't pick — this is the fundamental credit-assignment gap in sparse routing, and is why load-balancing losses (which use all $N$ router probabilities, not just the selected ones) are needed as an auxiliary signal encouraging exploration of underused experts. Without this term, experts that start out slightly ahead are the only ones ever selected, gradient never flows to the rest, and their weights never leave initialization — a degenerate local optimum that empirically occurs within the first few hundred training steps if $L_\text{aux}$ is omitted.

Honest tradeoffs. MoE is not a free lunch:

• Memory: total parameters must be stored (VRAM or across a cluster), so MoE trades compute-efficiency for memory-footprint, unlike dense models where the two scale together.
• Training instability: the discrete routing decision plus capacity dropping introduces noisy, discontinuous gradients; MoE training typically needs auxiliary losses, careful learning-rate schedules, and sometimes explicit "z-loss" terms to stabilize router logits from growing unboundedly large.
• Batch-size sensitivity: at small batch sizes, expert buffers are underfilled, hardware utilization drops, and the capacity-based token-dropping becomes a larger fraction of tokens — MoE benefits are much stronger at large batch/throughput regimes (training, or high-QPS serving) than at single-token, low-latency inference.
• Fine-tuning generalization: sparse experts can overfit narrow specializations seen in pretraining and can be harder to fine-tune on small, distribution-shifted downstream datasets than dense models of equal active-parameter size.

Consider one MoE layer with $N = 4$ experts, top-$k = 2$, and $T = 5$ tokens in the batch, with capacity $C = 2$ per expert (so total slots $= 8 \ge k \cdot T = 10$ is not satisfied — this intentionally forces an overflow to illustrate dropping). Suppose the router produces the following softmax scores (rows = tokens, columns = experts):

Tallying raw dispatch requests per expert: E1 gets {x3, x5} (2 requests), E2 gets {x1, x2, x4} (3 requests), E3 gets {x1, x4, x5} (3 requests), E4 gets {x2, x3} (2 requests). With capacity $C=2$, E1 and E4 fit their requests exactly. E2 and E3 each receive 3 requests but can only serve 2 — the lowest-scoring request at each overloaded expert is dropped:

• E2's requests are (x1: 0.688, x2: 0.667, x4: 0.714); the lowest score is x2 (0.667), so x2's E2 assignment is dropped.
• E3's requests are (x1: 0.312, x4: 0.286, x5: 0.750); the lowest score is x4 (0.286), so x4's E3 assignment is dropped.

Final dispatch:

Now check each token's final combine: x1 gets both its picks (E2 + E3) → fully served, normal weighted combine. x4 gets only E2 (its E3 pick was dropped) — practical implementations either renormalize over the surviving expert alone (weight $1.0$ on E2's output) or simply scale by the original weight and accept a partial combine; Switch/GShard-style systems use the former, so x4's output is just $E2(x4)$, with no contribution from E3. x2 lost its E2 assignment entirely but kept E4, so x2's output is just $E4(x2)$. No token loses both of its assignments in this example, so nobody falls back to the pure residual (identity) path — but the mechanism from the pseudocode's fallback step exists precisely for the case where that happens. This trace shows three real phenomena at once: uneven expert popularity (E2 and E3 are hot, E1 and E4 are comparatively cold — exactly what $L_\text{aux}$ penalizes over many batches), capacity-triggered token dropping (a systems-level approximation, not a bug), and the resulting effective $k$ for some tokens silently dropping from 2 to 1 mid-training — a subtlety often missed when reasoning about MoE purely in terms of the idealized top-$k$ formula.

The Mixture-of-Experts idea predates the transformer era by decades, originating in classical neural network ensembling research before being re-engineered for sparse, trillion-parameter-scale conditional computation.

MoE architectures underlie some of the largest deployed language models today. Mixtral 8x7B and Mixtral 8x22B (Mistral AI) use 8 experts per MoE layer with top-2 routing, achieving performance competitive with much larger dense models while activating only a fraction of total parameters per token. DeepSeek-MoE and DeepSeek-V2/V3 push this further with fine-grained experts (many small experts rather than few large ones) and shared always-active experts alongside routed ones, improving load balance and specialization simultaneously. Google's GShard and GLaM demonstrated multi-hundred-billion to trillion-parameter MoE translation and language models years before MoE became mainstream in open-source LLMs. Industry reporting has widely suggested that GPT-4 and similar frontier proprietary systems use MoE-style architectures to manage the tension between capacity and inference cost at scale. Beyond raw language modeling, MoE routing ideas appear in multi-modal models (routing different modalities to specialized experts), in multilingual translation systems (language-specific expert specialization emerges naturally from routing), and in efficient serving infrastructure, where expert-parallelism becomes a first-class dimension of model parallelism alongside tensor and pipeline parallelism. Understanding MoE is a natural bridge to the broader efficiency toolkit — techniques like Quantization, LoRA & Distillation attack the same cost-versus-capability tradeoff from the angle of compressing a fixed model rather than sparsely activating a large one, and both approaches are frequently combined in production systems. MoE routing decisions can themselves be seen as an extreme, hard form of the soft weighting performed by Attention Mechanism over tokens, applied instead over subnetworks, and both rely on the representational groundwork laid by the standard Transformer architecture.

Ask the tutor: How would you design the expert capacity factor and load-balancing coefficient differently for training versus low-latency production inference?