Beyond Dense Transformers: MoE, SSMs, Hybrids

Chapter 8 treated the dense block as the default. That default runs every token through every parameter, and it hides two walls a frontier builder eventually hits. One wall is in the feed-forward block: capacity and per-token compute are welded together, so the only way to know more is to compute more. The other is in attention: mixing every token with every other costs the square of the sequence length and leaves behind a key-value cache that grows without bound. A mixture-of-experts layer breaks the first wall by routing a token to only a few of its many experts. Against the second, a state-space model swaps attention's quadratic mixing for a linear recurrence. Hybrids keep both moves in the same stack. The sparse bet only pays off if the router can add parameters almost for free without starving experts, collapsing load, or making communication dominate the saved compute.

The two walls of the dense block

A dense feed-forward block runs every token through all of its parameters. That couples two things a frontier builder wants to separate: the model's capacity, how much it can know, and its per-token compute, how much each forward and backward step costs. Under that coupling, the only way to add knowledge is to add FLOPs, and FLOPs are the budget Chapter 5 spends. A model that wants more capacity than its compute budget can afford has no move inside the dense design. Call this the capacity wall.

The second wall is attention's own. Self-attention mixes every token with every other, so its cost grows with the square of sequence length, and the key-value cache it leaves behind grows linearly with context and dominates serving memory (Chapter 31). For long context that quadratic is the wall. The dense transformer therefore carries two separate inefficiencies: a feed-forward block that spends compute on capacity it may not need per token, and an attention block that spends compute and memory on a mixing pattern that may be denser than the task requires. The rest of this chapter takes the two walls in turn. MoE attacks the first, state-space models the second, and hybrids decline to choose.

Breaking the capacity wall: mixture-of-experts

Decoupling parameters from FLOPs

The mixture-of-experts answer is to decouple parameters from active FLOPs. An MoE layer holds $N$ expert feed-forward networks but routes each token to only $k$ of them, with $k \ll N$. The identity that matters: total parameters scale capacity and knowledge, while active parameters, the $k$ chosen experts plus the always-on parts of the model, scale per-token compute and step cost. Capacity becomes cheap relative to FLOPs, which is the whole frontier appeal. The same arithmetic cuts inference FLOPs too, though serving an MoE is its own problem and out of scope here. Figure 9.1 makes the split visible: as the expert count $N$ grows, total parameters climb while the active count that sets per-token FLOPs stays flat with $k$ fixed.

Sweep the expert count $N$ and watch total capacity climb while the active params that set per-token FLOPs stay almost flat.

import numpy as np, matplotlib.pyplot as plt

# MoE decouples total params (capacity) from active params (per-token FLOPs).
d, ffn = 4096, 14336           # hidden size and one expert's FFN width
k = 2                          # experts fired per token
per_expert = 2 * d * ffn       # up + down projection of one expert
N = np.arange(1, 65)           # number of experts in the layer
total  = N * per_expert                    # all experts are stored
active = k * per_expert + N * d            # k fired experts + router map

plt.figure(figsize=(5, 3))
plt.plot(N, total / 1e9,  label="total params (capacity)")
plt.plot(N, active / 1e9, label="active params (per-token FLOPs)")
plt.xlabel("experts N"); plt.ylabel("billions of params"); plt.legend()
plt.tight_layout(); plt.show()

print(f"N=64, k={k}: {total[-1]/1e9:.1f}B total but only {active[-1]/1e9:.2f}B active")
print(f"capacity-to-compute ratio: {total[-1]/active[-1]:.0f}x")

The new learned object is the router, also called the gate. It is a small linear map from a token's hidden state to one score per expert. A softmax or sigmoid turns scores into weights, the top-$k$ experts are selected, and their outputs are combined weighted by the renormalized gate values. In sketch form:

# router: a linear map h -> scores over N experts
scores = h @ W_gate                  # (tokens, N)
weights = softmax(scores)            # or sigmoid per-expert
top = topk(weights, k)               # indices and gate values
y = sum(g_i * Expert_i(h) for i, g_i in top)

This tiny map is the only genuinely new component, and it is the source of almost every MoE pathology. The whole rest of the design exists to keep the router in check. Figure 9.2 traces one token through the layer: the gate scores all experts, top-$k$ are selected and combined, while a shared expert (a later refinement) stays always-on in parallel.

The router and its pathologies

The first thing the router gets wrong on its own is balance. Left alone, the gate collapses onto a few favored experts in a winner-take-all loop: an expert that is chosen gets more gradient, improves, and is chosen more, while the rest starve and become dead parameters. Two families of fix exist. The auxiliary load-balancing loss, from Shazeer et al. (2017) (Shazeer et al. 2017) and refined in GShard (Lepikhin et al. 2020) and Switch (Fedus et al. 2021), adds a penalty term on the product of the fraction of tokens routed to each expert and the gate mass placed on it, pushing the distribution toward uniform. It works, but it injects a loss whose weight trades balance against task quality. The newer answer, loss-free or bias-based balancing from DeepSeek-V3 (DeepSeek-AI 2024), drops the auxiliary loss entirely. Instead it keeps a per-expert routing bias, nudged up for under-used experts and down for over-used ones between steps, steering selection without adding an interfering gradient (Wang et al. 2024). That sidesteps the balance-versus-quality tension the auxiliary loss creates.

