---
title: Mixture-of-Experts
module: Large Models & Agents
moduleNumber: 10
lessonNumber: 9
order: 1009
summary: >
  A mixture-of-experts layer replaces one feed-forward network with many and a
  router that sends each token to only a few of them, so the parameter count and
  the per-token compute become separate dials. We derive the gated output, sparse
  top-$k$ routing softmax, the load-balancing loss that stops the router from
  collapsing onto a single expert, and expert/token capacity with dropping, then
  work the dimension-annotated tensor shapes and FLOP arithmetic. We trace the
  architectures from the sparsely-gated LSTM through GShard, Switch Transformer,
  and Mixtral, cover distributed expert parallelism, and close on the training
  dynamics, failure modes, and serving costs of a sparse model.
topics: [Large Models & Agents]
sources:
  - book: Goodfellow
    ref: "Ch. 12 — conditional computation and mixtures of experts"
---

[Scaling laws](/deep-learning/architectures/the-transformer-architecture) say loss
keeps falling as a power law in parameter count $N$, but the
[Transformer accounting](/deep-learning/architectures/the-transformer-architecture)
ties every one of those $N$ parameters to per-token FLOPs: a dense forward pass
touches all of them. A **mixture-of-experts** layer breaks that tie. It holds many
feed-forward networks but routes each token through only a few, so the model can
carry an order of magnitude more parameters at the same arithmetic cost, the design
behind the largest [language models](/deep-learning/large-models-and-agents/large-language-models)
in deployment.[^shazeer]

## Conditional computation

A dense network applies the same weights to every input. **Conditional computation**
instead activates an input-dependent subset of the network, so the parameters a
token sees are a function of the token itself.[^gf-conditional]

> **Definition (Conditional computation).** A model whose effective function on
> input $x$ uses only a subset $\mathcal{A}(x) \subseteq \Theta$ of its parameters,
> selected by a data-dependent gate. Total capacity is $\abs{\Theta}$; the
> per-input cost depends only on $\abs{\mathcal{A}(x)}$, so the two can be scaled
> independently.

The two costs decouple. Write $N$ for the total parameter count and $k$
for the number of expert sub-networks each token activates out of $E$ total. Dense
computation is the special case $k = E$; sparse routing keeps $k$ fixed as $E$ grows.

$$
\underbrace{\text{FLOPs per token}}_{\text{dense}} \propto N,
\qquad
\underbrace{\text{FLOPs per token}}_{\text{sparse}} \propto \frac{k}{E}\,N \ll N
\;\;\text{when}\;\; k \ll E.
$$

A model with $E = 64$ experts and $k = 1$ active per token carries $64\times$ the
feed-forward parameters of its dense twin while doing the same per-token arithmetic.

$$
% caption: Dense (left) routes every token through all weights, so FLOPs track total
% parameters; sparse (right) activates $k$ of $E$ experts, holding FLOPs flat as $E$ grows.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  ex/.style={draw, black, minimum width=12mm, minimum height=7mm, align=center, inner sep=1pt},
  on/.style={draw=acc, text=acc, thick, minimum width=12mm, minimum height=7mm, align=center, inner sep=1pt, fill=acc!15},
  off/.style={draw, black, minimum width=12mm, minimum height=7mm, align=center, inner sep=1pt, fill=black!8}]
  \definecolor{acc}{HTML}{2348F2}
  % ---- dense ----
  \begin{scope}
    \node[black] (di) at (0,-1.0) {\texttt{token}};
    \node[on] (d1) at (-1.3,0.4) {\texttt{block}};
    \node[on] (d2) at (0.0,0.4) {\texttt{block}};
    \node[on] (d3) at (1.3,0.4) {\texttt{block}};
    \draw[->, acc, thick] (di) -- (d1);
    \draw[->, acc, thick] (di) -- (d2);
    \draw[->, acc, thick] (di) -- (d3);
    \node[black, anchor=south] at (0,1.1) {\texttt{all active}};
    \node[black, anchor=north, font=\footnotesize] at (0,-1.4) {\texttt{dense: FLOPs = c N}};
  \end{scope}
  % ---- sparse ----
  \begin{scope}[xshift=6.4cm]
    \node[black] (si) at (0,-1.0) {\texttt{token}};
    \node[off] (s1) at (-1.95,0.4) {\texttt{E1}};
    \node[on]  (s2) at (-0.65,0.4) {\texttt{E2}};
    \node[off] (s3) at (0.65,0.4) {\texttt{E3}};
    \node[off] (s4) at (1.95,0.4) {\texttt{E4}};
    \draw[->, acc, thick] (si) -- (s2);
    \draw[->, black, thick] (si) -- (s1);
    \draw[->, black, thick] (si) -- (s3);
    \draw[->, black, thick] (si) -- (s4);
    \node[acc, anchor=south] at (-0.65,1.1) {\texttt{1 active}};
    \node[black, anchor=north, font=\footnotesize] at (0,-1.4) {\texttt{sparse: FLOPs = c (k/E) N}};
  \end{scope}
\end{tikzpicture}
$$

## The mixture-of-experts layer

The layer that realizes this replaces a single feed-forward block with $E$ of them,
the **experts**, plus a small **gating network** (the **router**) that decides how
much each expert contributes.[^jacobs]

