Chapter 2 The Geometry of Self-Attention

Queries, keys and values

The transformer architecture introduced by Vaswani and colleagues in 2017 is built around attention: a way for each position in a sequence to gather information from all the others (Vaswani et al. 2017). Each token’s internal vector is projected three ways, into a query, a key and a value. The query asks; the keys advertise; the match between them decides how much of each value is gathered:

$\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V.$

Each row of $QK^{\mathsf T}$ holds one query’s dot products with every key; the softmax turns them into weights that sum to one; the result is a weighted average of the values. Multi-head attention runs several such operations in parallel, each with its own projections, so that different heads can attend to different kinds of relation, and many layers of this are stacked (Vaswani et al. 2017).

Why divide by the square root of the dimension

Calculation

For random queries and keys, dot products have a variance that grows with the dimension $d_k$ . The program measures the entropy of the attention weights over 16 keys with and without dividing by $\sqrt{d_k}$ .

import math, random

# Why divide by √d_k: dot products of random d-dimensional vectors have variance d,
# so without scaling the softmax over 16 keys becomes nearly one-hot as d grows.
random.seed(1)
def softmax_entropy(scores):
    m = max(scores)
    w = [math.exp(s - m) for s in scores]
    z = sum(w)
    return -sum((x / z) * math.log2(x / z) for x in w if x > 0)

for d in (16, 64, 256, 1024):
    raw, scaled = [], []
    for _ in range(200):
        q = [random.gauss(0, 1) for _ in range(d)]
        ks = [[random.gauss(0, 1) for _ in range(d)] for _ in range(16)]
        s = [sum(a * b for a, b in zip(q, k)) for k in ks]
        raw.append(softmax_entropy(s))
        scaled.append(softmax_entropy([x / math.sqrt(d) for x in s]))
    print(f"d_k = {d:4d}: attention entropy without scaling {sum(raw) / 200:4.2f} bits, "
          f"with 1/√d_k {sum(scaled) / 200:4.2f} bits (maximum 4.00)")
# d_k =   16: attention entropy without scaling 1.40 bits, with 1/√d_k 3.44 bits (maximum 4.00)
# d_k =   64: attention entropy without scaling 0.63 bits, with 1/√d_k 3.43 bits (maximum 4.00)
# d_k =  256: attention entropy without scaling 0.26 bits, with 1/√d_k 3.40 bits (maximum 4.00)
# d_k = 1024: attention entropy without scaling 0.16 bits, with 1/√d_k 3.41 bits (maximum 4.00)

Without the scaling, the weights become nearly one-hot as the dimension grows: the softmax saturates, almost all weight goes to one key, and learning slows because the gradients vanish. Dividing by $\sqrt{d_k}$ keeps the spread roughly constant. This was the reason the authors gave for the scaling (Vaswani et al. 2017).

A toy model of a question’s reach

In a trained model, the vectors of related concepts tend to lie in related directions. The toy model below has three “fields”, each a small cluster of keys in a 12-dimensional space, and asks what a single attention operation gathers for three queries: one pointing into one field, one pointing between two fields, and one pointing nowhere in particular.

Calculation

Scaled dot-product attention written from scratch, with keys in three clusters.

import math, random

random.seed(4)
D = 12                                    # embedding dimension
def unit(v):
    n = math.sqrt(sum(x * x for x in v))
    return [x / n for x in v]

# Three "fields", each a cluster of 4 keys around its own direction in a 12-d space.
fields = {"acoustics": 0, "control theory": 4, "Buddhist texts": 8}
keys, labels = [], []
for name, start in fields.items():
    centre = [1.0 if start <= i < start + 4 else 0.0 for i in range(D)]
    for _ in range(4):
        keys.append(unit([c + random.gauss(0, 0.25) for c in centre]))
        labels.append(name)
values = keys                             # let each key carry its own content

