Why this matters now

Everyone wants longer context — to pack in more knowledge, or to run agents over big workloads — and vendors have raced to ever-larger context windows. The problem is cost structure: the feed-forward part of the network grows linearly with sequence length, but attention is an all-to-all interaction and grows quadratically. So at short sequences FFN dominates the cost; at long sequences attention takes over.

Two parts of the toolkit reduce this. One is hybrids of local + global attention (Lecture 3): do full attention only every few layers, local windows elsewhere. The other is systems engineering — and a recurring theme of this course is that constant factors really matter.

The one core idea: associativity of multiplication

Standard attention computes ρ(Q Kᵀ) V, where ρ is the softmax (row-normalize). The Q Kᵀ term is the killer: it's n² · d, and n (context length) can be millions. Linear attention makes one lossy move — drop the softmax (treat ρ as the identity) — which lets you re-associate the matrix product:

# Softmax attention: the n² term is unavoidable
out = softmax(Q @ K.T) @ V       # cost ~ n² · d

# Drop softmax → re-associate: multiply K.T @ V FIRST
out = Q @ (K.T @ V)               # cost ~ n · d_k · d_v

Now the cost depends on n · d_k · d_v instead of n². The hidden dims are thousands, never millions — a far friendlier term. That's the whole trick: associativity changes which factor is quadratic.

The RNN duality

Something even nicer: that re-associated form Q (Kᵀ V) can be written incrementally. Sweep left to right, accumulate a fixed-size state S from the keys and values, then read it out with the query:

# Dense form (parallel — great for TRAINING):
out = Q @ (K.T @ V)

# Recurrent form (serial — great for INFERENCE):
S = 0
for t in sequence:
    S = S + k[t].outer(v[t])      # accumulate KV into a fixed-size state
    out[t] = q[t] @ S

The dense and recurrent forms are mathematically equivalent (the lossy step was only dropping the softmax). This is the best of both worlds: train with the parallel matmul form, infer with the RNN form whose state is fixed size (no growing KV cache). The catch is that linear attention alone is too weak — so it's the starting point, not the destination.

Mamba-2: add an input-dependent gate

Linear attention's weakness: it always passes the whole state forward. LSTMs taught us it's important to choose when to forget. Mamba-2 (the state-space-model family of Gu, Dao, et al.) adds a gate γₜ that modulates how much state carries forward:

# Linear attention:  S[t] = S[t-1] + k[t]·v[t]ᵀ
# Mamba-2:            S[t] = γ[t]·S[t-1] + k[t]·v[t]ᵀ     # γ[t] gates the carry-forward

Crucially, γₜ depends only on the current input xₜ — it's not state-dependent. That's what preserves the parallel/serial duality: as long as every term in the recurrence is input-dependent (no state dependence), you keep both the dense training form and the cheap recurrent inference form. (Mamba-2 also adds a small residual-like vₜ term, not core to the idea.) Nemotron-3 uses Mamba-2 as its lightweight layer interleaved with periodic full attention, getting good throughput at long context.

Gated DeltaNet: a second gate + a delta update

Push the recurrence further. Gated DeltaNet (used by Qwen-Next / Qwen 3.5) adds a second gate βₜ — a "write" gate (βₜ = 0 means "don't add the current input to the state," very LSTM-like) — and a smarter update direction:

S[t] = γ[t]·(I − β[t]·k[t]·k[t]ᵀ)·S[t-1] + β[t]·k[t]·v[t]ᵀ
#            └── project out the current key direction ──┘

The (I − β k kᵀ) term acts like a projector: when writing new information for key kₜ, it also erases previous content stored at that key — not just "add," but "clear and replace." This exact update was independently reinvented in fast-weight programming and test-time training from very different starting points. Qwen-Next/3.5 use a 3:1 gated-DeltaNet : full-attention hybrid — strong quality with much higher decoding throughput as context grows.

A different route: DeepSeek Sparse Attention (DSA)

Not all efficient attention is linear. DSA (DeepSeek v3.2) keeps full attention but only over a selected subset of tokens. A cheap indexer scores every position, takes the top-k, and full attention runs only on those:

This is not linear — the indexer still does all-to-all inner products — but the indexer can be made very cheap (low precision, low dimension), and the expensive full-attention step runs on a bounded top-k subset. A neat practical point: you don't pretrain with it. Train a normal transformer, then bolt on the indexer during the long-context extension stage (which you'd run anyway). DeepSeek v3.2 matches frontier models (Claude 4.5 Sonnet, Gemini 3) with much better prefill/decode scaling; GLM-5 adopted DSA too, with ablations showing little quality loss vs full attention even on hard long-context retrieval.

The fundamental trade-off of state-space models

