The Geometry of Attention Is a Choice of Kernel

· 16 min read

#attention#kernels#geometry#transformers#yat#jax

Part 6 of 6Attention Is a Kernel
  1. 1Attention is Explainable Because it is a Kernel
  2. 2What an MLP Knows, When It's a Kernel
  3. 3Cheap Attention: Linear-Time Kernel Approximation
  4. 4Why Attention Needs Q and K Projections
  5. 5The Kernel Between the Roles
  6. 6The Geometry of Attention Is a Choice of Kernelyou are here
Runnable JAX companionMeasuring Attention's Geometry in JAX/Flax NNXPrefer to read the code? This post has a hands-on JAX / Flax NNX implementation.Open the JAX companion

This series has read attention as a kernel smoother, and it has built the smoother for real, a transformer whose scores come from a genuine Mercer kernel with no softmax anywhere. Both posts asked how the mechanisms train. This one asks a different kind of question, the kind you ask about a country rather than an engine: what map does each mechanism draw?

Every attention head you have ever inspected was inspected at the tokens: a row of weights over the 128 positions of some window, a heatmap, maybe a name like “induction head” if the pattern is clean. But look at what the formula actually accepts. A head scores a query vector against the keys, and the query slot takes any vector in the space, not just the 128 the sequence happened to produce. The trained head defines a function over an entire continuous landscape, and the sequence only ever samples it at a handful of points. Which means there is a whole geography in there nobody draws: between the sampled queries stretches empty space, and everywhere in that space the head has already decided, silently, which key would take the top weight.

So draw it. Color every possible query by the key that wins it. Where do the borders run? Are the territories big or small, bounded or infinite? Does every key even get one?

The answers turn out not to depend on training, on the dataset, or on the head. They depend on exactly one design decision, the score function, and the two score functions this series has trained head to head draw maps so different they belong to different genres. Drawing those two maps, wedges of sky against pockets of land, is the whole business of this post. The trained transformers make only a brief appearance at the end, to confirm what the geometry decided before their first gradient step: they were never free to draw anything else.

A sequence is a solar system

Drawing that map runs straight into the standing obstacle of this whole blog: the space it lives in is the query space of one head, 48 dimensions in the models this series trains, and you cannot picture 48 dimensions. The rescue is the same physical picture the network posts have leaned on since the prototype days: things that pull.

The keys become bodies sitting in space. A query becomes a test particle dropped somewhere among them. Each body pulls on the particle with a strength given by the score function, and attention’s normalization does exactly what a physicist would do with several simultaneous pulls: it computes each body’s share of the total. The row of attention weights is not a list of 128 unrelated numbers; it is one measurement of a field, taken at the particle’s position, reported as fractions that sum to one. The output of the head is then the value vectors averaged by those fractions, which is why the first attention post could call attention a kernel smoother with a straight face.

And the geometry does not stop at the weights. A row of nonnegative fractions summing to one is a point on a simplex, and averaging the values by those fractions confines every output the head can ever produce to the convex hull of its value vectors; the attention row is nothing but the output’s barycentric coordinates inside that hull. Move the query and the output slides around inside the polytope; the key that owns the query is the vertex the output huddles nearest. So a map of who owns query space is simultaneously a map of the head’s outputs: it says which value vector anchors the answer, everywhere, including all the places no token of any actual sequence has yet stood.

attention objectphysical object
key kjk_ja body sitting in the space
query qqa test particle
score s(q,kj)s(q, k_j)the pull body jj exerts on the particle
normalized weight αj\alpha_jbody jj‘s share of the total pull
the attention rowthe particle’s sky: every body’s share at once
top-1 key argmaxjαj\arg\max_j \alpha_jthe body that owns the particle

The table has one deliberately loaded row. “The body that owns the particle” is a definition of territory: the set of queries a given key wins. Carve the whole space by ownership and you get the map this post is about. And the shape of that map is decided entirely by the force law, which is to say, by the score:

sdot(q,k)=qkdvs.κ(q,k)=(qk+b)2qk2+ε.s_{\text{dot}}(q, k) = \frac{q \cdot k}{\sqrt{d}} \qquad\text{vs.}\qquad \kappa(q, k) = \frac{(q \cdot k + b)^2}{\lVert q - k \rVert^2 + \varepsilon}.

The left law is the transformer’s original: a bilinear form, wrapped in an exponential and normalized, which is the whole softmax. The right law is the Yat kernel, a genuine Mercer kernel whose scores arrive nonnegative and need no exponential; its denominator is the same softened inverse-square that made prototype neurons into wells, now sitting inside attention.

And before any algebra, the physics already predicts the two genres of map. Fix a body kk and look at its pull as a function of where the particle stands. Under the kernel, the pull peaks at the body and dies off with distance: the body is a well, and a well should own the neighborhood around itself. Under the dot product, the pull is linear in qq: a tilted plane, rising forever in the direction of kk, indifferent to how far away you stand. The body is not a well at all; it is an infinite slope. And a slope does not own a neighborhood. It owns a direction: whoever’s plane tilts steepest along your bearing wins, whether you stand a step from the origin or a mile. Wells should carve land. Slopes should carve sky.

