How an attention head mixes information

In the word2vec article, we learned a table of word vectors from the contexts words appeared in during training. Words with related usage could acquire related representations. But the table has a limitation: looking up a word gives the same vector regardless of the sentence it is in. Context helped train the vector; the lookup does not consult the current context.

That is the polysemy problem we left with. In “the river bank” and “the money bank,” bank retrieves the same row. Its vector can reflect evidence from both uses, but it does not tell us which use this occurrence belongs to. The surrounding words supply that missing information.

We therefore need a step after the lookup: let each position gather information from the other positions in this sentence. An attention head computes a weighted mixture of the information they carry. The weights depend on the input, so different positions can gather different mixtures. The goal is a contribution that helps make the representation useful for this particular context.

Inside a transformer block, that contribution is projected and added to the representation already there. The head computes the mixture; the residual connection incorporates the update. This gives us a path from a shared word vector to a representation enriched by its current surroundings. The next article builds the block around these updates.

A standard transformer block contains several heads in its multi-head attention sublayer. Heads within a block run in parallel; transformer blocks run sequentially. The heads read the same input representations and compute separate mixtures, which are combined and projected into one attention update. Each subsequent block receives the previous block’s updated representations, rather than looking up the initial token vectors again.

Here we will build one head on small vectors: first watch information mix, then see how learned projections control what it gathers, and finally train a tiny model to use it.

You need vectors, dot products, and the idea that multiplying by a matrix transforms a vector. The weight-matrix article develops that last step. Our axes will be hand-chosen for legibility; real embeddings do not come with a labelled finance coordinate. We also treat each word as one token throughout this toy.

One forward pass, no training

Give four tokens four numbers each. The dimensions mean finance, water, structure, and generic content in this example only:

import numpy as np

emb = {
    #           finance water structure generic
    "the":   np.array([0.0, 0.0, 0.0, 1.0]),
    "river": np.array([0.0, 1.0, 0.0, 0.2]),
    "money": np.array([1.0, 0.0, 0.0, 0.2]),
    "bank":  np.array([0.5, 0.5, 0.8, 0.2]),
}

bank starts tied between finance and water. We want to compute a contribution at its position by mixing information from the sentence. With three tokens, the shape of that computation is:

The $v$ vectors carry information; the nonnegative $a$ weights sum to one. For now, use the embeddings themselves as those vectors. The open question is where the weights come from.

Our first rule is simple: dot the receiving token’s vector against each source vector, divide the scores by the square root of their dimension, and apply softmax. Softmax exponentiates the scores and divides each result by their sum. This makes larger scores receive larger shares while keeping the shares positive and summing to one.

Here is the reusable head. X is the sentence matrix, one token per row. We include three projection matrices now, but initially set all three to the identity, which leaves every vector unchanged.

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def attention(X, Wq, Wk, Wv, causal=False):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    scores = Q @ K.T / np.sqrt(K.shape[-1])
    if causal:
        future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
        scores = np.where(future, -np.inf, scores)
    A = softmax(scores)
    return A, A @ V

I = np.eye(4)
for tokens in [["the", "river", "bank"], ["the", "money", "bank"]]:
    X = np.stack([emb[t] for t in tokens])
    A, H = attention(X, I, I, I)
    print(tokens)
    print("  bank input:", X[-1])
    print("  weights:", A[-1].round(2))
    print("  mixture:", H[-1].round(2))

The result:

['the', 'river', 'bank']
  bank input: [0.5 0.5 0.8 0.2]
  weights: [0.26 0.31 0.43]
  mixture: [0.21 0.52 0.34 0.41]
['the', 'money', 'bank']
  bank input: [0.5 0.5 0.8 0.2]
  weights: [0.26 0.31 0.43]
  mixture: [0.52 0.21 0.34 0.41]

The same input vector at bank now receives different mixtures. river supplies water; money supplies finance. The mixture changes even though the weights are identical: the vectors being mixed have changed. Attention weights are only one part of the computation.

This is not yet the result of a complete transformer block. We have printed H, the head output, without an output projection or residual addition. Calling it bank’s new representation would skip those operations.

