Self-Attention: Scores, Allowed Positions, Weighted Values

Course role: v1 is a conceptual bridge with compact scalar exercises. This self-attention exercise implements one head. For cumulative learned multi-head projections, numerical boundaries, and end-to-end integration, continue with deep learning v2. Other v1 capstone code is a separate simplified implementation, not an import of this exercise.

What you will learn

Trace one row from query/key dot products to a weighted value, preserve gradients through stable softmax, and enforce the causal rule even when a future score is enormous. Prerequisites are embeddings, dot products, and the scalar Value computation graph.

1) Separate the three roles and their shapes

Each query asks which positions are relevant. Each key provides a vector to compare with the query. The corresponding value contains the information mixed into the output. Scores choose the mixture; values supply its contents.

For sequence length T and embedding width D, x has shape (T,D). The module projects every row through three Linear(D,D) layers to obtain Q,K,V, each (T,D). The score matrix is (T,T): row i is query position i, column j is key/value position j. The helper permits value width d_v to differ from query/key width d_k.

scores[i][j] = dot(Q[i],K[j]) / sqrt(d_k)
weights[i]  = softmax(allowed scores in row i)
attended[i] = sum_j weights[i][j] * V[j]
output[i]   = W_o(attended[i])

apply_attention returns the attended values; SelfAttention.forward also applies the supplied output projection. They are not interchangeable stages.

2) Work a complete two-position example

Set d_k=1, Q=[[0],[1]], K=[[0],[log(3)]], and V=[[2],[6]]. These are already-projected vectors; take the output projection as identity for this trace.

Query row	Raw scaled scores	Unmasked weights	Unmasked attended value
0	[0,0]	[0.5,0.5]	0.5*2+0.5*6=4
1	[0,log(3)]	[0.25,0.75]	0.25*2+0.75*6=5

With a causal rule, row 0 is allowed to use only position 0, so its full weight row is [1,0] and output is 2. Row 1 may use both positions and still outputs 5. Equal weights average values; weights need not select just one position.

Predict before moving on: if only V[1] changes from 6 to 10, the causal row-0 output remains 2, and row 1 becomes 0.25*2+0.75*10=8. This is a behavioral test of causality, not just a test of matrix shape.

3) Stable softmax with a computation graph

For finite allowed scores, choose m=max(s.data for s in scores). Compute exponentials of s-m with Value operations, sum them, and divide. For [1000,1001,1002] the shifted scores are [-2,-1,0]: no positive exponent can overflow, and at least one exponential is one. Common shifts cancel in the normalized ratio.

Read .data only to choose this numerical shift. Constructing Value(math.exp(s.data-m)) would detach each exponential from its score and break learning. (s-m).exp() preserves the graph. Treating the common maximum as a constant leaves the correct softmax derivative because softmax is invariant to a common shift. Small exponentials may underflow to zero.

As a gradient check, scores [0,0] and scalar values [2,6] give output 4. The derivative with respect to the first score is 0.5*(2-4)=-1, and with respect to the second it is 0.5*(6-4)=1. Nonzero gradients alone are a weaker test than these expected values.

4) Causality is exclusion before normalization

For autoregressive prediction, row i may use exactly j<=i. The exercise keeps a demonstration helper causal_mask(T) whose entries are 0.0 for allowed positions and -1e9 for future positions. Adding that finite bias is not a general implementation of the causal guarantee.

For scores [0,1e12], adding [0,-1e9] leaves [0,999000000000]; stable softmax still places essentially all weight on the future. Numerical stability does not imply causal correctness.

SelfAttention.forward(causal=True) instead normalizes only the allowed prefix and appends constant zero weights:

# scores is (T,T), with Value entries
weights = []
for i, row in enumerate(scores):
    allowed = softmax(row[:i + 1])
    weights.append(allowed + [Value(0.0) for _ in row[i + 1:]])

Every prefix contains the current position, so no row is fully masked. With causal=False, normalize the entire row. This preserves the finite-mask helper API for comparison while giving the actual attention module exact exclusion for finite scores.

Zeroing weights after full-row softmax is wrong without renormalization: on scores [0,0], row 0 would retain weight 0.5, giving value 1 instead of 2. It also leaves the normalization dependent on future scores. Exclude first.

5) Verify dependencies, not just shapes

With identity scalar projections, use inputs [[1],[1e12]]. The first causal output must be 1, regardless of the enormous second input. Backpropagating from that first output gives input gradients [1,0]: its current value matters, its future input does not. Future score entries may be computed in the dense matrix, but they are absent from the normalization path.

This tests three distinct properties: forbidden values cannot affect the output, allowed weights still sum to one, and no gradient from an earlier output flows to a later input. When checking gradients, start with fresh Value objects or reset accumulated gradients.

6) Multi-head attention and scope

Standard multi-head attention learns projections of the complete input for each head, arranges projected features into heads, mixes independently, concatenates, and applies an output projection. Splitting raw input coordinates before projection restricts what each head can see. This exercise implements a single head, not that full architecture. The original transformer paper specifies the multi-head construction.

Production tensor libraries support exact mask semantics; for example, PyTorch's attention documentation distinguishes additive biases from excluded positions. This scalar engine uses allowed-prefix normalization to avoid introducing infinities into its tiny autograd implementation.

7) Cost and next steps

Score computation and value mixing take O(T^2*D) time; dense projections add O(T*D^2). Intermediate storage is O(T^2+T*D), excluding parameters. Doubling length quadruples the dense score count, though a causal row uses only its prefix.

Before attempting the exercise, reproduce the two-row output table, then explain the finite-mask counterexample and the [1,0] gradient expectation. The checkpoint asks for transfer to new values. V2 extends this primitive into a cumulative implementation with learned multi-head projections and stronger integration checks.