Step-by-step answer
Core 60 min7 stepsImplement attention and verify it against PyTorch
Causal multi-head self-attention in about 40 lines of numpy, matched against PyTorch's scaled_dot_product_attention to within 1e-5 by a check script that also shows the permutation property, plus a measurement of what the √d_k scaling prevents. Without it, at d_k = 512, each softmax row puts 92% of its weight on one key.
Back to the exercise