NNX Scaled Dot-Product Attention

Turn query-key scores into a weighted average

Scaled dot-product attention lets each query combine a set of value vectors. Compare a query with every key, convert those scores into nonnegative weights that sum to one, then average the corresponding values using those weights.

This exercise implements the array calculation directly. There is no NNX module, parameter initialization, or random sampling. A larger attention layer can learn projections that produce the queries, keys and values; those projections are outside this exercise.

Shapes and the three steps

Queries q have shape (T_q, d_k), keys k have shape (T_k, d_k), and values v have shape (T_k, d_v). Key and value rows correspond: row j of each describes the same source position.

1. Compute scores = (q @ k.T) / sqrt(d_k), shaped (T_q, T_k). Each row compares one query with all keys.
2. Apply softmax along the key axis, axis=-1. Each query's weights sum to one across the T_k keys, up to floating-point roundoff.
3. Compute weights @ v, shaped (T_q, d_v). Each output row is a weighted average of the value rows.

Query and key sequence lengths may differ, as they often do in cross-attention. The value width may also differ from the key width. Only the query/key widths and the key/value row counts must match. Softmax over axis 0 would normalize across queries instead, which is not the attention rule specified here.

Why divide by the square root of the key width?

A dot product sums d_k products. Under a simple model with independent, zero-mean, unit-variance query and key components, its variance is d_k. Dividing by sqrt(d_k) gives variance one under those assumptions. Learned activations need not satisfy them, so scaling is not a guarantee that every score distribution has unit variance.

Large score differences can make softmax concentrate most weight on one key and make some gradients small. There is no particular feature width at which an unscaled calculation necessarily fails. The scale controls score magnitude; stable softmax separately avoids unnecessary overflow by subtracting a row maximum before exponentiation.

A complete numerical example

import jax
import jax.numpy as jnp

def sdpa(q, k, v):
    scores = (q @ k.T) / jnp.sqrt(q.shape[-1])
    weights = jax.nn.softmax(scores, axis=-1)
    return weights @ v

q = jnp.array([[jnp.log(2.0)], [0.0]], dtype=jnp.float32)
k = jnp.array([[1.0], [0.0]], dtype=jnp.float32)
v = jnp.array([[3.0, 0.0], [0.0, 6.0]], dtype=jnp.float32)
print(sdpa(q, k, v))  # [[2.0, 2.0], [1.5, 3.0]]

The first query produces scores [log(2), 0]. Exponentiation gives relative weights [2, 1], so softmax assigns [2/3, 1/3]. The value average is [2, 2]. The zero query scores both keys equally, so it returns the ordinary average [1.5, 3].

A matching key need not dominate the average. For example, scores [0.5, 0, 0, 0] give the first key only about 35.5% of the weight, with about 21.5% for each other key.

Scope and masks

There is one attention head, no mask, and no dropout. Every query can attend to every key. The visualization below shows allowed connections, not the numerical softmax weights. Its full pattern matches this exercise; the other patterns preview masked attention.

When adding a mask in another exercise, setting a forbidden score to zero does not remove it: exp(0) is still one. A usual approach replaces forbidden logits with negative infinity before softmax, with at least one allowed finite score in every row. An entirely masked row needs explicit handling to avoid undefined normalization. Do not add a mask to the function requested here.

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.