Self-Attention

Goal

Turn one sequence into context-aware token representations.

Shape Contract

• input x: (T, D)
• projected Q, K, V: each (T, D)
• scores: (T, T)
• weights: (T, T)
• output: (T, D)

Attention Equation

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

Row view:

• score row i compares token i to all tokens
• softmax row i is a probability distribution
• output row i is weighted sum of all V[j]

Causal Mask (Autoregressive)

• allow j <= i
• block j > i with -1e9
• add mask to scores before softmax

Numerical Stability

Stable softmax per row:

• m = max(row)
• exp(row - m)
• normalize by row sum

Module Build Order

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

Complexity

• time: O(T^2 * D)
• memory: O(T^2)

Failure Modes

• wrong softmax axis
• mask applied after softmax
• missing / sqrt(d_k)
• breaking Value graph by raw float ops