Neural Networks & LLM Components · Medium

Scaled Dot-Product Attention

Compute stable masked attention from query, key, and value matrices.

Task

Implement scaled_dot_product_attention(query: list[list[float]], key: list[list[float]], value: list[list[float]], mask: list[list[bool]]) -> list[list[float]]. Compute scaled query-key scores, apply a boolean validity mask and row-wise stable softmax, then return the weighted sum of value rows.

Requirements

query and key share a positive feature width d_k; key and value have the same number of rows.
Divide every query-key dot product by sqrt(d_k).
A True mask entry is valid and a False entry is blocked.
Compute softmax independently for each query using only valid scores and a stability shift.
Return an all-zero value vector when a query row is fully masked.
Preserve an empty query matrix by returning an empty list.

Example

query = [[1.0, 0.0]]
key = [[1.0, 0.0], [0.0, 1.0]]
value = [[10.0, 0.0], [0.0, 20.0]]
mask = [[True, True]]

scaled_dot_product_attention(query, key, value, mask)
# [[6.69761549, 6.60476901]]