Scaled Dot-Product Attention

medium
Unit 11 · Transformers0/2 to milestone

Attention is the core operation of the Transformer:

Attention(Q,K,V)=softmax(QK⊤dk+M)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V

Implement attention(Q, K, V, causal=False) where:

  • Q has shape (seq_q, d_k), K has shape (seq_k, d_k), V has shape (seq_k, d_v)
  • If causal=True, position ii may only attend to positions ≤i\le i (mask future positions with −∞-\infty before softmax). Assume seq_q == seq_k in this case.

Return a tuple (output, weights) where output has shape (seq_q, d_v) and weights has shape (seq_q, seq_k).

#transformers#attention#numpy3 visible + 1 hidden tests
How LLMs Work· 7 of 19
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
Scaled Dot-Product Attention · TechGradify