Long-context architectures: Exercises — AI-ML deep dive

This page tests everything in Long-context architectures: how attention scales, why sparse/linear/state-space/memory tricks each cut the cost, and how to compute the numbers by hand. Work each problem before opening its solution.

Every symbol used here is built in the parent note. Quick reminders as we go:

$n$ = number of tokens in the sequence.
$O (\cdot)$ = "grows no faster than", a way to say how the cost scales as $n$ gets huge — see Computational Complexity.
A comparison = one $q_{i} \cdot k_{j}$ dot product between two $d$ -dimensional vectors (one token looking at one other token). Each comparison itself hides an $O (d)$ multiply-add cost — we count comparisons and keep the $d$ factor in mind separately.
Per-head, per-layer. Every cost below ( $O (n^{2})$ , $O (n w + g n)$ , $O (n d_{ϕ} d_{v})$ , …) is for one attention head. Real models run $h$ heads in parallel, so the true layer cost is $h \times$ the per-head count. We hold $h$ fixed and count one head unless a problem says otherwise; multiply by $h$ for the full picture.

Figure 1 below is reused throughout — you'll count its lit cells directly in Exercises 2.1 and 3.1:

Figure 1. Left panel = full attention: every token (row) looks at every token (column), so all $n \times n$ cells are lit. Right panel = a sliding window of width $w$ : each row only lights the $w$ cells near the diagonal. Count the lit cells and you have counted the cost.

Level 1 — Recognition

Recall Solution 1.1

Full attention compares every token with every token, so the count is $n \times n = n^{2}$ . $400 0^{2} = 16, 000, 000 = 1.6 \times 1 0^{7} comparisons (per head).$ The scaling class is $O (n^{2})$ — quadratic. Double the tokens and the cost quadruples. With $h$ heads the layer does $h$ times this many.

Recall Solution 1.2

(a) full attention → $O (n^{2})$ (every pair).
(b) Longformer → $O (n w + g n)$ ( $w$ = window, $g$ = number of global tokens).
(c) linear attention → $O (n d_{ϕ} d_{v})$ ( $d_{ϕ}$ = feature-map dim, $d_{v}$ = value dim).
(d) S4 → $O (n d^{2})$ ( $d$ = fixed state dimension).

Notice: (b),(c),(d) are all linear in $n$ because their non- $n$ factors ( $w$ , $d_{ϕ} d_{v}$ , $d^{2}$ ) are constants that do not grow with the sequence. All four are per head; multiply by the head count $h$ for the whole layer.

Level 2 — Application

Recall Solution 2.1

Think of the right panel of Figure 1: regular tokens light only $w$ near-diagonal cells, global tokens light a whole row.

Regular tokens: there are $n - g = 12 - 2 = 10$ of them, each costing $w = 3$ : $10 \times 3 = 30.$ Global tokens: $g = 2$ of them, each costing $n = 12$ : $2 \times 12 = 24.$ Total $= 30 + 24 = 54$ comparisons per head.

Full attention (the left panel of Figure 1, all cells lit) would cost $n^{2} = 1 2^{2} = 144$ . So the sparse scheme uses $54/144 \approx 37.5%$ of the work. (Both counts scale by the same head factor $h$ , so the ratio is unchanged.)

Recall Solution 2.2

Quick recall of what $S$ and $Z$ are (built once, in the parent): $S = \sum_{j} ϕ (k_{j}) v_{j}$ is the value-weighted sum of key feature-maps, and $Z = \sum_{j} ϕ (k_{j})$ is the plain sum of key feature-maps (the normalizer). Both are computed one time over all keys, then reused by every query.

