Scaled Dot Product Attention

The core operation of every transformer.

scores  = Q @ K^T / sqrt(d_k)
weights = softmax(scores, over the last axis)
output  = weights @ V

Shapes: Q is (n, d_k), K is (m, d_k), V is (m, d_v). The output is (n, d_v) and the weights are (n, m). Return them as (output, weights).

The sqrt(d_k) divisor matters. Without it the dot products grow with the dimension, softmax saturates, and the gradients vanish.

mask is either None or a boolean array broadcastable to (n, m), where False marks a position that must not be attended to. Set those scores to -inf before the softmax so they come out as exactly zero weight. A causal mask, which is what makes a decoder autoregressive, is just np.tril(np.ones((n, n), dtype=bool)).

Use a numerically stable softmax: subtract the row maximum before exponentiating. Watch the fully masked row case, where every score is -inf and max is also -inf, so -inf - -inf gives nan.

Function signature

def scaled_dot_product_attention(Q, K, V, mask=None):

Example

Input

Q = np.array([[1.0, 0.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0]])
V = np.array([[1.0], [2.0]])
scaled_dot_product_attention(Q, K, V)

Output

(array([[1.33023845]]), array([[0.66976155, 0.33023845]]))

Reasoning

The scores are [1, 0] / sqrt(2) = [0.707, 0]. Softmax gives weights of about [0.67, 0.33]. The output is 0.67 * 1 + 0.33 * 2, about 1.33.

Constraints and notes

• Scale by sqrt(d_k).
• Masked scores become -inf before the softmax.
• A fully masked row must not produce nan.