Two refinements change what an "expert" is. Fine-grained experts slice each expert into smaller ones and raise $k$ proportionally, so the same active FLOPs buy far more routing combinations and sharper specialization. Shared experts are a few that every token always passes through, absorbing the common, domain-general computation so the routed experts can specialize instead of each re-learning the basics. DeepSeekMoE combines both (Dai et al. 2024), and the pair is now a common default for large open MoE models.

One systems fact intrudes on this otherwise clean picture. Routing is dynamic, but hardware wants fixed-shape buffers, so each expert is given a capacity cap, roughly capacity_factor * (tokens / experts). Tokens routed to an over-full expert are dropped, passed through on the residual or rerouted to a second choice. The capacity factor is a direct knob: set it low and you save memory and communication but drop more tokens, losing compute and adding training noise; set it high and you waste padding on under-full experts.

How the layer reached this form

Sparse experts predate the transformer era as a general idea, but the line that matters here begins with Shazeer et al. (2017) (Shazeer et al. 2017), who introduced top-$k$ gating and the auxiliary balancing loss inside an LSTM language model. GShard (Lepikhin et al., 2020) (Lepikhin et al. 2020) carried the layer into the transformer and into automatic sharding across devices. Switch Transformers (Fedus et al., 2021) (Fedus et al. 2021) pushed to the cheapest possible router, $k=1$, and to trillion-parameter scale, showing top-1 routing could train stably with the right stabilizers. GLaM (Du et al., 2022) (Du et al. 2022) made the efficiency case at scale, and Mixtral (Jiang et al., 2024) (Jiang et al. 2024) brought a strong open $k=2$ model into wide use.

Two threads then refined the layer. On routing, expert-choice (Zhou et al., 2022) (Zhou et al. 2022) inverted the question: instead of each token picking its experts, each expert picks its tokens, which makes per-expert load exact by construction and removes dropping, at the cost of some tokens getting no expert at all. BASE Layers (Lewis et al., 2021) (Lewis et al. 2021) framed routing as a balanced assignment problem outright. On specialization, DeepSeekMoE (Dai et al., 2024) (Dai et al. 2024) introduced fine-grained and shared experts, and DeepSeek-V3 (2024) (DeepSeek-AI 2024) replaced the auxiliary loss with bias-based balancing at frontier scale, with the standalone method documented separately by Wang et al. (2024) (Wang et al. 2024).

A parallel thread made MoE practical without paying for a from-scratch run. Sparse upcycling (Komatsuzaki et al., 2022) (Komatsuzaki et al. 2022) initializes each expert as a copy of a trained dense feed-forward block, adds a fresh router, and continues training; the experts diverge as routing specializes them. It buys strong quality for a fraction of from-scratch compute, with the failure mode that experts can stay near-identical if routing pressure is too weak, leaving MoE's cost without its benefit.

Stabilizing the router in practice

The router is the part to get right, and its instabilities have a specific cause. The top-$k$ selection is an argmax, it picks the highest-scoring experts, and argmax is discontinuous: a small change in scores flips a token's expert, making the loss landscape jagged and spike-prone, worse in low precision. The standard mitigations are narrow and worth naming. The router z-loss from ST-MoE (Zoph et al., 2022) (Zoph et al. 2022) penalizes large router logits to keep the gate numerically well-behaved; note this is distinct from the output-logit z-loss in the stability kit of Chapter 5. Keeping the router in fp32 under a bf16 or fp8 body (numeric precision formats, where fewer bits mean less memory but coarser arithmetic, defined in Chapter 10) costs little and removes a class of spikes. And early routing decisions can lock in before experts have differentiated, so a run is worth watching for that in its first phase. MoE loss curves are spikier than dense curves at the same scale; an unguarded router, with no z-loss and bf16 logits, can spike and fail to recover days into a run.

The failure modes are worth naming because each bites at a different time. Routing collapse funnels most tokens to a few experts and leaves the rest dead; it is exactly what load balancing exists to prevent, and it returns if balance is mis-weighted or disabled. Load imbalance and token dropping, short of collapse, overflows some experts' capacity and silently discards their tokens' compute while adding noise, and it stalls the slowest expert's device, a systems cost paid in Chapter 10. Expert under-training leaves rarely-chosen experts too few tokens to specialize, the exact risk mode of upcycling when routing pressure is too weak.

When the bet pays off

