Scaling dense language models incurs a brutal trade-off: every additional parameter increases both representational capacity and floating-point compute cost per token. Sparse Mixture-of-Experts (MoE) completely decouples parameter count from computational cost. By replacing dense feed-forward networks with an ensemble of specialized sub-networks (“experts”) and dynamically routing each token to only the top-1 or top-2 experts, MoE models (such as Mixtral 8x7B and GPT-4) unlock hundreds of billions of parameters at the inference speed of a fraction of their size.
1. The Philosophy of Conditional Computation
In a standard Transformer block, input tokens pass through identical Feed-Forward Network (FFN) parameters regardless of whether the token is a punctuation mark, a Python code snippet, or a quantum physics formula. Sparse MoE activates only a tiny, input-dependent sub-graph of the network for each token:
$mathbf{y} = sum_{i=1}^E G(mathbf{x})_i E_i(mathbf{x})$
Where $E$ is the total number of expert networks, $E_i(mathbf{x})$ is the output of the $i$-th expert FFN, and $G(mathbf{x}) in mathbb{R}^E$ is the sparse gating (routing) probability vector.
2. The Gating Mechanism & Top-$k$ Routing
To ensure sparsity, Noam Shazeer et al. (2017) introduced noisy Top-$k$ gating. Given token representation $mathbf{x} in mathbb{R}^d$ and trainable routing weight matrix $mathbf{W}_g in mathbb{R}^{d times E}$:
- Compute unnormalized routing logits with optional Gaussian exploration noise:
$H(mathbf{x})_i = (mathbf{x} mathbf{W}_g)_i + epsilon cdot text{Softplus}((mathbf{x} mathbf{W}_{text{noise}})_i), quad epsilon sim mathcal{N}(0, 1)$
- Select the indices of the top $k$ highest logit values (typically $k=2$ or $k=1$):
$text{TopK}(v, k)_i = begin{cases} v_i & text{if } v_i text{ is in top } k \ -infty & text{otherwise} end{cases}$
- Apply Softmax across the selected candidates to yield normalized routing weights:
$G(mathbf{x}) = text{Softmax}(text{TopK}(H(mathbf{x}), k))$
Because $G(mathbf{x})_i = 0$ for all but $k$ experts, only $k$ matrix multiplications are executed. In Mixtral 8x7B (8 experts, top-2 routing), while the total model contains 47 billion parameters, each token only encounters ~13 billion active parameters during forward execution.
3. The Pathology of Router Collapse & Expert Capacity
Left unconstrained, routing networks suffer from a fatal positive feedback loop known as Router Collapse. Early in training, a few experts randomly perform marginally better on common tokens. The router sends more tokens to them, making them update faster and improve further, while the remaining experts starve and never learn. Eventually, the entire model degenerates into a dense model using only 1 or 2 overworked experts.
To maintain parallel execution across distributed GPU hardware, each expert has a fixed Expert Capacity $C$: the maximum number of tokens it can process in a batch:
$C = text{Capacity Factor} times left( frac{text{Tokens per Batch}}{E} right)$
If the router sends more tokens to an expert than capacity $C$, the excess tokens are dropped (bypassed via residual connection without FFN computation), severely hurting generation quality.
4. Auxiliary Load Balancing Loss
To eliminate router collapse and ensure uniform token distribution across all $E$ experts, modern MoE models enforce an Auxiliary Load Balancing Loss $mathcal{L}_{text{aux}}$ added to the primary cross-entropy objective (Fedus et al., 2021):
$mathcal{L}_{text{aux}} = alpha cdot E sum_{i=1}^E f_i P_i$
Where for a batch of $T$ tokens:
- $f_i = frac{1}{T} sum_{t=1}^T mathbb{I}(text{token } t text{ routed to expert } i)$ is the empirical fraction of tokens routed to expert $i$.
- $P_i = frac{1}{T} sum_{t=1}^T G(mathbf{x}_t)_i$ is the average routing probability assigned to expert $i$.
Because $f$ and $P$ are non-negative vectors summing to 1, the dot product $sum f_i P_i$ is strictly minimized when both distributions are uniform ($f_i = 1/E, P_i = 1/E$), compelling the router to allocate workload with perfect balance across all hardware clusters.