The map under each law

A prediction that clean deserves to be embarrassed in public, so here it is on trial: six bodies you can drag, and every pixel of the plane a possible query, colored by the key that takes the top weight there, shaded darker the more of the row the winner holds.

The left map is all wedges, and every border between two wedges passes through the origin. That is not an accident of the six positions; it is a two-line theorem. The border between key jj‘s territory and key ll‘s is where their scores tie:

qkj=qkl    q(kjkl)=0,q \cdot k_j = q \cdot k_l \iff q \cdot (k_j - k_l) = 0,

a hyperplane through the origin, whatever the keys are. Each territory is an intersection of half-spaces through the origin, which is a convex cone: an infinite wedge of sky anchored at zero. Move a key and the wedges swing, but they can only ever swing; no border can bulge, close, or surround anything, because a hyperplane through the origin has no curvature to offer. The softmax on top changes none of this. Exponentials are monotone, so the winner under ese^{s} is the winner under ss; softmax decides how much of the row the winner takes, never who wins.

The right map is the other genre. The kernel’s tie surface, κ(q,kj)=κ(q,kl)\kappa(q, k_j) = \kappa(q, k_l), is a genuinely curved quartic, and near the bodies it does what curvature allows: it closes. Each body sits inside a bounded pocket of queries it owns outright, the map cartography calls a weighted Voronoi cell (Aurenhammer 1987), and dragging a body relocates its pocket like moving a planet relocates its gravitational capture zone.

Zoom out, though, and the kernel’s map plays a trick worth the slider. Far from all the bodies, along a direction uu, the kernel behaves like

κ(tu,k)=(tuk+b)2t22tuk+k2+ε    t    (uk)2,\kappa(t u, k) = \frac{(t\, u \cdot k + b)^2}{t^2 - 2t\, u \cdot k + \lVert k \rVert^2 + \varepsilon} \;\xrightarrow{\;t \to \infty\;}\; (u \cdot k)^2,

so at infinity the kernel’s sky is carved into cones too, by (ukj)2(u \cdot k_j)^2. Both laws agree that infinity belongs to directions. But look at the square: in the limit, the kernel’s far field cannot tell uu from u-u, a body from its antipode. The dot product’s sky is signed; the kernel’s sky, eventually, is not. Eventually is doing real work in that sentence: the sign survives in the next term of the expansion, 2(uk)((uk)2+b)/t2\,(u \cdot k)\,((u \cdot k)^2 + b)/t, which dies off only as 1/t1/t, so on the six bodies above the antipodal symmetry does not win a majority of directions until roughly radius twenty and does not saturate until near one hundred, far past the panel’s slider. What the slider does show is the structural point: away from the bodies the borders straighten and radiate, wedges again. Everything the two laws truly disagree about, they disagree about near the bodies, which is where every actual query of every actual sequence lives.

Length is a temperature, or an address

The wedge theorem has a corollary that sounds too strong to be true, so it deserves its own toy. Take any query and scale it: qtqq \to t\,q, with t>0t > 0. Every dot-product logit scales by the same tt, so the ranking of the keys is untouched, for every tt, on every ray, forever. Under the dot-product law, the length of a query cannot change who wins. What it changes is the softmax’s appetite: softmax(ts)\mathrm{softmax}(t \cdot s) at large tt piles the row onto the winner. Length is an inverse temperature. A long query is not somewhere else; it is the same place, said louder.

The kernel-attention post found softmax’s other gauge: adding a constant to every logit changes nothing, so the level of the scores is deleted information. Scaling is the gauge softmax keeps but spends on confidence. The kernel spends it on position:

Under the kernel, tt moves the particle. It passes body after body, and each handoff is visible as the row’s peak jumping; between bodies the total kernel mass, the row sum that survives as a channel softmax cannot represent, sags toward nothing. The same slider means “how sure” on one side of the panel and “where” on the other. That is the entire difference between the two attention laws, compressed into one degree of freedom.

The disenfranchised body

A query’s length was one freedom the two laws read differently. The arrangement of the keys themselves is the other, and it hides the theorem with teeth. Five bodies in a ring, a sixth in the middle, strictly inside the convex hull of the others: physics with wells says the middle body still owns its crater, because stand close enough to a well and it out-pulls everything, however the neighbors are arranged. What does physics with slopes say?

The middle key, sitting inside the hull, is a convex combination of the ring: k6=lλlklk_6 = \sum_l \lambda_l k_l with λl0\lambda_l \ge 0, λl=1\sum \lambda_l = 1. So for every query qq,

qk6=lλl(qkl)    maxl  qkl,q \cdot k_6 = \sum_l \lambda_l \,(q \cdot k_l) \;\le\; \max_l \; q \cdot k_l,

with equality only on ties. The middle body’s pull is an average of the others’ pulls, and an average never strictly beats its own maximum. Under the dot-product law, a key inside the others’ hull wins nothing, anywhere, for any query the network could ever compute. Its territory is not small; it is empty. And the converse holds too: outside the hull, a separating hyperplane hands it a whole infinite wedge. Ownership under softmax attention is a property of the hull, decided at the moment the keys are laid down.

