Module 2: Queries, keys, and one carefully chosen constant

What exactly is broken

Module 1 ended with two measured facts:

1. Self-focus. In Module 1's weight tables, every word's largest share was itself: "the" kept 57.6% of its attention, and bank kept 49.8% in one sentence and 52.5% in the other. The cause is structural: scores come from X @ X.T, so word i's score against itself is x_i . x_i, a sum of squares that another word can essentially never beat (Module 0, quiz question 3).
2. Symmetry. The score matrix always equals its own mirror image: S[i,j] and S[j,i] are the same dot product ("river looking at bank" and "bank looking at river" were both 0.540). Language is not like that. In "the river bank", "river" is enormously informative about "bank", while "bank" tells you little about "river".

Both facts have one root: in X @ X.T, a word's single embedding vector plays both roles, the one asking and the one being asked about. Whatever a word "looks for" is forced to be its own identity.

The fix: separate asking from advertising

Picture each word at a networking event. It wears a badge saying what it is (for "river": "I carry topic content"). And it carries a question about what it needs (for "bank": "I am an ambiguous place word; which topic am I in?"). Matching happens between one word's question and another word's badge, not between their identities. The question and the badge can be completely different from each other, and from the word's own content. That separation is the entire idea of queries and keys:

• the query $q_i$: what word $i$ is looking for,
• the key $k_j$: what word $j$ advertises,
• the score becomes $q_i·k_j$: how well $j$'s advertisement answers $i$'s question.

Where do queries and keys come from? Each is a linear remix of the word's embedding. Multiplying the embedding (4 numbers) by a 4x2 matrix produces a query with 2 numbers: each column of the matrix is a recipe saying how much of each embedding axis to pour into that slot of the query. Real models learn these matrices during training. We will instead wire them by hand, so that every number that comes out is explainable, and we only need two output slots (call them channels) to make the point.

Here is the wiring, exactly as the script defines it, comments included:

# Channel 1 is "topic": place words ASK for topic words on this channel
# (W_Q), and watery/money words ADVERTISE topic content on it (W_K).
# Channel 2 is "place": filler words ASK for place words, and place words
# ADVERTISE their place-ness. Rows are input axes (water, money, place,
# filler); columns are the two channels.
W_Q_WIRED = np.array([
    [0.0, 0.0],   # water-ness asks for nothing
    [0.0, 0.0],   # money-ness asks for nothing
    [2.0, 0.0],   # place words ask for TOPIC words   (channel 1)
    [0.0, 1.5],   # filler words ask for PLACE words  (channel 2)
])
W_K_WIRED = np.array([
    [1.0, 0.0],   # water-ness advertises "I carry topic"  (channel 1)
    [1.0, 0.0],   # money-ness advertises "I carry topic"  (channel 1)
    [0.0, 1.0],   # place-ness advertises "I am a place"   (channel 2)
    [0.0, 0.0],   # filler advertises nothing
])

Read W_Q_WIRED row by row (rows are the embedding axes): the only ways to acquire a question are to have place-ness (row 3: place words ask for topic, with gain 2.0) or filler-ness (row 4: filler words ask for places, gain 1.5). And in W_K_WIRED: the ways to earn an advertisement are water-ness or money-ness (both advertise "topic") or place-ness (advertises "place"). Asking and advertising now run on different axes: a word can seek something it does not itself have.

Watch the wired head think

The script computes Q = X @ W_Q_WIRED and K = X @ W_K_WIRED for "the river bank"; here is every word's query and key (cross-check: results.log lines 54-57):

  queries and keys for "the river bank" (2 channels: topic, place):
    the     q = [ 0.00,  1.50]   k = [ 0.00,  0.00]
    river   q = [ 0.60,  0.00]   k = [ 0.90,  0.30]
    bank    q = [ 1.80,  0.00]   k = [ 0.60,  0.90]

Let's go deeper and build two of these by hand from the wiring above, because every row of that table is just one embedding pushed through one matrix. The inputs are the embedding rows: the [0, 0, 0, 1], river [0.9, 0, 0.3, 0], bank [0.3, 0.3, 0.9, 0].

A key is x @ W_K_WIRED: each axis of the embedding flows through its own row of the wiring, and the partial results add up, exactly like Module 1's slow-motion blend. River's key, axis by axis:

River's badge reads: "I carry topic content (0.90, from my water-ness), and I am a little place-like (0.30, from my own place-ness)." Note which path each number took: river's place badge comes from river's place axis, not from its water axis. The wiring routes each quality uniformly for every word; facts about an individual word stay in that word's embedding row.

