Queries against keys

Scaled dot-product attention compares each query with every key in the same batch item. The feature dimension is shared, but the query and key sequence lengths can differ:

q:      (batch, q_len, d)
k:      (batch, k_len, d)
scores: (batch, q_len, k_len)

scores[b, i, j] = dot(q[b, i], k[b, j]) / sqrt(d)

Each output entry is a similarity score for one query/key pair. These are scores before masking or softmax, not probabilities.

Align the feature axes

For these matrix operands, @ sums over the last axis of the left operand and the second-to-last axis of the right. Transpose the last two key axes so both contracted dimensions have size d:

import torch

q = torch.tensor([[[1.0, 0.0], [0.0, 2.0]]])
k = torch.tensor([[[1.0, 1.0], [2.0, 0.0], [0.0, 3.0]]])
dots = q @ k.transpose(-2, -1)
print(dots.tolist())  # [[[1.0, 2.0, 0.0], [2.0, 0.0, 6.0]]]
print(list(dots.shape))  # [1, 2, 3]

For strided tensors, this transpose shares storage rather than copying key values. It still creates a view, and the resulting layout can affect the following multiplication's implementation and cost.

Negative axis numbers name the matrix axes even if an additional batch or head axis is inserted before them. Those leading dimensions must still be compatible for the intended multiplication.

Why divide by the square root of d?

Under a simple model where query and key components are independent, zero-mean and unit-variance, their dot product has variance d. Dividing by sqrt(d) gives variance one under those assumptions.

This scaling helps prevent larger feature dimensions from producing increasingly spread-out scores solely through summation. Large score gaps can make softmax highly concentrated and reduce some gradients, but saturation is not inevitable. Learned queries and keys need not satisfy the assumptions, so scaling does not guarantee a fixed score variance throughout training.

Attention combines values using learned weights

Attention lets each output position combine information from source positions. It uses three sets of vectors:

• A query for each output position.
• A key for each source position, compared with the query to get a score.
• A value for each source position, containing the information to combine.

Scaled dot-product attention has three steps: compute query-key scores, turn those scores into weights with softmax, then take a weighted sum of the values.

$$A=\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}+M\right),\qquad O=AV.$$

Here, d_k is the width of each query and key vector. The mask M is zero for allowed positions and negative infinity for blocked positions.

For a single head, Q has shape (query_length, key_width), K has shape (source_length, key_width), and V has shape (source_length, value_width). The weights A have shape (query_length, source_length). Softmax runs over source positions, so each query gets its own set of weights that sum to one. The output has shape (query_length, value_width).

For weights [0.25, 0.75] and scalar values [2, 6], the output is 0.25 * 2 + 0.75 * 6 = 5. Both values contribute; attention does not simply choose the value with the largest score.

Mask before softmax

A causal mask blocks future positions. A padding mask blocks positions added to make sequences the same length. Setting a blocked score to negative infinity makes its exponential zero, so it receives zero weight. Setting it to zero does not work: exp(0) = 1.

Each row needs at least one allowed, finite score. Applying ordinary softmax to an all-negative-infinity row produces undefined values. If you instead zero weights after softmax, the remaining weights no longer sum to one unless you normalise them again.

Why divide by the square root of the key width?

If query and key components are independent, with mean zero and variance one, the variance of their dot product is d_k. Dividing by sqrt(d_k) brings it back to one. This helps keep softmax from concentrating nearly all the weight on one source position, where its derivatives become small. Scaling helps under these assumptions; it does not prevent extreme weights for every possible input.

In self-attention, queries, keys, and values come from the same sequence. In cross-attention, queries come from one sequence and keys and values come from another.