> **Definition (Mixture-of-experts layer).** Given $E$ expert networks
> $E_1, \dots, E_E$ (each a feed-forward block $\mathbb{R}^d \to \mathbb{R}^d$) and a
> gating network $g : \mathbb{R}^d \to \Delta^{E-1}$ producing nonnegative weights
> $g_i(x)$ that sum to $1$, the layer output is the gated combination
> $$
> y = \sum_{i=1}^{E} g_i(x)\, E_i(x).
> $$
> When $g$ is **dense** every term is computed; when $g$ is **sparse** ($g_i(x) = 0$
> for most $i$) only the active experts run.

The experts are ordinary [position-wise FFNs](/deep-learning/architectures/the-transformer-architecture):
each is $E_i(x) = W_2^{(i)}\,\sigma(W_1^{(i)} x)$ with its own weights
$W_1^{(i)} \in \mathbb{R}^{d_{\text{ff}} \times d}$ and
$W_2^{(i)} \in \mathbb{R}^{d \times d_{\text{ff}}}$, exactly the two-layer FFN the
Transformer already uses, replicated $E$ times. The router is a single linear map
$W_g \in \mathbb{R}^{E \times d}$ followed by a softmax. In a Transformer, the MoE
layer drops in where the FFN sublayer sat, leaving attention untouched.

### Dimension walk-through

Trace one token $x \in \mathbb{R}^d$ through the layer. With model width $d = 4096$,
expert hidden width $d_{\text{ff}} = 4d = 16384$, and $E = 8$ experts routed top-$2$:

$$
\begin{aligned}
x &\in \mathbb{R}^{d} & &(d = 4096) \\
h = W_g\, x &\in \mathbb{R}^{E} & &(E = 8 \text{ router logits}) \\
\mathcal{T} = \argtop_2(h) &\subset \{1,\dots,E\} & &(\abs{\mathcal{T}} = 2) \\
g_i = \softmax(h_{\mathcal{T}})_i &\in [0,1],\ i \in \mathcal{T} & &(\textstyle\sum_{i\in\mathcal{T}} g_i = 1) \\
E_i(x) = W_2^{(i)}\sigma(W_1^{(i)} x) &\in \mathbb{R}^{d} & &(\text{two matmuls per active } i) \\
y = \textstyle\sum_{i\in\mathcal{T}} g_i\, E_i(x) &\in \mathbb{R}^{d}. & &
\end{aligned}
$$

The router costs $d E = 4096 \cdot 8 \approx 33\text{k}$ multiply-adds, negligible
against the $2 \cdot 2 d\, d_{\text{ff}} = 2 \cdot 2 \cdot 4096 \cdot 16384 \approx
0.27\text{B}$ per active expert. Routing top-$2$ runs two experts, so per-token FFN
cost is $\approx 0.54\text{B}$ MACs, independent of how many experts the layer holds.
Adding experts grows the parameter store, not this number.

$$
% caption: An MoE layer: the router scores the token $x\in\mathbb{R}^d$, two experts
% fire, and their outputs combine by the gate weights $g_2,g_4$ into the output.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  ex/.style={draw, black, minimum width=13mm, minimum height=7mm, align=center, inner sep=1pt, fill=black!8},
  on/.style={draw=acc, text=acc, thick, minimum width=13mm, minimum height=7mm, align=center, inner sep=1pt, fill=acc!15},
  rt/.style={draw=acc, text=acc, thick, minimum width=16mm, minimum height=7mm, align=center, inner sep=1pt},
  op/.style={draw, black, circle, inner sep=0pt, minimum size=6mm}]
  \definecolor{acc}{HTML}{2348F2}
  \node[black] (x) at (0,0) {\texttt{token x}};
  \node[rt] (r) at (0,1.4) {\texttt{router g}};
  \draw[->, thick] (x) -- (r);
  % experts
  \node[ex] (e1) at (-3.0,3.0) {\texttt{E1}};
  \node[on] (e2) at (-1.0,3.0) {\texttt{E2}};
  \node[ex] (e3) at (1.0,3.0) {\texttt{E3}};
  \node[on] (e4) at (3.0,3.0) {\texttt{E4}};
  \draw[->, black, thick] (r) -- (e1);
  \draw[->, acc, thick] (r) -- (e2);
  \draw[->, black, thick] (r) -- (e3);
  \draw[->, acc, thick] (r) -- (e4);
  \node[text=acc, font=\scriptsize] at (-2.0,1.95) {$g_2$};
  \node[text=acc, font=\scriptsize] at (2.0,1.95) {$g_4$};
  % combine
  \node[op] (sum) at (0,4.6) {$+$};
  \draw[->, acc, thick] (e2) -- (sum);
  \draw[->, acc, thick] (e4) -- (sum);
  \draw[->, black, thick, dashed] (e1) -- (sum);
  \draw[->, black, thick, dashed] (e3) -- (sum);
  \node[black] (y) at (0,5.8) {\texttt{output y}};
  \draw[->, thick] (sum) -- (y);
  \node[black, anchor=west, font=\footnotesize] at (3.6,4.6) {\texttt{weighted sum}};
\end{tikzpicture}
$$

## Sparse top-$k$ gating

A dense gate that runs every expert defeats the purpose. The fix is to keep only the
$k$ highest-scoring experts and zero the rest before the softmax, so the combination
has at most $k$ nonzero terms.[^shazeer]

