Chapter 26 A Map of Efficient, Sparse, and Alternative Architectures
26.1 Which resource has become expensive?
Section titled “26.1 Which resource has become expensive?”The distant-key example asks a model to use a color near the beginning of a long prefix. Chapter 25 changed the block’s position, normalization and feed-forward operations. It did not remove the need to carry information across that prefix. Making the prefix longer can increase both the context stored during generation and the work needed to compare positions. Making the feed-forward network wider introduces another expense: more parameters to store and use.
These are three different problems. A smaller collection of stored keys and values need not reduce the number of query positions. Fewer allowed attention pairs need not reduce the number of feed-forward parameters. Using only some experts for one token need not make the other experts disappear from memory. We will locate each proposal by the resource it changes before judging its name.
Recall the local objects from Chapter 17. A query Q asks which earlier positions to use; keys K provide comparison vectors; values V provide the vectors being mixed. An attention head has its own such computation. During sequential generation, earlier K and V can be retained rather than recomputed. Chapter 30 will follow that cache through time. Here we count what an architecture would require it to contain.
26.2 Eight query heads need not mean eight KV heads
Section titled “26.2 Eight query heads need not mean eight KV heads”Example 26.1 — Count the stored numbers first. Assign one layer eight query heads, a head width of 64 coordinates, and a retained prefix of 1,024 positions. Each stored coordinate occupies two bytes. We count only keys and values for one request, with no padding, temporary buffer or metadata.
With eight separate KV heads, one position contains numbers. The leading 2 is for K and V, not for bytes. Multiplying by 1,024 positions and then two bytes gives bytes. Multiplying by two for K and V again would double-count them: they are already included in the 1,024 numbers per position.
Table 26.1 — Same queries, different stored context; derived bytes per layer
| Arrangement | Query heads | KV heads | Stored numbers per position | Bytes for 1,024 positions |
|---|---|---|---|---|
| Separate K/V for every query head | 8 | 8 | 1,024 | 2,097,152 = 2 MiB |
| Four query heads share each K/V pair | 8 | 2 | 256 | 524,288 = 0.5 MiB |
| All query heads share one K/V pair | 8 | 1 | 128 | 262,144 = 0.25 MiB |
Here one MiB is bytes. The first arrangement is multi-head attention (MHA). Sharing all K/V projections gives multi-query attention (MQA); using groups in between gives grouped-query attention (GQA). In the middle row, heads 1–4 share one KV head and heads 5–8 share the other. Their queries remain distinct, so sharing a key does not force identical attention weights.
The storage saving follows from the count. A speedup or quality claim does not. The model still produces eight query-head outputs; its projections and learned behavior have changed. One cannot usually turn an existing MHA checkpoint into a quality-equivalent GQA model by editing a config field. Conversion or training requires evaluation. MQA and GQA are also compatible with a dense feed-forward network: “shared KV” and “sparse parameters” answer different questions.
26.3 Choosing experts saves use, not necessarily residence
Section titled “26.3 Choosing experts saves use, not necessarily residence”Suppose a large feed-forward stage has many alternative transformations, and a small learned router chooses which ones receive a token. Each alternative is an expert: a set of numerical transformations, not a person or a guaranteed subject specialist. A mixture of experts (MoE) combines the selected outputs using routing weights. This is conditional parameter use.
Example 26.2 — Read a sparse model card. Our assigned card counts 20 million shared parameters, including routing and all other shared components, plus four experts of 10 million each. It selects exactly two experts for each token. Total parameters are million. Parameters participating for that token are million. In this card, each parameter set is counted once; this is not a FLOP count.
At two bytes per weight, all weights occupy an ideal 120 MB, where MB means one million bytes. Selected weights represent 80 MB. If all experts stay on the accelerator, 120 MB is still needed for weights, before activations and caches. Loading only selected experts from another device exchanges residence for transfers and scheduling. Across a batch, different tokens can collectively select all four experts.
A router that sends nearly everything to expert 1 can leave other devices idle and overflow expert 1’s assigned capacity. Balancing or capacity rules address that problem, but may add objectives, movement or handling for excess tokens. The original Switch paper studies top-1 routing; this chapter’s top-2 count is an original illustration. Neither count alone predicts quality. A card saying “40 million active” is incomplete without total weights, routing policy, batch conditions and runtime placement.
26.4 Removing comparisons also removes direct paths
Section titled “26.4 Removing comparisons also removes direct paths”Another expense is the list of position pairs. In a causal sequence of six positions, full attention permits query-key pairs per head. The triangular count includes each position itself; future positions are excluded.
Example 26.3 — A distant clue outside the window. Allow each position to attend only to itself and the immediately preceding position. The first position has one allowed pair; each of the other five has two. There are pairs. At position 6, the direct connections reach only positions 5 and 6. A red key at position 1 is outside them.
Information can travel through intermediate positions across layers: with this window, one layer moves information forward at most one position. Reaching position 6 from position 1 requires at least five layers of such links. That is an available path, not a promise that training stores and preserves the color along it. A designated global position could shorten paths, but someone must specify which positions are global and preserve causality when generating.
This is sparse attention: omit some position pairs according to a pattern. Longformer’s local/global patterns are a research example; its bidirectional pattern is not a ready-made causal mask for our decoder. “Efficient attention” is broader. An implementation may calculate the same full attention with fewer memory transfers, as FlashAttention will do in Chapter 31, without deleting any logical pair. Sparse connectivity and efficient execution should not share one unlabeled score.
26.5 Summarizing a prefix instead of storing every position
Section titled “26.5 Summarizing a prefix instead of storing every position”Can the distant prefix be compressed into a fixed-size state? Consider an intentionally simple alternative to softmax attention. Each arriving item has a nonnegative scalar key weight and scalar value . A query has weight 1. We retain the weighted-value sum S and the key-weight sum Z. Start both at zero.
Example 26.4 — Three items, two running sums. Assign keys and values , paired in that order. After the first item, S is and Z is 1, so the output is 4. After the second, add to S and 2 to Z: the output is . After the third, the output is .
Table 26.2 — Derived recurrent state
| Item processed | Contribution to S | S | Z | Output S/Z |
|---|---|---|---|---|
| First | 4 | 4 | 1 | 4 |
| Second | 2 | 6 | 3 | 2 |
| Third | 3 | 9 | 4 | 2.25 |
The general update for this scalar example is
The subscript i means “after item i.” S stores value times weight; Z stores weight alone, which explains the denominator. The output is defined only when Z is positive. Nothing requires saving the full list to obtain this weighted average. In vector forms of linear attention, a chosen feature map of Q and K permits related sums; their size depends on feature and value dimensions. For fixed dimensions, the recurrent state’s size does not grow with prefix length.
There is a price. Swapping the first and third key-value pairs leaves final S and Z unchanged. This scalar state cannot tell which of those two items came first. Nor does it compute ordinary softmax attention: we deliberately changed the similarity rule. Other feature maps and positional mechanisms can retain other distinctions, but “linear” alone cannot certify exact recall of every old token.
A state-space model or another recurrent route also updates a state as tokens arrive. Our weighted sum already retains an old value and adds a weighted new value. A schematic scalar update is : h is the one-number retained state, x the incoming scalar, a the multiplier of the old state and b the multiplier of the input. Setting a to 1, b to the current key weight and x to the current value recovers the S update above. Changing a can instead attenuate older contributions. This scalar illustration is not the complete Mamba architecture. Selective state-space models such as Mamba make some update quantities depend on the input. That can change which information survives. It still needs evidence on distant-key tasks and changed-policy tasks; fixed state is neither an external document store nor guaranteed durable memory.
26.6 Read the map across independent axes
Section titled “26.6 Read the map across independent axes”- A list of numbers8 query heads; 8 KV heads
- An operationGroup queries 1–4 and 5–8
- A list of numbers2 KV heads: 2 MiB → 0.5 MiB
Distinct queries still compute distinct mixtures.
- Stored model numbers20M shared + 4 × 10M experts
- An operationOne token selects two experts
- Stored model numbers40M active; 60M total
All loaded experts still occupy 120 MB at two bytes per weight.
- An operationFull causal: 1+2+3+4+5+6 = 21
- An operationCurrent + previous: 1+5×2 = 11
- A list of numbersAt position 6: only positions 5 and 6
The clue at position 1 needs a multi-layer path or another connection.
Figure 26.1 follows three reductions and their corresponding lost or retained obligations. Each arrow names an actual count or information constraint. No arrow means that one route universally replaces another.
Table 26.3 — Questions to write beside a model card
| Card statement | What it locates | What remains to check |
|---|---|---|
| Dense Transformer with GQA | Shared KV, no routed FFN experts | KV dimensions, quality, supported runtime |
| Attention model with routed FFN experts | Sparse parameter use | Total/active count definitions, routing and transfers |
| Local attention | Sparse token interactions | Window, global positions, causal paths and task quality |
| Kernel linear attention | Different interaction rule with aggregate state | Feature/state size, numerical stability and recall |
| Selective recurrent/state-space blocks | State updates rather than a full KV history in those blocks | State capacity, training behavior and any hybrid attention layers |
A hybrid can occupy several rows. “Dense/sparse” must say whether it refers to weights used per token or position-pair connectivity. “Attention/alternative” must say whether all blocks or only some blocks use that route. Active parameters, stored bytes and arithmetic operations belong in separate columns.
The companion scalar program reproduces the counts in data/part-v/mechanisms.json; it does not implement or benchmark these industrial architectures. The model-card annotation is complete when a reader can mark the sharing, routing and context representation, identify missing evidence, and decline to convert an active-parameter claim into a memory guarantee. Even after selecting a route, training needs room for much more than weights. That is the next resource ledger.
26.7 Summary
Section titled “26.7 Summary”Architectural changes target different expenses. MQA/GQA share K/V; MoE chooses parameter subsets; sparse attention removes pairs; linear and recurrent routes summarize history differently. Every saving has a counting convention and an information or systems obligation. A useful comparison keeps them separate and then evaluates the intended task.
26.8 Exercises
Section titled “26.8 Exercises”- Double Example 26.1’s retained positions and use four KV heads. Calculate bytes per layer. Does retaining eight query heads multiply that answer by eight again?
- The MoE card now selects one expert. Calculate active parameters and ideal resident weight bytes if all experts remain loaded. Explain why these are different questions.
- In Example 26.3, allow the current position and two previous positions. Count all permitted pairs for six positions and the minimum number of layers connecting position 1 to position 6.
- Append key 2 and value 6 to Example 26.4. Compute the new state and output. Diagnose an implementation that divides S by the number of items.
- A card advertises “sparse, 40M active, long context.” Write three questions needed before concluding that the red-key task fits in memory and succeeds.
26.9 Reference answers
Section titled “26.9 Reference answers”- The count is bytes, or 2 MiB. K and V are stored per KV head; the eight distinct queries do not create eight copies of that cache.
- Active parameters become million. Total parameters stay 60 million, so resident weights still need an ideal 120 MB at two bytes each. Selecting fewer weights per token does not unload the others.
- The row counts are 1, 2, 3, 3, 3, 3, totaling 15. Each layer can advance information at most two positions. A distance of five requires at least three layers, for example 1→3→5→6. This establishes connectivity only.
- S becomes ; Z becomes ; the output is 3.5. Dividing by four items would give 5.25 and ignore unequal weights. Z counts weight, not records.
- Ask what is sparse and how many total weights must reside; ask for KV/state dimensions, precision and concurrent context lengths; ask for evidence on this distant-clue distribution at the requested length. A permitted context length does not answer the quality question.
26.10 References
Section titled “26.10 References”- Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019. Sections 2.4–3 explain shared keys and values in MQA.
- Ainslie et al., GQA, 2023. Section 2.2 places grouped queries between MHA and MQA and describes model conversion/uptraining.
- Fedus et al., Switch Transformers, 2022. Section 2 explains routed experts and capacity concerns; its top-1 design differs from our assigned example.
- Beltagy et al., Longformer, 2020. Section 3 presents local/global attention patterns; our causal window is an original comparison.
- Katharopoulos et al., Transformers are RNNs, 2020. Sections 3.2–3.4 derive aggregate-state linear attention.
- Gu and Dao, Mamba, v2. Sections 3.1–3.3 connect selective state updates, compression and implementation costs.