Self-Attention: Trace One Row

• Queries and keys choose weights; values supply the information mixed.
• Score row i belongs to query i; column j belongs to key/value j.
• This v1 exercise implements one head with scalar autograd.

Shapes and projections

• x,Q,K,V: (T,D) in the module.
• scores: (T,T) with divisor sqrt(d_k).
• apply_attention: weights (T,T) times values (T,d_v).
• SelfAttention also applies its output projection.

Two-position trace

Q=[[0],[1]], K=[[0],[log(3)]], V=[[2],[6]]

• Unmasked row 0: weights [0.5,0.5], output 4.
• Row 1: weights [0.25,0.75], output 5.
• Causal row 0: weights [1,0], output 2.

Stable softmax keeps the graph

• Subtract the maximum allowed score before exponentiating.
• [1000,1001,1002] -> [-2,-1,0].
• Use (s-m).exp(), not Value(exp(s.data-m)).
• With scores [0,0] and values [2,6], score gradients are [-1,1].

A penalty is not exclusion

• Helper mask uses 0 and -1e9 for demonstration.
• Future score 1e12 remains dominant after that penalty.
• Actual causal forward normalizes only row[:i+1].
• Append constant zero weights for future positions.

Check output and gradient independence

With identity scalar projections and inputs [1,1e12]:

• First causal output must be 1.
• Its input gradients must be [1,0].
• Zeroing after full softmax can lose normalization and retain future dependence.

Cost and scope

• Scores and mixing: O(T^2*D) time.
• Dense projections add O(T*D^2).
• Score storage grows fourfold when sequence length doubles.
• V2 develops cumulative learned multi-head attention.