> **Definition (Sparse top-$k$ gating).** Let $h(x) = W_g x \in \mathbb{R}^{E}$ be
> the router logits. Define $\TopK(h, k)$ to keep the $k$ largest
> entries and set the rest to $-\infty$, then take the softmax over the survivors:
> $$
> g_i(x) = \frac{\exp\!\parens{h_i(x)}}{\sum_{j \in \mathcal{T}} \exp\!\parens{h_j(x)}}
> \quad\text{for } i \in \mathcal{T}, \qquad g_i(x) = 0 \text{ otherwise},
> $$
> where $\mathcal{T} = \argtop_k h(x)$ is the index set of the $k$
> largest logits. The renormalization over $\mathcal{T}$ keeps $\sum_i g_i(x) = 1$.

### A worked routing example

Take $E = 4$ experts and top-$2$ routing. Suppose the router produces logits
$h = (2.0,\ 0.5,\ 2.6,\ -0.3)$ for one token. The two largest are $h_3 = 2.6$ and
$h_1 = 2.0$, so $\mathcal{T} = \{1, 3\}$. Softmax over just those two:

$$
g_3 = \frac{e^{2.6}}{e^{2.6} + e^{2.0}} = \frac{13.46}{13.46 + 7.39} = 0.646,
\qquad
g_1 = \frac{e^{2.0}}{e^{2.6} + e^{2.0}} = 0.354,
$$

with $g_2 = g_4 = 0$. The token's output is $y = 0.646\, E_3(x) + 0.354\, E_1(x)$.
Renormalizing over the survivors (rather than the full softmax over all four) matters:
the full softmax would put $\approx 0.10$ on the dropped experts, and folding that mass
back onto the top two is what keeps the gate a valid convex combination.

### Why gate on the softmax, not a hard argmax

The gate values multiply the expert outputs, so gradients reach the router through the
$g_i$. If the router just picked an expert and passed the output through unweighted,
$\partial y / \partial W_g = 0$ and the routing decision would never learn. Weighting by
$g_i$ makes the selection differentiable in the magnitude even though the discrete choice
of $\mathcal{T}$ is not. The original sparsely-gated layer adds tunable Gaussian noise to
the logits before the top-$k$,

$$
h_i(x) = (W_g x)_i + \varepsilon_i \cdot \softplus\!\parens{(W_{\text{noise}} x)_i},
\qquad \varepsilon_i \sim \mathcal{N}(0,1),
$$

which spreads tokens across experts during training (a token near a decision boundary
lands on either side across steps) and gives the top-$k$ selection a smooth expected
behaviour.[^shazeer]

$$
% caption: Top-$k$ routing ($k=2$): the router scores all four experts, keeps the
% two highest logits, and renormalizes the gate over just those before combining.
\begin{tikzpicture}[>=stealth, font=\scriptsize, x=1.0cm, y=1.0cm]
  \definecolor{acc}{HTML}{2348F2}
  % logit bars
  \draw[->, thick] (0,0) -- (5.6,0) node[right, font=\footnotesize] {\texttt{experts}};
  \draw[->, thick] (0,0) -- (0,3.2) node[above, font=\footnotesize] {\texttt{router logit}};
  % kept (top-2): outlined accent
  \draw[draw=acc, thick, fill=acc!15] (0.4,0) rectangle (1.1,2.6);
  \draw[draw=acc, thick, fill=acc!15] (1.6,0) rectangle (2.3,2.0);
  % dropped: outlined black
  \draw[draw=black, thick, fill=black!8] (2.8,0) rectangle (3.5,1.0);
  \draw[draw=black, thick, fill=black!8] (4.0,0) rectangle (4.7,0.6);
  \node[acc, anchor=south, font=\footnotesize] at (0.75,2.6) {\texttt{E2}};
  \node[acc, anchor=south, font=\footnotesize] at (1.95,2.0) {\texttt{E4}};
  \node[black, anchor=south, font=\footnotesize] at (3.15,1.0) {\texttt{E1}};
  \node[black, anchor=south, font=\footnotesize] at (4.35,0.6) {\texttt{E3}};
  \node[acc, anchor=west, font=\footnotesize] at (2.6,2.6) {\texttt{kept (top-2)}};
  \node[black, anchor=west, font=\footnotesize] at (3.7,1.6) {\texttt{dropped}};
\end{tikzpicture}
$$

Setting $k = 1$ routes each token to a single expert, the **Switch** simplification;
$k = 2$ is the common choice in GShard and Mixtral, giving the gate at least two
gradients to weigh against each other per token.[^fedus][^lepikhin] The figure below
contrasts the two: top-$1$ passes the whole token to one expert; top-$2$ splits it
across two by the renormalized gate.