To preview the addition without introducing another matrix yet, suppose the output projection is identity and the stream has the same four coordinates. Adding the river mixture to the original gives [0.71, 1.02, 1.14, 0.61]; adding the money mixture gives [1.02, 0.71, 1.14, 0.61]. Both keep an additive contribution from the original and acquire a different context contribution. This isolates the residual operation; normalization and general output projections belong to the block article.

A plain average also separates these two sentences: the river sentence’s mean has water 0.50 and finance 0.17. So this first experiment demonstrates information mixing, not a task that requires selective attention. Our hand-designed vectors made even a simple mixture useful. Next we will see what controls that mixture and where this first matching rule falls short.

Why it mixed at all

Trace just the bank row in “the river bank.” Its unscaled dot products with the, river, and itself are [0.20, 0.54, 1.18]. Since the vectors have four dimensions, dividing by $\sqrt{4}=2$ gives [0.10, 0.27, 0.59]. Softmax turns those scores into approximately [0.26, 0.31, 0.43].

Then the head computes:

mixture = 0.26 × the + 0.31 × river + 0.43 × bank
        ≈ [0.21, 0.52, 0.34, 0.41]

The displayed arithmetic is rounded; the program keeps full precision. The river term contributes water to the mixture. The largest weight happens to be on bank itself: its self-dot-product is largest in this constructed space. This is a property of these numbers, not a requirement that attention preserve identity by attending to itself. In a block, the residual path separately carries the existing representation forward.

All rows do the same operation at once. With identity projections:

scores = X @ X.T / np.sqrt(X.shape[-1])  # (n, n)
A = softmax(scores)                      # (n, n): one distribution per row
H = A @ X                               # (n, d): one mixture per row

A[i, j] is the share of source position j in the mixture received at position i. It is a distribution over positions, not over next-token vocabulary choices.

When raw similarity asks the wrong question

The first rule says “gather from vectors similar to mine.” But useful information need not resemble the receiving token. A structure can need an identity supplied by a name.

To see this cleanly, switch to a separate, two-dimensional teaching space. Its axes are structure and Paris identity. This is not a slice of the four-dimensional bank example:

# rows: the, Eiffel, tower
X2 = np.array([[0., 0.], [0., 1.], [1., 0.]])
I2 = np.eye(2)
A, H = attention(X2, I2, I2, I2)
print(A[-1].round(4))  # [0.2483 0.2483 0.5035]
print(H[-1].round(4))  # [0.5035 0.2483]

tower · Eiffel = 0. Under raw similarity, Eiffel scores exactly like the zero-vector filler. Some Paris content still enters because softmax gives both positions positive weight, but the matching rule has no way to prefer Eiffel over the filler.

For our chosen task, we want tower to gather identity information. That is a limitation of the matching rule, not a failure of weighted mixing. We need to change what tower matches against while keeping its original vector available.

Producing a query

Each token produces one query, one key, and one value vector per attention head, per transformer block. Within one head, the same learned matrix Wq transforms every position: q_i = x_i @ Wq, or Q = X @ Wq for the whole sentence. Separate projections Wk and Wv produce that token’s key and value. These are computed vectors, rather than additional vocabulary embedding lookups. A token in a block with 12 heads therefore has 12 queries, 12 keys, and 12 values. The query is a view used for matching; the original representation remains available to the rest of the block.

“What information could this position use from other tokens?” is a useful way to interpret a query. For bank, a hypothetical head might match water-related evidence or finance-related evidence; for tower, it might match a name that supplies identity. More precisely, the query determines which keys produce large dot-product scores. These scores become weights for mixing values. Producing a query is the matching setup; it has not yet gathered information or updated the token’s representation.

A query reads the token’s full current representation through Wq; it need not preserve raw similarity or keep the same number of coordinates. A large model can project thousands of representation coordinates into a much smaller query/key space. Queries and keys within a head must have the same width so their dot products are defined. Their coordinates do not have to be named concepts, and a head is not programmed with a literal question.

Transform the vector used to ask the question. A swap matrix maps tower = [1, 0] to a query [0, 1], which matches Eiffel’s identity coordinate:

