Beyond attention
Why attention gets expensive on long text, and the designs that trade exact recall for linear cost.
Read first:Self-attentionKV cache
Step 1 of 6· Illustrative tokens and gates; growth rates are real with constant factors ignored, and the cache figures come from the Mistral and Jamba papers
In full attention, every token compares itself with every earlier token, so doubling the length of a text roughly quadruples the work, and the memory for past tokens keeps growing as a model writes. Several designs aim to avoid this. Sliding-window attention looks only at recent tokens. Linear attention rewrites the maths so the past folds into a fixed-size summary. State-space models such as Mamba keep a running state and let each new token decide what to keep or forget. All three grow in proportion to the length rather than its square, but a fixed-size memory cannot hold everything, so hybrids such as Jamba keep a few full-attention layers.
Why it matters for your product
Context length is a cost line. Quadratic compute and a growing key-value cache are why long prompts are slower and pricier to serve; as Mistral's authors note, a growing cache means higher latency and lower throughput. Jamba's authors report a 128GB key-value cache for LLaMA-2 7B at 256K tokens, against 4GB for their hybrid. For long-document products, weigh these designs, or simply retrieve less text, before paying for the square.
For engineersShow the maths
Cost per layer: full attention ∝ n²·d, sliding window ∝ n·W·d, linear attention and SSMs ∝ n
n is the sequence length, d the head size and W the window. Only full attention squares the length.
Worked example: At n = 32,768 with W = 4,096, a rolling cache holds 4,096 positions instead of 32,768: 8 times less memory, as Mistral's authors report.
Sᵢ = Sᵢ₋₁ + φ(kᵢ)·vᵢᵀ, zᵢ = zᵢ₋₁ + φ(kᵢ), yᵢ = φ(qᵢ)ᵀSᵢ / φ(qᵢ)ᵀzᵢ
S is a running d × d summary of every key and value so far and z a running sum of keys. Each new token adds to them, and reading them costs the same at any length.
Worked example: With a head size of d = 64, S holds 64 × 64 = 4,096 numbers, whether the text is a thousand tokens long or a million.
hₜ = (1 − gₜ)·hₜ₋₁ + gₜ·xₜ
The simplest case of Mamba's selection: a gate g, computed from the current token, blends the old state with the new input. Near 1 it overwrites; near 0 it keeps the past.
Worked example: With h = 0.8 and x = 0.2, a gate of 0.1 gives 0.74 (mostly kept) and a gate of 0.9 gives 0.26 (mostly replaced).
Linear-time layers compress the past into a fixed state, so exact recall of details far back gets harder: Jamba's authors found a pure Mamba model struggled with in-context learning where their hybrid did not. Sliding windows can also miss a fact that lies beyond their stacked reach.
The words you will hear
Quadratic cost
Primary sources
- 01Attention Is All You NeedVaswani et al. · 2017
- 02Longformer: The Long-Document TransformerBeltagy et al. · 2020
- 03Transformers are RNNs: Fast Autoregressive Transformers with Linear AttentionKatharopoulos et al. · 2020
- 04Mamba: Linear-Time Sequence Modeling with Selective State SpacesGu, Dao · 2023
- 05Mistral 7BJiang et al. · 2023
- 06Jamba: A Hybrid Transformer-Mamba Language ModelLieber et al. · 2024
Connected ideas
Generating text
KV cache
Why a model keeps the keys and values of earlier tokens instead of recomputing them, and what that memory costs.
ExploreSpeed and cost
FlashAttention
The same attention maths, done in tiles so the big score table never touches slow memory.
ExploreInside the transformer
Self-attention
How every word looks at every other word to work out what it means here.
ExploreGenerating text
Prefill and decode
Why the first word of an answer takes a moment and the rest stream out at a steady pace.
ExploreWant this working on your data?
We design and build the systems these ideas power: retrieval, agents, voice and the models behind them. Start with a free discovery call.