$$
% caption: Top-$1$ (left) sends the token to a single expert with gate $1$; top-$2$
% (right) splits it across two experts weighted by $g_a,g_b$ with $g_a+g_b=1$.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  ex/.style={draw, black, minimum width=11mm, minimum height=6.5mm, align=center, inner sep=1pt, fill=black!8},
  on/.style={draw=acc, text=acc, thick, minimum width=11mm, minimum height=6.5mm, align=center, inner sep=1pt, fill=acc!15},
  rt/.style={draw=acc, text=acc, thick, minimum width=13mm, minimum height=6.5mm, align=center, inner sep=1pt}]
  \definecolor{acc}{HTML}{2348F2}
  % ---- top-1 ----
  \begin{scope}
    \node[black] (t) at (0,0) {\texttt{token}};
    \node[rt] (r) at (0,1.2) {\texttt{router}};
    \draw[->, thick] (t) -- (r);
    \node[ex] (a1) at (-1.6,2.7) {\texttt{Ea}};
    \node[on] (b1) at (0.0,2.7) {\texttt{Eb}};
    \node[ex] (c1) at (1.6,2.7) {\texttt{Ec}};
    \draw[->, acc, thick] (r) -- (b1) node[midway, right, text=acc, font=\footnotesize] {\texttt{g=1}};
    \draw[->, black, thick, dashed] (r) -- (a1);
    \draw[->, black, thick, dashed] (r) -- (c1);
    \node[black, anchor=north, font=\footnotesize] at (0,-0.45) {\texttt{top-1}};
  \end{scope}
  % ---- top-2 ----
  \begin{scope}[xshift=6.2cm]
    \node[black] (t2) at (0,0) {\texttt{token}};
    \node[rt] (r2) at (0,1.2) {\texttt{router}};
    \draw[->, thick] (t2) -- (r2);
    \node[on] (a2) at (-1.6,2.7) {\texttt{Ea}};
    \node[on] (b2) at (0.0,2.7) {\texttt{Eb}};
    \node[ex] (c2) at (1.6,2.7) {\texttt{Ec}};
    \draw[->, acc, thick] (r2) -- (a2) node[midway, left, text=acc, font=\scriptsize] {$g_a$};
    \draw[->, acc, thick] (r2) -- (b2) node[midway, right, text=acc, font=\scriptsize] {$g_b$};
    \draw[->, black, thick, dashed] (r2) -- (c2);
    \node[black, anchor=north, font=\footnotesize] at (0,-0.45) {\texttt{top-2}};
  \end{scope}
\end{tikzpicture}
$$

## The load-balancing problem

Nothing in the gated objective forces the router to use all the experts. Early in
training a few experts get slightly more gradient, the router prefers them, they
improve faster, and the preference compounds.

> **Definition (Routing collapse).** A degenerate state in which the router sends
> almost all tokens to a small subset of experts. The unused experts receive no
> gradient and stay near initialization, so the layer's effective capacity falls to
> that of the few active experts and the extra parameters are wasted.

The standard remedy is an **auxiliary loss** that rewards balanced assignment. For a
batch $\mathcal{B}$ of $T$ tokens and $E$ experts, let $f_i$ be the fraction of
tokens routed to expert $i$ and $P_i$ the average gate probability mass it received.

> **Definition (Load-balancing loss).** With
> $$
> f_i = \frac{1}{T}\sum_{x \in \mathcal{B}} \mathbb{1}\!\brackets{i \in \argtop_k h(x)},
> \qquad
> P_i = \frac{1}{T}\sum_{x \in \mathcal{B}} g_i(x),
> $$
> the auxiliary loss is
> $$
> \mathcal{L}_{\text{aux}} = \alpha\, E \sum_{i=1}^{E} f_i\, P_i,
> $$
> added to the task loss with a small coefficient $\alpha$. It is minimized when
> $f_i = P_i = 1/E$ for all $i$, i.e. a uniform split.

The product $f_i P_i$ pairs the (non-differentiable) routing count $f_i$ with the
differentiable probability $P_i$, so gradients flow into the router through $P_i$
while $f_i$ scales the penalty toward whichever experts are overloaded.[^fedus] The
$P_i$ term nudges the router to shave probability off experts that are already crowded;
the $f_i$ multiplier makes that push proportional to how crowded each is. At the uniform
optimum the sum equals $E \cdot E \cdot (1/E)^2 = 1$, scaled by $\alpha$.

For example, with $E = 4$, a fully collapsed router that sends every token to
expert $1$ has $f = (1,0,0,0)$ and $P \approx (1,0,0,0)$, so
$\mathcal{L}_{\text{aux}} = \alpha \cdot 4 \cdot 1 = 4\alpha$. A uniform router has
$f = P = (\tfrac14,\dots)$, giving $\alpha \cdot 4 \cdot 4 \cdot \tfrac{1}{16} = \alpha$.
The collapsed state costs $4\times$ more, and the gradient points away from it.

> **Theorem (Uniform minimizer of $\mathcal{L}_{\text{aux}}$).** Over the simplex
> $\sum_i P_i = 1$ with $f_i$ held at its induced values, the convex surrogate
> $E \sum_i P_i^2$ is minimized uniquely at $P_i = 1/E$.

> **Proof.** Minimize $\sum_i P_i^2$ subject to $\sum_i P_i = 1$. The Lagrangian
> $\sum_i P_i^2 - \lambda\parens{\sum_i P_i - 1}$ has $\partial / \partial P_i =
> 2 P_i - \lambda = 0$, so $P_i = \lambda/2$ is constant across $i$; the constraint
> forces $P_i = 1/E$. The Hessian $2I$ is positive definite, so the stationary point
> is the global minimum, with value $E \cdot E \cdot (1/E)^2 = 1$. $\qed$

The coefficient $\alpha$ is small, $10^{-2}$ in Switch, so the balancing pressure never
dominates the task loss; it only breaks the symmetry that would otherwise let the router
collapse. Too large an $\alpha$ forces uniform routing regardless of the token
distribution, which discards the specialization the experts could learn.