A query is the same recipe through the other matrix. Bank's, axis by axis:

Bank asks, loudly, for topic words because it is a strongly place-flavored word: exactly the "ambiguous place word seeking context" behavior we wanted. The other four rows of the log check the same way; two are worth reading aloud. Bank's key is [0.60, 0.90]: its water and money each advertise as topic (0.3 + 0.3 = 0.60) and its 0.9 place-ness is its place badge. And "the", pure filler, gets its whole query from the filler row of W_Q_WIRED (1.0 x [0.0, 1.5] = [0.00, 1.50], asking for places) and its whole key from the filler row of W_K_WIRED (1.0 x [0.0, 0.0] = [0.00, 0.00]): there is no badge for pure filler.

Scoring is now every query against every key (cross-check: results.log lines 58-63):

  score matrix Q @ K.T:
                the   river    bank
    the       0.000   0.450   1.350
    river     0.000   0.540   0.360
    bank      0.000   1.620   1.080
  symmetric? False  (bank->river = 1.620, river->bank = 0.360)

Both Module 1 shortfalls are visibly dead. Symmetry is gone: bank gives river a score of 1.620 (bank's topic-question of 1.8 meets river's topic-badge of 0.9), while river gives bank only 0.360; the capture's symmetric? False line is the run flagging exactly that. And self-focus is gone: the biggest entry in bank's row is now "river" (1.620), not itself (1.080). Bank's question is not "who looks like me?" anymore; it is "who carries topic?", and it does not itself carry much.

The weights, after a division we will justify in a moment (cross-check: results.log lines 64-71):

  weights after dividing by sqrt(2) and softmaxing each row:
                the   river    bank
    the       0.201   0.276   0.522
    river     0.266   0.390   0.344
    bank      0.159   0.500   0.341
  bank's row,  raw vs wired: [0.185, 0.317, 0.498] -> [0.159, 0.500, 0.341]
  'the's row,  raw vs wired: [0.576, 0.212, 0.212] -> [0.201, 0.276, 0.522]

Let's go deeper on the two comparison lines, because they are the module's headline. Under raw attention, bank kept 49.8% of itself and gave "river" 31.7%; under the wired head it takes half its update from "river" (50.0%) and keeps only 34.1%. And "the", which raw attention left staring at itself (57.6%), now spends 52.2% of its attention on "bank": the place word its question was wired to find. The same two matrices, four numbers each, produced both behavior changes.

Values: the third face, and the full pipeline

One more role remains. Once bank decides to take 50% of its update from "river", what does river hand over? In the wired head, the answer is river's raw embedding: the blend works exactly as in Module 1, only with better weights. Here is the head, as the script defines it:

def wired_attention(X):
    """Attention with the hand-wired W_Q and W_K (values stay raw).

    Keeping the values equal to the raw embeddings (formally, W_V is the
    identity matrix) is a hand-wiring choice: it keeps every output
    readable on our four named axes while Q and K do the interesting work.
    """
    # Each word's question, and each word's advertisement.
    Q = X @ W_Q_WIRED
    K = X @ W_K_WIRED

    # scores[i, j] = how well word j's advertisement answers word i's
    # question. No longer symmetric: asking and advertising differ.
    scores = Q @ K.T

    # Scale into softmax's responsive range (measured in [S5]).
    d_k = W_K_WIRED.shape[1]
    weights = softmax(scores / np.sqrt(d_k))

    return weights @ X, weights, scores

In a real head, the hand-over is a third linear remix, the value: $V=X{W}_V$, with its own learned matrix. What a word contributes to the blend can then differ from its identity too (a word might advertise "topic" but contribute only its topic-relevant content). Our wired head simply chooses "contribute your whole embedding", which is W_V = identity, so the outputs stay readable on the named axes. The payoff: here are bank's three lives in the river sentence, side by side on the named axes:

Read the water column downward: balanced, then leaning (3 to 1 over money), then decisive (better than 5 to 1). Same blend machinery, better weights, stronger disambiguation (cross-check: results.log lines 46 and 72). To see where each drop comes from, rebuild Module 1's slow-motion table with the wired recipe [0.159, 0.500, 0.341] (bank's row of the wired weight table above) in place of the raw one, and the bottom row falls out.

The complete pipeline, which is also exactly the shape of the code above:

Follow the boxes: the sentence becomes X (one embedding row per word); three remixes produce Q (queries), K (keys), and V (values); queries meet keys in scores = Q @ K.T, which is then divided by sqrt(d_k) (where $d_k$ is the length of each query and key vector, 2 in this head); softmax turns each row of scores into shares; and the blend edge carries the values into the final weighted average, output = weights @ V, one new context-aware row per word. In symbols, this whole diagram is the formula you will find in every paper and textbook:

$$\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}(Q,K,V)=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

You can now read it the way you read the code: $Q{K}^{\top }$ is "every question meets every advertisement", the division keeps softmax sane (next section), softmax makes each row a recipe, and multiplying by $V$ executes the blend.

The constant: why divide by sqrt(d_k)?

One piece of the formula is unexplained: the division by $\sqrt{d_k}$, where $d_k$ is just the length of each query and key vector (2 in the wired head, 64 in GPT-2). The usual one-line justification is "it stops the scores from getting too big", which is true but vague, and this is a workshop, so we measure precisely what happens.

First, why would scores grow with $d_k$ at all? A dot product of length-$d_k$ vectors is a sum of $d_k$ terms. When queries and keys come from learned matrices applied to arbitrary inputs, those terms behave like random numbers of either sign: some cancel, some reinforce. The sum's typical size therefore grows, but slower than $d_k$ itself; the classic random-walk result says it grows like $\sqrt{d_k}$. We do not ask you to take that on faith. The script draws random queries and keys at four sizes and measures the spread of the resulting scores, along with what that spread does to softmax; here is its inner loop, the part that writes the capture below:

    stds, max_unscaled, max_scaled = [], [], []
    for d in dims:
        dots, mx_raw, mx_scl = [], [], []
        for _ in range(TRIALS):
            # One query meeting a row of 8 keys, exactly like one row of
            # the score matrix in a real head of dimension d.
            q = rng.standard_normal(d)
            keys = rng.standard_normal((ROW, d))
            scores = keys @ q

            dots.extend(scores.tolist())
            mx_raw.append(softmax(scores).max())
            mx_scl.append(softmax(scores / np.sqrt(d)).max())

Predict. Apply Module 0's lesson before looking. If typical scores really grow like $\sqrt{d_k}$, they arrive 8 times larger at d = 64 and 16 times larger at d = 256, and the "x 5" demo already pushed a top share from 0.498 to 0.899. What should the average top share do as d grows without the division, and what should dividing by $\sqrt{d_k}$ do to it?

The measurements (cross-check: results.log lines 75-80):

[S5] why scores are divided by sqrt(d_k)
  random q, k entries drawn from N(0,1); 8 keys per row; 2000 rows per size
  d =   4: dot-product std =   1.97  (sqrt(d) =  2.00) | avg top weight: unscaled = 0.522, scaled = 0.338
  d =  16: dot-product std =   4.02  (sqrt(d) =  4.00) | avg top weight: unscaled = 0.750, scaled = 0.356
  d =  64: dot-product std =   8.01  (sqrt(d) =  8.00) | avg top weight: unscaled = 0.879, scaled = 0.362
  d = 256: dot-product std =  16.03  (sqrt(d) = 16.00) | avg top weight: unscaled = 0.941, scaled = 0.366

Let's go deeper, column by column. The std column confirms the random-walk story with striking precision: the measured spread of the scores is 1.97 at $d=4$ (prediction: 2.00), 8.01 at $d=64$ (prediction: 8.00), 16.03 at $d=256$ (prediction: 16.00). Typical scores really do grow like $\sqrt{d_k}$.

The unscaled column shows what that does downstream. Module 0 measured that inflating scores sharpens softmax (the "x 5" demo: top share went from 0.498 to 0.899). At $d=256$, scores arrive 16 times larger than at $d=1$, and the average largest share in an 8-word row has climbed to 0.941, heading for 1.0: one word soaks up everything. That is winner-take-all attention: the blend stops blending, context reduces to a single word, and (a practical aside beyond this workshop's scope) a softmax pinned at its extreme is also notoriously hard to train. In the figure, that is the red curve climbing toward 1.0.

The scaled column is the fix working: divide the scores by $\sqrt{d_k}$, the exact factor by which they grew, and the average top weight stays essentially flat (0.338 to 0.366, the blue curve) no matter the dimension. The division is not a hack; it is cancelling a measured growth law so that a head behaves the same at $d_k=2$ as at $d_k=256$. That is the entire story of the "magic" constant in the formula, and of the / np.sqrt(d_k) in the code.

Next: Module 3: The real thing