Kernelization: The Licence to Move the Parentheses

Part 3 of 4 on how a transformer remembers. KV cache · Linear attention · → Kernelization · Recurrent memory

The last post moved a pair of parentheses and got a constant-size cache out of it. But it did that by quietly assuming the score is a plain dot product, when in real attention it’s $\exp(q\cdot k)$. This post is about that assumption: what exactly has to be true for the parentheses to move, why the exponential almost qualifies, and what people do about the almost.

The whole thing reduces to one move: get $q$ and $k$ to stand in separate columns, so neither ever needs to know about the other.

A kernel is an inner product wearing a disguise

One line:

\[\kappa(x,y) \;=\; \phi(x)^{\top}\phi(y)\]

That’s it. A kernel is an ordinary inner product in some (possibly much bigger) space, written as a complicated-looking function of the original inputs. The trivial case is $\kappa(x,y)=x^\top y$ with $\phi(x)=x$ — no lift at all. The instructive case is one dimension up:

\[\kappa(x,y)=(1+xy)^2 = 1+2xy+x^2y^2 = \begin{bmatrix}1\\ \sqrt2\,x\\ x^2\end{bmatrix}^{\!\top}\begin{bmatrix}1\\ \sqrt2\,y\\ y^2\end{bmatrix}\]

The left side tangles $x$ and $y$ into a quadratic. The right side is a perfectly ordinary dot product of two 3-vectors. They are identically equal. Check it channel by channel:

The dashed line in that panel is the whole point. The left column is a function of $x$ alone; the right column is a function of $y$ alone; they touch exactly once, in the final multiply-and-sum. That separation is the licence. It’s what lets you compute $\sum_j \phi(k_j)v_j^{\top}$ before any query exists — and a sum you can compute in advance is a sum you can store in fixed size.

Choosing a kernel is choosing what attention looks like

Write attention as a kernel-weighted average and it is, near enough, Nadaraya–Watson kernel regression:

\[o_i \;=\; \frac{\sum_j \kappa(q_i,k_j)\,v_j}{\sum_j \kappa(q_i,k_j)}\]

Numerator and denominator depend on nothing but $\kappa$. So the choice of kernel is the choice of attention’s behaviour — not a detail of it.

Which reframes the exponential in softmax. It isn’t a nonlinearity someone tossed in for flavour. It is a deliberately chosen kernel that happens to satisfy every requirement at once: strictly positive (so weights are weights), sharply peaked (so retrieval is decisive), normalisable (so the denominator never crosses zero) — and, it turns out, a legitimate inner product.

The exponential kernel’s φ really does exist

No trick required; the power series does the work. Substituting $z=q^\top k$:

\[e^{q^\top k}=\sum_{n=0}^{\infty}\frac{(q^\top k)^n}{n!}\]

Each term $(q^\top k)^n$ is the inner product $\langle q^{\otimes n},k^{\otimes n}\rangle$ of two $n$-th tensor powers. Concatenate the orders and you have an honest — if infinite — feature map. Here it is written out for $d=2$:

So the exponential is a valid kernel and $q$ can be factored out — it just needs infinitely many dimensions to do it, and truncating to order $M$ costs $\binom{d+M}{M}$ channels. At $d=64$ that’s 2145 channels by order 2 and tens of millions by order 6, per head. Valid, and unusable.

Random features: sample the infinite inner product

Performer’s answer is to estimate it instead. Draw $\omega\sim\mathcal N(0,I)$ and define

\[f_\omega(x)=\exp\!\Big(\omega^{\top}x-\tfrac{\|x\|^2}{2}\Big)\]

Then there’s an exact expectation identity, straight out of the Gaussian moment generating function $\mathbb E[e^{\omega^\top t}]=e^{|t|^2/2}$ with $t=q+k$:

\[\mathbb E_\omega\big[f_\omega(q)\,f_\omega(k)\big] \;=\; e^{q^{\top}k}\]

Sample $m$ of them, set $\phi(x)=\tfrac{1}{\sqrt m}[f_{\omega_1}(x),\dots,f_{\omega_m}(x)]$, and the inner product becomes a Monte Carlo estimate of the exact one. $m$ is now a dimension you pick.

Where it all lands

Three posts, one chain, and every link is load-bearing:

Kernelization issues the licence — if $\kappa$ factors as $\phi(q)^\top\phi(k)$, then $q$ and $k$ are separable. Associativity cashes it in — separate them and the parentheses move, and an $O(N)$ cache collapses into a $d\times d$ state. Fixed capacity sends the bill — $d\times d$ holds $d$ things, so past $d$ tokens the system has no choice but to learn what to throw away. That’s the last post: erase, decay, and the five-year argument about how to forget.

The practical epilogue is a little deflating and worth saying anyway. Almost nobody actually approximates $\exp$. Explicit expansions blow up combinatorially; random features are unbiased but noisy at $O(1/\sqrt m)$, which is one reason Performer-style methods aren’t always stable on long sequences. So the field mostly shrugged and reached for a cheap non-negative map like ELU(x)+1 — which doesn’t even pretend to approximate softmax, and only guarantees the two properties you actually need: non-negative, and normalisable.

That’s not laziness, it’s a budget decision. If you’re going to lose retrieval fidelity either way, lose it cheaply and spend the compute somewhere it buys more — which turned out to be the forgetting rules in the next post.

---

Separate the variables, and everything else is bookkeeping.

References

• Linear attention via feature maps — Katharopoulos et al., Transformers are RNNs (arXiv:2006.16236)
• Random features for kernels — Rahimi & Recht, Random Features for Large-Scale Kernel Machines (NIPS 2007)
• Applying them to softmax attention — Choromanski et al., Rethinking Attention with Performers (FAVOR+, arXiv:2009.14794)
• A survey of the design space — Tay et al., Efficient Transformers: A Survey (arXiv:2009.06732)