A second mechanism bounds the damage when balance still fails. Each expert is given a
fixed **capacity**, and tokens beyond it are dropped past the layer through the
residual connection.

> **Definition (Capacity factor and token dropping).** Each expert accepts at most
> $$
> C = \ceil*{\,c \cdot \frac{T}{E}\,}
> $$
> tokens per batch, where $c \ge 1$ is the **capacity factor** ($C = T/E$ slots when
> $c = 1$). Tokens routed to an already-full expert are **dropped**: they skip the
> experts and pass through unchanged via the residual. Larger $c$ drops fewer tokens
> at the cost of padding compute and memory.

Capacity has to be fixed ahead of time because the expert weight tensors are allocated
as dense $[E, C, d]$ buffers for the batched matmul, so the buffer size cannot depend on
the data-dependent routing counts. With $T = 4096$ tokens, $E = 8$ experts, and $c = 1.25$,
each expert holds $C = \ceil{1.25 \cdot 512} = 640$ slots. A perfectly uniform batch fills
$512$ of them and wastes $128$ per expert to padding; a skewed batch overflows the crowded
experts and drops their excess. The capacity factor provides headroom against skew,
at the cost of padded compute.

$$
% caption: Load balancing: tokens spread evenly under the auxiliary loss (left); under
% collapse (right) one expert overflows its capacity and the excess tokens are dropped.
\begin{tikzpicture}[>=stealth, font=\scriptsize, x=1.0cm, y=1.0cm]
  \definecolor{acc}{HTML}{2348F2}
  \definecolor{red}{HTML}{C0392B}
  % ---- balanced ----
  \begin{scope}
    \draw[->, thick] (0,0) -- (4.4,0) node[right, font=\footnotesize] {\texttt{experts}};
    \draw[->, thick] (0,0) -- (0,2.8) node[above, font=\footnotesize] {\texttt{tokens}};
    \draw[acc, thick, dashed] (0,1.7) -- (4.0,1.7) node[right, text=acc, font=\footnotesize] {\texttt{capacity}};
    \draw[draw=acc, thick, fill=acc!15] (0.3,0) rectangle (1.0,1.4);
    \draw[draw=acc, thick, fill=acc!15] (1.3,0) rectangle (2.0,1.5);
    \draw[draw=acc, thick, fill=acc!15] (2.3,0) rectangle (3.0,1.3);
    \draw[draw=acc, thick, fill=acc!15] (3.3,0) rectangle (4.0,1.45);
    \node[black, anchor=north, font=\footnotesize] at (2.0,-0.25) {\texttt{balanced}};
  \end{scope}
  % ---- collapsed ----
  \begin{scope}[xshift=6.2cm]
    \draw[->, thick] (0,0) -- (4.4,0) node[right, font=\footnotesize] {\texttt{experts}};
    \draw[->, thick] (0,0) -- (0,2.8) node[above, font=\footnotesize] {\texttt{tokens}};
    \draw[acc, thick, dashed] (0,1.7) -- (4.0,1.7) node[right, text=acc, font=\footnotesize] {\texttt{capacity}};
    % overflowing expert
    \draw[draw=acc, thick, fill=acc!15] (0.3,0) rectangle (1.0,1.7);
    \draw[draw=red, thick, fill=red!15] (0.3,1.7) rectangle (1.0,2.5);
    \draw[draw=black, thick, fill=black!8] (1.3,0) rectangle (2.0,0.4);
    \draw[draw=black, thick, fill=black!8] (2.3,0) rectangle (3.0,0.3);
    \draw[draw=black, thick, fill=black!8] (3.3,0) rectangle (4.0,0.35);
    \node[red, anchor=west, font=\footnotesize] at (1.05,2.2) {\texttt{dropped}};
    \node[black, anchor=north, font=\footnotesize] at (2.0,-0.25) {\texttt{collapsed}};
  \end{scope}
\end{tikzpicture}
$$

An alternative inverts the decision. Token-choice routing has each token pick its
experts; **expert-choice** routing has each expert pick its top-$C$ tokens, which
guarantees perfect balance by construction (every expert fills exactly its capacity)
at the price that some tokens may be chosen by many experts and others by none.[^zhou]

## Architectures

The same gated layer appears in a line of progressively simpler and larger models.

**Sparsely-gated MoE (2017).** The first practical instance interleaved an MoE layer
between stacked LSTM layers in a language and translation model, with up to $131{,}072$
experts and noisy top-$k$ gating ($k = 4$), reaching $137$B parameters at the compute
of a far smaller dense model.[^shazeer]

**GShard (2021).** Ported the MoE layer into the Transformer, replacing every other
FFN sublayer with a top-$2$ MoE, and added the sharding and capacity machinery to
train a $600$B-parameter multilingual translation model across thousands of
devices.[^lepikhin]

**Switch Transformer (2022).** Simplified routing to $k = 1$: each token goes to
exactly one expert. This halves the routing and communication cost of top-$2$, and the
paper shows that with the right capacity factor and auxiliary loss, top-$1$ matches or
beats top-$2$ quality, scaling to $1.6$T parameters.[^fedus]

