The query term cancels in the softmax; the key-norm term cancels too when all key norms are equal. Keeping only the dot product is a modeling choice: a compatibility score that learned representations can shape freely.
For unit-variance entries, \operatorname{Var}(\mathbf{q}^\top\mathbf{k}) = d — scores grow with vector length for no semantic reason. Hence the Transformer’s scoring function:
Random queries and keys, attention over 64 candidates, entropy of the weight distribution as d grows:
Uniform over 64 keys ≈ 4.2 nats; scaled scores hold ≈ 3.7 nats at every d.
Unscaled: below 0.2 nats by d = 512 — top weight above 0.9, from vector length alone.
Saturation kills the gradient
The softmax Jacobian \partial \boldsymbol{\alpha} / \partial \mathbf{a} = \mathrm{diag}(\boldsymbol{\alpha}) - \boldsymbol{\alpha}\boldsymbol{\alpha}^\top tends to zero as \boldsymbol{\alpha} approaches one-hot:
for d in (4, 64, 512): _, jac_unscaled = attention_stats(d, scaled=False) _, jac_scaled = attention_stats(d, scaled=True)print(f'd = {d:3d}: Jacobian norm {jac_unscaled:.3f} (unscaled), 'f'{jac_scaled:.3f} (scaled)')
d = 4: Jacobian norm 0.258 (unscaled), 0.178 (scaled)
d = 64: Jacobian norm 0.218 (unscaled), 0.184 (scaled)
d = 512: Jacobian norm 0.103 (unscaled), 0.186 (scaled)
Scaled: constant at every dimension. Unscaled: decaying — about half by d = 512, still falling.
One division by \sqrt{d} removes the problem. Part of the definition, not a tuning trick.
Two reasons to mask
Padding — variable-length sequences share a minibatch:
Dive into Deep Learning
Learn to code <blank>
Hello world <blank> <blank>
Causality — a language model computes all positions in parallel, but query t must not see keys beyond t.
Same fix for both: overwrite invalid scores before the softmax.
Not with -10^6 (overflows float16), not with -\infty (a fully masked query → NaN). Use the dtype’s most negative finite value.
masked_softmax
def masked_softmax(X, valid_lens):"""Perform softmax operation by masking elements on the last axis."""# X: 3D tensor, valid_lens: 1D or 2D tensorif valid_lens isNone:return jax.nn.softmax(X, axis=-1) shape = X.shapeif valid_lens.ndim ==1: valid_lens = jnp.repeat(valid_lens, shape[1])else: valid_lens = valid_lens.reshape(-1) mask = jnp.arange(shape[-1])[None, :] < valid_lens[:, None]# Most negative finite score: exactly zero weight after the softmax,# at any precision, without the NaN risk of literal -inf X = jnp.where(mask, X.reshape(-1, shape[-1]), jnp.finfo(X.dtype).min)return jax.nn.softmax(X.reshape(shape), axis=-1)
Masking in action
Valid lengths 2 and 3 — weights beyond the prefix are exactly zero, rows still sum to one:
Query t sees keys 1, \ldots, t: valid lengths are just (1, 2, \ldots, n). The pattern is lower triangular:
scores = jax.random.normal(jax.random.key(2), (1, 6, 6))causal_lens = jnp.arange(1, 7)[None, :] # query t sees keys 1..td2l.show_heatmaps(masked_softmax(scores, causal_lens)[None], xlabel='Keys', ylabel='Queries')
On the attention side, this mask is the key difference between reading a sequence and being trainable, in parallel, to generate one.
Composing masks
The general interface is a boolean tensor — entry (i, j): may query i see key j? Requirements compose by AND, broadcasting does the bookkeeping: padding (\textrm{batch}, 1, \textrm{keys}) ∧ causal (1, \textrm{queries}, \textrm{keys}).
The composite must still leave every query at least one valid key — masks that are harmless alone can intersect to an empty row.
Batched attention
Attention over a minibatch is two batched matrix multiplications and one softmax:
Scoring, masking, dropout on the weights, value pooling — a dozen lines the rest of the book reuses:
class DotProductAttention(nnx.Module):"""Scaled dot product attention."""def__init__(self, dropout, rngs=None): rngs = nnx.Rngs(dropout=0) if rngs isNoneelse rngsself.dropout = nnx.Dropout(dropout, rngs=rngs)# Shape of queries: (batch_size, no. of queries, d)# Shape of keys: (batch_size, no. of key-value pairs, d)# Shape of values: (batch_size, no. of key-value pairs, value dimension)# Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)def__call__(self, queries, keys, values, valid_lens=None): d = queries.shape[-1]# Swap the last two dimensions of keys with keys.swapaxes(1, 2) scores = queries@(keys.swapaxes(1, 2)) / math.sqrt(d) attention_weights = masked_softmax(scores, valid_lens)# NNX idiom: return (output, weights); PyTorch stores weights on selfreturnself.dropout(attention_weights) @ values, attention_weights
Masked attention, verified
Valid lengths (2, 6): the stored weights vanish beyond the second and sixth key.
From alignment to attention
Where all of this came from
2014 machine translation: one fixed vector between encoder and decoder RNNs — long sentences don’t fit. Bahdanau, Cho & Bengio: let the decoder state query all encoder states, one fresh summary per output token — “jointly learning to align and translate”.
Learned weights behave like soft alignments — monotone, except where the languages reorder words.
Additive scoring
Decoder state and encoder states had different sizes, so the original score was a tiny MLP:
Dot products won on hardware: one matmul vs. an n \times m \times h tensor of activations. Learned projections + dot product give the metric back — exactly the Transformer’s form.
Recap
Scaled dot product is the scoring function: \mathbf{q}^\top\mathbf{k}/\sqrt{d} keeps score variance at 1, so the softmax neither saturates nor starves its gradient — we measured both.
Masking handles padding and causality with one primitive: overwrite invalid scores with the dtype’s most negative finite value.
