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 — a near-deterministic weighting, from vector length alone.
Saturation reduces the softmax 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.
Division by \sqrt{d} keeps this scale independent of dimension.
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) weights = jax.nn.softmax(X, axis=-1)# A fully masked query (no valid key) would be a uniform average over# invalid values; zero those rows so no padded position leaks through weights = jnp.where(mask.any(-1, keepdims=True), weights, 0.0)return weights.reshape(shape)
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, so the valid lengths are (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')
The mask permits parallel training over all positions without exposing future tokens.
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:
The implementation combines scoring, masking, dropout on the weights, and value pooling:
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.
Attention as learned alignment
Machine translation
In 2014 neural machine translation, one fixed vector connected the encoder and decoder RNNs, restricting the representation of long sentences. 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 are efficient on current hardware: one matrix multiplication instead of an n \times m \times h tensor of activations. Learned projections followed by a dot product give the metric back, which gives the form used in the Transformer.
Recap
Scaled dot product is the standard scoring function: \mathbf{q}^\top\mathbf{k}/\sqrt{d} keeps score variance at 1, so the softmax avoids dimension-induced saturation and small score gradients.
Masking handles padding and causality with one primitive: overwrite invalid scores with the dtype’s most negative finite value.
DotProductAttention uses BMM, masked softmax, dropout, and a second BMM.
Additive attention introduced learned soft alignment for translation.