The kernel’s answer is the well’s answer, and it comes with its own guarantee. At its own location the kernel’s pull is κ(k,k)=(k2+b)2/ε\kappa(k, k) = (\lVert k \rVert^2 + b)^2 / \varepsilon: as the softening ε\varepsilon shrinks, that peak grows without bound while every rival’s pull at that spot stays finite. With softening below the spacing of the keys, every body owns at least the ground it stands on, hull or no hull. And ε\varepsilon is not a fixed constant in the trained models; each head learns its own, which means each head learns how fiercely bodies defend their home ground.

The theorem has one boundary worth drawing before it goes to work: softmax weights are never zero, so an interior key still receives shares and still feeds the average. Disenfranchisement is about winning, the top of the row, not about mattering at all. But the top of the row is precisely what sharp heads route by, and it is exactly what the maps in this post have been coloring.

A coda: the trained models were never free

Everything above is exact in any dimension, indifferent to the data, and settled before training begins, which makes it a checkable kind of arrogance. So it was checked: the two character-level Shakespeare transformers from the kernel-attention post retrained with every head’s trained query and key vectors exported, and the theorems replayed on them (provenance in the scope note; the companion walks the code).

The wedge corollary made the strongest claim, that no rescaling of any query can ever change a softmax winner, and the trained model honors it to the digit: rescale every real query by anything from a half to four, and across all six layers, all four heads, all scorable rows, exactly 0 winners move in 15,240 comparisons. Twelve thousand steps of gradient descent could not buy softmax a meaning for magnitude, because the geometry never put one on sale. The kernel model, whose lengths are addresses, re-elects 12 percent of its rows at half scale and 23 percent at quadruple:

The territory maps behave like their genres too. Scrub the census below through training and ownership gets carved out of noise, but the two laws carve in character: the softmax model’s rows are far sharper (mean normalized entropy 0.49 against 0.72), yet it spreads its crowns over more distinct winners, 72 of 128 against the kernel’s 54, because the kernel’s diffuse rows keep re-electing the same few keys, visible as the horizontal annexation runs in its map. Concentration of a row and concentration of a map are different geometric quantities. And the hull theorem’s day in court is quiet by necessity: in 48 dimensions a cloud of 128 keys is essentially all hull (a linear program finds 0 interior keys of 3,072 in either model), so the plane’s absolute disenfranchisement relaxes, in a trained head, into exactly this census of who wins the queries that actually occur.

Scope: the coda’s numbers come from one seed-0 run per variant of scripts/yat_attention.py on Kaggle with geometry telemetry (bundles kgl_blog-attngeom-qk and kgl_blog-attngeom-goat, each variant at the per-variant learning rate swept in the kernel-attention post), measured on one fixed 128-character validation window; the analysis is scripts/export_attention_geometry_viz.py, a local replay of the exported vectors and maps. Quality and entropy comparisons of these same architectures, with seeds, live in the kernel-attention post and are not re-litigated here. The theorems are exact in any dimension; the census statistics are one window of one seed at one scale.

The choice of gravity

Two attention mechanisms, trained on the same text, holding within about a percent of each other in quality, and underneath: two different physical universes. One carves query space into signed wedges of sky that all meet at the origin, cannot let magnitude mean position, and holds elections a key inside the hull can never win. The other carves bounded neighborhoods around bodies that always own their ground, reads a query’s length as an address, and only agrees with the first law about what happens at infinity, up to a sign it cannot see.

None of that was trained. All of it was chosen, silently, the moment the score function was written down. That is what it means to say the kernel is the choice of gravity: before a single gradient step, the force law has already decided what kind of map the model is allowed to draw, and training only gets to place the bodies on it.

Cite as

Bouhsine, T. (). The Geometry of Attention Is a Choice of Kernel. Records of the !mmortal Data Scientist. https://tahabouhsine.com/blog/the-geometry-of-attention/

BibTeX
@misc{bouhsine2026thegeometryofattention,
  author       = {Bouhsine, Taha},
  title        = {The Geometry of Attention Is a Choice of Kernel},
  year         = {2026},
  month        = {jul},
  howpublished = {\url{https://tahabouhsine.com/blog/the-geometry-of-attention/}},
  note         = {Blog post, Records of the !mmortal Data Scientist}
}

References

  1. Vaswani, A., et al. (2017). Attention Is All You Need. arXiv:1706.03762
  2. Tsai, Y.-H. H., et al. (2019). Transformer Dissection: An Unified Understanding for Transformer's Attention via the Lens of Kernel. EMNLP-IJCNLP 2019.
  3. Aurenhammer, F. (1987). Power Diagrams: Properties, Algorithms and Applications. SIAM Journal on Computing 16(1), 78-96.
  4. Bouhsine, T. (2026). A Universal Reproducing Kernel Hilbert Space from Polynomial Alignment and IMQ Distance. arXiv:2605.03262