> **Definition (Expert and model parallelism).** Experts are partitioned across
> devices (**expert parallelism**): each device holds a few experts, and routed
> tokens are sent to the device owning their expert by an **all-to-all** exchange,
> then returned. This is composed with the usual data and **model parallelism** that
> splits each large tensor across devices, so an MoE layer's communication is two
> all-to-all shuffles per layer rather than the dense all-reduce.

The two all-to-all shuffles are the crux of the distributed cost: the first
(**dispatch**) sends each token to the device holding its expert; the expert runs
locally; the second (**combine**) returns the weighted output to the token's home
device. Because routing is data-dependent, the volume each device sends is uneven, and a
single overloaded expert stalls the collective while the others wait. This is the
communication reason to care about balance, on top of the wasted-capacity reason.

$$
% caption: Expert parallelism across two devices: dispatch all-to-all sends each token
% to its expert's device, experts run locally, and combine all-to-all returns outputs.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  tok/.style={draw, black, minimum width=8mm, minimum height=6mm, align=center, inner sep=1pt, fill=black!8},
  ex/.style={draw=acc, text=acc, thick, minimum width=10mm, minimum height=6.5mm, align=center, inner sep=1pt, fill=acc!15},
  dev/.style={draw, black, thick, inner sep=4mm}]
  \definecolor{acc}{HTML}{2348F2}
  % device 0 tokens
  \node[tok] (t1) at (0,0) {\texttt{t1}};
  \node[tok] (t2) at (1.1,0) {\texttt{t2}};
  % device 1 tokens
  \node[tok] (t3) at (4.6,0) {\texttt{t3}};
  \node[tok] (t4) at (5.7,0) {\texttt{t4}};
  % experts
  \node[ex] (ea) at (0,2.3) {\texttt{E1}};
  \node[ex] (eb) at (1.1,2.3) {\texttt{E2}};
  \node[ex] (ec) at (4.6,2.3) {\texttt{E3}};
  \node[ex] (ed) at (5.7,2.3) {\texttt{E4}};
  % dispatch arrows (all-to-all)
  \draw[->, acc, thick] (t1) -- (ea);
  \draw[->, acc, thick] (t2) -- (ec);
  \draw[->, acc, thick] (t3) -- (eb);
  \draw[->, acc, thick] (t4) -- (ed);
  \node[black, anchor=south, font=\footnotesize] at (2.85,2.6) {\texttt{dispatch all-to-all}};
  \node[black, anchor=east, font=\footnotesize] at (-0.55,1.15) {\texttt{device 0}};
  \node[black, anchor=west, font=\footnotesize] at (6.25,1.15) {\texttt{device 1}};
  \draw[black, dashed] (2.85,-0.6) -- (2.85,3.0);
\end{tikzpicture}
$$

The forward pass of one MoE layer makes the routing, dispatch, and combine explicit.

```algorithm
caption: $\textsc{MoELayer}(X, k)$ — sparse top-$k$ mixture-of-experts forward pass
input batch $X = (x_1, \dots, x_T)$; experts $E_1, \dots, E_E$; capacity $C$
$\mathit{load}_i \gets 0$ for each expert $i$ // tokens assigned so far
for each token $x_t$ in $X$ do
  $h \gets W_g\, x_t$ // router logits
  $\mathcal{T} \gets \argtop_k(h)$ // the $k$ highest-scoring experts
  $g \gets \softmax(h_{\mathcal{T}})$ // renormalize over the survivors
  $y_t \gets 0$
  for each expert $i$ in $\mathcal{T}$ do
    if $\mathit{load}_i < C$ then // expert still has a free slot
      $y_t \gets y_t + g_i \cdot E_i(x_t)$
      $\mathit{load}_i \gets \mathit{load}_i + 1$
    else
      $y_t \gets y_t + x_t / k$ // dropped: pass through residual
return $(y_1, \dots, y_T)$ and the batch loads for $\mathcal{L}_{\text{aux}}$
```

In practice the per-token loop is a single batched gather/scatter: tokens are sorted
by their chosen expert, dispatched in one all-to-all, processed as dense matrix
multiplies inside each expert, and scattered back. The dense expert matmul operates on
the padded $[E, C, d]$ buffer, so its shape and FLOP count are fixed at compile time
regardless of the realized routing.

## Practical issues

Sparse models trade their FLOP savings for a set of engineering and statistics
problems the dense model never had.

**Training instability.** The router's argmax is discontinuous, so small logit
changes flip a token between experts and produce loss spikes; the **ST-MoE** recipe
adds a **router $z$-loss**, $\mathcal{L}_z = \frac{1}{T}\sum_x \parens{\log\sum_i e^{h_i(x)}}^2$,
which penalizes large logits and keeps the softmax in a numerically stable range,
the single most effective stabilizer reported.[^zoph] The $\log\sum e^{h_i}$ term is
the log-partition of the router softmax; squaring and minimizing it pulls the logits
toward a smaller magnitude, so a single flip moves the gate values less and the loss
surface stays smoother.

**Fine-tuning.** Sparse models overfit downstream tasks faster than dense ones of
equal quality, because each expert sees only a fraction of the (already small)
fine-tuning set; ST-MoE recommends fine-tuning a subset of parameters and a smaller
auxiliary-loss weight.[^zoph]

**Expert specialization.** Trained experts do specialize, though rarely along
human-legible lines; in Mixtral the routing is largely uniform across experts and only
weakly correlated with topic or syntax, so the gain comes from added capacity more
than from clean division of labor.[^jiang]