• Sparse versus dense. MoE buys more capacity per training FLOP and per active FLOP, but pays in memory, since all experts are stored and sharded, in communication, since expert parallelism needs an all-to-all (every device exchanges data with every other, Chapter 10), in fragility, and in a harder serving story. At small scale or under tight memory, dense wins. The crossover favors MoE as total scale grows.
• top-1 versus top-k. A router with $k=1$ is cheapest and simplest but more brittle and less expressive (Fedus et al. 2021); $k=2$ is the common sweet spot (Jiang et al. 2024); higher $k$ approaches dense cost and erodes the advantage.
• Capacity factor. Lower is cheaper but drops more tokens and adds noise; higher wastes padding. There is no single right value, and it is often raised at evaluation and inference, where dropping a token is unacceptable.

Breaking the context wall: state-space models

A fixed state instead of an addressable past

State-space models answer the other wall, attention's quadratic. An SSM maintains a fixed-size hidden state and updates it token by token through a linear recurrence, where each step's state feeds into the next, then reads an output from it:

$$h_t = A\, h_{t-1} + B\, x_t, \qquad y_t = C\, h_t$$

Because the state is fixed in size, the cost is linear in sequence length and there is no growing key-value cache to serve, only a constant-size state. The price is that all of a token's past must be summarized into that fixed state, where attention keeps every past token addressable. The design question is whether a learned, input-dependent recurrence can compress the past well enough to compete. Mamba's contribution was to make the recurrence selective, letting $A$, $B$, and $C$ depend on the input so the model can choose what to keep and what to forget, recovering much of what content-based attention does while staying linear (Gu and Dao 2023).

From S4 to Mamba

The state-space line is younger as a serious attention competitor. The structured state-space model S4 (Gu et al., 2021) (Gu et al. 2021) showed a carefully parametrized linear recurrence could handle very long sequences, and Mamba (Gu and Dao, 2023) (Gu and Dao 2023) made the recurrence selective and hardware-efficient, bringing it within reach of language modeling at scale. Mamba-2 (Dao and Gu, 2024) (Dao and Gu 2024) tied state-space models back to attention through a duality that let them reuse attention's hardware.

Mamba is not the only linear line. Linear attention reaches the same shape from the other side: drop the softmax from attention, and the past can be folded into a fixed-size running state, updated token by token like an SSM. Gated DeltaNet is the current refinement, adding a gated forgetting rule to the Mamba-2 recurrence (Yang et al. 2024), and MiniMax's lightning attention is the same family run at production scale. For years the practical objection to all of these was tooling: the recurrence needs a hardware-efficient scan to be fast, and that path was thinner than the dense-transformer one. The hybrid systems below are what retired the objection.

Recall is the catch

Attention keeps every past token exactly addressable and excels at precise recall, at quadratic cost and a growing key-value cache. An SSM is linear in length with a constant-size state, but must compress the past, and tends to lag attention on tasks that need exact long-range lookup. Figure 9.4 is the cost side of that trade: the two curves part ways as context grows, which is what makes the recall penalty worth paying at long context. Hybrids, the subject of the next section, exist precisely because neither pure form dominates.

Neither alone: the hybrid stack

In practice the frontier does not choose one or the other. Hybrid stacks interleave a minority of full-attention layers with a majority of SSM or linear-recurrence layers, so a few global-mixing layers preserve exact recall while the linear layers carry the bulk of the sequence cheaply. MoE and SSM are also orthogonal: a hybrid can put its feed-forward blocks behind an MoE router and its sequence-mixing behind a state-space recurrence at the same time. Jamba (Lieber et al., 2024) (Lieber et al. 2024) interleaved Mamba and attention layers and added MoE on top, demonstrating that the three ideas compose into one stack. What Jamba demonstrated, the 2025 generation shipped at scale: MiniMax-01 interleaved lightning attention with softmax attention at 456B total parameters in January 2025 (MiniMax 2025), Qwen3-Next runs Gated DeltaNet against gated full attention at roughly three linear layers per attention layer (Qwen Team 2025), and Nemotron-H does the same with Mamba layers (NVIDIA 2025). Figure 9.6 shows the three ideas composing: the sequence-mixing slot alternates SSM and attention, while the feed-forward slot is an MoE layer.

The few attention layers in a hybrid still carry a key-value cache, so a hybrid does not escape attention's costs entirely; it pays them on a minority of layers instead of all of them. Figure 9.7 lays out the threads of this chapter as a genealogy: a routing line, a specialization line, the upcycling shortcut, and the separate state-space line, all converging on the hybrid stack.

Two trade-offs span the whole chapter and frame the rest of the book's view of these architectures.

• Attention versus state-space. Attention keeps every past token exactly addressable and excels at precise recall, at quadratic cost and a growing key-value cache. An SSM is linear in length with a constant-size state, but must compress the past, and tends to lag attention on tasks that need exact long-range lookup. Hybrids exist precisely because neither pure form dominates (Lieber et al. 2024).
• Train and serve complexity. Every choice here adds operational surface: a router to tune, balance to monitor, expert placement to lay out, and a memory and communication profile that Chapter 10 must absorb. The quality win is real but never free engineering.

Further reading