Long-context and efficient attention

"They handed me a window long enough to remember the whole shift, then a compute budget that could barely afford the last five seconds. Efficiency is just me learning which seconds to forget."

A Budget-Conscious AI Agent

Why quadratic attention breaks on long sensor windows

Self-attention builds an \(L \times L\) score matrix, where \(L\) is the number of tokens. For queries, keys, and values of width \(d\), computing \(QK^\top\), softmaxing it, and multiplying by \(V\) costs on the order of

What this means concretely: double the window and you quadruple the work. Why it bites sensing harder than language is a matter of raw arithmetic. A sentence is maybe 30 tokens; a page is 500. Now take an industrial accelerometer sampled at \(20\,\mathrm{kHz}\). A ten-second capture is \(200{,}000\) raw samples. Even after patchification collapses, say, 64 samples into one token, you are left with roughly \(3{,}000\) tokens, and the score matrix has \(3{,}000^2 \approx 9\) million entries per head, per layer. Push the window to a full minute of high-rate data and the matrix crosses a billion entries; the memory alone exceeds a commodity GPU before a single gradient flows. When does this matter? Precisely when the phenomenon you care about is slow or intermittent relative to the sample rate: a bearing fault that rings once per revolution, a cardiac arrhythmia that appears in one beat out of a thousand, a machine that drifts out of tolerance over an hour. Downsample to shorten the sequence and you blur the very event you are hunting.

Sparse and local attention: exploit the structure of signals

Full attention assumes any token might be decisively relevant to any other. Physical signals rarely honor that assumption. The value of a sample is dominated by its recent neighbors (local temporal correlation), plus a sparse set of far-away anchors (a periodic beat, a shift-start marker, a global summary). Sparse attention hard-codes this prior by letting each token attend to only a chosen subset of the others, dropping the cost from \(O(L^2)\) toward \(O(L)\).

The workhorse pattern is sliding-window (local) attention: token \(i\) attends only to the \(w\) tokens on either side, costing \(O(L w)\) with \(w \ll L\). This is attention's answer to the dilated receptive field you met in Chapter 14, and like a stack of dilated convolutions, stacking windowed-attention layers grows the effective receptive field linearly with depth, so distant tokens still communicate indirectly. To restore true global reach cheaply, designs such as Longformer and BigBird add two ingredients: dilated windows that skip tokens to cover more span at fixed cost, and a handful of global tokens that every position attends to and that attend to everyone, giving an \(O(L)\) information highway. A single \([\text{CLS}]\)-style global token is often enough to route a whole-window classification decision. The design question is always the same: what is the physical range of dependency in this signal, and how few links can honor it?

Linear attention: replace the softmax, not just the pattern

Sparse patterns still compute an exact softmax over the links they keep. A different family attacks the math itself. The bottleneck is that softmax forces you to form \(QK^\top\) explicitly before you can normalize. Linear attention replaces the softmax kernel with a feature map \(\phi(\cdot)\) so that the similarity between query \(i\) and key \(j\) factorizes as \(\phi(q_i)^\top \phi(k_j)\). Once it factorizes, the associativity of matrix multiplication lets you reorder the computation:

The left side builds the \(L \times L\) matrix; the right side first builds a \(d \times d\) matrix and never materializes anything of size \(L^2\). The cost falls to \(O(L d^2)\), linear in sequence length. Performer approximates the true softmax kernel with random features so the results stay close to exact attention, while simpler schemes (for example \(\phi(x) = \mathrm{elu}(x) + 1\)) trade a little fidelity for speed and, usefully, admit a recurrent form that streams one token at a time. That streaming form is why linear attention is attractive for the always-on, never-ending inputs of Chapter 60. The catch, which you should treat as a working rule rather than a footnote, is that linear attention often underperforms exact attention when the task hinges on sharp, precise matches between two specific distant tokens, because the low-rank feature map smears fine distinctions. Reach for it when context is long and the signal is smooth; be skeptical when a single needle in the haystack decides the label.

You do not implement any of this by hand. PyTorch exposes the fused, memory-efficient kernels behind one function, and it selects FlashAttention automatically when the inputs qualify. Listing 15.4 shows the whole thing.

import torch
import torch.nn.functional as F

# 8 heads, a long sensor window of L=4096 tokens, head width 64
B, H, L, d = 2, 8, 4096, 64
q = torch.randn(B, H, L, d, device="cuda", dtype=torch.float16)
k = torch.randn(B, H, L, d, device="cuda", dtype=torch.float16)
v = torch.randn(B, H, L, d, device="cuda", dtype=torch.float16)

# One call: fused, IO-aware, exact softmax attention.
# is_causal=True masks the future so a token never attends forward in time.
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape)   # torch.Size([2, 8, 4096, 64]); no 4096x4096 matrix ever stored

Choosing an efficient-attention strategy

The three families answer different questions, and Table 15.4 lines them up. Full attention is your accuracy reference and the right choice whenever the post-patchification sequence fits in memory, which, thanks to FlashAttention, now covers many sensor tasks people assumed needed approximation. Reach for sparse or local attention when the signal has a clear locality structure and a few known long-range anchors, because you can encode that prior directly and cheaply. Reach for linear attention when the context is very long, the dynamics are smooth, and you want a streaming-friendly recurrent form. And when the sequence is truly enormous and the dependencies are long and diffuse, the honest answer is often to leave attention behind entirely for a state-space model, the linear-time architecture that Chapter 16 is devoted to.

Two disciplines cut across all four rows. First, the cheapest efficient attention is the token you never create: aggressive but sensible patchification (Section 15.2) shrinks \(L\) before any attention runs, and halving \(L\) quarters the full-attention bill for free. Second, whatever you choose, evaluate it under the leakage-safe protocol this book insists on, splitting by recording or subject rather than by window, so a long context does not quietly let the model memorize a device it will also be tested on; the machinery for that lives in Chapter 65.

What's Next

In Section 15.5, we turn from the length axis to the channel axis and confront a design fork that quietly decides both accuracy and cost: should each sensor channel be modeled on its own (channel-independent), or should attention mix channels so the model can learn cross-sensor structure (channel-mixing)? The efficient-attention budget you just built is exactly what that choice spends.