swap = np.array([[0., 1.], [1., 0.]])
A, H = attention(X2, swap, I2, I2)
print(A[-1].round(4))  # [0.2483 0.5035 0.2483]
print(H[-1].round(4))  # [0.2483 0.5035]
print(X2 @ swap)
# [[0. 0.]    the's query
#  [1. 0.]    Eiffel's query
#  [0. 1.]]   tower's query

The same matrix acts on every row. Different input vectors give different queries; we did not need a separate matrix per word. The token representation and its query are now different objects: the representation supplies the input, and the query is one transformed view used for matching.

At the first layer, the input includes embeddings and any positional signal. Deeper layers operate on representations already updated by earlier layers. The projection rule is shared; its inputs need not be static word vectors.

What is query/key space, and how does it relate to subspaces?

Query/key space is the coordinate space where a head compares queries with keys. If a head uses 128 coordinates for each, both live in $\mathbb R^{128}$, and their dot product produces a matching score there. The learned maps Wq and Wk translate the much wider token representations into these matching coordinates. Each head has its own maps; the coordinates need not have named semantic meanings.

For example, a 12,288-coordinate input can be multiplied by a 12,288 × 128 query matrix to produce a 128-coordinate query. The key matrix has the same shape and produces 128-coordinate keys from the source representations. The diagram uses bank and river to show this relationship; these widths illustrate a possible model configuration, rather than the sizes of our toy.

This connects directly to row space and null space: what the matrix reads and ignores. That article uses column vectors, while our code uses row vectors. To apply its notation, set $M=W_q^\top$, so $q^\top=Mx^\top$. Now the same subspace picture applies:

• In the 12,288-dimensional input space, the row space of $M$ contains the directions it reads. Its null space contains directions that contribute zero to the query. Adding a vector from this null space to the input leaves the query unchanged.
• In the 128-dimensional output space, the column space of $M$ contains every query the linear map can produce. It may fill all of $\mathbb R^{128}$ or be a smaller subspace if the map’s rank is below 128. See column space: which outputs can the matrix reach?.

The query map can read at most 128 independent input directions in this example, leaving at least 12,160 independent input directions in its null space. Those directions are ignored by this query map; the residual stream still carries the input, and other heads or projections can use information this map ignores. The input directions read by Wq and Wk can also differ, even though their outputs are compared using the same 128 matching coordinates.

The 128-dimensional matching space is therefore a separate output space, while the relevant input subspace lies inside the original 12,288-dimensional space. In attention, “projection” usually means a learned linear map: it can stretch and mix directions and is not generally an orthogonal projection onto a plane inside the original space.

What information does the query map lose?

Make this concrete with a separate 3D → 2D example. Using the column-vector notation above, choose:

This matrix reads the two sums $x_1+x_3$ and $x_2+x_3$. For the vector $n$, both sums cancel:

So $n$ is in the null space. Adding any multiple of it leaves the query unchanged: $M(e+tn)=Me+tMn=Me$. For example:

The query map cannot distinguish these two inputs. What it loses is the amount of input along the null-space direction $(1,1,-1)$: increasing the first two coordinates together while decreasing the third by the same amount leaves both output sums unchanged. It does not simply discard the third coordinate; that coordinate contributes to both sums. The arithmetic identifies an ignored direction; identifying any semantic information encoded along it would require experiments on a trained model.

If the keys are held fixed, these identical queries also produce identical attention scores and weights. The input change could still affect this token’s key or value, another head, or the residual stream. “Lost” here means unavailable through this query map, rather than erased from the model’s entire representation.

One query can combine several matching criteria

One query can carry several matching preferences at once. If queries and keys have 128 coordinates, their scaled dot product is:

All 128 coordinate pairs contribute to one compatibility score for this query/key pair. The query is compared with every allowed key, producing a separate score for each source position. Softmax converts those scores into one attention distribution. This is the scaled dot-product computation in the Transformer paper, §3.2.1.

