Multi-Head Self-Attention Forward Pass

hard
Checkpoint challenge
Solve this to show you know the material and jump straight to Unit 22. Everything before it counts as cleared.

Implement a full multi-head self-attention layer for one sequence X of shape (T, d):

  1. project: Q = X @ Wq, K = X @ Wk, V = X @ Wv (each d × d)
  2. split the last dimension into h heads of size d_h = d / h
  3. per head: softmax ⁣(QK⊤dh+M)V\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}} + M\right)V, where M is 0 or −∞-\infty from an optional causal mask
  4. concatenate heads and apply Wo

Implement mha(X, Wq, Wk, Wv, Wo, h, causal=False) returning the (T, d) output.

#attention#transformers2 visible + 1 hidden tests
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
Multi-Head Self-Attention Forward Pass · TechGradify