def attention(q, scale=3.0):
    """Softmax of scaled dot products. With unit vectors, 'scale' stands in for the
    learned length of the projected queries and keys divided by √d_k."""
    s = [scale * sum(a * b for a, b in zip(q, k)) for k in keys]
    m = max(s)
    w = [math.exp(x - m) for x in s]
    z = sum(w)
    return [x / z for x in w]

def report(name, q):
    w = attention(q)
    per = {f: sum(wi for wi, l in zip(w, labels) if l == f) for f in fields}
    out = [sum(wi * v[i] for wi, v in zip(w, values)) for i in range(D)]
    print(f"{name:34s} weight by field: " + ", ".join(f"{f} {per[f]:.2f}" for f in fields)
          + f" | weight entropy {-sum(x * math.log2(x) for x in w):.2f} bits")
    return out

q1 = unit([1.0 if i < 4 else 0.0 for i in range(D)])                       # one field only
q2 = unit([1.0 if i < 4 else (1.0 if 4 <= i < 8 else 0.0) for i in range(D)])  # two fields
q3 = unit([1.0] * D)                                                        # everything at once
for n, q in (("query inside acoustics", q1), ("query spanning acoustics + control", q2),
             ("query with no direction", q3)):
    report(n, q)
# query inside acoustics             weight by field: acoustics 0.87, control theory 0.07, Buddhist texts 0.07 | weight entropy 2.69 bits
# query spanning acoustics + control weight by field: acoustics 0.42, control theory 0.49, Buddhist texts 0.10 | weight entropy 3.31 bits
# query with no direction            weight by field: acoustics 0.26, control theory 0.32, Buddhist texts 0.42 | weight entropy 3.53 bits

Attention weight gathered from each field by three queries in the toy model. A query inside one field stays there; a query between two fields draws on both; a query with no direction spreads its weight over everything.

The three cases have intuitive readings. A query that lies inside one field gathers almost only from it. A query with components along two fields draws on both in roughly equal measure, and the weighted average it returns is a mixture that neither cluster contains on its own. A query with no particular direction spreads its weight most evenly of all, and returns something close to the average of everything: the attention equivalent of a generic answer.

Note

This is a single attention operation on hand-made vectors, an illustration of the geometry, not a model of how a trained transformer combines ideas. Real models stack dozens of layers with many heads, learned projections and other components, and how they compose information is an active research field (Elhage et al. 2021).

The bell

Source text

善待问者如撞钟，叩之以小者则小鸣，叩之以大者则大鸣，待其从容，然后尽其声。(佚名, n.d.)

One who is good at answering questions is like a bell being struck: struck lightly, it sounds lightly; struck heavily, it sounds fully; and given time, it sounds to its full extent.

(Record of Learning, in the Book of Rites)

The Record of Learning describes the teacher, not the student, as the bell; the passage is about how a good teacher answers in proportion to the question. Read beside the toy model, it gives an image for what the query vector does: the response depends on how, and how hard, the bell is struck.

Hypothesis · not established

Not established. Prompts that explicitly name concepts from two or more distant fields and the relation to be examined between them lead a language model to produce combinations that do not appear together in its training data more often than prompts confined to one field. The hypothesis predicts that, for matched topics, independent raters searching the literature will judge the central combination in an output to be “not previously published” more often after the cross-field prompt. It would be refuted if, in a blinded comparison of at least fifty topics, the proportions did not differ.

References

Elhage, Nelson, Neel Nanda, Catherine Olsson, et al. 2021. A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread. https://transformer-circuits.pub/2021/framework/index.html.

Vaswani, Ashish, Noam Shazeer, Niki Parmar, et al. 2017. “Attention Is All You Need.” Advances in Neural Information Processing Systems 30.

佚名. n.d. 礼记·学记.

Fan, Y. (2026). The Algorithmic Mirror: Self-Attention Geometry, Ashby's Requisite Variety and Cognitive Projection in Information Theory and East Asian Philosophy of Mind [算法之镜] (Version 1.0.0). 5D Library. https://doi.org/10.5281/zenodo.23723557 · Licensed under Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International (CC BY-NC-ND 4.0).