Skip to content

Chapter 26 A Map of Efficient, Sparse, and Alternative Architectures

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 2×8×64=1,0242\times8\times64=1{,}024 numbers. The leading 2 is for K and V, not for bytes. Multiplying by 1,024 positions and then two bytes gives 1,024×1,024×2=2,097,1521{,}024\times1{,}024\times2=2{,}097{,}152 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

ArrangementQuery headsKV headsStored numbers per positionBytes for 1,024 positions
Separate K/V for every query head881,0242,097,152 = 2 MiB
Four query heads share each K/V pair82256524,288 = 0.5 MiB
All query heads share one K/V pair81128262,144 = 0.25 MiB

Here one MiB is 220=1,048,5762^{20}=1{,}048{,}576 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 20+4×10=6020+4\times10=60 million. Parameters participating for that token are 20+2×10=4020+2\times10=40 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 1+2+3+4+5+6=211+2+3+4+5+6=21 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 1+5×2=111+5\times2=11 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 kik_i and scalar value viv_i. 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 [1,2,1][1,2,1] and values [4,1,3][4,1,3], paired in that order. After the first item, S is 1×4=41\times4=4 and Z is 1, so the output is 4. After the second, add 2×1=22\times1=2 to S and 2 to Z: the output is 6/3=26/3=2. After the third, the output is 9/4=2.259/4=2.25.

Table 26.2 — Derived recurrent state

Item processedContribution to SSZOutput S/Z
First4414
Second2632
Third3942.25

The general update for this scalar example is

Si=Si−1+kivi,Zi=Zi−1+ki,yi=Si/Zi.S_i=S_{i-1}+k_i v_i,\qquad Z_i=Z_{i-1}+k_i,\qquad y_i=S_i/Z_i.

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 hi=aihi−1+bixih_i=a_i h_{i-1}+b_i x_i: 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.

Share stored context
  1. A list of numbers8 query heads; 8 KV heads
  2. An operationGroup queries 1–4 and 5–8
  3. A list of numbers2 KV heads: 2 MiB → 0.5 MiB

Distinct queries still compute distinct mixtures.

Route parameters
  1. Stored model numbers20M shared + 4 × 10M experts
  2. An operationOne token selects two experts
  3. Stored model numbers40M active; 60M total

All loaded experts still occupy 120 MB at two bytes per weight.

Restrict position pairs
  1. An operationFull causal: 1+2+3+4+5+6 = 21
  2. An operationCurrent + previous: 1+5×2 = 11
  3. 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 — Three savings leave different obligations. Original book diagram · CC BY-SA 4.0.

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 statementWhat it locatesWhat remains to check
Dense Transformer with GQAShared KV, no routed FFN expertsKV dimensions, quality, supported runtime
Attention model with routed FFN expertsSparse parameter useTotal/active count definitions, routing and transfers
Local attentionSparse token interactionsWindow, global positions, causal paths and task quality
Kernel linear attentionDifferent interaction rule with aggregate stateFeature/state size, numerical stability and recall
Selective recurrent/state-space blocksState updates rather than a full KV history in those blocksState 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.

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.

  1. 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?
  2. 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.
  3. 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.
  4. 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.
  5. A card advertises “sparse, 40M active, long context.” Write three questions needed before concluding that the red-key task fits in memory and succeeds.
  1. The count is 2×2,048×4×64×2=2,097,1522\times2{,}048\times4\times64\times2=2{,}097{,}152 bytes, or 2 MiB. K and V are stored per KV head; the eight distinct queries do not create eight copies of that cache.
  2. Active parameters become 20+10=3020+10=30 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.
  3. 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.
  4. S becomes 9+2×6=219+2\times6=21; Z becomes 4+2=64+2=6; the output is 3.5. Dividing by four items would give 5.25 and ignore unequal weights. Z counts weight, not records.
  5. 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.
  • 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.