insideout Get Pro

16 · Linear attention and hybrid models

Pro

Transformers, lesson 16

Linear attention and hybrid models

Most layers keep a fixed-size state, and a few keep full attention.

Full attention compares each new token with every token before it, so its work and its KV cache grow with the length. Linear attention keeps a fixed-size state instead, and hybrid models mix the two. You need "Attention as a soft lookup"; the cache sizes follow "How big the cache gets", in Inference and Serving.

By the end, you can:

  1. Rewrite attention without the softmax as a running sum, and count what each token costs.
  2. Show why a fixed-size state loses exact recall of one token far back.
  3. Update a state with a decay gate and with the delta rule, and say what each fixes.
  4. Size a hybrid model's KV cache, and say why only its attention layers grow one.

Part 1. Attention as a running sum

Rewrite attention without the softmax as a running sum, and count what each token costs.

Full attention scores each new query against every earlier key and blends every earlier value. Token 100,000 does 100,000 times the work of token 1, and every key and value stays in the cache. Drop the softmax, and the same sum can be done in another order.

Example 1.1

Two keys, two orders

Keys (1, 0) and (1, 1) hold values (2, 0) and (0, 4). Query (1, 2) weights each value by its score q · k, with no softmax, and divides by the total. Work it out directly, then by folding the keys and values into one grid first.

  1. Score, blend, divide(1 × (2, 0) + 3 × (0, 4)) ÷ 4 = (0.5, 3)

    q · k₁ = 1 and q · k₂ = 3, which add up to 4.

  2. Fold firstS = k₁v₁ᵀ + k₂v₂ᵀ = [[2, 4], [0, 4]]

    Each key times its value is a 2 × 2 grid: [[2, 0], [0, 0]] plus [[0, 4], [0, 4]]. Add the keys too: z = (2, 1).

  3. Read with the queryqᵀS ÷ (q · z) = (2, 12) ÷ 4

    (1, 2) times the grid's columns gives (2, 12), and q · z = 4.

Answer(0.5, 3) both ways, the second from just the grid and the key total.

Scoring first adds up (q · kᵢ) × vᵢ. Regrouped, that's q times the sum of every kᵢvᵢᵀ. That sum doesn't involve the query, so you can keep it as you go: each token adds its kvᵀ, and its key and value can be thrown away.

Full attentionkeeps every key and valuekv…token 1token t2 × t × d numbers a head: 25.6 million at t = 100,000Linear attentionkeeps one statetoken t+ k vᵀd × d numbers16,384 at any t(d = 128)
Figure 1Full attention keeps a key and a value for every token, so its cache grows with t. Linear attention adds each token into one d × d state and keeps nothing else.

Linear attention

Attention whose score splits into a query part times a key part, φ(q) · φ(k), so it runs as a running sum. The softmax's e^(q · k) can't split. The state is a d × d grid per head. Katharopoulos and colleagues (2020) used φ(x) = elu(x) + 1, which keeps scores positive.

S ← S + k vᵀ, z ← z + k, out = qᵀS ÷ (q · z)

Cost per token changes shape too. For heads of d = 128, attention at token t does about 4 × t × 128 operations; the state costs about 4 × 128² = 65,536 at any t. They cross at t = 128, and at token 100,000 attention does 780 times the work.

Your turn

A third token, with key (0, 1) and value (4, 2), joins the example's state. What does query (1, 1) read?

k vᵀ = [[0, 0], [4, 2]], so S becomes [[2, 4], [4, 6]] and z becomes (2, 2). qᵀS = (6, 10) and q · z = 4. (2, 3.33) leaves z alone, (0.5, 3.5) adds v kᵀ in place of k vᵀ, and (6, 10) skips the divide.

Part 2. What a fixed state can't hold

Show why a fixed-size state loses exact recall of one token far back.

Keep going with Pro

You've done the first part. Pro opens the rest of this lesson, and every other one.

Still to come
  1. What a fixed state can't hold
  2. Gates that decide what to keep
  3. Hybrid models and their cache
See Pro plansI have Pro: sign in
Next lessonTraining a small GPT from zeroTrain and sample a tiny GPT in the browser, with every line in view. Start lesson 17

Transformers overviewSpotted a mistake? Tell us.