**Sparse upcycling.** Rather than train an MoE from scratch, copy a trained dense
FFN into $E$ identical experts, add a fresh router, and continue training; the copies
diverge under the load-balancing pressure, recovering most of the gain at a fraction
of the from-scratch cost.

**Inference routing cost.** At serve time the FLOP savings are real but the full
parameter set must be resident in memory, and the all-to-all dispatch adds latency and
makes throughput sensitive to how evenly a batch happens to route. The binding
constraint shifts from compute to memory and communication.

**Mixtral (2024).** A modern decoder-only LLM, Mixtral 8×7B, places $E = 8$ experts
with top-$2$ routing in every layer. It holds $47$B total parameters but activates only
$13$B per token, matching or beating a $70$B dense model at the inference cost of a
$13$B one.[^jiang]

## Dense vs sparse: the accounting

The three routing regimes differ only in $k$, but that one number sets active
parameters, communication, and the headline parameter count.

| Model | Routing | Active params / token | Total params | Key idea |
| --- | --- | --- | --- | --- |
| Dense Transformer | all FFNs | $N$ | $N$ | one FFN per layer, every weight used |
| Switch Transformer | top-$1$ | $\approx N / E$ | up to $1.6$T | one expert per token; cheapest routing |
| GShard / Mixtral | top-$2$ | $\approx 2N / E$ | $600$B / $47$B | two experts per token; more gate signal |

The pattern is constant per-token cost with growing capacity: each row adds experts
(total params) without adding much active compute, exactly the decoupling
conditional computation promised. The figure below sizes the trade for one FFN sublayer:
the dense block uses all its parameters every token; the sparse block stores $E$ times
as many but touches only $k$ of them.

$$
% caption: Dense FFN uses all its parameters per token (params = FLOPs); the MoE stores
% $E$ experts but activates $k$, so params scale with $E$ while FLOPs scale with $k$.
\begin{tikzpicture}[>=stealth, font=\scriptsize, x=1.0cm, y=1.0cm]
  \definecolor{acc}{HTML}{2348F2}
  % ---- dense: params bar and FLOP bar equal ----
  \begin{scope}
    \node[black, anchor=south, font=\footnotesize] at (0.85,3.1) {\texttt{dense FFN}};
    \draw[draw=black, thick, fill=black!8] (0.2,0) rectangle (0.8,1.0);
    \draw[draw=acc, thick, fill=acc!15] (1.0,0) rectangle (1.6,1.0);
    \node[black, anchor=north, font=\footnotesize] at (0.5,-0.15) {\texttt{params}};
    \node[acc, anchor=north, font=\footnotesize] at (1.3,-0.15) {\texttt{FLOPs}};
  \end{scope}
  % ---- sparse: params tall (E experts), FLOPs short (k active) ----
  \begin{scope}[xshift=4.2cm]
    \node[black, anchor=south, font=\footnotesize] at (0.85,3.1) {\texttt{MoE (E=8, k=2)}};
    \draw[draw=black, thick, fill=black!8] (0.2,0) rectangle (0.8,3.0);
    \draw[draw=acc, thick, fill=acc!15] (1.0,0) rectangle (1.6,0.75);
    \node[black, anchor=north, font=\footnotesize] at (0.5,-0.15) {\texttt{params}};
    \node[acc, anchor=north, font=\footnotesize] at (1.3,-0.15) {\texttt{FLOPs}};
    \node[black, anchor=west, font=\footnotesize] at (1.8,2.6) {\texttt{store 8 experts}};
    \node[acc, anchor=west, font=\footnotesize] at (1.8,0.55) {\texttt{run 2}};
  \end{scope}
\end{tikzpicture}
$$

## Fine-grained and shared experts

Goodfellow describes the classic gated mixture; the design that ships in the largest open models since makes two modifications to the top-$k$ layer, both aimed at getting more specialization out of the same active-parameter budget.[^dai-deepseekmoe]

**Fine-grained experts.** Instead of $E$ experts each of hidden width $d_{\text{ff}}$ routed top-$2$, split each expert into $m$ slices of width $d_{\text{ff}}/m$, giving $mE$ small experts routed top-$2m$. The active-parameter count is unchanged — the same total hidden width fires per token — but the router now chooses from a much larger menu, so each token can assemble a more precise combination of specialists. The intuition is combinatorial: routing $8$ of $64$ small experts gives far more distinct expert-subsets than routing $2$ of $8$ large ones, and finer partitions let experts specialize more sharply.

**Shared experts.** A fine-grained MoE also keeps one or two **shared experts** that _every_ token passes through, alongside the routed ones. The shared expert absorbs the common computation that all tokens need — the syntax and general structure that would otherwise be redundantly relearned by every routed expert — so the routed experts are freed to specialize on what actually differs between tokens. This removes the redundancy that makes plain MoE experts specialize only weakly.

