Scaled Dot-Product Attention - Transformer Architecture

Scaled Dot-Product Attention

What you'll build

The core operation of every Transformer: each query takes a weighted average of the values, with weights given by how well it matches each key. Return both the output and the attention weights.

Background

Attention is a differentiable lookup table. A query q_i is compared with every key k_j by a dot product; a softmax turns those scores into a probability distribution over positions; the output is the matching mixture of value vectors. Dividing by sqrt(d_k) keeps the scores at unit scale: for random vectors with unit-variance entries, q · k has variance d_k, and without the scaling the softmax saturates and its gradients vanish.

Formula

S   = Q @ K.T / sqrt(d_k)      # (seq_len, seq_len) scores
A   = softmax(S, axis=-1)      # each row sums to 1
out = A @ V                    # (seq_len, d_v)

• Q: queries, K: keys, V: values; d_k = Q.shape[-1]
• Row i of A holds query i's weights over all keys

Inputs

• Q: float array (seq_len, d_k)
• K: float array (seq_len, d_k)
• V: float array (seq_len, d_v)

Expected Return Format

A tuple (output, attention_weights):

• output: (seq_len, d_v)
• attention_weights: (seq_len, seq_len), non-negative, rows sum to 1

Edge cases and conventions

• Use a numerically stable softmax: subtract each row's max before np.exp.
• No mask here (every query sees every key). Masking is the Causal Attention Mask problem.

Constraints

NumPy only, fully vectorised (no Python loops over positions).

Follow-up

• Why sqrt(d_k) and not d_k? What happens to the gradients of a saturated softmax?
• Time and memory are O(seq_len^2 · d) and O(seq_len^2). Which later problems (Online Softmax, KV Cache) attack each?

Input: Q = K = np.array([[1.0, 0.0], [0.0, 1.0]]), V = np.array([[1.0, 2.0], [3.0, 4.0]])
Output: output = [[1.6605, 2.6605], [2.3395, 3.3395]], attention_weights = [[0.6698, 0.3302], [0.3302, 0.6698]]

Input: Q = np.zeros((3, 4)), any K, V = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
Output: every output row = [3.0, 4.0]

import numpy as np

def scaled_dot_product_attention(Q, K, V):
    """
    Compute scaled dot-product attention.

    Args:
        Q: Queries (seq_len, d_k)
        K: Keys (seq_len, d_k)
        V: Values (seq_len, d_v)

    Returns:
        (output, attention_weights):
            output: (seq_len, d_v)
            attention_weights: (seq_len, seq_len), each row sums to 1
    """
    # Your code here
    pass