Linear attention & kernelized approximations

Switch queries. Both numbers move together, every time, because they're computing the same thing two different ways — one rescans all four keys from scratch, the other looks up a state that was built before any query even ran.

Drop softmax and score a query against a key with a plain dot product instead — no per-row normalization needed. The double sum every query requires then factors, by ordinary associativity:

• $q_i$ — the query vector for position $i$, the one being scored against every key.
• $k_j$ — the key vector for position $j$, one of the keys a query is compared against.
• $v_j$ — the value vector for position $j$, the payload retrieved when key $j$ matches.
• $i$ — the index of the query currently being scored.
• $j$ — the index ranging over every key (and its paired value) in the sequence.

1. The sums don't depend on the query

   The parenthesized sums on the right don't mention $i$ at all — they're the same for every query, so computing them once (a single pass over the keys, $O(n)$) and reusing them turns each query into one $O(1)$ dot product instead of an $O(n)$ rescan.

2. Real layers add a feature map for non-negativity

   Real linear-attention layers apply a positive feature map $\phi$ (commonly $\text{elu}(x)+1$, where ELU is short for Exponential Linear Unit) to $q$ and $k$ before the dot product, so the "attention weights" stay non-negative the way softmax's do — this chapter uses the identity map to keep every number exact and checkable by hand, but the factoring trick itself doesn't depend on that choice at all.

The bold curve is the same quadratic from last chapter. The dashed one — $2n$, not $n^2$ — is what scoring every query the factored way actually costs: one pass to build the running state, one more to evaluate every query against it.

Four keys $k_j$ and values $v_j$, scored against query $q_0=(1,0)$:

1. The naive way: rescan every key

   Scores, one per key:

   • $q_0\cdot k_1 = (1,0)\cdot(2,1) = 2$
   • $q_0\cdot k_2 = (1,0)\cdot(-1,2) = -1$
   • $q_0\cdot k_3 = (1,0)\cdot(1,-1) = 1$
   • $q_0\cdot k_4 = (1,0)\cdot(0,3) = 0$

   Weighted sum $= 2(3)+(-1)(-2)+1(1)+0(4) = 6+2+1+0 = 9$. Normalizer $= 2-1+1+0 = 2$. Output $= 9/2 = 4.5$.

2. Build the state once, before any query

   Building $s=\sum_j k_jv_j$, one key at a time:

   • $k_1v_1 = (2,1)(3) = (6,3)$
   • $k_2v_2 = (-1,2)(-2) = (2,-4)$
   • $k_3v_3 = (1,-1)(1) = (1,-1)$
   • $k_4v_4 = (0,3)(4) = (0,12)$

   Summing: $s = (6+2+1+0,\ 3-4-1+12) = (9,\,10)$.

   Building $z=\sum_j k_j$: $(2,1)+(-1,2)+(1,-1)+(0,3) = (2+(-1)+1+0,\ 1+2+(-1)+3) = (2,\,5)$. Neither $s$ nor $z$ mentions $q_0$.

3. Every query is now one dot product

   • $q_0 \cdot s = (1,0)\cdot(9,10) = 9$.
   • $q_0 \cdot z = (1,0)\cdot(2,5) = 2$.

   Output $= 9/2 = 4.5$ — exactly the naive answer, with the four keys never rescanned.