Transformers from First PrinciplesPost 10 of 18 · 18 of 18 written
  1. How to use this guide
  2. 01The machinery: how a neural network computes and learns
  3. 02Making deep networks trainable
  4. 03Turning text into numbers
  5. 04Attention from first principles
  6. 05The Transformer block
  7. 06Position: telling the model where each token is
  8. 07Architecture families and training objectives
  9. 08Pretraining at scale
  10. 09Inference and serving
  11. 10Efficiency variants: MoE, long context and beyond attention
  12. 11Post-training: from text predictor to assistant
  13. 12Adapting models: prompting, fine-tuning, LoRA and distillation
  14. 13Embeddings, retrieval and RAG
  15. 14Beyond text: vision, speech and multimodal models
  16. 15Inside the model: interpretability
  17. 16Evaluation, failure modes and security
  18. 17From model to system
  19. 18Fluency kit

10. Efficiency variants: MoE, long context and beyond attention

Transformers · 10 of 184 min readNov 2026

Two costs shape architecture research: dense models spend every parameter on every token, and attention’s cost grows with the square of context length. Mixture-of-experts attacks the first by activating only part of the network per token; sparse attention, linear attention and state space models attack the second. Each wins something and gives something up, and recent models often combine them.

Mixture of experts (MoE)

Replace each dense MLP with E smaller expert MLPs plus a router. For each token, the router scores the experts and sends the token to the top k (commonly 1–8 out of 8–256):

y=∑i∈TopK(xWr)gi(x) Ei(x)y = \sum_{i \in \mathrm{TopK}(x W_r)} g_i(x)\, E_i(x)

Here Eᵢ is expert i, W_r the router weights, and gᵢ the normalized router weights over the selected experts.

  • Compute follows active parameters; memory follows total parameters. Mixtral 8x7B has about 47B parameters in total but uses about 13B per token: it runs at roughly 13B-model cost but needs memory for 47B (approximate).
  • Load balancing: routers tend to collapse onto a few favourite experts. An auxiliary loss (or a bias adjustment) pushes toward even use; each expert has a capacity limit, and overflow tokens may be dropped or rerouted.
  • Fine-grained and shared experts: many small experts plus a few always-on shared experts (DeepSeekMoE style) improve specialization.
  • Serving: experts are spread across GPUs (expert parallelism) with all-to-all communication. At small batch sizes the memory needs make MoE less attractive.

Experts don’t specialize in human topics like “biology”; routing follows mostly token-level and syntactic patterns. Don’t over-interpret them.

The quadratic attention problem

Full attention over 1M tokens means 10¹² scores per head per layer. Three families of fixes:

Sparse and local attention. Each token attends to a subset: a sliding window of nearby tokens, a few global tokens, and strided or random links (Longformer, BigBird, and sliding-window layers in several recent open models). Information still travels far by hopping across layers. Cheap and effective, but some long-range lookups are lost.

Linear attention. Replace softmax(QKᵀ)V with φ(Q)(φ(K)ᵀV) for some feature map φ. Multiplying right to left makes the cost linear in T, and the running sum φ(K)ᵀV is a fixed-size state, so at inference the model behaves like an RNN. The weakness: a fixed-size state can’t store arbitrary details, so exact recall and copying from long contexts suffer.

Exact attention, distributed. Keep full attention but split the sequence across GPUs (ring attention, context parallelism) and use FlashAttention-style kernels. Many long-context models actually work this way.

State space models (SSMs) and Mamba

SSMs process sequences through a linear recurrence on a hidden state, an idea borrowed from control theory:

ht=Aˉ ht−1+Bˉ xt,yt=C hth_t = \bar{A}\, h_{t-1} + \bar{B}\, x_t, \qquad y_t = C\, h_t

Because the recurrence is linear, training can be parallelized (as a convolution or a parallel scan), while inference is a cheap recurrence with constant memory per step - no growing KV cache. Mamba (Gu & Dao, 2023) made the parameters depend on the input (“selective”), letting the model decide per token what to keep or forget, which closed much of the quality gap with transformers.

The trade-off is the same as for linear attention: a fixed-size state compresses the past. Pure SSMs are weaker at precise retrieval and copying from context, which is exactly what attention does best. Hence hybrids: mostly SSM or linear-attention layers with a few full-attention layers interleaved (Jamba and several later models), aiming for near-transformer quality with far less memory at long context.

Comparison

Mechanism Training cost vs T Inference memory per sequence Strength Weakness
Full attention O(T²), with fast kernels KV cache grows with T exact recall from anywhere in context cost and memory at long T
Sliding window O(T · w) bounded by the window w cheap, strong locally distant details only via hops
Linear attention O(T) constant fast on long sequences lossy memory, weaker recall
SSM (Mamba) O(T) constant fast, streaming lossy memory, weaker copying
Hybrid mostly O(T) small KV cache for the few attention layers balances both more complex to build and serve

Other efficiency ideas worth recognizing

  • Multi-token prediction: extra heads predict several future tokens during training; this can improve learning and enables faster decoding.
  • Early exit and adaptive compute: stop at an earlier layer when the model is already confident.
  • Pruning and distillation: remove parameters, or train a smaller student (Part 12).

Questions you should now be able to answer

  • An MoE has 400B total and 20B active parameters. What does that imply for FLOPs per token, GPU memory and batch-size needs?
  • Why does an MoE router need a load-balancing loss?
  • Why does linear attention turn into an RNN at inference time?
  • What kind of task would expose an SSM’s weakness relative to a transformer?
  • Why might a hybrid keep only a few full-attention layers?