Query-Key Attention Scores
~12 mincode completion
Implement attention_scores(Q, K) that returns the scaled score matrix.
Examples
2 queries, 2 keys, identity-like: scores on diagonal
- Input
- attention_scores([[1, 0], [0, 1]], [[1, 0], [0, 1]])
- Output
- [[0.70711, 0], [0, 0.70711]]
1 query attending to 3 keys
- Input
- attention_scores([[1, 1]], [[1, 0], [0, 1], [1, 1]])
- Output
- [[0.70711, 0.70711, 1.41421]]
d_k=3 scaling applied correctly
- Input
- attention_scores([[2, 0, 1]], [[1, 0, 0]])
- Output
- [[1.1547]]
Hints
Hint 1
Use a matrix product rather than nested loops, and check which operand transposes.
Hint 2
Do not forget to sqrt dk scaling. That step is easy to skip.
Requirements
Q: Query matrix, shape (T_q, d_k)K: Key matrix, shape (T_k, d_k)Return Score matrix, shape (T_q, T_k): Q @ K.T / sqrt(d_k)
Constraints
Allowed library: NumPy only
Time limit: 200 ms, Memory: 64 MB
Where this shows up
~12 min
••••••••••••••••
8 employers weight this skill
4 frontier labs, 3 AI product companies, 1 enterprise vendor. Top match scores 93.
Python
import numpy as np
def attention_scores(Q: np.ndarray, K: np.ndarray) -> np.ndarray:
"""
Compute scaled dot-product attention scores.
Args:
Q: Query matrix, shape (T_q, d_k)
K: Key matrix, shape (T_k, d_k)
Returns:
Score matrix, shape (T_q, T_k): Q @ K.T / sqrt(d_k)
"""
# YOUR CODE HERE
pass