Why not go fully RNN? Expressive power. Softmax attention's all-to-all connection is extremely powerful and easy to train. State-space models historically shared the LSTM's weaknesses, but caught on because the duality (RNN form ⇄ dense matmul form) finally solved their hardware-efficiency problem. What remains is the state-vs-context trade-off: a finite-size state must compress everything it has seen, so for a fixed (small) state, a long context inevitably loses information. Make the state as big as the context and you're back to paying large costs. No free lunch — yet.

What an MoE is

Conceptually an MoE is just a more efficient MLP. Replace the single FFN with several FFNs ("experts") plus a router that picks which expert(s) handle each token. With 4 experts you have 4× the FFN parameters, but each forward/backward pass touches only one (or k) of them — so you pay roughly one expert's worth of FLOPs.

Mental model: increase parameters without increasing FLOPs. Routing is at the token level — every token gets routed — and routers are deliberately naive (a single matmul; they don't "understand" the token).

Why everyone ships MoEs now

Empirically, holding total compute fixed and increasing the number of sparse parameters just makes models better. Fedus et al. (Switch Transformer, 2022) showed test loss dropping as you add experts (active params fixed), and better performance per unit of training compute. OLMoE reproduced it: an MoE trains roughly 2× faster than its dense counterpart to the same quality. DeepSeek's MoEs match or beat dense models with far fewer active parameters — which is what determines inference cost.

The fourth row matters: experts are natural chunks you can place on different devices (expert parallelism), giving another way to shard a model too big for one device — at the cost of shipping activations between devices (comms).

Most big open models past a certain size are MoEs. In the West: Llama 4, GPT-OSS. Much of the MoE research and training momentum is in China — Qwen, DeepSeek, miniCPM pioneered and popularized the recipes (e.g. Qwen 1.5-MoE's 2.7B-active model beating 7B dense models). Note: MoE-ifying the attention block has been tried but is much harder; essentially everyone applies MoE to the FFN only.

Design axis 1 — routing

You never activate all experts; you always pick a Top-k. Who chooses?

Token choice (standard)

Each token picks its favorite k experts. Lower validation loss and better downstream scores than the alternatives (OLMoE).

Expert choice

Each expert picks its favorite tokens. Trains fine, but generally harder to get working well.

Global assignment

Solve a linear-assignment problem for the optimal token↔expert matching. Elegant and optimal, but too expensive — not seen at scale.

The standard router is trivially simple — an inner product between the token and a learned vector per expert, then Top-k + softmax for the gates:

s = softmax(u @ W_router)      # score per expert (a single matmul)
g = topk(s, k)                  # keep the k largest; others = 0
out = u + Σ_i  g[i] · expert_i(u)  # gated sum over selected experts

Variants exist: hashing tokens to experts (no learning — surprisingly works, a common baseline), RL routing (the "natural" bandit framing — Bengio 2013 — but high variance/overhead, rarely used), and linear assignment (optimal, too costly). Learned Top-k inner-product routing won because simple heuristics make it just work.

Design axis 2 — shared + fine-grained experts

DeepSeekMoE's now-standard idea: cut experts into smaller fine-grained chunks, and make a few of them shared experts that are always on (they bypass the router and process every token). The shared expert absorbs the common processing all tokens need, freeing the routed experts to specialize.

DeepSeek's ablations show both interventions help (finer-grained experts + a shared expert), with large gains on knowledge tasks (TriviaQA, NaturalQuestions). OLMoE agrees fine-grained helps but found shared experts didn't help much — so the shared-expert point is mildly contested. Either way, the DeepSeekMoE design (shared + fine-grained) became the standard everyone copies, much as Llama set the dense-transformer template.

Design axis 3 — training (the hard part)

Sparsity is what makes MoEs hard. During training only k experts are active, so gating is non-differentiable and you never observe the counterfactual experts you didn't pick — a bandit-flavored problem. Three families of solutions: RL on the router (works, but high variance — Clark 2020), stochastic perturbations (inject noise to tie-break/explore — the original Shazeer MoE; Fedus's multiplicative version was later dropped), and heuristics. Practice uses heuristics.

The fix is a load-balancing auxiliary loss added to the modeling loss. The Switch Transformer form multiplies, per expert, the fraction of tokens dispatched (fᵢ) by the soft probability mass the router assigned (Pᵢ):

L_balance = N · Σ_i  f_i · P_i
# derivative w.r.t. P_i is f_i → popular experts get pushed down,
# proportional to how many tokens they grabbed.

You can't read the objective off directly, but its gradient is clear: push down probability mass on over-popular experts. DeepSeek adds a second, per-device balancing loss (so machines holding different experts stay evenly utilized), and DeepSeek-V3 moves toward aux-loss-free balancing via a per-expert bias updated by online learning — though some auxiliary loss still remains to prevent extreme imbalance.

MoE systems, stability, and fine-tuning

Structured-sparse matmuls

Multiple experts on one GPU shouldn't be many tiny matmuls. Use block-diagonal / structured sparsity (natively supported in hardware) to do them as one big efficient multiply — a hardware/architecture co-design that fits MoEs naturally.

Comms trick (Nemotron-3)

Expert parallelism ships activations between devices. Down-project the residual to a lower dimension before the all-to-all collective, cutting communication without the full downside of a smaller hidden dimension. (Shared experts aren't communicated, so they can stay full-width.)

Token dropping

If one expert's queue overflows, old infra would silently drop tokens (sending zeros), creating weird cross-user stochasticity. Modern "dropless" frameworks (MegaBlocks, etc.) fixed this.

Router stability

The router adds another softmax (exp + division = danger). Use fp32 for the router and a z-loss (Lecture 3) — OLMoE shows z-loss removes spiky router training curves. (Barret Zoph's ST-MoE work catalogued these.)

Fine-tuning

MoEs have so many parameters they overfit badly on small fine-tuning sets (huge train/val gap vs dense). Mitigate by fine-tuning only attention or non-MoE FFN layers — or, bitter-lesson style, just use a lot more fine-tuning data and retrain the experts.

Upcycling (a fading trick worth knowing)

Want an MoE but already have a dense model? Upcycle: copy the dense MLP into many expert copies, add a randomly-initialized router, and train. Stochastic routing makes the copies specialize into real experts. miniCPM (2.4B → 13.4B) and the first Qwen-MoE (1.8B → 2.7B-active) showed near-free wins this way. It's faded now — if you're committing to an MoE, you just train the MoE from scratch rather than converting a dense model.

Case study: the DeepSeek MoE evolution

The DeepSeek papers are well worth reading; their MoE design is the field's template.

Two more V3 ideas worth knowing:

Multi-head latent attention (MLA)

Instead of projecting the hidden state directly to Q/K/V, first project to a low-dimensional latent c, then derive Q/K/V from c. You KV-cache only the small c — big memory savings. The one wrinkle: it conflicts with RoPE, so a few non-latent dimensions carry position separately.

Multi-token prediction (MTP)

Predict several future tokens at once. Statistical upside (better foresight) plus a systems upside: it gives you a built-in speculative decoder to speed up inference (more in the inference lecture).

Golden rules & takeaways

• Associativity is the key to linear attention: drop softmax, compute Q(KᵀV), and gain an exact RNN duality — parallel for training, fixed-state recurrent for inference.
• Everything that works at scale is a hybrid — interleave cheap layers (linear attention / Mamba-2 / gated DeltaNet) with periodic full attention; pure-RNN degrades, especially on long-context retrieval.
• Don't over-index on asymptotics. FlashAttention (constant factors) and DSA (cheap indexer + bounded Top-k) win big without changing the O(n²).
• MoE = more parameters at the same FLOPs. Token-choice Top-k routing with shared + fine-grained experts is the standard.
• The only thing standing between you and a working MoE is the load-balancing loss — without it, experts collapse to a handful. Otherwise, pump gradients straight through the non-differentiable Top-k.
• Respect the systems. Expert parallelism, structured-sparse matmuls, comms-saving down-projection, fp32 router + z-loss — MoE quality is as much systems as deep learning.
• Top-k selection + balancing losses is a reusable primitive — you've now seen it in DSA, MoE routing, and tokenizer-free models (H-Nets).

Next: scaling laws (and then systems — kernels, parallelism, inference — where FlashAttention and expert parallelism get their full treatment).

References

Primary source is the lecture itself. Key works named:

1. Dao et al. (2022) — FlashAttention
2. Katharopoulos et al. (2020) — Transformers are RNNs (linear attention)
3. Dao & Gu (2024) — Mamba-2 (state-space duality)
4. Yang et al. (2024) — Gated DeltaNet
5. Shazeer et al. (2017) — Outrageously Large Neural Networks (sparsely-gated MoE)
6. Fedus et al. (2022) — Switch Transformers
7. Zoph et al. (2022) — ST-MoE (MoE stability, z-loss)
8. Dai et al. (2024) — DeepSeekMoE (shared + fine-grained experts)
9. Muennighoff et al. (2024) — OLMoE (open MoE, controlled ablations)
10. DeepSeek-AI (2024) — DeepSeek-V2 (multi-head latent attention)
11. DeepSeek-AI (2024) — DeepSeek-V3 (aux-loss-free balancing, MTP)
12. Komatsuzaki et al. (2022) — Sparse Upcycling