Scaled Dot-Product Attention
~25 mincode completion
Implement scaled_dot_product_attention(Q, K, V).
Examples
Q=K=V identity: uniform attention, output is weighted avg of V
- Input
- scaled_dot_product_attention([[1, 0], [0, 1]], [[1, 0], [0, 1]], [[1, 0], [0, 1]])
- Output
- [[0.66976, 0.33024], [0.33024, 0.66976]]
Hints
Hint 1
Subtract the row max before exponentiating to keep the result stable.
Hint 2
Do not forget to scaling by sqrt dk. That step is easy to skip.
Requirements
Q: Query matrix of shape (n_q, d_k)K: Key matrix of shape (n_k, d_k)V: Value matrix of shape (n_k, d_v)Return Output matrix of shape (n_q, d_v).
Constraints
Allowed library: NumPy only
Time limit: 200 ms, Memory: 64 MB
Where this shows up
~25 min
••••••••••••••••
8 employers weight this skill
4 frontier labs, 3 AI product companies, 1 enterprise vendor. Top match scores 91.
Python
import numpy as np
def scaled_dot_product_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray) -> np.ndarray:
"""
Compute scaled dot-product attention.
Args:
Q: Query matrix of shape (n_q, d_k)
K: Key matrix of shape (n_k, d_k)
V: Value matrix of shape (n_k, d_v)
Returns:
Output matrix of shape (n_q, d_v).
"""
# YOUR CODE HERE
pass