Implement multi-head attention from scratch

CodingIntermediate⏱ 3 hoursDeep Learningattentionimplementationcausal masknumerical stability

Problem

Implement multi-head self-attention with a causal mask, and verify the mathematics.

Check that the attention matrix rows sum to one and that the causal mask is respected.
Verify permutation equivariance numerically when the positional encoding is removed.
Demonstrate softmax saturation by removing the $1/\sqrt{d_k}$ scaling.
Verify that rotary embeddings make the logits depend only on the position difference.

Implement it, then check the implementation against something you can verify independently.

Solve it here5 tests · Python runs in your browser

⌘/Ctrl + ↵ to run · limit 30s · draft saved in this browser

5 tests will run against your code. The first run downloads the Python runtime — several megabytes, once per browser.

Opening this before you have a real attempt on paper is the fastest way to feel like you understand something you cannot reproduce.

Problem	Field	Type	Difficulty	Time
When does attention dominate the MLP?	Deep Learning	Numerical	Beginner	30 min
Attention is permutation equivariant	Deep Learning	Proof	Intermediate	30 min
Implement backpropagation from scratch run	Deep Learning	Coding	Intermediate	3 hours
Implement kernel ridge regression run	Classical ML	Coding	Intermediate	1 hour
Implement power iteration and verify the convergence rate run	Foundations	Coding	Intermediate	1 hour