From recurrence to attention

RNNs process a sequence one step at a time, which makes it hard for information from far-back tokens to influence the current step (it has to survive being squeezed through the hidden state at every intermediate timestep). Attention takes a completely different approach: every position can directly look at — and pull information from — every other position in a single operation, regardless of distance.

Queries, keys, and values

Attention reframes "look up relevant information" as a soft, differentiable dictionary lookup:

• Query (Q) — "what am I looking for?" — one vector per position.
• Key (K) — "what do I contain?" — one vector per position, compared against every query.
• Value (V) — "what do I actually return if selected?" — the payload that gets mixed together based on how well queries matched keys.

The full formula:

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

Step by step

1. Similarity scores. Q @ K^T computes, for every pair of positions (i, j), the dot product between query i and key j — a raw measure of how relevant position j is to position i. This produces a T × T matrix of scores.

2. Scaling. Divide every score by √d_k. Without this, as d_k grows, dot products tend to grow large in magnitude (more terms summed), which pushes softmax into a regime where gradients vanish (one input dominates completely, everything else gets essentially zero weight). Scaling keeps the scores in a range where softmax stays well-behaved regardless of dimensionality.

3. Softmax. Turn each row of scores into a probability distribution over the other positions — the classic softmax(x)_i = exp(x_i - max(x)) / Σ exp(x_j - max(x)) (subtracting the row max before exponentiating avoids overflow, and doesn't change the result since it cancels in the ratio).

4. Weighted sum of values. Multiply the T × T softmax weights by V (T × d_v) — each output position becomes a weighted blend of every value vector, weighted by how much attention that position paid to each source position.

Causal masking

For autoregressive models (predicting the next token), a position must never attend to positions that come after it — that would leak future information the model wouldn't actually have at generation time. Causal masking enforces this by setting the score for any (i, j) pair with j > i to -infinity before the softmax step — exp(-inf) = 0, so those positions get exactly zero attention weight after normalization, without needing any special-casing in the weighted-sum step.

if causal and j > i:
    score[i][j] = -infinity

Apply the mask row by row, right after computing raw scores and before dividing by √d_k or running softmax — the order (mask, then scale, then softmax) is a common source of subtle bugs if reversed carelessly, though masking before or after scaling produces the same result since -inf scaled by a positive constant is still -inf. What must not be reordered is doing the mask after softmax — by then the invalid positions already have nonzero probability mass baked in.

Putting it together

Your exercise builds both SCORES (the intermediate softmax weight matrix — useful for inspecting what the model is attending to) and ATTN (the final output after applying those weights to V). Getting the shapes right matters: Q/K share d_k, V can have a different dimension d_v, and the output of attention always has shape T × d_v — the same shape as V, just recombined according to the learned attention pattern.