Self-Attention

Self-attention is the mechanism that lets each token decide which previous tokens matter most for its next representation.

This module turns the transformer core equation into implementation steps you can debug.

Learning goals

By the end, you should be able to:

• reason about all attention tensor shapes
• compute scaled dot-product attention row by row
• apply causal masking correctly for autoregressive language models
• implement numerically stable softmax
• map the equation to concrete functions and classes

Notation and shapes

We use:

• T: sequence length
• D: embedding dimension (d_k = d_v = D in this module)

Input:

• x: (T, D)

Learned projections:

• W_q, W_k, W_v, W_o: each Linear(D, D)

Projected sequences:

• Q = W_q(x): (T, D)
• K = W_k(x): (T, D)
• V = W_v(x): (T, D)

Core equation

For each token position i:

1. Compare query Q[i] to every key K[j].
2. Normalize those scores into probabilities.
3. Take a weighted sum of value vectors V[j].

Elementwise form:

score[i, j] = (Q[i] dot K[j]) / sqrt(D)
weights[i, :] = softmax(score[i, :])
attended[i] = sum_j weights[i, j] * V[j]
output[i] = W_o(attended[i])

Matrix form:

Attention(Q, K, V) = softmax((Q K^T) / sqrt(D)) V

Why scaling by sqrt(D)

Without scaling, dot products grow with dimension and push softmax into extreme values too early, making training unstable.

Dividing by sqrt(D) keeps score magnitudes in a healthier range.

Causal masking

In autoregressive language modeling, token i cannot look at future tokens j > i.

Mask rule:

mask[i, j] = 0.0      if j <= i
mask[i, j] = -1e9     if j > i

Apply mask before softmax:

masked_score[i, j] = score[i, j] + mask[i, j]

-1e9 drives future-position probability to approximately zero after softmax.

Softmax stability (non-negotiable)

Naive exp(score) overflows for large scores.

Stable softmax:

m = max(scores)
probs = exp(scores - m) / sum(exp(scores - m))

Subtracting the row maximum does not change probabilities, only numerical range.

Worked micro-trace (shape-first)

Assume T = 3, D = 2.

After projections:

• Q: 3 vectors of length 2
• K: 3 vectors of length 2
• V: 3 vectors of length 2

For token i = 1:

1. Compute three dot products with keys K[0], K[1], K[2].
2. Divide each by sqrt(2).
3. If causal:
   • keep positions 0 and 1
   • mask position 2 with -1e9
4. Softmax the 3 scores -> probabilities summing to 1.
5. Weighted sum of the three V[j] vectors -> one length-2 vector.
6. Apply W_o -> final output vector at position 1.

Repeat for each token position.

Implementation map for this pack

You implement attention in five pieces:

1. softmax(scores)
2. attention_scores(Q, K, d_k)
3. causal_mask(seq_len)
4. apply_attention(weights, V)
5. SelfAttention.forward(x, causal=True)

Think of the flow as:

x -> Q/K/V -> scores -> (+mask) -> weights -> attended -> output

Complexity

For sequence length T and dim D:

• score matrix construction: O(T^2 * D)
• weight application: O(T^2 * D)
• memory for scores/weights: O(T^2)

This quadratic T^2 cost is why long-context attention is expensive.

Common bugs and fixes

1. Wrong softmax axis

• Bug: softmax over all matrix entries.
• Fix: apply softmax per row scores[i, :].

1. Mask applied after softmax

• Bug: future tokens still get probability mass.
• Fix: add mask to scores before softmax.

1. Missing scaling

• Bug: peaky/unreliable distributions.
• Fix: divide by sqrt(d_k) in attention_scores.

1. Breaking autograd graph

• Bug: converting too much to raw floats.
• Fix: keep core math in Value ops; only use .data for max selection.

1. Shape mismatch

• Bug: transposed assumptions in dot products or weighted sums.
• Fix: re-check each function contract with (T, D) expectations.

Debugging checklist

Before moving to transformer blocks, verify:

• each softmax row sums to approximately 1.0
• causal mask has large negatives above diagonal
• attention_scores returns (T, T)
• apply_attention returns (T, D)
• SelfAttention.forward output shape matches input shape (T, D)

If these hold, your attention primitive is ready for multi-head composition.