128 coordinates does not mean 128 semantic features. A feature might correspond to a direction spread across many coordinates, and several features can use overlapping coordinates. Ten features could be distributed across 128 dimensions. A larger collection, such as 1,000 possible features, could in principle use overlapping directions, especially when only a few features are active at once. But 128 dimensions provide at most 128 linearly independent directions: they cannot independently preserve 1,000 arbitrary feature strengths through a linear encoding.

Toy Models of Superposition demonstrates how toy networks can represent more sparse features than dimensions, with interference between overlapping representations. This supports a representational possibility; it does not establish that a particular query contains 1,000 identifiable features or tell us what they are.

The dot product combines the evidence into one score, rather than returning a separate answer for every feature. Each head produces one attention distribution and one value mixture per receiving token. Multiple heads let different matching criteria produce separate distributions and mixtures; a single head can still combine several criteria within its own score.

Try three hypothetical matching directions

Return to the four-coordinate bank vectors, this time projecting into a three-dimensional query/key space. We hand-name its axes finance, water, and structure, and keep the key projection fixed. This is a new teaching construction, not a claim about the coordinates or behavior of a trained head.

Choose Find water evidence, Find finance evidence, or Find structure evidence to try a possible interpretation of bank’s query. Each preset changes the displayed Wq; the widget computes the resulting queries, keys, scores, and attention weights. Use Rotate view to inspect the three axes, or open Adjust Wq coefficients to combine or oppose directions.

Notice that the query need not identify a word’s intended sense on its own. In this construction, bank starts with the same input and produces the same query in both sentences. The available keys and values supply the differing context. Whether any of these hypothetical matching patterns helps a task is something training and evaluation must establish.

The embedding table and the projection matrices are separate

The embedding table E stores one starting vector per vocabulary token. The projection matrices Wq, Wk, and Wv transform a position’s current representation into its query, key, and value. Each head normally has its own projections in each transformer block; the token embedding table supplies the starting representations below the stack.

How many tables and vectors are there?

The architecture below expands 2 heads processing 3 tokens in one standard self-attention sublayer. The embedding table is outside attention. Its selected rows become X; both heads read that same full input, using separate learned projection matrices. Within each head, Q, K, and V are computed matrices with one vector per token row. Follow the highlighted bank rows to count its queries, keys, and values across the two heads.

For n token positions and h heads in one standard self-attention sublayer:

So one token has h queries, h keys, and h values in a block, while the block has 3h learned Q/K/V projection matrices and one output projection. A stack of N such blocks normally has its own projections in every block: 3Nh Q/K/V projection matrices and N output projections. These counts describe the standard multi-head arrangement, omit biases, and count logical matrices; implementations often pack several projections into a single parameter tensor.

Only the lookup is indexed by token ID. The projection matrices have input-coordinate rows and output-coordinate columns, so every token uses the same projection matrices within a head. Later blocks receive updated representations from the previous block, without repeating the embedding lookup. The block article adds the normalization and feed-forward details to this outline.

“Initial” describes where the lookup happens when processing a sentence. Before training from scratch, the embedding table usually starts with random numbers. Training learns those vectors jointly with the projections and the rest of the model; there is usually no separate stage that teaches the embeddings a “generic meaning” first. After training, the lookup retrieves a learned vector. It is context-independent because it does not inspect the current sentence. The same token ID retrieves the same row, even though later representations depend on context.

Our first examples use hand-set vectors so we can follow the arithmetic. The training example below learns the embedding table as well as the projections.

Three roles: matching and carrying information

The query is one of three transformed views:

Queries and keys meet in Q @ K.T. Both sides can be learned: the model can transform how a receiver asks and how a source is found. The two projections need not be equal, so the score from i to j can differ from the score from j to i.

Values have a different job. Keep the queries and keys fixed and swap only the value coordinates:

A, H = attention(X2, swap, I2, swap)
print(A[-1].round(4))  # [0.2483 0.5035 0.2483] -- unchanged
print(H[-1].round(4))  # [0.5035 0.2483]        -- different content

Q and K control the mixture’s weights; V controls its content. A large weight on Eiffel does not by itself tell us which direction the output moves. We must also inspect Eiffel’s value. Our swap is hand-set to make that independence visible; training will learn the transformations together.