First map the query: $ϕ (q_{2}) = [1, 0, 1]$ (that's $[1, x_{1}, x_{2}]$ with $x = [0, 1]$ ).

Numerator = $ϕ (q_{2}) \cdot S = 1 \cdot 10 + 0 \cdot 4 + 1 \cdot 5 = 15$ . Denominator = $ϕ (q_{2}) \cdot Z = 1 \cdot 4 + 0 \cdot 2 + 1 \cdot 2 = 6$ . $Attention_{2} = \frac{15}{6} = 2.5.$ Because $S$ and $Z$ were built once for all queries, this query cost only two length-3 dot products — no scan over the keys at all. That is the whole point of kernel-style linear attention.

Level 3 — Analysis

Recall Solution 3.1

Picture the right panel of Figure 1 widening: as $w$ grows the lit band thickens until it fills the whole left panel. Set the costs equal: $n w = n^{2} \Rightarrow w = n = 4096.$ So the window only reaches full-attention cost when it is as wide as the whole sequence — i.e. when the band fills the grid and it stops being a window. For any $w < n$ the sliding window is strictly cheaper, and the savings ratio is $w / n$ . At $w = 512$ you pay $512/4096 = 1/8$ of full attention.

Recall Solution 3.2

Count in "comparisons", where each comparison is one $d$ -dimensional dot product (the same unit as full attention), so the hidden $O (d)$ factor is common to every term and cancels in the ratio.

Per token the cost is: local $w = 512$ comparisons, plus memory attention over the $k = 64$ retrieved values (64 comparisons), plus ANN search $lo g_{2} M$ (each a $d$ -dim similarity): $lo g_{2} (1 0^{6}) \approx 19.93 \approx 20.$ Per-token cost $\approx 512 + 64 + 20 = 596$ comparisons. Over $n = 50, 000$ tokens: $50, 000 \times 596 = 2.98 \times 1 0^{7} comparisons .$ Full attention: $50, 00 0^{2} = 2.5 \times 1 0^{9}$ . Ratio $\approx 2.98 \times 1 0^{7} /2.5 \times 1 0^{9} \approx 1.2%$ .

Dominant term: the local window ( $512$ ) dwarfs both retrieval ( $64$ ) and search ( $20$ ). The retrieval machinery is almost free; the sliding window is where the money goes — which tells you that is the knob to shrink first.

Level 4 — Synthesis

Recall Solution 4.1

Where does $d$ enter? The attention score matrix stores one scalar similarity per token pair — the $d$ -dimensional $q, k$ vectors are dotted down to a single number, so the score-matrix size is $n \times n$ and is independent of $d$ . Likewise the S4 running state is a length- $s$ vector of numbers, also independent of $d$ . So for "numbers stored as running attention/state", $d$ does not appear — it would only matter if we counted the $Q, K, V$ activation tensors ( $n \times d$ each), which both methods share and which are linear, not quadratic. That's why the headline comparison drops $d$ .

(a) Full attention stores an $n \times n$ score matrix (per head): $100, 00 0^{2} = 1.0 \times 1 0^{10} scores.$ (b) An S4 layer carries a fixed-size hidden state $h_{k}$ of dimension $s = 256$ — it does not grow with $n$ . So it holds on the order of $256$ numbers as the running state. $reduction = \frac{1 0 ^{10}}{256} \approx 3.9 \times 1 0^{7} .$ That's a ~39-million-fold shrink in the attention-matrix memory.

Bookkeeping caveat: S4 also stores fixed model parameters — the state matrices $\overset{ˉ}{A}, \overset{ˉ}{B}, C$ and the convolution kernel (order $s$ to a few $\times 1 0^{3}$ numbers). These are constant in $n$ : negligible next to the $1 0^{10}$ scores and unchanged whether the sequence is 1k or 100k tokens, but they are still part of the true footprint and shouldn't be forgotten. The dominant, $n$ -dependent term is what the reduction factor captures.

The catch the parent flags: that fixed 256-vector must compress all history — anything it can't hold is gone, unlike attention which keeps every token addressable. Compare with the RNN hidden state (same fixed-size idea) and the Information Bottleneck view (a fixed state is a bottleneck).

Recall Solution 4.2

Linear beats full when $c_{lin} < c_{full}$ : $n d_{ϕ} d_{v} < n^{2} \Rightarrow d_{ϕ} d_{v} < n \Rightarrow n > d_{ϕ} d_{v} .$ So the break-even point is $n^{*} = d_{ϕ} d_{v}$ . $n^{*} = 256 \times 64 = 16, 384 tokens.$ Reading it: below ~16k tokens, full attention is actually cheaper here — the feature-map constants outweigh the quadratic term. Linear attention only wins once the sequence is long enough that $n$ overtakes the constant $d_{ϕ} d_{v}$ . This is why "linear = always faster" is false at short context. (Both sides scale by the same head count $h$ , so $h$ cancels and does not move the break-even.)

Level 5 — Mastery

Recall Solution 5.1

Path (A) — score then weight, term by term:

$ϕ (q) \cdot ϕ (k_{1}) = 1 \cdot 1 + 2 \cdot 0 = 1$ , times $v_{1} = 3 \Rightarrow 3$ .
$ϕ (q) \cdot ϕ (k_{2}) = 1 \cdot 0 + 2 \cdot 1 = 2$ , times $v_{2} = 5 \Rightarrow 10$ .
Sum $= 3 + 10 = 13$ .

Path (B) — precompute $S = \sum_{j} ϕ (k_{j}) v_{j}$ first: $S = [1, 0] \cdot 3 + [0, 1] \cdot 5 = [3, 5] .$ Then $ϕ (q) \cdot S = 1 \cdot 3 + 2 \cdot 5 = 3 + 10 = 13.$

Both give $13$ . They must — it's just the distributive/associative law of the dot product. Path (B) is the whole efficiency win: build $S$ once, then every query is a single fixed-size dot product.

Recall Solution 5.2

A defensible choice is memory-augmented + local window (Memorizing-Transformer style). Quantitative backing (all per-head; multiply by head count $h$ uniformly):

Full attention is out: $200, 00 0^{2} = 4 \times 1 0^{10}$ scores/layer — infeasible.
Sparse (Longformer): with $w = 512$ , $g = 64$ globals, cost $\approx n w + g n = 200, 000 (512 + 64) = 1.152 \times 1 0^{8}$ — feasible, but purely-positional windows can miss a clause cited 150k tokens earlier.
S4/SSM: cheapest memory (fixed 256-state, ~ $3.9 \times 1 0^{7} \times$ smaller matrix from Ex 4.1), but that fixed state compresses away rare specific facts — bad for "quote clause 14.3 exactly".
Memory-augmented: local $w = 512$ for fluency + ANN retrieval ( $lo g_{2} M \approx 20$ ) that addresses individual earlier facts by content, cost from the Ex 3.2 logic $\approx n (w + k + lo g_{2} M) \approx 200, 000 \times 596 = 1.192 \times 1 0^{8}$ — same ballpark as sparse but with content-based recall of exact clauses.

Decision: memory-augmented + local window wins for this task, because a legal assistant must quote exact clauses from anywhere in the document. It costs $\approx 1.192 \times 1 0^{8}$ comparisons/head — essentially the same as sparse ( $1.152 \times 1 0^{8}$ ) and vastly below full attention ( $4 \times 1 0^{10}$ , a ~336× saving) — yet, unlike sparse's fixed positional window or S4's lossy fixed state, its content-addressed retrieval can pull a clause cited 150k tokens earlier.

Biggest risk: retrieval quality. If the ANN index returns the wrong neighbors (stale or poorly-embedded memory), the model confidently attends to irrelevant clauses — a failure sparse attention cannot have because its pattern is fixed and positional. Mitigation: periodically re-embed the memory bank and evict stale entries so retrieved keys stay faithful to the current model.

(There is no single "correct" family — the mastery skill is matching the cost structure and failure mode to the task: exact-recall legal work rewards addressable memory; smooth long-range reasoning rewards SSMs.)

Recall Self-test recap (cloze)

Full attention scales as ::: $O (n^{2})$ per head Longformer total cost with window $w$ and $g$ globals ::: $O (n w + g n)$ Linear-attention break-even sequence length ::: $n^{*} = d_{ϕ} d_{v}$ The linear-attention denominator plays the role of ::: softmax normalization ANN memory search costs ::: $O (lo g M)$ , not $O (M)$ An S4 hidden state has ::: fixed dimension, independent of $n$ Every per-head cost must be multiplied by ::: the number of heads $h$