Skip to content

← Attention 101 step 2 of 5

Hard End-to-End

Self-Attention Layer

Self-attention lets each position in a sequence read information from other positions in that same sequence. It first transforms each input row into three vectors: a query for comparing positions, a key for being compared, and a value for the information to combine.

The query-key scores determine a weighted average of the value vectors. The output at a position therefore depends on both its query and the values available across the sequence.

Your task

Implement self_attention(X, W_Q, W_K, W_V) for floating-point matrices:

  • X: shape (seq_len, d_model), one input row per position.
  • W_Q, W_K, W_V: each has shape (d_model, d_k).

All dimensions are positive. The projected width d_k can differ from the input width d_model. Compute:

Q = X @ W_Q
K = X @ W_K
V = X @ W_V
scores = Q @ K.T / sqrt(d_k)
weights = softmax(scores, axis=-1)
output = weights @ V

Q, K and V have shape (seq_len, d_k). Scores have shape (seq_len, seq_len): rows are queries and columns are keys. Normalize each score row over the keys using a numerically stable softmax. Return the weighted values with shape (seq_len, d_k).

This is one attention head with no mask. Each position can attend to every position, including itself and later positions. There are no biases, dropout, residual connections, normalization layers or output projection.

Example

Let X = [[1], [2]] and all three weights be [[1]]. Then Q, K and V are all equal to X, and the score rows are [1, 2] and [2, 4]. Their softmax weights are approximately [0.268941, 0.731059] and [0.119203, 0.880797]. Mixing the values 1 and 2 gives output [[1.731059], [1.880797]].

Both positions favor the second value, but the second query favors it more strongly. The weights sum to one separately for each query.