The full head has the familiar form $H=\operatorname{softmax}(QK^\top/\sqrt{d_k})V$, as specified in the original Transformer paper, §3.2. The diagram follows its operations and dimensions, then shows the output projection used by the surrounding block:

d_v need not equal d. A head returns a mixture in value coordinates. The block’s output projection will map it into the coordinates of the representation we are updating.

Try the two-dimensional example below. Start with Raw query, switch to Swap query, then toggle Swap values. The numbers are computed from the matrices above; changing values leaves the weights untouched.

What the weight actually measures

A score compares one query/key pair. A weight compares that score with every other allowed score in the same row:

Adding another key can change a_ij without changing q_i or k_j. Likewise the weight is not a runtime word-co-occurrence count. Training changes the parameters; inference computes fresh scores and mixtures using those parameters. Position can also affect scores when positional information is supplied—we have deliberately left it out of this toy.

Dot products, cosine similarity, and scaling

We have used a dot product to compare each query with each key, then divided by $\sqrt{d_k}$. What does that division change? Compare three operations on the same vectors:

A dot product includes vector length

For nonzero vectors, the dot product can be written as:

It depends on both lengths and the angle between them. Take $q=(1,0)$ and $k=(1,0)$: their dot product is $1$. Make both vectors ten times longer, giving $q=(10,0)$ and $k=(10,0)$, and the dot product becomes $100$. The directions have not changed.

Hold the angle fixed in the explorer and change either length. Then hold both lengths fixed and change the angle. The vectors are labelled $a$ and $b$ here, but the same arithmetic applies to a query and a key.

Length multiplies the signed angular contribution: making a vector longer makes a positive dot product more positive and a negative one more negative. Perpendicular vectors still give zero. A longer key therefore does not automatically receive a higher score.

Cosine similarity removes the lengths

Dividing by the product of the two lengths cancels those factors:

This is equivalent to taking the dot product after normalizing each vector to unit length. For nonzero vectors it lies between $-1$ and $1$; it is undefined if either vector is zero.

Scaling preserves the influence of lengths

In our head, $d_k$ is the number of coordinates in each query and key. Every pair in that head uses the same divisor $\sqrt{d_k}$, regardless of the vectors’ actual lengths.

For example, a head with $d_k=64$ divides every dot product by $8$:

These are illustrative scores for 64-dimensional vectors, separate from the two-dimensional examples above. Multiplying either vector by ten still multiplies its scaled dot product by ten. Dividing all scores by the same positive constant preserves their ordering and rescales their differences; it does not make the vectors unit length or bound the scores to $[-1,1]$.

Why divide by the square root of the dimension?

A dot product adds one product per coordinate:

To see how its scale can grow, assume all query and key components are independent random variables with mean zero and variance one. Each product then has mean zero and variance one:

Variances add for independent terms, so:

The expected dot product remains zero. What grows is its spread: positive and negative fluctuations become larger. At $d_k=64$, the variance is $64$ and the standard deviation is $8$. Dividing by $8$ brings the standard deviation back to $1$ because dividing a random variable by a constant divides its variance by the constant’s square:

This is the motivation given in the original Transformer paper, §3.2.1. It explains the choice of divisor under simplifying assumptions; learned queries and keys need not remain independent or have unit component variance, so a trained head’s score variance is not guaranteed to be one.

Scaling changes how sharply softmax chooses

Consider one query with raw scores $[4,2,0]$ against three keys in a four-dimensional head. Dividing by $\sqrt{4}=2$ gives $[2,1,0]$:

Key 1 still wins, but receives a smaller share. Softmax responds to differences between scores: for two allowed keys, the weight ratio is $a_1/a_2=\exp(s_1-s_2)$. Here scaling reduces that score gap from $2$ to $1$, changing the ratio from $e^2$ to $e$. Adding the same constant to every score would leave the weights unchanged.

Large score gaps can make softmax nearly one-hot. Its derivatives with respect to its input scores then become small, which can weaken the learning signal through the attention weights. Scaling helps control this effect as query/key dimension increases. It does not prevent a trained head from learning sharp attention when useful.

