Next-gen architectures: MoE, state-space models, linear & ring attention

Standard attention costs grow with the square of the sequence length, and dense models pay for every parameter on every token. This stage covers the four main ways around that: route each token to a few experts, replace attention with a recurrent state, make attention linear, or spread one long sequence across many GPUs.

Source: Suraj Sharma, X post, stage 3 · MoE routing at frontier scale and linear attention (KDA) are covered in depth in Module 05. This page links to it and focuses on what's new.

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings, try the demo, and read the reality check.

1. What this stage covers

• Learn: MoE routing, SSMs (S4, Mamba), linear attention and ring attention.
• Practice: write a sparse MoE layer with a load-balancing loss from scratch and benchmark its throughput against a dense Transformer.
• Why (the post's claim): quadratic attention is being replaced, so you need to understand the alternatives from first principles. Section 4 tests that claim.

2. Concepts, explained

Mixture-of-experts routing and the auxiliary load-balancing loss

An MoE layer replaces the Transformer's single feed-forward block with E expert feed-forward blocks and a router. The router is a linear layer that gives each token a score per expert. The token goes to its top-k experts, and their outputs are summed, weighted by the router's gate values. Shazeer et al. (2017, "Outrageously Large Neural Networks") made this work at scale with a noisy top-k gate between LSTM layers. Switch Transformer (Fedus, Zoph & Shazeer, 2021) simplified it to top-1: one expert per token, less compute and less communication, and reported up to 7× faster pre-training than a dense T5 baseline at the same compute.

The failure mode is collapse. Early in training, a few experts get slightly better, so the router sends them more tokens, so they get better still. The rest starve. Switch's fix is an auxiliary loss added to the training loss:

f_i = fraction of tokens in the batch whose top-1 choice is expert i     (hard counts)
P_i = mean router probability for expert i over the batch               (soft, differentiable)

L_aux = α · E · Σ_i f_i · P_i          α ≈ 0.01 in the Switch paper

uniform routing:     f_i = P_i = 1/E   →  E · Σ f_i P_i = 1   (the minimum)
total collapse:      f_1 = P_1 ≈ 1     →  E · Σ f_i P_i ≈ E

Why both f and P? The counts f aren't differentiable, but the probabilities P are. Multiplying them means the gradient pushes down the probability of experts that are already overloaded. The factor E keeps the loss scale the same whatever the number of experts.

Capacity factor. For efficient batched kernels each expert gets a fixed number of slots: capacity = (tokens / E) × capacity_factor. Tokens routed to a full expert are dropped: they skip the layer and pass through on the residual connection only. A capacity factor above 1 leaves headroom for imbalance but wastes slots and memory when routing is balanced. The demo below shows both effects.

What Module 05 adds. The auxiliary loss is a trade-off: too weak and experts collapse, too strong and it fights the language-modelling loss. Newer models often drop it in favour of a per-expert bias that only affects selection. Module 05's routing deep dive explains that approach and Kimi K3's Quantile Balancing for 896 experts, with its own demo. Read it after this section.

State-space models: S4 and Mamba

A state-space model comes from control theory, which is why it feels familiar to roboticists. A continuous input u(t) drives a hidden state x(t) of size N, and the output reads from the state:

continuous:   x'(t) = A x(t) + B u(t)        y(t) = C x(t)
discretised:  x_k = Ā x_{k-1} + B̄ u_k        y_k = C x_k      (Ā, B̄ from A, B and a step size Δ)

as a convolution (when A, B, C don't depend on the input):
  y = K * u,   K = (C B̄, C Ā B̄, C Ā² B̄, …)

The same layer has two forms. For training, compute the kernel K once and apply it as a long convolution over the whole sequence, which parallelises well (via FFT). For generation, run the recurrence: one fixed-size state update per token, no growing cache.

S4 (Gu, Goel & Ré, 2021) made this work for long sequences. Two ingredients: initialising A from HiPPO theory, so the state stores a compressed history of the input, and a structured parameterisation of A that makes computing K fast. S4 did very well on the Long Range Arena benchmark, including tasks with sequences of thousands of steps that Transformers struggled with.

Mamba (Gu & Dao, 2023) fixed S4's weakness on language. In S4, A, B, C and Δ are the same for every token, so the model can't decide to remember this token and ignore that one. Mamba makes B, C and Δ functions of the current input. This is the "selective" SSM. A large Δ means "reset and take in this token"; a small Δ means "ignore it and keep the state". The cost: with input-dependent parameters there's no single kernel K, so the convolution trick is gone. Mamba instead uses a hardware-aware parallel scan that keeps the expanded state in fast on-chip memory. The paper reports about 5× higher inference throughput than Transformers and a 3B model that matches Transformers twice its size. A follow-up, Mamba-2 (Dao & Gu, 2024), showed a close mathematical link between selective SSMs and a form of linear attention.

Linear attention

Softmax attention computes softmax(QKᵀ)V, which needs the full L×L score matrix. Katharopoulos et al. (2020, "Transformers are RNNs") replaced the softmax with a feature map φ and used associativity:

softmax attention:   o_t = Σ_{s≤t} exp(q_t·k_s) v_s / Σ_{s≤t} exp(q_t·k_s)        O(L²)
linear attention:    o_t = φ(q_t)ᵀ S_t / φ(q_t)ᵀ z_t
                     S_t = S_{t-1} + φ(k_t) v_tᵀ     (a d×d matrix state)
                     z_t = z_{t-1} + φ(k_t)                                       O(L)

The causal version is literally an RNN with a matrix-valued state. That is the same shape as an SSM, which is why the two families have merged. Plain linear attention only adds to its state and never forgets, so it degrades on long text. Modern variants add decay (gating) and the delta rule, which overwrites old associations instead of piling new ones on top. Kimi Delta Attention is one such variant. Module 05's length deep dive explains it and why K3 still keeps one full-attention layer in four.

Ring attention

Ring attention (Liu, Zaharia & Abbeel, 2023) attacks a different problem. It doesn't change what attention computes. It keeps exact softmax attention and spreads the memory across devices.

1. Split a sequence of length L into D blocks, one per device. Each device keeps its own query block.
2. Each device computes attention between its queries and the key/value block it currently holds, using the blockwise "online softmax" trick (keep a running max and running sum, as in FlashAttention), so partial results can be combined exactly.
3. Each device passes its K/V block to the next device in the ring and receives one from the previous device.
4. After D steps every query block has seen every K/V block. The send for the next step overlaps with the compute of the current one, so if each block is big enough, communication is hidden.

per-device activation memory  ∝ L / D        (one block of Q, K, V at a time)
total attention compute       ∝ L²           (unchanged: it's still exact attention)
max context                   ≈ D × (what one device can hold)

The paper's title says "near-infinite". The honest reading: context grows linearly with the number of devices, and the paper reports sequences of millions of tokens. The total compute is still quadratic. Ring attention is a training and long-prefill tool for clusters. It does nothing for a single GPU.

3. The practice project

(a) Laptop / Colab version: a Switch MoE in a tiny GPT

1. Start from a small character-level GPT (nanoGPT-style): 4 layers, d_model = 256, context 256, on Tiny Shakespeare. Train it dense as the baseline.
2. Swap the feed-forward block in every other layer for the MoE layer below, with E = 8 experts, top-1, each expert the same size as the dense feed-forward block. Active compute per token stays roughly the same; total parameters in those layers grow about 8×.
3. Train three variants for the same number of steps: dense; MoE with α = 0; MoE with α = 0.01.
4. Measure: validation loss against steps; tokens per second (forward + backward, after warm-up, with torch.cuda.synchronize() around the timer); per-expert load histogram every 200 steps; fraction of dropped tokens.
5. Repeat at capacity factor 1.0 and 1.5.

import torch, torch.nn as nn, torch.nn.functional as F

class SwitchMoE(nn.Module):
    def __init__(s, d, h, E, cf=1.25):
        super().__init__()
        s.E, s.cf = E, cf
        s.router = nn.Linear(d, E, bias=False)
        s.w1 = nn.Parameter(torch.randn(E, d, h) * d**-0.5)
        s.w2 = nn.Parameter(torch.randn(E, h, d) * h**-0.5)

    def forward(s, x):                       # x: (B, T, d)
        B, T, d = x.shape
        x = x.reshape(-1, d)
        p = F.softmax(s.router(x), -1)        # (N, E)
        gate, idx = p.max(-1)                # top-1
        f = F.one_hot(idx, s.E).float().mean(0)
        aux = s.E * (f * p.mean(0)).sum()    # Switch loss, multiply by α outside
        cap = int(s.cf * x.shape[0] / s.E)
        out = torch.zeros_like(x)
        for e in range(s.E):                 # simple loop, fine for E=8
            sel = (idx == e).nonzero().squeeze(-1)[:cap]   # drop overflow
            h = F.gelu(x[sel] @ s.w1[e])
            out[sel] = gate[sel, None] * (h @ s.w2[e])
        return out.reshape(B, T, d), aux

What to expect, honestly. With α = 0 you should see the load histogram skew, sometimes badly, and dropped tokens rise. With α = 0.01 the loads should stay much flatter. On loss per step, the MoE will usually match or beat the dense baseline, since it has more parameters at similar compute. On tokens per second, the MoE will very likely be slower than dense at this size. The Python loop, the gather/scatter and the small per-expert matmuls dominate. That's a real finding, not a bug: MoE throughput wins need grouped-matmul kernels and large batches.

(b) The full version from the post

Benchmarking MoE throughput "against a dense transformer" properly means a GPT-2-scale or larger model, many experts (16–64), and expert parallelism: experts live on different GPUs and tokens move with an all-to-all collective. Use an existing kernel library (for example MegaBlocks-style grouped/block-sparse matmuls) rather than the loop above. Compare against two baselines: a dense model with the same active parameters (same FLOPs per token) and one with the same total parameters. Report tokens/s, loss at fixed tokens, peak memory and all-to-all time from a profiler. Hardware: one multi-GPU node is enough to see the communication cost. Cost scales with tokens trained; estimate it with the 6 × N_active × D FLOPs rule from Day 25 before renting anything.

(c) Pitfalls

• Comparing at the wrong baseline. "MoE beats dense" means nothing unless you say whether dense matches active or total parameters.
• Router in low precision. Switch kept the router computation in float32 for stability. Do the same if you train in bf16.
• Measuring the first steps. Warm up, synchronise the GPU, and time many steps.
• Ignoring dropped tokens. A good loss with 15% of tokens skipping the layer isn't the model you think it is.
• Loss weight too high. A large α flattens the loads but forces the router to ignore content. Watch the main loss, not just the histogram.
• For SSMs: the naive sequential scan in Python is extremely slow. Use the reference Mamba kernels, or stick to S4's convolution form for learning.

4. Reality check

5. In this course

• Module 05 · width: routing 896 experts, why auxiliary losses are being replaced, and Quantile Balancing (with a demo).
• Module 05 · length: hybrid linear + global attention, Kimi Delta Attention, why no positional encoding is needed.
• Module 05 · serving: what a mixed fixed-state / growing-cache model means for inference memory.
• 06 Day 13 · Attention and Day 14 · The Annotated Transformer: the baseline everything here modifies.
• 06 Day 1 and Day 2 · LSTM: recurrent state, the ancestor of SSMs and linear attention.
• 06 Day 24 · GPipe: another way to split one model across devices, useful context for ring attention.
• 06 Day 25 · Scaling laws: the FLOPs accounting you need to compare MoE with dense fairly.

6. Why it matters for robotics

7. Self-check

8. Courses & content

Start with the Hugging Face MoE guide: it maps directly onto the practice project. Then do the Annotated S4 for the SSM half.

Full course map for all 12 stages →

9. Papers & docs

• Shazeer et al. (2017), Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Fedus, Zoph & Shazeer (2021), Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Gu, Goel & Ré (2021), Efficiently Modeling Long Sequences with Structured State Spaces (S4)
• Gu & Dao (2023), Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• Katharopoulos, Vyas, Pappas & Fleuret (2020), Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Liu, Zaharia & Abbeel (2023), Ring Attention with Blockwise Transformers for Near-Infinite Context