$$
% caption: Fine-grained plus shared experts. Every token passes through a shared
% expert (always on); the router additionally picks several fine-grained experts from
% a large pool, so active compute is unchanged but the routed menu is far larger.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  sh/.style={draw=green, text=green, thick, minimum width=11mm, minimum height=6.5mm, inner sep=1pt, fill=green!14},
  on/.style={draw=acc, text=acc, thick, minimum width=8mm, minimum height=6.5mm, inner sep=1pt, fill=acc!15},
  off/.style={draw, black, minimum width=8mm, minimum height=6.5mm, inner sep=1pt, fill=black!8}]
  \definecolor{acc}{HTML}{2348F2}
  \definecolor{green}{HTML}{1F9D4D}
  \node[black] (t) at (0,0) {\texttt{token}};
  \node[sh] (s) at (-2.6,1.6) {\texttt{shared}};
  \node[on] (e1) at (-1.1,1.6) {\texttt{f2}};
  \node[off] (e2) at (0.0,1.6) {\texttt{f5}};
  \node[on] (e3) at (1.1,1.6) {\texttt{f7}};
  \node[off] (e4) at (2.2,1.6) {\texttt{...}};
  \draw[->, green, thick] (t) -- (s);
  \draw[->, acc, thick] (t) -- (e1);
  \draw[->, acc, thick] (t) -- (e3);
  \draw[->, black, thick, dashed] (t) -- (e2);
  \node[green, anchor=west, font=\scriptsize] at (-2.6,2.4) {\texttt{always on}};
  \node[acc, anchor=west, font=\scriptsize] at (0.2,2.4) {\texttt{routed (fine-grained)}};
\end{tikzpicture}
$$

These two changes, on top of the auxiliary-loss and capacity machinery derived above, are why modern MoE LLMs push the active-to-total ratio far lower than Mixtral's $13/47$ while keeping quality: a large pool of finely-sliced routed experts, a shared expert absorbing common work, and the same $z$-loss and load-balancing machinery. The layer is the same $y = \sum_i g_i(x)\, E_i(x)$; only the granularity and a fixed always-on term changed.

## Takeaways

- **Conditional computation** activates an input-dependent subset of weights, so
  total capacity $N$ and per-token FLOPs ($\propto \tfrac{k}{E}N$) become separate
  dials.
- A **mixture-of-experts layer** replaces one FFN with $E$ experts and a router,
  outputting $y = \sum_i g_i(x)\, E_i(x)$; it drops in where the Transformer FFN sat.
- **Sparse top-$k$ gating** keeps the $k$ largest router logits, renormalizes the
  softmax over them, and zeroes the rest, so only $k$ experts run per token.
- Without intervention the router **collapses** onto a few experts; the
  **load-balancing loss** $\alpha E \sum_i f_i P_i$ is minimized at the uniform split
  $f_i = P_i = 1/E$.
- A **capacity factor** $c$ caps tokens per expert and **drops** the overflow through
  the residual; **expert-choice** routing balances by construction.
- The architectures simplify and grow: sparsely-gated **LSTM** MoE, **GShard**
  (top-$2$ Transformer, $600$B), **Switch** (top-$1$, $1.6$T); experts are sharded by
  **expert parallelism** with two all-to-all shuffles per layer.
- Sparse models need a **router $z$-loss** for stability, need careful fine-tuning,
  and shift the serving bottleneck from compute to **memory and communication**.
- **Mixtral 8×7B** ($47$B total, $13$B active, top-$2$) is the canonical modern MoE,
  matching a $70$B dense model at a fraction of the inference compute.
- **Modern designs:** **fine-grained experts** (many thin experts, larger routed
  menu, same active budget) and **shared experts** (one always-on expert absorbing
  common computation) let modern MoE LLMs push the active-to-total ratio far below
  Mixtral's while keeping quality.

[^dai-deepseekmoe]: **Dai et al.**, _DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models_, 2024 — fine-grained expert segmentation plus isolated shared experts, raising specialization at a fixed active-parameter budget.
[^jacobs]: **Jacobs, Jordan, Nowlan & Hinton**, _Adaptive Mixtures of Local Experts_, Neural Computation 1991 — the original mixture-of-experts: a gating network softly partitions the input space among specialist sub-networks trained jointly.
[^shazeer]: **Shazeer et al.**, _Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer_, ICLR 2017 — sparse noisy top-$k$ gating between LSTM layers, the load-balancing loss, and the first MoE at hundreds of billions of parameters.
[^lepikhin]: **Lepikhin et al.**, _GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding_, ICLR 2021 — top-$2$ MoE Transformer with capacity, automatic sharding, and all-to-all dispatch for a $600$B-parameter translation model.
[^fedus]: **Fedus, Zoph & Shazeer**, _Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity_, JMLR 2022 — top-$1$ routing, the differentiable load-balancing loss, and scaling to $1.6$T parameters.
[^zoph]: **Zoph et al.**, _ST-MoE: Designing Stable and Transferable Sparse Expert Models_, 2022 — the router $z$-loss for training stability and a fine-tuning recipe for sparse models.
[^zhou]: **Zhou et al.**, _Mixture-of-Experts with Expert Choice Routing_, NeurIPS 2022 — inverts routing so each expert selects its top tokens, guaranteeing perfect load balance without an auxiliary loss.
[^jiang]: **Jiang et al.**, _Mixtral of Experts_, 2024 — a decoder-only LLM with $8$ experts and top-$2$ routing per layer, $47$B total / $13$B active parameters, matching a $70$B dense model.
[^gf-conditional]: **Goodfellow**, _Deep Learning_, §12.4.3 — conditional computation and mixtures of experts: gating networks that activate input-dependent sub-networks to decouple capacity from per-example cost.