Attention¶

The Problem¶

The model has a 16-dimensional vector x for the current token. But language depends on context: the right next character depends on what came before. The letter after "qua" is likely "r" or "l" — but the model is processing one token at a time.

How does the current token look back at previous tokens to gather context?

The Intuition: Library Lookup¶

The Code Setup (Lines 117–122)¶

x = rmsnorm(x)
q = linear(x, state_dict[f'layer{li}.attn_wq'])    # query: what am I looking for?
k = linear(x, state_dict[f'layer{li}.attn_wk'])    # key: what do I contain?
v = linear(x, state_dict[f'layer{li}.attn_wv'])    # value: what's my information?
keys[li].append(k)       # store key for future tokens to see
values[li].append(v)     # store value for future tokens to see

Each token produces:

• A query \(\mathbf{q}\) (16-dim): "What information do I need?"
• A key \(\mathbf{k}\) (16-dim): "Here's what I contain."
• A value \(\mathbf{v}\) (16-dim): "Here's my actual information."

All three are produced by different linear transformations of the same input x.

Scaled Dot-Product Attention (Lines 127–131)¶

attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5
               for t in range(len(k_h))]
attn_weights = softmax(attn_logits)
head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h)))
            for j in range(head_dim)]

Step 1: Compute Attention Scores¶

The dot product measures similarity between the current query and each past key. Divide by \(\sqrt{d_k}\) (the scaling factor) to keep values from getting too large.

Step 2: Convert to Weights (Softmax)¶

Step 3: Weighted Sum of Values¶

Blend the value vectors, weighted by how relevant each past token is.

The KV Cache¶

keys[li].append(k)
values[li].append(v)

Causal Masking (Built-In)¶

Notice we only attend to past tokens — the ones already in the keys list. The current token can't see the future because future keys haven't been appended yet. This is causal masking achieved by the sequential processing order, rather than an explicit mask matrix.

Attention as a Flow¶

flowchart TD
    X["x (current token)"] --> Q["Q = linear(x, Wq)"]
    X --> K["K = linear(x, Wk)"]
    X --> V["V = linear(x, Wv)"]
    K --> CACHE["KV Cache<br>(all past K,V)"]
    V --> CACHE
    Q --> SCORE["Score = Q · Kᵢ / √d"]
    CACHE --> SCORE
    SCORE --> SM["Softmax → weights α"]
    SM --> BLEND["Output = Σ αᵢ · Vᵢ"]
    CACHE --> BLEND