Why not use cosine similarity instead?

Cosine similarity would change what the matching rule can use. Standard scaled dot-product attention retains the influence of both query/key directions and lengths. With keys fixed, increasing a query’s length by a positive factor scales its whole score row and can sharpen its attention distribution. Changing one key’s length affects that key’s score, with the effect depending on its sign.

Cosine-based attention is a possible alternative, but normalizing queries and keys removes these particular effects of their lengths. The fixed $\sqrt{d_k}$ divisor instead adjusts the expected score scale as the head width changes.

Finally, keep vector normalization separate from weight normalization:

The divisor scales compatibility scores. Softmax turns each score row into nonnegative mixing weights that sum to one. Neither operation requires each query or key to have unit length.

Make it learn

Return to the bank task, but discard our hand-designed axes. The model now learns eight-dimensional embeddings and the three projection matrices. A classifier reads only the mixture at the final bank position and predicts water (0) or finance (1). It is a tiny supervised classifier, not a next-token language model or a complete transformer block.

Follow a sentence from lookup to an update

The widget uses a two-dimensional version of this classifier to keep every vector and parameter visible. Choose a sentence, then use Next stage to follow sentence → embedding lookup → Q/K/V → attention mixture → prediction → loss → gradients → parameter updates. At the gradient stage, the parameters have not changed yet. Apply update takes one SGD step; Next sentence continues with the updated parameters and the other example.

Compare the embedding table with the projection matrices, then compare those stored parameters with the vectors computed from them. The gradient and update stages show all parameter groups together, including the embedding row that received no gradient because its word was absent. Use Train 20 steps to alternate both examples, or Reset training to return to the same random initialization.

The runnable eight-dimensional model

import torch
import torch.nn as nn

vocab = {"the": 0, "river": 1, "money": 2, "bank": 3}
data = [(["the", "river", "bank"], 0),
        (["the", "money", "bank"], 1),
        (["river", "bank"], 0),
        (["money", "bank"], 1)]

def encode(tokens):
    return torch.tensor([vocab[t] for t in tokens])

class OneHeadSense(nn.Module):
    def __init__(self, vocab_size, d=8):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, d)
        self.Wq, self.Wk, self.Wv = (
            nn.Linear(d, d, bias=False) for _ in range(3))
        self.cls = nn.Linear(d, 2)
        self.d = d

    def forward(self, ids, return_attn=False):
        X = self.emb(ids)
        Q, K, V = self.Wq(X), self.Wk(X), self.Wv(X)
        A = torch.softmax(Q @ K.T / self.d**0.5, dim=-1)
        H = A @ V
        logits = self.cls(H[-1])
        return (logits, A) if return_attn else logits

torch.manual_seed(0)
model = OneHeadSense(len(vocab))
opt = torch.optim.Adam(model.parameters(), lr=0.05)
loss_fn = nn.CrossEntropyLoss()
for step in range(300):
    for tokens, label in data:
        logits = model(encode(tokens)).unsqueeze(0)
        loss = loss_fn(logits, torch.tensor([label]))
        opt.zero_grad()
        loss.backward()
        opt.step()

for tokens, _ in data[:2]:
    logits, A = model(encode(tokens), return_attn=True)
    print(tokens, "prediction:", logits.argmax().item())
    print({t: round(w, 2) for t, w in zip(tokens, A[-1].tolist())})

In the checked seed-0 CPU run, the predictions are 0 and 1, and the rounded attention rows are:

{'the': 0.0, 'river': 1.0, 'bank': 0.0}
{'the': 0.0, 'money': 1.0, 'bank': 0.0}

This run learned strongly selective routing. That is an observed solution, not something the four-example task forces: a uniform mixture also contains enough information to classify the examples. The printed 0.0 values are rounded, not proof of exactly zero softmax weights. There is no held-out evaluation here; this demonstrates fitting a tiny task.

