Lecture 04 · Part 1

Attention Alternatives, Mixture of Experts

Tatsu covers two advanced architecture ideas layered on top of the basic transformer. First, attention alternatives that trade the quadratic cost of softmax attention for near-linear cost — linear attention and its RNN duality, gated variants like Mamba-2 and gated delta net used in hybrids, and DeepSeek's sparse attention. Second, mixture of experts, which buys more parameters without more FLOPs by routing each token to a few experts, tamed in practice by top-k routing and load-balancing losses.

▶ Lecture video Lecture slides (PDF) ↗

Why attention gets expensive, and the basic toolkit

Everyone wants longer context — more knowledge packed in, agents operating over more state — and if you plot context window size against release date on a log scale, there’s a clear rush by the top vendors toward millions of tokens. The problem is cost. Split a transformer layer’s compute into the feed-forward part and the attention part and watch what happens as the sequence grows: feed-forward starts large but grows linearly; attention is an all-to-all interaction between positions, so it grows quadratically and quickly outpaces everything else. At short sequence lengths feed-forward dominates; at long ones, attention is the whole problem.

There are two cheap ways to fight this before changing the math. The first is to combine local attention with occasional global attention — if you do full global attention only once every eight layers and keep the rest local, you have controlled the cost while keeping some long-range mixing. The second is pure systems engineering, and Tatsu is emphatic that this is underappreciated: “constant factors really, really matter.”

But if you want 5–10 million tokens, constant factors aren’t enough. You want a genuinely milder dependence on sequence length.

Linear attention and the RNN in disguise

Standard attention is $\text{Attn}(Q,K,V) = \rho(QK^\top)V$ , with $Q, K \in \mathbb{R}^{n \times d_k}$ and $V \in \mathbb{R}^{n \times d_v}$ . The $QK^\top$ product is what costs $n^2 d_k$ , and $n$ is the millions-long context. Now forget the softmax $\rho$ for a moment — pretend it’s the identity. Then multiplication is associative, and you can move the parentheses:

$(QK^\top)V = Q(K^\top V)$

That single re-grouping changes which factor is quadratic. Instead of $n^2 d_k + n^2 d_v$ you pay $2 n d_v d_k$ . The key and value dimensions can be large in principle, but they’re usually thousands to tens of thousands — nobody has a million hidden coordinates — so depending on $d_v d_k$ instead of $n^2$ is a huge win.

The re-grouped form has a second gift. The right-hand product, read incrementally as you sweep left to right, looks exactly like an RNN with a fixed-size state $S$ :

Code walkthrough — The duality: parallel for training, serial for inference

The dense form and this loop compute the same thing:

S = torch.zeros(d_k, d_v)          # fixed-size state — this is what makes RNNs nice
outputs = []
for t in range(n):                 # sweep the context left to right
    S = S + torch.outer(k[t], v[t])   # accumulate KV into the state
    outputs.append(q[t] @ S)          # read out with the query

At inference you carry $S$ forward one token at a time — fixed memory, no growing KV cache. At training you use the dense parallel matmul form instead. Linear attention lets you have it both ways; a plain RNN gives you only the serial form.

Unlike the lossy softmax-dropping step, this equivalence between the linear form and the recurrent form is exact. Minimax M1, a large high-performance open model, uses a 7-to-1 hybrid — seven linear-attention layers per one full softmax layer — and stays competitive with models like O3 and DeepSeek-R1 while scaling much more gently in context length. No one has yet proven out fully linear attention at scale; everything that works is a hybrid.

Making the recurrence expressive: gating, and a sparse alternative

Plain linear attention feels too naive, so people elaborate the recurrence — and the rule of thumb is that as long as the new terms are input-dependent (functions of $x_t$ ) and not state-dependent, the parallel/serial duality survives.

Mamba-2 adds a gate. Linear attention always passes the whole state forward; Mamba-2 modulates it, $S_t = \gamma_t S_{t-1} + k_t v_t^\top$ with $\gamma_t = f(x_t)$ , borrowing the LSTM lesson that knowing when to forget matters. (It also adds a residual pass-through $v_t^\top D$ to the output, which isn’t core.) Nemotron-3 uses a roughly 3-to-1 Mamba-2 / attention hybrid and gets good throughput at long context.

Gated delta net pushes further with a second gate $\beta_t$ — a “no-input” gate that, at $\beta_t = 0$ , refuses to write the current token into the state — plus a delta update:

$S_t = \gamma_t\left(I - \beta_t k_t k_t^\top\right)S_{t-1} + \beta_t k_t v_t^\top$

The blue $(I - \beta_t k_t k_t^\top)$ acts like a projector that erases previous information in the direction of the current key before writing the new key in — not just adding, but clearing out. The same update has been reinvented from meta-learning least-squares, fast weight programming, and test-time training. Qwen 3.5 and Qwen Next use a 3-to-1 gated delta net / attention hybrid with strong decoding throughput as context grows.

A controlled study (ByteDance Seed and UC Santa Cruz) sweeps the hybrid ratio: at low ratios of RNN-style layers there’s basically no hit, but past a point long-context retrieval and QA degrade, ending badly at a pure RNN. The all-to-all softmax connection is expressive in a way a finite state can’t fully replace.

A completely different attack is DeepSeek Sparse Attention (DSA). Instead of attending to all tokens, a lightweight “lightning indexer” scores each preceding token, $I_{t,s} = \sum_j w_{t,j}\,\text{ReLU}(q_{t,j}\cdot k_s)$ , and full attention runs only on the top- $k$ selected tokens. This isn’t linear time — the indexer still does all-to-all inner products — but the indexer can be tiny (low precision, low dimension), and the expensive full attention runs on a small bounded subset. Better still, you can bolt it on after dense short-context pretraining, during the long-context extension stage. DeepSeek v3.2 and GLM-5 both do this and lose almost nothing versus full attention. The top- $k$ selection idea returns immediately in the second half of the lecture.

Mixture of experts: more parameters, same FLOPs

Conceptually MoEs don’t change the game — a mixture of experts is just a more efficient MLP. Take the feed-forward network and replace it with several feed-forward networks plus a selector that picks which one(s) each input uses. If you have four experts but activate one per token, you have 4× the FFN parameters but pay only one FFN’s worth of compute per forward or backward pass.

Why is this everywhere now? Because for fixed compute, more sparse parameters just keeps helping. Fedus et al. (Switch Transformer, 2022) show test loss dropping monotonically from 1 expert to 256 experts at constant active parameters; OlMoE shows an MoE training roughly 2× faster than its dense counterpart to the same loss; and on the activated-parameter axis that governs inference cost, released MoEs like DeepSeek-V2 sit well above dense models of the same active size. MoEs also open a new parallelism axis — each expert is a natural chunk you can place on its own device (expert parallelism) — at the cost of shipping activations between devices, a comms/compute trade-off explored in the systems lecture.

Designing the router and the experts

Three things vary: the routing function, the expert sizes, and the training procedure.

For routing, almost everyone does token-choice top- $k$ : each token picks its favorite experts. The alternatives — expert-choice (each expert picks tokens), global assignment via a linear-assignment solver, hashing, and RL-learned routes — all exist, but token-choice gets lower validation loss than expert-choice (per OlMoE), hashing is a common baseline that’s never deployed, global assignment is too expensive at scale, and RL is the theoretically “correct” bandit framing (Bengio 2013) that loses on gradient variance and overhead. The standard DeepSeek v1–v2 router is barely more than a softmax: