Scaled Dot-Product Attention

Choose a mixture of values for each query

Scaled dot-product attention compares each query with a set of keys. Softmax converts those scores into nonnegative weights that sum to one across keys. The output for that query is the weighted average of the corresponding value vectors:

scores  = Q @ K.T / sqrt(d)
weights = softmax(scores, axis=-1)
output  = weights @ V

Keys determine the scores; values supply the vectors to combine. Each key must have a corresponding value, but a value can have a different feature count from a query or key.

This exercise receives queries, keys and values directly. It includes no learned projections, masking, dropout or output projection. In a self-attention layer these inputs usually come from projections of the same sequence. In cross-attention, queries come from one sequence and keys and values from another. The calculation here supports either.

Follow the shapes

All inputs are two-dimensional, representing one sequence per input and one attention head:

Apply softmax along the last axis so each query distributes weight across its keys. Normalizing along axis zero would distribute weight across queries instead, which is a different operation.

Why divide by the square root of the feature count?

Suppose query and key components are independent, have mean zero and variance one, and the component products are independent across features. A dot product sums d products, so its variance is d. Dividing the dot product by sqrt(d) gives variance one under these assumptions.

Large score differences can make softmax concentrate almost all its weight on one key, making its derivatives small. Scaling helps control that effect. Learned queries and keys need not satisfy the assumptions, so the scale is not a guarantee of unit variance or healthy gradients. Dividing by d would shrink the scores by an additional factor of sqrt(d) compared with the required formula.

A numerical example

Use two queries, three keys and values with two features:

import jax
import jax.numpy as jnp

q = jnp.eye(3)[:2]                          # (2, 3)
k = jnp.eye(3)                              # (3, 3)
v = jnp.array([[1., 0.], [0., 1.], [1., 1.]])
scores = q @ k.T / jnp.sqrt(3.0)           # (2, 3)
weights = jax.nn.softmax(scores, axis=-1)
out = weights @ v                         # (2, 2)

assert jnp.allclose(weights.sum(axis=-1), jnp.ones(2))
assert jnp.allclose(
    weights[0], jnp.array([0.471083, 0.264458, 0.264458]), atol=1e-6
)
assert jnp.allclose(
    out, jnp.array([[0.735541, 0.528917], [0.528917, 0.735541]]),
    atol=1e-6,
)

Query zero has score 1/sqrt(3) for key zero and zero for the other keys. Its first output feature is 0.471083 + 0.264458, because values zero and two have a 1 in that feature. Its second output feature is 0.264458 + 0.264458, from values one and two.

If all scores in a query row are equal, its output is the mean of the value vectors. Softmax is still necessary: multiplying unnormalized scores by values does not give the specified weighted average.

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.