What changed during training? The embedding table, Q/K/V projections, and classifier: 4×8 + 3×8×8 + (8×2+2) = 242 parameters. The weights A are activations, recomputed from each input, not another table of trainable parameters. Training adjusts the fixed transformations so their input-dependent mixtures help the classifier lower its loss.

Heads and blocks train together

The toy has one head. In standard full-model transformer training, all attention heads in all transformer blocks train together. A forward pass computes predictions and a loss. Backpropagation carries gradients from that loss through the blocks, through each head’s query/key and value routes, and back to the embedding table. The optimizer then updates all trainable parameter groups, including the output projections and FFNs. Each parameter receives its own gradient; this does not require every scalar to change on every step.

There is no separate training loop for each head. Heads within one block compute in parallel, while blocks pass updated representations forward sequentially. Backpropagation follows those dependencies in reverse. Our loop’s model.parameters() includes the embeddings, three projections, and classifier together; a full transformer’s optimizer likewise includes the parameters across the whole trainable model.

Can two small matrices recover identifiable features?

The standalone feature-recovery experiment trains only two 9 × 3 query and value projections on synthetic data. It tests which known features the maps recover, checks held-out predictions, and intervenes on relevant and irrelevant input directions. Fixed keys and output targets deliberately anchor the feature meanings so they can be identified.

Attention learns the data, not the world

The companion train_biased_dataset.py assigns arbitrary labels to six context words, including opposite labels for lawyer and attorney. In a checked seed-0 run, it fits those assignments:

The result is more interesting than “the model always attends to the clue.” It solves the two classes with different attention patterns. Very small weights, other values, and the classifier all matter; a rounded heatmap alone is not an explanation of the prediction.

The narrow conclusion is that this supervised toy can learn arbitrary label assignments. It does not establish that a language model reproduces every bias in its data or cannot learn useful generalizations. To study those questions we would need a different experiment and evaluation.

Causal masking: a decoder only looks left

So far every position can read every other position. During left-to-right language-model training, the representation used to predict the next token must not see that future token. A causal mask changes which source positions are allowed.

The function already includes the necessary step, just before softmax:

future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
scores = np.where(future, -np.inf, scores)

Exponentiating negative infinity gives zero, so the forbidden positions get zero weight. Setting their scores to zero would not work: exp(0) = 1 would still give them a share. Every row retains at least its own position.

The toggle below compares the resulting access patterns for another sentence:

For “the river bank,” the last row is unchanged: bank can already see the entire sequence. For “the bank river,” causal bank cannot receive anything from the future river position. The source restriction changed; the projection-and-mixing mechanism did not.

A mask specifies access, not a full positional encoding. Our position-free toy still has limitations; the block article explains where positional information enters. For generation, previously computed keys and values can be reused under causal attention—the KV-cache article follows that consequence. Cross-attention instead takes its queries and its key/value pairs from different sequences.

Watch real heads

The hand-set axes made arithmetic readable. In a trained model, inspect both where information is gathered and what contribution is produced.

BertViz visualizes attention patterns head by head. AttentionViz additionally supports exploring query/key geometry and patterns across inputs. exBERT is another exploration interface. These are companion tools, not evidence that a particular head must reproduce our toy.

Try two sentences containing different uses of bank. Inspect several heads and layers, then compare the corresponding hidden representations. A heatmap displays mixing coefficients; the values and output projection determine the vector actually contributed. Treat a proposed role such as “disambiguation head” as a hypothesis to test across inputs and with interventions.

Where this leaves us

We built a contribution in three stages: match, weight, mix. The initial bank experiment showed information flowing through a weighted sum. The Eiffel example separated matching from raw input similarity, and changing values showed that the same weights can carry different content. Training supplied the parameters instead of our hand choosing them.

The next step is to incorporate that contribution. In the simplest same-width case with identity output projection, tower = [1, 0] plus its swap-query mixture [0.2483, 0.5035] gives [1.2483, 0.5035]. The mixture is one contribution to an existing representation. A real block handles the output projection, multiple heads, normalization, and another update from the FFN.

Inside a transformer block follows those updates through the complete computation.

Optional: the mathematics behind the head

Theoretical and mathematical foundations

Reverse-engineer a head you trained